diff --git a/.env.example b/.env.example index 9cde796..bf3130b 100644 --- a/.env.example +++ b/.env.example @@ -34,6 +34,7 @@ OPENAI_API_KEY=sk-pro ANTHROPIC_API_KEY=your_anthropic_api_key_here GOOGLE_API_KEY=your_google_api_key_here OPENROUTER_API_KEY=sk-or +CEREBRAS_API_KEY=your_cerebras_api_key_here # Banner system configuration diff --git a/backend/modules/llm/litellm_caller.py b/backend/modules/llm/litellm_caller.py index abfa1a0..b81ef58 100644 --- a/backend/modules/llm/litellm_caller.py +++ b/backend/modules/llm/litellm_caller.py @@ -57,6 +57,8 @@ def _get_litellm_model_name(self, model_name: str) -> str: return f"anthropic/{model_id}" elif "google" in model_config.model_url: return f"google/{model_id}" + elif "cerebras" in model_config.model_url: + return f"cerebras/{model_id}" else: # For custom endpoints, use the model_id directly return model_id @@ -95,10 +97,12 @@ def _get_model_kwargs(self, model_name: str, temperature: Optional[float] = None os.environ["ANTHROPIC_API_KEY"] = api_key elif "google" in model_config.model_url: os.environ["GOOGLE_API_KEY"] = api_key + elif "cerebras" in model_config.model_url: + os.environ["CEREBRAS_API_KEY"] = api_key # Set custom API base for non-standard endpoints if hasattr(model_config, 'model_url') and model_config.model_url: - if not any(provider in model_config.model_url for provider in ["openrouter", "api.openai.com", "api.anthropic.com"]): + if not any(provider in model_config.model_url for provider in ["openrouter", "api.openai.com", "api.anthropic.com", "api.cerebras.ai"]): kwargs["api_base"] = model_config.model_url # Handle extra headers with environment variable expansion @@ -339,4 +343,4 @@ def _format_rag_metadata(self, metadata) -> str: summary_parts.append(f" • ... and {remaining} more document(s)") summary_parts.append(f" **Retrieval Method:** {metadata.retrieval_method}") - return "\n".join(summary_parts) \ No newline at end of file + return "\n".join(summary_parts) diff --git a/config/overrides/llmconfig.yml b/config/overrides/llmconfig.yml index c529249..5e6d164 100644 --- a/config/overrides/llmconfig.yml +++ b/config/overrides/llmconfig.yml @@ -38,4 +38,9 @@ models: model_url: "https://api.openai.com/v1/chat/completions" model_name: "gpt-4.1-nano" api_key: "${OPENAI_API_KEY}" - compliance_level: "External" \ No newline at end of file + compliance_level: "External" + cerebras-gpt-oss-120b: + model_url: "https://api.cerebras.ai/v1" + model_name: "gpt-oss-120b" + api_key: "${CEREBRAS_API_KEY}" + compliance_level: "External"