Big update to llm crate and configuration of same

llm/Cargo.toml

@@ -21,3 +21,4 @@ toml = { workspace = true }
 xdg = { workspace = true }
 derive_builder = { workspace = true }
 anyhow = "1.0.82"
+chrono = "0.4.38"

llm/src/config.rs

@@ -6,6 +6,7 @@ use crate::vendor::{ollama, openai::compat};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Config {
+    pub default_system_message: Option<String>,
     pub providers: Providers,
 }
 

llm/src/main.rs

@@ -1,39 +1,36 @@
 use clap::Parser;
 use itertools::Itertools;
-use llm::{
-    vendor::{self, openai},
-    Config,
-};
+use llm::{vendor, Config};
 
 #[derive(Debug, clap::Parser)]
 #[command(version, about, long_about = None)]
 struct App {
+    #[clap(short, long)]
+    provider: vendor::Provider,
+
     #[command(subcommand)]
     command: Commands,
 }
 
 #[derive(Debug, clap::Subcommand)]
 enum Commands {
-    Completion(Completion),
+    Completion(vendor::Completion),
     ListModels,
 }
 
-#[derive(Debug, clap::Args)]
-struct Completion {
-    #[clap(short, long)]
-    provider: vendor::Provider,
-
-    #[clap(short, long)]
-    system_message: Option<String>,
-
-    user_message: String,
-}
-
-impl Completion {
-    async fn run(self, config: &Config) -> Result<openai::compat::Response, anyhow::Error> {
-        self.provider
-            .completion(config, self.system_message, self.user_message)
-            .await
+impl App {
+    async fn run(self, config: &Config) -> Result<(), anyhow::Error> {
+        match self.command {
+            Commands::Completion(c) => {
+                let response = self.provider.completion(config, c).await?;
+                println!("{}", response.content());
+            }
+            Commands::ListModels => {
+                let models = self.provider.list_models(config).await?;
+                println!("{}", models.iter().join("\n"));
+            }
+        }
+        Ok(())
     }
 }
 
@@ -43,21 +40,5 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let app = App::parse();
     let config = Config::read()?;
-    match app.command {
-        Commands::Completion(c) => {
-            let response = c.run(&config).await?;
-            println!("{}", response.content());
-        }
-        Commands::ListModels => {
-            let models: Vec<_> = vendor::ollama::list_models().await?.into();
-            println!(
-                "{:?}",
-                models
-                    .iter()
-                    .map(|m| (m.name(), m.human_size()))
-                    .collect_vec()
-            );
-        }
-    }
-    Ok(())
+    app.run(&config).await
 }

llm/src/vendor/groq.rs

@@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize};
 
 use super::openai::{self, compat::Response};
 
-const GROQ_CHAT_API: &str = "https://api.groq.com/openai/v1/chat/completions";
+pub(crate) const GROQ_CHAT_API: &str = "https://api.groq.com/openai/v1";
 
 const MIXTRAL: &str = "mixtral-8x7b-32768";
 const LLAMA3_70B: &str = "llama3-70b-8192";
@@ -59,15 +59,9 @@ impl Model {
 
 pub async fn completion(
     api_key: String,
-    system_message: Option<String>,
+    model: String,
+    system_message: String,
     user_message: String,
 ) -> Result<Response, anyhow::Error> {
-    openai::compat::completion(
-        GROQ_CHAT_API,
-        api_key,
-        Model::default(),
-        system_message,
-        user_message,
-    )
-    .await
+    openai::compat::completion(GROQ_CHAT_API, api_key, model, system_message, user_message).await
 }

llm/src/vendor/mod.rs

@@ -2,10 +2,23 @@ use anyhow::anyhow;
 
 use crate::Config;
 
+use self::openai::compat;
+
 pub mod groq;
 pub mod ollama;
 pub mod openai;
 
+#[derive(Debug, Clone, clap::Args)]
+pub struct Completion {
+    #[clap(short, long)]
+    model: Option<String>,
+
+    #[clap(short, long)]
+    system_message: Option<String>,
+
+    user_message: String,
+}
+
 #[derive(Debug, Clone, Copy, clap::ValueEnum)]
 #[value(rename_all = "lower")]
 pub enum Provider {
@@ -18,27 +31,86 @@ impl Provider {
     pub async fn completion(
         &self,
         config: &Config,
-        system_message: Option<String>,
-        user_message: String,
+        comp: Completion,
     ) -> Result<openai::compat::Response, anyhow::Error> {
+        let model = comp
+            .model
+            .clone()
+            .or_else(|| self.default_model(config))
+            .ok_or_else(|| {
+                anyhow!(
+                    "No model specified and no default configured for {}",
+                    self.name()
+                )
+            })?;
+
+        let system_message = comp
+            .system_message
+            .clone()
+            .or_else(|| config.default_system_message.clone())
+            .ok_or_else(|| anyhow!("No default system message configured"))?;
+
+        match self {
+            Provider::Groq => {
+                let provider = &config.providers.groq;
+                let api_key = provider
+                    .get_api_key()
+                    .await?
+                    .ok_or_else(|| anyhow!("no api key?"))?;
+                groq::completion(api_key, model, system_message, comp.user_message).await
+            }
+            Provider::OpenAi => {
+                let provider = &config.providers.openai;
+                let api_key = provider
+                    .get_api_key()
+                    .await?
+                    .ok_or_else(|| anyhow!("no api key?"))?;
+                openai::completion(api_key, model, system_message, comp.user_message).await
+            }
+            Provider::Ollama => {
+                let model = comp
+                    .model
+                    .ok_or_else(|| anyhow!("No model specified for Ollama"))?;
+                ollama::completion(model, system_message, comp.user_message).await
+            }
+        }
+    }
+
+    pub async fn list_models(&self, config: &Config) -> anyhow::Result<Vec<String>> {
         match self {
             Provider::Groq => {
                 let provider = &config.providers.groq;
                 let api_key = provider
                     .get_api_key()
                     .await?
                     .ok_or_else(|| anyhow!("no api key?"))?;
-                groq::completion(api_key, system_message, user_message).await
+                compat::list_models(groq::GROQ_CHAT_API, api_key).await
             }
             Provider::OpenAi => {
                 let provider = &config.providers.openai;
                 let api_key = provider
                     .get_api_key()
                     .await?
                     .ok_or_else(|| anyhow!("no api key?"))?;
-                openai::completion(api_key, system_message, user_message).await
+                openai::list_models(api_key).await
             }
-            Provider::Ollama => ollama::completion(system_message, user_message).await,
+            Provider::Ollama => ollama::list_models().await,
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        match self {
+            Provider::Groq => "Groq",
+            Provider::OpenAi => "OpenAI",
+            Provider::Ollama => "Ollama",
+        }
+    }
+
+    fn default_model(&self, config: &Config) -> Option<String> {
+        match self {
+            Provider::Groq => config.providers.groq.default_model(),
+            Provider::OpenAi => config.providers.openai.default_model(),
+            Provider::Ollama => config.providers.ollama.default_model(),
         }
     }
 }

llm/src/vendor/ollama/list.rs

@@ -25,8 +25,6 @@ impl Default for Host {
     }
 }
 
-struct Model(String);
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct ListModelsResponse {
     models: Vec<LocalModel>,
@@ -94,11 +92,13 @@ impl LocalModel {
     }
 }
 
-pub async fn list_models() -> anyhow::Result<ListModelsResponse> {
-    let resp = reqwest::Client::new()
+pub async fn list_models() -> anyhow::Result<Vec<String>> {
+    let resp: ListModelsResponse = reqwest::Client::new()
         .get(format!("{OLLAMA_API}/tags"))
         .send()
+        .await?
+        .json()
         .await?;
 
-    Ok(resp.json().await?)
+    Ok(resp.models.into_iter().map(|m| m.name).collect())
 }