Merge pull request #110 from ktpedre/cerebras

Add support for Cerebras to litellm_caller

Author: garland3

.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

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)
+        return "\n".join(summary_parts)

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"
+    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"