Merge branch 'aiscriptdev:main' into feat-integrate-minijinja-as-template-engine

luojiyin1987 · web-flow · commit 40c71ef48aea · 2025-03-30T10:41:31.000+08:00
diff --git a/README.md b/README.md
@@ -113,6 +113,28 @@ AIScript excels in these scenarios:
 
 Check out the [examples](./examples) directory for more sample code.
 
+## Supported AI Models
+
+AIScript supports the following AI models:
+
+- [x] OpenAI ((uses `OPENAI_API_KEY` environment variable by default))
+- [x] DeepSeek
+- [ ] Anthropic
+
+Configuration by `project.toml`:
+
+```toml
+# use OpenAI
+[ai.openai]
+api_key = "YOUR_API_KEY"
+model = "gpt-3.5-turbo"
+
+# or use DeepSeek
+[ai.deepseek]
+api_key = "YOUR_API_KEY"
+model = "deepseek-chat"
+```
+
 ## Roadmap
 
 See our [roadmap](https://aiscript.dev/guide/contribution/roadmap) for upcoming features and improvements.
diff --git a/aiscript-directive/src/validator/format.rs b/aiscript-directive/src/validator/format.rs
@@ -54,7 +54,7 @@ mod uscc {
     });
 
     static USCC_BASE_CHARS: LazyLock<HashMap<char, u8>> = LazyLock::new(|| {
-        let mut base_chars = HashMap::with_capacity(17);
+        let mut base_chars = HashMap::with_capacity(31);
 
         [
             '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G',
diff --git a/aiscript-runtime/src/config/mod.rs b/aiscript-runtime/src/config/mod.rs
@@ -3,6 +3,7 @@ use std::{env, fmt::Display, fs, ops::Deref, path::Path, sync::OnceLock};
 use auth::AuthConfig;
 use serde::Deserialize;
 
+use aiscript_vm::AiConfig;
 use db::DatabaseConfig;
 pub use sso::{SsoConfig, get_sso_fields};
 
@@ -64,6 +65,8 @@ impl AsRef<str> for EnvString {
 
 #[derive(Debug, Deserialize, Default)]
 pub struct Config {
+    #[serde(default)]
+    pub ai: Option<AiConfig>,
     #[serde(default)]
     pub database: DatabaseConfig,
     #[serde(default)]
@@ -116,9 +119,9 @@ impl Config {
         }
     }
 
-    pub fn load(path: &str) -> &Config {
+    pub fn load() -> &'static Config {
         CONFIG.get_or_init(|| {
-            Config::new(path).unwrap_or_else(|e| {
+            Config::new("project.toml").unwrap_or_else(|e| {
                 eprintln!("Error loading config file: {}", e);
                 Config::default()
             })
diff --git a/aiscript-runtime/src/endpoint.rs b/aiscript-runtime/src/endpoint.rs
@@ -500,8 +500,13 @@ impl Future for RequestProcessor {
                     let redis_connection = self.endpoint.redis_connection.clone();
                     let handle: JoinHandle<Result<ReturnValue, VmError>> =
                         task::spawn_blocking(move || {
-                            let mut vm =
-                                Vm::new(pg_connection, sqlite_connection, redis_connection);
+                            let ai_config = Config::load().ai.clone();
+                            let mut vm = Vm::new(
+                                pg_connection,
+                                sqlite_connection,
+                                redis_connection,
+                                ai_config,
+                            );
                             if let Some(fields) = sso_fields {
                                 vm.inject_sso_instance(fields);
                             }
diff --git a/aiscript-vm/src/ai/agent.rs b/aiscript-vm/src/ai/agent.rs
@@ -8,7 +8,6 @@ use openai_api_rs::v1::{
         ChatCompletionMessage, ChatCompletionMessageForResponse, ChatCompletionRequest, Content,
         MessageRole, Tool, ToolCall, ToolChoiceType, ToolType,
     },
-    common::GPT3_5_TURBO,
     types::{self, FunctionParameters, JSONSchemaDefine},
 };
 use tokio::runtime::Handle;
@@ -278,6 +277,8 @@ pub async fn _run_agent<'gc>(
     mut agent: Gc<'gc, Agent<'gc>>,
     args: Vec<Value<'gc>>,
 ) -> Value<'gc> {
+    use super::default_model;
+
     let message = args[0];
     let debug = args[1].as_boolean();
     let mut history = Vec::new();
@@ -288,11 +289,11 @@ pub async fn _run_agent<'gc>(
         tool_calls: None,
         tool_call_id: None,
     });
-    let mut client = super::openai_client();
+    let mut client = super::openai_client(state.ai_config.as_ref());
     loop {
         let mut messages = vec![agent.get_instruction_message()];
         messages.extend(history.clone());
-        let mut req = ChatCompletionRequest::new(GPT3_5_TURBO.to_string(), messages);
+        let mut req = ChatCompletionRequest::new(default_model(state.ai_config.as_ref()), messages);
         let tools = agent.get_tools();
         if !tools.is_empty() {
             req = req
diff --git a/aiscript-vm/src/ai/mod.rs b/aiscript-vm/src/ai/mod.rs
@@ -4,13 +4,84 @@ mod prompt;
 use std::env;
 
 pub use agent::{Agent, run_agent};
-use openai_api_rs::v1::api::OpenAIClient;
+use openai_api_rs::v1::{api::OpenAIClient, common::GPT3_5_TURBO};
 pub use prompt::{PromptConfig, prompt_with_config};
 
+use serde::Deserialize;
+
+const DEEPSEEK_API_ENDPOINT: &str = "https://api.deepseek.com/v1";
+const DEEPSEEK_V3: &str = "deepseek-chat";
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum AiConfig {
+    #[serde(rename = "openai")]
+    OpenAI(ModelConfig),
+    #[serde(rename = "anthropic")]
+    Anthropic(ModelConfig),
+    #[serde(rename = "deepseek")]
+    DeepSeek(ModelConfig),
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelConfig {
+    pub api_key: String,
+    pub model: Option<String>,
+}
+
+impl AiConfig {
+    pub(crate) fn take_model(&mut self) -> Option<String> {
+        match self {
+            Self::OpenAI(ModelConfig { model, .. }) => model.take(),
+            Self::Anthropic(ModelConfig { model, .. }) => model.take(),
+            Self::DeepSeek(ModelConfig { model, .. }) => model.take(),
+        }
+    }
+
+    pub(crate) fn set_model(&mut self, m: String) {
+        match self {
+            Self::OpenAI(ModelConfig { model, .. }) => model.replace(m),
+            Self::Anthropic(ModelConfig { model, .. }) => model.replace(m),
+            Self::DeepSeek(ModelConfig { model, .. }) => model.replace(m),
+        };
+    }
+}
+
 #[allow(unused)]
-pub(crate) fn openai_client() -> OpenAIClient {
-    OpenAIClient::builder()
-        .with_api_key(env::var("OPENAI_API_KEY").unwrap().to_string())
-        .build()
-        .unwrap()
+pub(crate) fn openai_client(config: Option<&AiConfig>) -> OpenAIClient {
+    match config {
+        None => OpenAIClient::builder()
+            .with_api_key(env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set"))
+            .build()
+            .unwrap(),
+        Some(AiConfig::OpenAI(model_config)) => {
+            let api_key = if model_config.api_key.is_empty() {
+                env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set")
+            } else {
+                model_config.api_key.clone()
+            };
+            OpenAIClient::builder()
+                .with_api_key(api_key)
+                .build()
+                .unwrap()
+        }
+        Some(AiConfig::DeepSeek(ModelConfig { api_key, .. })) => OpenAIClient::builder()
+            .with_endpoint(DEEPSEEK_API_ENDPOINT)
+            .with_api_key(api_key)
+            .build()
+            .unwrap(),
+        Some(AiConfig::Anthropic(_)) => unimplemented!("Anthropic API not yet supported"),
+    }
+}
+
+pub(crate) fn default_model(config: Option<&AiConfig>) -> String {
+    match config {
+        None => GPT3_5_TURBO.to_string(),
+        Some(AiConfig::OpenAI(ModelConfig { model, .. })) => {
+            model.clone().unwrap_or(GPT3_5_TURBO.to_string())
+        }
+        Some(AiConfig::DeepSeek(ModelConfig { model, .. })) => {
+            model.clone().unwrap_or(DEEPSEEK_V3.to_string())
+        }
+        Some(AiConfig::Anthropic(_)) => unimplemented!("Anthropic API not yet supported"),
+    }
 }
diff --git a/aiscript-vm/src/ai/prompt.rs b/aiscript-vm/src/ai/prompt.rs
@@ -1,9 +1,11 @@
 use openai_api_rs::v1::common::GPT3_5_TURBO;
 use tokio::runtime::Handle;
 
+use super::{AiConfig, ModelConfig, default_model};
+
 pub struct PromptConfig {
     pub input: String,
-    pub model: Option<String>,
+    pub ai_config: Option<AiConfig>,
     pub max_tokens: Option<i64>,
     pub temperature: Option<f64>,
     pub system_prompt: Option<String>,
@@ -13,27 +15,42 @@ impl Default for PromptConfig {
     fn default() -> Self {
         Self {
             input: String::new(),
-            model: Some(GPT3_5_TURBO.to_string()),
+            ai_config: Some(AiConfig::OpenAI(ModelConfig {
+                api_key: Default::default(),
+                model: Some(GPT3_5_TURBO.to_string()),
+            })),
             max_tokens: Default::default(),
             temperature: Default::default(),
             system_prompt: Default::default(),
         }
     }
 }
 
+impl PromptConfig {
+    fn take_model(&mut self) -> String {
+        self.ai_config
+            .as_mut()
+            .and_then(|config| config.take_model())
+            .unwrap_or_else(|| default_model(self.ai_config.as_ref()))
+    }
+
+    pub(crate) fn set_model(&mut self, model: String) {
+        if let Some(config) = self.ai_config.as_mut() {
+            config.set_model(model);
+        }
+    }
+}
+
 #[cfg(feature = "ai_test")]
 async fn _prompt_with_config(config: PromptConfig) -> String {
     return format!("AI: {}", config.input);
 }
 
 #[cfg(not(feature = "ai_test"))]
 async fn _prompt_with_config(mut config: PromptConfig) -> String {
-    use openai_api_rs::v1::{
-        chat_completion::{self, ChatCompletionRequest},
-        common::GPT3_5_TURBO,
-    };
-
-    let mut client = super::openai_client();
+    use openai_api_rs::v1::chat_completion::{self, ChatCompletionRequest};
+    let mut client = super::openai_client(config.ai_config.as_ref());
+    let model = config.take_model();
 
     // Create system message if provided
     let mut messages = Vec::new();
@@ -57,13 +74,7 @@ async fn _prompt_with_config(mut config: PromptConfig) -> String {
     });
 
     // Build the request
-    let mut req = ChatCompletionRequest::new(
-        config
-            .model
-            .take()
-            .unwrap_or_else(|| GPT3_5_TURBO.to_string()),
-        messages,
-    );
+    let mut req = ChatCompletionRequest::new(model, messages);
 
     if let Some(max_tokens) = config.max_tokens {
         req.max_tokens = Some(max_tokens);
diff --git a/aiscript-vm/src/lib.rs b/aiscript-vm/src/lib.rs
@@ -16,6 +16,7 @@ use std::collections::HashMap;
 use std::fmt::Display;
 use std::ops::Deref;
 
+pub use ai::AiConfig;
 use aiscript_arena::Collect;
 use aiscript_arena::Mutation;
 pub(crate) use aiscript_lexer as lexer;
diff --git a/aiscript-vm/src/vm/mod.rs b/aiscript-vm/src/vm/mod.rs
@@ -6,6 +6,7 @@ pub use state::State;
 
 use crate::{
     ReturnValue, Value,
+    ai::AiConfig,
     ast::ChunkId,
     builtins, stdlib,
     string::{InternedString, InternedStringSet},
@@ -35,7 +36,7 @@ impl Display for VmError {
 
 impl Default for Vm {
     fn default() -> Self {
-        Self::new(None, None, None)
+        Self::new(None, None, None, None)
     }
 }
 
@@ -48,13 +49,15 @@ impl Vm {
         pg_connection: Option<PgPool>,
         sqlite_connection: Option<SqlitePool>,
         redis_connection: Option<redis::aio::MultiplexedConnection>,
+        ai_config: Option<AiConfig>,
     ) -> Self {
         let mut vm = Vm {
             arena: Arena::<Rootable![State<'_>]>::new(|mc| {
                 let mut state = State::new(mc);
                 state.pg_connection = pg_connection;
                 state.sqlite_connection = sqlite_connection;
                 state.redis_connection = redis_connection;
+                state.ai_config = ai_config;
                 state
             }),
         };
diff --git a/aiscript-vm/src/vm/state.rs b/aiscript-vm/src/vm/state.rs
@@ -15,7 +15,7 @@ use sqlx::{PgPool, SqlitePool};
 
 use crate::{
     NativeFn, OpCode, ReturnValue, Value,
-    ai::{self, PromptConfig},
+    ai::{self, AiConfig, PromptConfig},
     ast::{ChunkId, Visibility},
     builtins::BuiltinMethods,
     module::{ModuleKind, ModuleManager, ModuleSource},
@@ -110,6 +110,7 @@ pub struct State<'gc> {
     pub pg_connection: Option<PgPool>,
     pub sqlite_connection: Option<SqlitePool>,
     pub redis_connection: Option<redis::aio::MultiplexedConnection>,
+    pub ai_config: Option<AiConfig>,
 }
 
 unsafe impl Collect for State<'_> {
@@ -152,6 +153,7 @@ impl<'gc> State<'gc> {
             pg_connection: None,
             sqlite_connection: None,
             redis_connection: None,
+            ai_config: None,
         }
     }
 
@@ -1011,15 +1013,21 @@ impl<'gc> State<'gc> {
                 let result = match value {
                     // Simple string case
                     Value::String(s) => {
-                        let input = s.to_str().unwrap().to_string();
-                        ai::prompt_with_config(PromptConfig {
-                            input,
+                        let mut config = PromptConfig {
+                            input: s.to_str().unwrap().to_string(),
                             ..Default::default()
-                        })
+                        };
+                        if let Some(ai_cfg) = &self.ai_config {
+                            config.ai_config = Some(ai_cfg.clone());
+                        }
+                        ai::prompt_with_config(config)
                     }
                     // Object config case
                     Value::Object(obj) => {
                         let mut config = PromptConfig::default();
+                        if let Some(ai_cfg) = &self.ai_config {
+                            config.ai_config = Some(ai_cfg.clone());
+                        }
                         let obj_ref = obj.borrow();
 
                         // Extract input (required)
@@ -1037,7 +1045,7 @@ impl<'gc> State<'gc> {
                         if let Some(Value::String(model)) =
                             obj_ref.fields.get(&self.intern(b"model"))
                         {
-                            config.model = Some(model.to_str().unwrap().to_string());
+                            config.set_model(model.to_str().unwrap().to_string());
                         }
 
                         // Extract max_tokens (optional)
diff --git a/aiscript/src/main.rs b/aiscript/src/main.rs

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ mod uscc {`
`54`	`54`	`});`
`55`	`55`
`56`	`56`	`static USCC_BASE_CHARS: LazyLock<HashMap<char, u8>> = LazyLock::new(\|\| {`
`57`		`- let mut base_chars = HashMap::with_capacity(17);`
	`57`	`+ let mut base_chars = HashMap::with_capacity(31);`
`58`	`58`
`59`	`59`	`[`
`60`	`60`	`'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G',`