feat: add `--sync-models` cli option (#1114)

1 files changed, 81 insertions, 37 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 8813930..341470f 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -12,7 +12,7 @@ use self::session::Session;
 
 use crate::client::{
     create_client_config, list_client_types, list_models, ClientConfig, MessageContentToolCalls,
-    Model, ModelType, OPENAI_COMPATIBLE_PLATFORMS,
+    Model, ModelType, ProviderModels, OPENAI_COMPATIBLE_PROVIDERS,
 };
 use crate::function::{FunctionDeclaration, Functions, ToolResult};
 use crate::rag::Rag;
@@ -24,7 +24,7 @@ use anyhow::{anyhow, bail, Context, Result};
 use indexmap::IndexMap;
 use inquire::{list_option::ListOption, validator::Validation, Confirm, MultiSelect, Select, Text};
 use parking_lot::RwLock;
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
 use simplelog::LevelFilter;
 use std::collections::{HashMap, HashSet};
@@ -64,6 +64,8 @@ const CLIENTS_FIELD: &str = "clients";
 
 const SERVE_ADDR: &str = "127.0.0.1:8000";
 
+const SYNC_MODELS_URL: &str = "https://cdn.jsdelivr.net/gh/sigoden/aichat/models.yaml";
+
 const SUMMARIZE_PROMPT: &str =
     "Summarize the discussion briefly in 200 words or less to use as a prompt for future context.";
 const SUMMARY_PROMPT: &str = "This is a summary of the chat history as a recap: ";
@@ -140,6 +142,7 @@ pub struct Config {
     pub serve_addr: Option<String>,
     pub user_agent: Option<String>,
     pub save_shell_history: bool,
+    pub sync_models_url: Option<String>,
 
     pub clients: Vec<ClientConfig>,
 
@@ -214,6 +217,7 @@ impl Default for Config {
             serve_addr: None,
             user_agent: None,
             save_shell_history: true,
+            sync_models_url: None,
 
             clients: vec![],
 
@@ -240,9 +244,12 @@ impl Config {
     pub fn init(working_mode: WorkingMode, info_flag: bool) -> Result<Self> {
         let config_path = Self::config_file();
         let mut config = if !config_path.exists() {
-            match env::var(get_env_name("platform")) {
-                Ok(v) => Self::load_dynamic(&v)?,
-                Err(_) => {
+            match env::var(get_env_name("provider"))
+                .ok()
+                .or_else(|| env::var(get_env_name("platform")).ok())
+            {
+                Some(v) => Self::load_dynamic(&v)?,
+                None => {
                     if *IS_STDOUT_TERMINAL {
                         create_config_file(&config_path)?;
                     }
@@ -417,6 +424,10 @@ impl Config {
         }
     }
 
+    pub fn models_override_file() -> PathBuf {
+        Self::local_path("models-override.json")
+    }
+
     pub fn state(&self) -> StateFlags {
         let mut flags = StateFlags::empty();
         if let Some(session) = &self.session {
@@ -1362,23 +1373,12 @@ impl Config {
         let temp_file = temp_file(&format!("-rag-{}", rag.name()), ".txt");
         tokio::fs::write(&temp_file, &document_paths.join("\n"))
             .await
-            .with_context(|| {
-                format!(
-                    "Failed to write current document paths to '{}'",
-                    temp_file.display()
-                )
-            })?;
+            .with_context(|| format!("Failed to write to '{}'", temp_file.display()))?;
         let editor = config.read().editor()?;
         edit_file(&editor, &temp_file)?;
-        let new_document_paths =
-            tokio::fs::read_to_string(&temp_file)
-                .await
-                .with_context(|| {
-                    format!(
-                        "Failed to read new document paths from '{}'",
-                        temp_file.display()
-                    )
-                })?;
+        let new_document_paths = tokio::fs::read_to_string(&temp_file)
+            .await
+            .with_context(|| format!("Failed to read '{}'", temp_file.display()))?;
         let new_document_paths = new_document_paths
             .split('\n')
             .filter_map(|v| {
@@ -1535,12 +1535,7 @@ impl Config {
                 &agent_config_path,
                 "# see https://github.com/sigoden/aichat/blob/main/config.agent.example.yaml\n",
             )
-            .with_context(|| {
-                format!(
-                    "Failed to write to agent config file at '{}'",
-                    agent_config_path.display()
-                )
-            })?;
+            .with_context(|| format!("Failed to write to '{}'", agent_config_path.display()))?;
         }
         let editor = self.editor()?;
         edit_file(&editor, &agent_config_path)?;
@@ -1860,6 +1855,50 @@ impl Config {
             .collect()
     }
 
+    pub fn sync_models_url(&self) -> String {
+        self.sync_models_url
+            .clone()
+            .unwrap_or_else(|| SYNC_MODELS_URL.into())
+    }
+
+    pub async fn sync_models(url: &str, abort_signal: AbortSignal) -> Result<()> {
+        let content = abortable_run_with_spinner(fetch(url), "Fetching models.yaml", abort_signal)
+            .await
+            .with_context(|| format!("Failed to fetch '{url}'"))?;
+        println!("✓ Fetched '{url}'");
+        let list = serde_yaml::from_str::<Vec<ProviderModels>>(&content)
+            .with_context(|| "Failed to parse models.yaml")?;
+        let models_override = ModelsOverride {
+            version: env!("CARGO_PKG_VERSION").to_string(),
+            list,
+        };
+        let models_override_data =
+            serde_json::to_string_pretty(&models_override).with_context(|| "Failed to serde {}")?;
+
+        let model_override_path = Self::models_override_file();
+        ensure_parent_exists(&model_override_path)?;
+        std::fs::write(&model_override_path, models_override_data)
+            .with_context(|| format!("Failed to write to '{}'", model_override_path.display()))?;
+        println!("✓ Updated '{}'", model_override_path.display());
+        Ok(())
+    }
+
+    pub fn loal_models_override() -> Result<Vec<ProviderModels>> {
+        let model_override_path = Self::models_override_file();
+        let err = || {
+            format!(
+                "Failed to load models at '{}'",
+                model_override_path.display()
+            )
+        };
+        let content = read_to_string(&model_override_path).with_context(err)?;
+        let models_override: ModelsOverride = serde_json::from_str(&content).with_context(err)?;
+        if models_override.version != env!("CARGO_PKG_VERSION") {
+            bail!("Incompatible version")
+        }
+        Ok(models_override.list)
+    }
+
     pub fn render_options(&self) -> Result<RenderOptions> {
         let theme = if self.highlight {
             let theme_mode = if self.light_theme { "light" } else { "dark" };
@@ -2175,17 +2214,17 @@ impl Config {
     }
 
     fn load_dynamic(model_id: &str) -> Result<Self> {
-        let platform = match model_id.split_once(':') {
+        let provider = match model_id.split_once(':') {
             Some((v, _)) => v,
             _ => model_id,
         };
-        let is_openai_compatible = OPENAI_COMPATIBLE_PLATFORMS
+        let is_openai_compatible = OPENAI_COMPATIBLE_PROVIDERS
             .into_iter()
-            .any(|(name, _)| platform == name);
+            .any(|(name, _)| provider == name);
         let client = if is_openai_compatible {
-            json!({ "type": "openai-compatible", "name": platform })
+            json!({ "type": "openai-compatible", "name": provider })
         } else {
-            json!({ "type": platform })
+            json!({ "type": provider })
         };
         let config = json!({
             "model": model_id.to_string(),
@@ -2323,6 +2362,9 @@ impl Config {
         if let Some(Some(v)) = read_env_bool(&get_env_name("save_shell_history")) {
             self.save_shell_history = v;
         }
+        if let Some(v) = read_env_value::<String>(&get_env_name("sync_models_url")) {
+            self.sync_models_url = v;
+        }
     }
 
     fn load_functions(&mut self) -> Result<()> {
@@ -2502,6 +2544,12 @@ pub struct MacroVariable {
     pub default: Option<String>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelsOverride {
+    pub version: String,
+    pub list: Vec<ProviderModels>,
+}
+
 #[derive(Debug, Clone)]
 pub struct LastMessage {
     pub input: Input,
@@ -2581,12 +2629,8 @@ fn create_config_file(config_path: &Path) -> Result<()> {
     );
 
     ensure_parent_exists(config_path)?;
-    std::fs::write(config_path, config_data).with_context(|| {
-        format!(
-            "Failed to write to config file at '{}'",
-            config_path.display()
-        )
-    })?;
+    std::fs::write(config_path, config_data)
+        .with_context(|| format!("Failed to write to '{}'", config_path.display()))?;
     #[cfg(unix)]
     {
         use std::os::unix::prelude::PermissionsExt;