Support ai config (#42)

Author: Folyd

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()
             })

aiscript-runtime/src/endpoint.rs

@@ -499,8 +499,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);
                             }

aiscript-vm/src/ai/mod.rs

@@ -7,6 +7,21 @@ pub use agent::{Agent, run_agent};
 use openai_api_rs::v1::api::OpenAIClient;
 pub use prompt::{PromptConfig, prompt_with_config};
 
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub struct AiConfig {
+    pub openai: Option<ModelConfig>,
+    pub anthropic: Option<ModelConfig>,
+    pub deepseek: Option<ModelConfig>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelConfig {
+    pub api_key: String,
+    pub model: Option<String>,
+}
+
 #[allow(unused)]
 pub(crate) fn openai_client() -> OpenAIClient {
     OpenAIClient::builder()

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;

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
             }),
         };

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,
         }
     }
 

aiscript/src/main.rs

@@ -48,7 +48,7 @@ enum Commands {
 #[tokio::main]
 async fn main() {
     dotenv::dotenv().ok();
-    Config::load("project.toml");
+    let config = Config::load();
 
     let cli = AIScriptCli::parse();
     match cli.command {
@@ -69,7 +69,12 @@ async fn main() {
                 let sqlite_connection = aiscript_runtime::get_sqlite_connection().await;
                 let redis_connection = aiscript_runtime::get_redis_connection().await;
                 task::spawn_blocking(move || {
-                    let mut vm = Vm::new(pg_connection, sqlite_connection, redis_connection);
+                    let mut vm = Vm::new(
+                        pg_connection,
+                        sqlite_connection,
+                        redis_connection,
+                        config.ai.clone(),
+                    );
                     vm.run_file(path);
                 })
                 .await // must use await to wait for the thread to finish