chore: support more LLMs provider

ldclabs · Feb 10, 2025 · 35bba07 · 35bba07
1 parent 02f1655
commit 35bba07
Show file tree

Hide file tree

Showing 8 changed files with 39 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/agents/anda_bot/nitro_enclave/Config.toml b/agents/anda_bot/nitro_enclave/Config.toml
@@ -7,10 +7,13 @@ token_ledgers = []
 [llm]
 # DeepSeek API key
 deepseek_api_key = "sk-0**1a2b"
+deepseek_endpoint = "https://api.deepseek.com"
+deepseek_model = "deepseek-chat"
 cohere_api_key = "Y4f**CXd"
 cohere_embedding_model = "embed-multilingual-v3.0"
 # Or use OpenAI
 openai_api_key = ""
+openai_endpoint = ""
 openai_embedding_model = ""
 openai_completion_model = ""
 

diff --git a/agents/anda_bot/src/config.rs b/agents/anda_bot/src/config.rs
@@ -14,12 +14,18 @@ pub struct Llm {
     #[serde(default)]
     pub deepseek_api_key: String,
     #[serde(default)]
+    pub deepseek_endpoint: String,
+    #[serde(default)]
+    pub deepseek_model: String,
+    #[serde(default)]
     pub cohere_api_key: String,
     #[serde(default)]
     pub cohere_embedding_model: String,
     #[serde(default)]
     pub openai_api_key: String,
     #[serde(default)]
+    pub openai_endpoint: String,
+    #[serde(default)]
     pub openai_embedding_model: String,
     #[serde(default)]
     pub openai_completion_model: String,

diff --git a/agents/anda_bot/src/main.rs b/agents/anda_bot/src/main.rs
@@ -376,12 +376,30 @@ fn connect_model(cfg: &config::Llm) -> Result<Model, BoxError> {
                     .embedding_model(&cfg.cohere_embedding_model),
             ),
             Arc::new(
-                deepseek::Client::new(&cfg.deepseek_api_key)
-                    .completion_model(deepseek::DEEKSEEK_V3),
+                deepseek::Client::new(
+                    &cfg.deepseek_api_key,
+                    if cfg.deepseek_endpoint.is_empty() {
+                        None
+                    } else {
+                        Some(cfg.deepseek_endpoint.clone())
+                    },
+                )
+                .completion_model(if cfg.deepseek_model.is_empty() {
+                    deepseek::DEEKSEEK_V3
+                } else {
+                    &cfg.deepseek_model
+                }),
             ),
         ))
     } else {
-        let cli = openai::Client::new(&cfg.openai_api_key);
+        let cli = openai::Client::new(
+            &cfg.openai_api_key,
+            if cfg.deepseek_endpoint.is_empty() {
+                None
+            } else {
+                Some(cfg.deepseek_endpoint.clone())
+            },
+        );
         Ok(Model::new(
             Arc::new(cli.embedding_model(&cfg.openai_embedding_model)),
             Arc::new(cli.completion_model(&cfg.openai_completion_model)),

diff --git a/anda_engine/Cargo.toml b/anda_engine/Cargo.toml
@@ -3,7 +3,7 @@ name = "anda_engine"
 description = "Agents engine for Anda -- a framework for AI agent development."
 repository = "https://github.com/ldclabs/anda/tree/main/anda_engine"
 publish = true
-version = "0.3.5"
+version = "0.3.6"
 edition.workspace = true
 keywords.workspace = true
 categories.workspace = true

diff --git a/anda_engine/src/extension/attention.rs b/anda_engine/src/extension/attention.rs
@@ -434,7 +434,7 @@ mod tests {
         dotenv::dotenv().ok();
 
         let api_key = std::env::var("DEEPSEEK_API_KEY").expect("DEEPSEEK_API_KEY is not set");
-        let client = Client::new(&api_key);
+        let client = Client::new(&api_key, None);
         let model = client.completion_model(DEEKSEEK_V3);
         let attention = Attention::default();
         let res = attention

diff --git a/anda_engine/src/model/deepseek.rs b/anda_engine/src/model/deepseek.rs
@@ -38,17 +38,17 @@ impl Client {
     ///
     /// # Returns
     /// Configured DeepSeek client instance
-    pub fn new(api_key: &str) -> Self {
+    pub fn new(api_key: &str, endpoint: Option<String>) -> Self {
         Self {
-            endpoint: DEEKSEEK_API_BASE_URL.to_string(),
+            endpoint: endpoint.unwrap_or_else(|| DEEKSEEK_API_BASE_URL.to_string()),
             http: reqwest::Client::builder()
                 .use_rustls_tls()
                 .https_only(true)
                 .http2_keep_alive_interval(Some(Duration::from_secs(25)))
                 .http2_keep_alive_timeout(Duration::from_secs(15))
                 .http2_keep_alive_while_idle(true)
                 .connect_timeout(Duration::from_secs(10))
-                .timeout(Duration::from_secs(60))
+                .timeout(Duration::from_secs(120))
                 .user_agent(APP_USER_AGENT)
                 .default_headers({
                     let mut headers = reqwest::header::HeaderMap::new();
@@ -332,7 +332,7 @@ mod tests {
         println!("Character path: {}", character_path);
         let character = std::fs::read_to_string(character_path).expect("Character file not found");
         let character = Character::from_toml(&character).expect("Character should parse");
-        let client = Client::new(&api_key);
+        let client = Client::new(&api_key, None);
         let now = Instant::now();
         let model = client.completion_model(DEEKSEEK_V3);
         let req = character.to_request("I am Yan, glad to see you".into(), Some("Yan".into()));

diff --git a/anda_engine/src/model/openai.rs b/anda_engine/src/model/openai.rs
@@ -90,9 +90,9 @@ impl Client {
     ///
     /// # Arguments
     /// * `api_key` - OpenAI API key for authentication
-    pub fn new(api_key: &str) -> Self {
+    pub fn new(api_key: &str, endpoint: Option<String>) -> Self {
         Self {
-            endpoint: OPENAI_API_BASE_URL.to_string(),
+            endpoint: endpoint.unwrap_or_else(|| OPENAI_API_BASE_URL.to_string()),
             http: reqwest::Client::builder()
                 .use_rustls_tls()
                 .https_only(true)