[goose-llm] add providerConfig param for exposed LLM functions (#2491)

crates/goose-llm/README.md

@@ -31,26 +31,19 @@ Structure:
 │               └── goose_llm.kt   ← auto-generated bindings
 ```
 
-Create Kotlin bindings:
-```
+#### Create Kotlin bindings:
+
+```bash
 # run from project root directory
 cargo build -p goose-llm 
 
 cargo run --features=uniffi/cli --bin uniffi-bindgen generate --library ./target/debug/libgoose_llm.dylib --language kotlin --out-dir bindings/kotlin
 ```
 
-Download jars in `bindings/kotlin/libs` directory (only need to do this once):
-```
-pushd bindings/kotlin/libs/
-curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlin/kotlin-stdlib/1.9.0/kotlin-stdlib-1.9.0.jar
-curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlinx/kotlinx-coroutines-core-jvm/1.7.3/kotlinx-coroutines-core-jvm-1.7.3.jar
-curl -O https://repo1.maven.org/maven2/net/java/dev/jna/jna/5.13.0/jna-5.13.0.jar
-popd
-```
 
+#### Kotlin -> Rust: run example
 
-Compile & Run usage example from Kotlin -> Rust:
-```
+```bash
 pushd bindings/kotlin/
 
 kotlinc \
@@ -68,3 +61,19 @@ java \
 popd
 ```
 
+You will have to download jars in `bindings/kotlin/libs` directory (only the first time):
+```bash
+pushd bindings/kotlin/libs/
+curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlin/kotlin-stdlib/1.9.0/kotlin-stdlib-1.9.0.jar
+curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlinx/kotlinx-coroutines-core-jvm/1.7.3/kotlinx-coroutines-core-jvm-1.7.3.jar
+curl -O https://repo1.maven.org/maven2/net/java/dev/jna/jna/5.13.0/jna-5.13.0.jar
+popd
+```
+
+#### Python -> Rust: generate bindings, run example
+
+```bash
+cargo run --features=uniffi/cli --bin uniffi-bindgen generate --library ./target/debug/libgoose_llm.dylib --language python --out-dir bindings/python
+
+DYLD_LIBRARY_PATH=./target/debug python bindings/python/usage.py
+```

crates/goose-llm/examples/simple.rs

@@ -3,6 +3,7 @@ use std::vec;
 use anyhow::Result;
 use goose_llm::{
     completion,
+    extractors::generate_tooltip,
     types::completion::{
         CompletionRequest, CompletionResponse, ExtensionConfig, ToolApprovalMode, ToolConfig,
     },
@@ -13,8 +14,12 @@ use serde_json::json;
 #[tokio::main]
 async fn main() -> Result<()> {
     let provider = "databricks";
-    // let model_name = "goose-claude-3-5-sonnet"; // sequential tool calls
-    let model_name = "goose-gpt-4-1"; // parallel tool calls
+    let provider_config = json!({
+        "host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
+        "token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
+    });
+    // let model_name = "goose-gpt-4-1"; // parallel tool calls
+    let model_name = "claude-3-5-haiku";
     let model_config = ModelConfig::new(model_name.to_string());
 
     let calculator_tool = ToolConfig::new(
@@ -92,18 +97,26 @@ async fn main() -> Result<()> {
     ] {
         println!("\n---------------\n");
         println!("User Input: {text}");
-        let messages = vec![Message::user().with_text(text)];
-        let completion_response: CompletionResponse = completion(CompletionRequest {
-            provider_name: provider.to_string(),
-            model_config: model_config.clone(),
-            system_preamble: system_preamble.to_string(),
-            messages: messages,
-            extensions: extensions.clone(),
-        })
+        let messages = vec![
+            Message::user().with_text("Hi there!"),
+            Message::assistant().with_text("How can I help?"),
+            Message::user().with_text(text),
+        ];
+        let completion_response: CompletionResponse = completion(CompletionRequest::new(
+            provider.to_string(),
+            provider_config.clone(),
+            model_config.clone(),
+            system_preamble.to_string(),
+            messages.clone(),
+            extensions.clone(),
+        ))
         .await?;
         // Print the response
         println!("\nCompletion Response:");
         println!("{}", serde_json::to_string_pretty(&completion_response)?);
+
+        let tooltip = generate_tooltip(provider, provider_config.clone().into(), &messages).await?;
+        println!("\nTooltip: {}", tooltip);
     }
 
     Ok(())

crates/goose-llm/src/completion.rs

@@ -29,8 +29,12 @@ pub fn print_messages(messages: Vec<Message>) {
 pub async fn completion(req: CompletionRequest) -> Result<CompletionResponse, CompletionError> {
     let start_total = Instant::now();
 
-    let provider = create(&req.provider_name, req.model_config)
-        .map_err(|_| CompletionError::UnknownProvider(req.provider_name.to_string()))?;
+    let provider = create(
+        &req.provider_name,
+        req.provider_config.clone(),
+        req.model_config.clone(),
+    )
+    .map_err(|_| CompletionError::UnknownProvider(req.provider_name.to_string()))?;
 
     let system_prompt = construct_system_prompt(&req.system_preamble, &req.extensions)?;
     let tools = collect_prefixed_tools(&req.extensions);

crates/goose-llm/src/extractors/session_name.rs

@@ -1,9 +1,8 @@
-use crate::message::Message;
 use crate::model::ModelConfig;
-use crate::providers::base::Provider;
-use crate::providers::databricks::DatabricksProvider;
+use crate::providers::create;
 use crate::providers::errors::ProviderError;
 use crate::types::core::Role;
+use crate::{message::Message, types::json_value_ffi::JsonValueFfi};
 use anyhow::Result;
 use indoc::indoc;
 use serde_json::{json, Value};
@@ -50,7 +49,20 @@ fn build_system_prompt() -> String {
 
 /// Generates a short (≤4 words) session name
 #[uniffi::export(async_runtime = "tokio")]
-pub async fn generate_session_name(messages: &[Message]) -> Result<String, ProviderError> {
+pub async fn generate_session_name(
+    provider_name: &str,
+    provider_config: JsonValueFfi,
+    messages: &[Message],
+) -> Result<String, ProviderError> {
+    // Use OpenAI models specifically for this task
+    let model_name = if provider_name == "databricks" {
+        "goose-gpt-4-1"
+    } else {
+        "gpt-4.1"
+    };
+    let model_cfg = ModelConfig::new(model_name.to_string()).with_temperature(Some(0.0));
+    let provider = create(provider_name, provider_config.into(), model_cfg)?;
+
     // Collect up to the first 3 user messages (truncated to 300 chars each)
     let context: Vec<String> = messages
         .iter()
@@ -75,10 +87,6 @@ pub async fn generate_session_name(messages: &[Message]) -> Result<String, Provi
     let system_prompt = build_system_prompt();
     let user_msg_text = format!("Here are the user messages:\n{}", context.join("\n"));
 
-    // Instantiate DatabricksProvider with goose-gpt-4-1
-    let model_cfg = ModelConfig::new("goose-gpt-4-1".to_string()).with_temperature(Some(0.0));
-    let provider = DatabricksProvider::from_env(model_cfg)?;
-
     // Use `extract` with a simple string schema
     let schema = json!({
         "type": "object",

crates/goose-llm/src/extractors/tooltip.rs

@@ -1,9 +1,9 @@
 use crate::message::{Message, MessageContent};
 use crate::model::ModelConfig;
-use crate::providers::base::Provider;
-use crate::providers::databricks::DatabricksProvider;
+use crate::providers::create;
 use crate::providers::errors::ProviderError;
 use crate::types::core::{Content, Role};
+use crate::types::json_value_ffi::JsonValueFfi;
 use anyhow::Result;
 use indoc::indoc;
 use serde_json::{json, Value};
@@ -54,7 +54,20 @@ fn build_system_prompt() -> String {
 /// Generates a tooltip summarizing the last two messages in the session,
 /// including any tool calls or results.
 #[uniffi::export(async_runtime = "tokio")]
-pub async fn generate_tooltip(messages: &[Message]) -> Result<String, ProviderError> {
+pub async fn generate_tooltip(
+    provider_name: &str,
+    provider_config: JsonValueFfi,
+    messages: &[Message],
+) -> Result<String, ProviderError> {
+    // Use OpenAI models specifically for this task
+    let model_name = if provider_name == "databricks" {
+        "goose-gpt-4-1"
+    } else {
+        "gpt-4.1"
+    };
+    let model_cfg = ModelConfig::new(model_name.to_string()).with_temperature(Some(0.0));
+    let provider = create(provider_name, provider_config.into(), model_cfg)?;
+
     // Need at least two messages to summarize
     if messages.len() < 2 {
         return Err(ProviderError::ExecutionError(
@@ -128,10 +141,6 @@ pub async fn generate_tooltip(messages: &[Message]) -> Result<String, ProviderEr
         rendered.join("\n")
     );
 
-    // Instantiate the provider
-    let model_cfg = ModelConfig::new("goose-gpt-4-1".to_string()).with_temperature(Some(0.0));
-    let provider = DatabricksProvider::from_env(model_cfg)?;
-
     // Schema wrapping our tooltip string
     let schema = json!({
         "type": "object",

crates/goose-llm/src/providers/databricks.rs

@@ -26,66 +26,83 @@ pub const _DATABRICKS_KNOWN_MODELS: &[&str] = &[
     "databricks-claude-3-7-sonnet",
 ];
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum DatabricksAuth {
-    Token(String),
+fn default_timeout() -> u64 {
+    60
 }
 
-impl DatabricksAuth {
-    pub fn token(token: String) -> Self {
-        Self::Token(token)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DatabricksProviderConfig {
+    pub host: String,
+    pub token: String,
+    #[serde(default)]
+    pub image_format: ImageFormat,
+    #[serde(default = "default_timeout")]
+    pub timeout: u64, // timeout in seconds
+}
+
+impl DatabricksProviderConfig {
+    pub fn new(host: String, token: String) -> Self {
+        Self {
+            host,
+            token,
+            image_format: ImageFormat::OpenAi,
+            timeout: default_timeout(),
+        }
+    }
+
+    pub fn from_env() -> Self {
+        let host = get_env("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST");
+        let token = get_env("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN");
+        Self::new(host, token)
     }
 }
 
 #[derive(Debug)]
 pub struct DatabricksProvider {
-    client: Client,
-    host: String,
-    auth: DatabricksAuth,
+    config: DatabricksProviderConfig,
     model: ModelConfig,
-    image_format: ImageFormat,
+    client: Client,
+}
+
+impl DatabricksProvider {
+    pub fn from_env(model: ModelConfig) -> Self {
+        let config = DatabricksProviderConfig::from_env();
+        DatabricksProvider::from_config(config, model)
+            .expect("Failed to initialize DatabricksProvider")
+    }
 }
 
 impl Default for DatabricksProvider {
     fn default() -> Self {
+        let config = DatabricksProviderConfig::from_env();
         let model = ModelConfig::new(DATABRICKS_DEFAULT_MODEL.to_string());
-        DatabricksProvider::from_env(model).expect("Failed to initialize Databricks provider")
+        DatabricksProvider::from_config(config, model)
+            .expect("Failed to initialize DatabricksProvider")
     }
 }
 
 impl DatabricksProvider {
-    pub fn from_env(model: ModelConfig) -> Result<Self> {
-        let host = get_env("DATABRICKS_HOST")?;
-        let api_key = get_env("DATABRICKS_TOKEN")?;
-
+    pub fn from_config(config: DatabricksProviderConfig, model: ModelConfig) -> Result<Self> {
         let client = Client::builder()
-            .timeout(Duration::from_secs(600))
+            .timeout(Duration::from_secs(config.timeout))
             .build()?;
 
         Ok(Self {
-            client,
-            host,
-            auth: DatabricksAuth::token(api_key),
+            config,
             model,
-            image_format: ImageFormat::OpenAi,
+            client,
         })
     }
 
-    async fn ensure_auth_header(&self) -> Result<String> {
-        match &self.auth {
-            DatabricksAuth::Token(token) => Ok(format!("Bearer {}", token)),
-        }
-    }
-
     async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
-        let base_url = Url::parse(&self.host)
+        let base_url = Url::parse(&self.config.host)
             .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
         let path = format!("serving-endpoints/{}/invocations", self.model.model_name);
         let url = base_url.join(&path).map_err(|e| {
             ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
         })?;
 
-        let auth_header = self.ensure_auth_header().await?;
+        let auth_header = format!("Bearer {}", &self.config.token);
         let response = self
             .client
             .post(url)
@@ -189,7 +206,13 @@ impl Provider for DatabricksProvider {
         messages: &[Message],
         tools: &[Tool],
     ) -> Result<ProviderCompleteResponse, ProviderError> {
-        let mut payload = create_request(&self.model, system, messages, tools, &self.image_format)?;
+        let mut payload = create_request(
+            &self.model,
+            system,
+            messages,
+            tools,
+            &self.config.image_format,
+        )?;
         // Remove the model key which is part of the url with databricks
         payload
             .as_object_mut()

crates/goose-llm/src/providers/factory.rs

@@ -2,14 +2,28 @@ use std::sync::Arc;
 
 use anyhow::Result;
 
-use super::{base::Provider, databricks::DatabricksProvider, openai::OpenAiProvider};
+use super::{
+    base::Provider,
+    databricks::{DatabricksProvider, DatabricksProviderConfig},
+    openai::{OpenAiProvider, OpenAiProviderConfig},
+};
 use crate::model::ModelConfig;
 
-pub fn create(name: &str, model: ModelConfig) -> Result<Arc<dyn Provider>> {
+pub fn create(
+    name: &str,
+    provider_config: serde_json::Value,
+    model: ModelConfig,
+) -> Result<Arc<dyn Provider>> {
     // We use Arc instead of Box to be able to clone for multiple async tasks
     match name {
-        "openai" => Ok(Arc::new(OpenAiProvider::from_env(model)?)),
-        "databricks" => Ok(Arc::new(DatabricksProvider::from_env(model)?)),
+        "openai" => {
+            let config: OpenAiProviderConfig = serde_json::from_value(provider_config)?;
+            Ok(Arc::new(OpenAiProvider::from_config(config, model)?))
+        }
+        "databricks" => {
+            let config: DatabricksProviderConfig = serde_json::from_value(provider_config)?;
+            Ok(Arc::new(DatabricksProvider::from_config(config, model)?))
+        }
         _ => Err(anyhow::anyhow!("Unknown provider: {}", name)),
     }
 }

crates/goose-llm/src/providers/openai.rs

@@ -3,6 +3,7 @@ use std::{collections::HashMap, time::Duration};
 use anyhow::Result;
 use async_trait::async_trait;
 use reqwest::Client;
+use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
 
 use super::{
@@ -20,82 +21,112 @@ use crate::{
 pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o";
 pub const _OPEN_AI_KNOWN_MODELS: &[&str] = &["gpt-4o", "gpt-4.1", "o1", "o3", "o4-mini"];
 
+fn default_timeout() -> u64 {
+    60
+}
+
+fn default_base_path() -> String {
+    "v1/chat/completions".to_string()
+}
+
+fn default_host() -> String {
+    "https://api.openai.com".to_string()
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OpenAiProviderConfig {
+    pub api_key: String,
+    #[serde(default = "default_host")]
+    pub host: String,
+    #[serde(default)]
+    pub organization: Option<String>,
+    #[serde(default = "default_base_path")]
+    pub base_path: String,
+    #[serde(default)]
+    pub project: Option<String>,
+    #[serde(default)]
+    pub custom_headers: Option<HashMap<String, String>>,
+    #[serde(default = "default_timeout")]
+    pub timeout: u64, // timeout in seconds
+}
+
+impl OpenAiProviderConfig {
+    pub fn new(api_key: String) -> Self {
+        Self {
+            api_key,
+            host: default_host(),
+            organization: None,
+            base_path: default_base_path(),
+            project: None,
+            custom_headers: None,
+            timeout: 600,
+        }
+    }
+
+    pub fn from_env() -> Self {
+        let api_key = get_env("OPENAI_API_KEY").expect("Missing OPENAI_API_KEY");
+        Self::new(api_key)
+    }
+}
+
 #[derive(Debug)]
 pub struct OpenAiProvider {
-    client: Client,
-    host: String,
-    base_path: String,
-    api_key: String,
-    organization: Option<String>,
-    project: Option<String>,
+    config: OpenAiProviderConfig,
     model: ModelConfig,
-    custom_headers: Option<HashMap<String, String>>,
+    client: Client,
+}
+
+impl OpenAiProvider {
+    pub fn from_env(model: ModelConfig) -> Self {
+        let config = OpenAiProviderConfig::from_env();
+        OpenAiProvider::from_config(config, model).expect("Failed to initialize OpenAiProvider")
+    }
 }
 
 impl Default for OpenAiProvider {
     fn default() -> Self {
+        let config = OpenAiProviderConfig::from_env();
         let model = ModelConfig::new(OPEN_AI_DEFAULT_MODEL.to_string());
-        OpenAiProvider::from_env(model).expect("Failed to initialize OpenAI provider")
+        OpenAiProvider::from_config(config, model).expect("Failed to initialize OpenAiProvider")
     }
 }
 
 impl OpenAiProvider {
-    pub fn from_env(model: ModelConfig) -> Result<Self> {
-        let api_key: String = get_env("OPENAI_API_KEY")?;
-        let host: String =
-            get_env("OPENAI_HOST").unwrap_or_else(|_| "https://api.openai.com".to_string());
-        let base_path: String =
-            get_env("OPENAI_BASE_PATH").unwrap_or_else(|_| "v1/chat/completions".to_string());
-        let organization: Option<String> = get_env("OPENAI_ORGANIZATION").ok();
-        let project: Option<String> = get_env("OPENAI_PROJECT").ok();
-        let custom_headers: Option<HashMap<String, String>> = get_env("OPENAI_CUSTOM_HEADERS")
-            .or_else(|_| get_env("OPENAI_CUSTOM_HEADERS"))
-            .ok()
-            .map(parse_custom_headers);
-        // parse get_env("OPENAI_TIMEOUT") to u64 or set default to 600
-        let timeout_secs = get_env("OPENAI_TIMEOUT")
-            .ok()
-            .and_then(|s| s.parse::<u64>().ok())
-            .unwrap_or(600);
+    pub fn from_config(config: OpenAiProviderConfig, model: ModelConfig) -> Result<Self> {
         let client = Client::builder()
-            .timeout(Duration::from_secs(timeout_secs))
+            .timeout(Duration::from_secs(config.timeout))
             .build()?;
 
         Ok(Self {
-            client,
-            host,
-            base_path,
-            api_key,
-            organization,
-            project,
+            config,
             model,
-            custom_headers,
+            client,
         })
     }
 
     async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
-        let base_url = url::Url::parse(&self.host)
+        let base_url = url::Url::parse(&self.config.host)
             .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
-        let url = base_url.join(&self.base_path).map_err(|e| {
+        let url = base_url.join(&self.config.base_path).map_err(|e| {
             ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
         })?;
 
         let mut request = self
             .client
             .post(url)
-            .header("Authorization", format!("Bearer {}", self.api_key));
+            .header("Authorization", format!("Bearer {}", self.config.api_key));
 
         // Add organization header if present
-        if let Some(org) = &self.organization {
+        if let Some(org) = &self.config.organization {
             request = request.header("OpenAI-Organization", org);
         }
 
         // Add project header if present
-        if let Some(project) = &self.project {
+        if let Some(project) = &self.config.project {
             request = request.header("OpenAI-Project", project);
         }
 
-        if let Some(custom_headers) = &self.custom_headers {
+        if let Some(custom_headers) = &self.config.custom_headers {
             for (key, value) in custom_headers {
                 request = request.header(key, value);
             }
@@ -198,14 +229,3 @@ impl Provider for OpenAiProvider {
         Ok(ProviderExtractResponse::new(data, model, usage))
     }
 }
-
-fn parse_custom_headers(s: String) -> HashMap<String, String> {
-    s.split(',')
-        .filter_map(|header| {
-            let mut parts = header.splitn(2, '=');
-            let key = parts.next().map(|s| s.trim().to_string())?;
-            let value = parts.next().map(|s| s.trim().to_string())?;
-            Some((key, value))
-        })
-        .collect()
-}

crates/goose-llm/src/providers/utils.rs

@@ -19,12 +19,22 @@ struct OpenAIErrorResponse {
     error: OpenAIError,
 }
 
-#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
+#[derive(Debug, Copy, Clone, Serialize, Deserialize, Default)]
 pub enum ImageFormat {
+    #[default]
     OpenAi,
     Anthropic,
 }
 
+/// Timeout in seconds.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub struct Timeout(u32);
+impl Default for Timeout {
+    fn default() -> Self {
+        Timeout(60)
+    }
+}
+
 /// Convert an image content into an image json based on format
 pub fn convert_image(image: &ImageContent, image_format: &ImageFormat) -> Value {
     match image_format {

crates/goose-llm/src/types/completion.rs

@@ -11,17 +11,64 @@ use crate::types::json_value_ffi::JsonValueFfi;
 use crate::{message::Message, providers::Usage};
 use crate::{model::ModelConfig, providers::errors::ProviderError};
 
-// Lifetimes are not supported in Uniffi, cause other languages don't have them
-// https://github.com/mozilla/uniffi-rs/issues/1526#issuecomment-1528851837
-#[derive(uniffi::Record)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CompletionRequest {
     pub provider_name: String,
+    pub provider_config: serde_json::Value,
     pub model_config: ModelConfig,
     pub system_preamble: String,
     pub messages: Vec<Message>,
     pub extensions: Vec<ExtensionConfig>,
 }
 
+impl CompletionRequest {
+    pub fn new(
+        provider_name: String,
+        provider_config: serde_json::Value,
+        model_config: ModelConfig,
+        system_preamble: String,
+        messages: Vec<Message>,
+        extensions: Vec<ExtensionConfig>,
+    ) -> Self {
+        Self {
+            provider_name,
+            provider_config,
+            model_config,
+            system_preamble,
+            messages,
+            extensions,
+        }
+    }
+}
+
+#[uniffi::export]
+pub fn create_completion_request(
+    provider_name: &str,
+    provider_config: JsonValueFfi,
+    model_config: ModelConfig,
+    system_preamble: &str,
+    messages: Vec<Message>,
+    extensions: Vec<ExtensionConfig>,
+) -> CompletionRequest {
+    CompletionRequest::new(
+        provider_name.to_string(),
+        provider_config.into(),
+        model_config,
+        system_preamble.to_string(),
+        messages,
+        extensions,
+    )
+}
+
+uniffi::custom_type!(CompletionRequest, String, {
+    lower: |tc: &CompletionRequest| {
+        serde_json::to_string(&tc).unwrap()
+    },
+    try_lift: |s: String| {
+        Ok(serde_json::from_str(&s).unwrap())
+    },
+});
+
 // https://mozilla.github.io/uniffi-rs/latest/proc_macro/errors.html
 #[derive(Debug, Error, uniffi::Error)]
 #[uniffi(flat_error)]
@@ -149,7 +196,7 @@ uniffi::custom_type!(ToolConfig, String, {
 
 // — Register the newtypes with UniFFI, converting via JSON strings —
 
-#[derive(Debug, Clone, Serialize, uniffi::Record)]
+#[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
 pub struct ExtensionConfig {
     name: String,
     instructions: Option<String>,

crates/goose-llm/tests/extract_session_name.rs

@@ -4,15 +4,32 @@ use goose_llm::extractors::generate_session_name;
 use goose_llm::message::Message;
 use goose_llm::providers::errors::ProviderError;
 
-#[tokio::test]
-async fn test_generate_session_name_success() -> Result<(), ProviderError> {
-    // Load .env for Databricks credentials
+fn should_run_test() -> Result<(), String> {
     dotenv().ok();
-    let has_creds =
-        std::env::var("DATABRICKS_HOST").is_ok() && std::env::var("DATABRICKS_TOKEN").is_ok();
-    if !has_creds {
-        println!("Skipping generate_session_name test – Databricks creds not set");
-        return Ok(());
+    if std::env::var("DATABRICKS_HOST").is_err() {
+        return Err("Missing DATABRICKS_HOST".to_string());
+    }
+    if std::env::var("DATABRICKS_TOKEN").is_err() {
+        return Err("Missing DATABRICKS_TOKEN".to_string());
+    }
+    Ok(())
+}
+
+async fn _generate_session_name(messages: &[Message]) -> Result<String, ProviderError> {
+    let provider_name = "databricks";
+    let provider_config = serde_json::json!({
+        "host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
+        "token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
+    });
+
+    generate_session_name(provider_name, provider_config.into(), messages).await
+}
+
+#[tokio::test]
+async fn test_generate_session_name_success() {
+    if should_run_test().is_err() {
+        println!("Skipping...");
+        return;
     }
 
     // Build a few messages with at least two user messages
@@ -22,7 +39,10 @@ async fn test_generate_session_name_success() -> Result<(), ProviderError> {
         Message::user().with_text("What’s the weather in New York tomorrow?"),
     ];
 
-    let name = generate_session_name(&messages).await?;
+    let name = _generate_session_name(&messages)
+        .await
+        .expect("Failed to generate session name");
+
     println!("Generated session name: {:?}", name);
 
     // Should be non-empty and at most 4 words
@@ -34,20 +54,23 @@ async fn test_generate_session_name_success() -> Result<(), ProviderError> {
         "Name must be 4 words or less, got {}: {}",
         word_count,
         name
-    );
-
-    Ok(())
+    )
 }
 
 #[tokio::test]
 async fn test_generate_session_name_no_user() {
+    if should_run_test().is_err() {
+        println!("Skipping 'test_generate_session_name_no_user'. Databricks creds not set");
+        return;
+    }
+
     // No user messages → expect ExecutionError
     let messages = vec![
         Message::assistant().with_text("System starting…"),
         Message::assistant().with_text("All systems go."),
     ];
 
-    let err = generate_session_name(&messages).await;
+    let err = _generate_session_name(&messages).await;
     assert!(
         matches!(err, Err(ProviderError::ExecutionError(_))),
         "Expected ExecutionError when there are no user messages, got: {:?}",

crates/goose-llm/tests/extract_tooltip.rs

@@ -6,13 +6,32 @@ use goose_llm::providers::errors::ProviderError;
 use goose_llm::types::core::{Content, ToolCall};
 use serde_json::json;
 
-#[tokio::test]
-async fn test_generate_tooltip_simple() -> Result<(), ProviderError> {
-    // Skip if no Databricks creds
+fn should_run_test() -> Result<(), String> {
     dotenv().ok();
-    if std::env::var("DATABRICKS_HOST").is_err() || std::env::var("DATABRICKS_TOKEN").is_err() {
-        println!("Skipping simple tooltip test – Databricks creds not set");
-        return Ok(());
+    if std::env::var("DATABRICKS_HOST").is_err() {
+        return Err("Missing DATABRICKS_HOST".to_string());
+    }
+    if std::env::var("DATABRICKS_TOKEN").is_err() {
+        return Err("Missing DATABRICKS_TOKEN".to_string());
+    }
+    Ok(())
+}
+
+async fn _generate_tooltip(messages: &[Message]) -> Result<String, ProviderError> {
+    let provider_name = "databricks";
+    let provider_config = serde_json::json!({
+        "host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
+        "token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
+    });
+
+    generate_tooltip(provider_name, provider_config.into(), messages).await
+}
+
+#[tokio::test]
+async fn test_generate_tooltip_simple() {
+    if should_run_test().is_err() {
+        println!("Skipping...");
+        return;
     }
 
     // Two plain-text messages
@@ -21,7 +40,9 @@ async fn test_generate_tooltip_simple() -> Result<(), ProviderError> {
         Message::assistant().with_text("I'm fine, thanks! How can I help?"),
     ];
 
-    let tooltip = generate_tooltip(&messages).await?;
+    let tooltip = _generate_tooltip(&messages)
+        .await
+        .expect("Failed to generate tooltip");
     println!("Generated tooltip: {:?}", tooltip);
 
     assert!(!tooltip.trim().is_empty(), "Tooltip must not be empty");
@@ -29,16 +50,13 @@ async fn test_generate_tooltip_simple() -> Result<(), ProviderError> {
         tooltip.len() < 100,
         "Tooltip should be reasonably short (<100 chars)"
     );
-    Ok(())
 }
 
 #[tokio::test]
-async fn test_generate_tooltip_with_tools() -> Result<(), ProviderError> {
-    // Skip if no Databricks creds
-    dotenv().ok();
-    if std::env::var("DATABRICKS_HOST").is_err() || std::env::var("DATABRICKS_TOKEN").is_err() {
-        println!("Skipping tool‐based tooltip test – Databricks creds not set");
-        return Ok(());
+async fn test_generate_tooltip_with_tools() {
+    if should_run_test().is_err() {
+        println!("Skipping...");
+        return;
     }
 
     // 1) Assistant message with a tool request
@@ -57,7 +75,9 @@ async fn test_generate_tooltip_with_tools() -> Result<(), ProviderError> {
 
     let messages = vec![tool_req_msg, tool_resp_msg];
 
-    let tooltip = generate_tooltip(&messages).await?;
+    let tooltip = _generate_tooltip(&messages)
+        .await
+        .expect("Failed to generate tooltip");
     println!("Generated tooltip (tools): {:?}", tooltip);
 
     assert!(!tooltip.trim().is_empty(), "Tooltip must not be empty");
@@ -65,5 +85,4 @@ async fn test_generate_tooltip_with_tools() -> Result<(), ProviderError> {
         tooltip.len() < 100,
         "Tooltip should be reasonably short (<100 chars)"
     );
-    Ok(())
 }

crates/goose-llm/tests/providers_extract.rs

@@ -25,8 +25,8 @@ impl ProviderType {
 
     fn create_provider(&self, cfg: ModelConfig) -> Result<Arc<dyn Provider>> {
         Ok(match self {
-            ProviderType::OpenAi => Arc::new(OpenAiProvider::from_env(cfg)?),
-            ProviderType::Databricks => Arc::new(DatabricksProvider::from_env(cfg)?),
+            ProviderType::OpenAi => Arc::new(OpenAiProvider::from_env(cfg)),
+            ProviderType::Databricks => Arc::new(DatabricksProvider::from_env(cfg)),
         })
     }
 }