From 75c9a65feb3c734258c5798b721f66893950dafa Mon Sep 17 00:00:00 2001
From: danijerez <danieljerezgarrido@gmail.com>
Date: Sat, 26 Oct 2024 16:35:38 +0200
Subject: [PATCH] add custom provider cerebras

---
 README.md                                        |  1 +
 src/libs/OpenAI/CustomProviders.cs               | 16 +++++++++++++++-
 .../OpenAI.IntegrationTests/CustomProvider.cs    |  1 +
 .../OpenAI.IntegrationTests/Tests.Helpers.cs     |  9 ++++++++-
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index cb46082d..8d04e8e7 100644
--- a/README.md
+++ b/README.md
@@ -187,6 +187,7 @@ using var api = CustomProviders.Perplexity("API_KEY");
 using var api = CustomProviders.SambaNova("API_KEY");
 using var api = CustomProviders.Mistral("API_KEY");
 using var api = CustomProviders.Codestral("API_KEY");
+using var api = CustomProviders.Cerebras("API_KEY");
 using var api = CustomProviders.Ollama();
 using var api = CustomProviders.LmStudio();
 ```
diff --git a/src/libs/OpenAI/CustomProviders.cs b/src/libs/OpenAI/CustomProviders.cs
index a65ed53a..f75434c8 100755
--- a/src/libs/OpenAI/CustomProviders.cs
+++ b/src/libs/OpenAI/CustomProviders.cs
@@ -75,7 +75,12 @@ public static class CustomProviders
     /// 
     /// </summary>
     public const string LmStudioBaseUrl = "http://localhost:1234/v1";
-    
+
+    /// <summary>
+    /// https://inference-docs.cerebras.ai/openai
+    /// </summary>
+    public const string CerebrasBaseUrl = "https://api.cerebras.ai/v1";
+
     /// <summary>
     /// Creates an API to use for GitHub Models: https://github.com/marketplace/models
     /// </summary>
@@ -210,4 +215,13 @@ public static OpenAiApi LmStudio(Uri? baseUri = null)
     {
         return new OpenAiApi(baseUri: baseUri ?? new Uri(LmStudioBaseUrl));
     }
+
+    /// <summary>
+    /// Create an API to use for Cerebras.
+    /// </summary>
+    /// <returns></returns>
+    public static OpenAiApi Cerebras(string apiKey)
+    {
+        return new OpenAiApi(apiKey, baseUri: new Uri(CerebrasBaseUrl));
+    }
 }
\ No newline at end of file
diff --git a/src/tests/OpenAI.IntegrationTests/CustomProvider.cs b/src/tests/OpenAI.IntegrationTests/CustomProvider.cs
index ebb50059..687d7677 100644
--- a/src/tests/OpenAI.IntegrationTests/CustomProvider.cs
+++ b/src/tests/OpenAI.IntegrationTests/CustomProvider.cs
@@ -18,4 +18,5 @@ public enum CustomProvider
     Mistral,
     Codestral,
     Hyperbolic,
+    Cerebras
 }
\ No newline at end of file
diff --git a/src/tests/OpenAI.IntegrationTests/Tests.Helpers.cs b/src/tests/OpenAI.IntegrationTests/Tests.Helpers.cs
index 5769f2af..379aaf06 100755
--- a/src/tests/OpenAI.IntegrationTests/Tests.Helpers.cs
+++ b/src/tests/OpenAI.IntegrationTests/Tests.Helpers.cs
@@ -135,7 +135,14 @@ internal static (OpenAiApi Api, string Model) GetAuthorizedChatApi(CustomProvide
             
             return pair;
         }
-        
+        if (customProvider == CustomProvider.Cerebras)
+        {
+            return (CustomProviders.Cerebras(apiKey:
+                Environment.GetEnvironmentVariable("CEREBRAS_API_KEY") ??
+                throw new AssertInconclusiveException("CEREBRAS_API_KEY environment variable is not found.")),
+                model ?? "llama3.1-70b");
+        }
+
         var apiKey =
             Environment.GetEnvironmentVariable("OPENAI_API_KEY") ??
             throw new AssertInconclusiveException("OPENAI_API_KEY environment variable is not found.");