feat: Add Cerebras AI

efugier · Dec 5, 2024 · 0c2bea1 · 0c2bea1
1 parent 4a9e169
commit 0c2bea1
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -255,6 +255,11 @@ api_key = "<yet_another_api_key>"
 url = "https://api.anthropic.com/v1/messages"
 default_model = "claude-3-opus-20240229"
 version = "2023-06-01"  # anthropic API version, see https://docs.anthropic.com/en/api/versioning
+
+[cerebras]
+api_key = "<your_api_key>"
+default_model = "llama3.1-70b"
+url = "https://api.cerebras.ai/v1/chat/completions"
 ```
 
 `prompts.toml`

diff --git a/src/config/api.rs b/src/config/api.rs
@@ -21,6 +21,7 @@ pub enum Api {
     Mistral,
     Openai,
     AzureOpenai,
+    Cerebras,
 }
 
 impl FromStr for Api {
@@ -34,6 +35,7 @@ impl FromStr for Api {
             "mistral" => Ok(Api::Mistral),
             "groq" => Ok(Api::Groq),
             "anthropic" => Ok(Api::Anthropic),
+            "cerebras" => Ok(Api::Cerebras),
             _ => Err(()),
         }
     }
@@ -48,6 +50,7 @@ impl ToString for Api {
             Api::Mistral => "mistral".to_string(),
             Api::Groq => "groq".to_string(),
             Api::Anthropic => "anthropic".to_string(),
+            Api::Cerebras => "cerebras".to_string(),
             v => panic!(
                 "{:?} is not implemented, use one among {:?}",
                 v,
@@ -178,6 +181,17 @@ impl ApiConfig {
             timeout_seconds: None,
         }
     }
+
+    pub(super) fn cerebras() -> Self {
+        ApiConfig {
+            api_key_command: None,
+            api_key: None,
+            url: String::from("https://api.cerebras.ai/v1/chat/completions"),
+            default_model: Some(String::from("llama3.1-70b")),
+            version: None,
+            timeout_seconds: None,
+        }
+    }
 }
 
 pub(super) fn api_keys_path() -> PathBuf {
@@ -192,6 +206,7 @@ pub(super) fn generate_api_keys_file() -> std::io::Result<()> {
     api_config.insert(Api::Mistral.to_string(), ApiConfig::mistral());
     api_config.insert(Api::Groq.to_string(), ApiConfig::groq());
     api_config.insert(Api::Anthropic.to_string(), ApiConfig::anthropic());
+    api_config.insert(Api::Cerebras.to_string(), ApiConfig::cerebras());
 
     // Default, should override one of the above
     api_config.insert(Prompt::default().api.to_string(), ApiConfig::default());

diff --git a/src/text/api_call.rs b/src/text/api_call.rs
@@ -49,7 +49,7 @@ pub fn post_prompt_and_get_answer(
         .expect("Unable to initialize HTTP client");
 
     let prompt_format = match prompt.api {
-        Api::Ollama | Api::Openai | Api::AzureOpenai | Api::Mistral | Api::Groq => {
+        Api::Ollama | Api::Openai | Api::AzureOpenai | Api::Mistral | Api::Groq | Api::Cerebras => {
             PromptFormat::OpenAi(OpenAiPrompt::from(prompt.clone()))
         }
         Api::Anthropic => PromptFormat::Anthropic(AnthropicPrompt::from(prompt.clone())),
@@ -61,9 +61,15 @@ pub fn post_prompt_and_get_answer(
         .header("Content-Type", "application/json")
         .json(&prompt_format);
 
+    // https://stackoverflow.com/questions/77862683/rust-reqwest-cant-make-a-request
+    let request = match prompt.api {
+        Api::Cerebras => request.header("User-Agent", "CUSTOM_NAME/1.0"),
+        _ => request,
+    };
+
     // Add auth if necessary
     let request = match prompt.api {
-        Api::Openai | Api::Mistral | Api::Groq => request.header(
+        Api::Openai | Api::Mistral | Api::Groq | Api::Cerebras => request.header(
             "Authorization",
             &format!("Bearer {}", &api_config.get_api_key()),
         ),
@@ -81,7 +87,7 @@ pub fn post_prompt_and_get_answer(
 
     let response_text: String = match prompt.api {
         Api::Ollama => handle_api_response::<OllamaResponse>(request.send()?),
-        Api::Openai | Api::AzureOpenai | Api::Mistral | Api::Groq => {
+        Api::Openai | Api::AzureOpenai | Api::Mistral | Api::Groq | Api::Cerebras => {
             handle_api_response::<OpenAiResponse>(request.send()?)
         }
         Api::Anthropic => handle_api_response::<AnthropicResponse>(request.send()?),