Merge pull request #66 from vithursant/vithu-dev/cerebras-api

feat: Add Cerebras API via llama-index

Author: oulianov

.env.example

@@ -1,6 +1,7 @@
 MISTRAL_API_KEY=""
 OPENAI_API_KEY=""
 GROK_API_KEY=""
+CEREBRAS_API_KEY=""
 DISABLE_LLM="False"
 
 # AWS credentials

agent/config.py

@@ -17,6 +17,8 @@
     },
     "GROQ": {"groq:gemma-7b-it"},
     "ANTHROPIC": {"anthropic:claude-3-haiku-20240307"},
+    "CEREBRAS": {"cerebras:llama3.1-8b",
+                 "cerebras:llama3.1-70b"}
 }
 
 

agent/llm.py

@@ -32,3 +32,8 @@ def get_client(model_str):
         from llama_index.llms.bedrock import Bedrock
 
         return Bedrock(model=model_name)
+    elif provider == "cerebras":
+        from llama_index.llms.cerebras import Cerebras
+
+        return Cerebras(model=model_name)
+        

eval/game.py

@@ -47,7 +47,10 @@ def verify_provider_name(self):
             assert (
                 os.environ.get("MISTRAL_API_KEY") is not None
             ), "Mistral API key not set"
-
+        if self.model.startswith("cerebras"):
+            assert (
+                os.environ.get("CEREBRAS_API_KEY") is not None
+            ), "Cerebras API key not set"
 
 class Player1(Player):
     def __init__(

requirements.txt

@@ -53,3 +53,4 @@ llama-index-llms-anthropic
 llama-index-llms-groq
 llama-index-llms-ollama
 llama-index-llms-bedrock
+llama-index-llms-cerebras