diff --git a/llm/client.py b/llm/client.py index 9f21798..e27fc8e 100644 --- a/llm/client.py +++ b/llm/client.py @@ -19,15 +19,39 @@ def __init__(self, model: str = "gpt-4.1", temperature: float = 0.2): self.temperature = temperature self.client = None - api_key = os.environ.get("OPENAI_API_KEY") - if OpenAI and api_key: - self.client = OpenAI(api_key=api_key) - self.enabled = True - else: + # Select provider via environment variable. Supported: "openai" (default), "cerebras". + provider = os.environ.get("LLM_PROVIDER", "openai").lower() + + if not OpenAI: self.enabled = False - if not OpenAI: - logger.warning("openai package not installed; LLM calls are mocked") + logger.warning("openai package not installed; LLM calls are mocked") + return + + if provider == "cerebras": + api_key = os.environ.get("CEREBRAS_API_KEY") + base_url = os.environ.get("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1") + if api_key: + try: + self.client = OpenAI(base_url=base_url, api_key=api_key) + self.enabled = True + except Exception: + logger.exception("failed to initialize Cerebras LLM client") + self.enabled = False + else: + self.enabled = False + logger.warning("CEREBRAS_API_KEY not set; LLM calls are mocked") + else: + # default: OpenAI + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + try: + self.client = OpenAI(api_key=api_key) + self.enabled = True + except Exception: + logger.exception("failed to initialize OpenAI LLM client") + self.enabled = False else: + self.enabled = False logger.warning("OPENAI_API_KEY not set; LLM calls are mocked") def _mock_response(self, prompt: str) -> str: