From de95745d3af20cbe20f5c2403c0f5e2377b49895 Mon Sep 17 00:00:00 2001 From: CSP386 Date: Sat, 7 Mar 2026 15:39:42 +0000 Subject: [PATCH] fix: use provider_urls config for embedding driver base_url - Add base_url_override parameter to create_embedding_driver() - Pass URL from config.provider_urls to embedding driver for ollama, openai - Add unit tests for URL override functionality Fixes embedding fallback when using remote Ollama provider --- crates/openfang-kernel/src/kernel.rs | 9 ++-- crates/openfang-runtime/src/embedding.rs | 69 +++++++++++++++++++----- 2 files changed, 62 insertions(+), 16 deletions(-) diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index 88e0baf44..bc8a07dec 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -765,7 +765,8 @@ impl OpenFangKernel { if let Some(ref provider) = config.memory.embedding_provider { // Explicit config takes priority — use the configured embedding model let api_key_env = config.memory.embedding_api_key_env.as_deref().unwrap_or(""); - match create_embedding_driver(provider, configured_model, api_key_env) { + let url = config.provider_urls.get(provider).map(|s| s.as_str()); + match create_embedding_driver(provider, configured_model, api_key_env, url) { Ok(d) => { info!(provider = %provider, model = %configured_model, "Embedding driver configured from memory config"); Some(Arc::from(d)) @@ -781,7 +782,8 @@ impl OpenFangKernel { } else { configured_model.as_str() }; - match create_embedding_driver("openai", model, "OPENAI_API_KEY") { + let url = config.provider_urls.get("openai").map(|s| s.as_str()); + match create_embedding_driver("openai", model, "OPENAI_API_KEY", url) { Ok(d) => { info!("Embedding driver auto-detected: OpenAI"); Some(Arc::from(d)) @@ -798,7 +800,8 @@ impl OpenFangKernel { } else { configured_model.as_str() }; - match create_embedding_driver("ollama", model, "") { + let url = config.provider_urls.get("ollama").map(|s| s.as_str()); + match create_embedding_driver("ollama", model, "", url) { Ok(d) => { info!("Embedding driver auto-detected: Ollama (local)"); Some(Arc::from(d)) diff --git a/crates/openfang-runtime/src/embedding.rs b/crates/openfang-runtime/src/embedding.rs index 2fd414c29..893f25630 100644 --- a/crates/openfang-runtime/src/embedding.rs +++ b/crates/openfang-runtime/src/embedding.rs @@ -179,6 +179,7 @@ pub fn create_embedding_driver( provider: &str, model: &str, api_key_env: &str, + base_url_override: Option<&str>, ) -> Result, EmbeddingError> { let api_key = if api_key_env.is_empty() { String::new() @@ -186,18 +187,22 @@ pub fn create_embedding_driver( std::env::var(api_key_env).unwrap_or_default() }; - let base_url = match provider { - "openai" => OPENAI_BASE_URL.to_string(), - "groq" => GROQ_BASE_URL.to_string(), - "together" => TOGETHER_BASE_URL.to_string(), - "fireworks" => FIREWORKS_BASE_URL.to_string(), - "mistral" => MISTRAL_BASE_URL.to_string(), - "ollama" => OLLAMA_BASE_URL.to_string(), - "vllm" => VLLM_BASE_URL.to_string(), - "lmstudio" => LMSTUDIO_BASE_URL.to_string(), - other => { - warn!("Unknown embedding provider '{other}', using OpenAI-compatible format"); - format!("https://{other}/v1") + let base_url = if let Some(url) = base_url_override { + url.to_string() + } else { + match provider { + "openai" => OPENAI_BASE_URL.to_string(), + "groq" => GROQ_BASE_URL.to_string(), + "together" => TOGETHER_BASE_URL.to_string(), + "fireworks" => FIREWORKS_BASE_URL.to_string(), + "mistral" => MISTRAL_BASE_URL.to_string(), + "ollama" => OLLAMA_BASE_URL.to_string(), + "vllm" => VLLM_BASE_URL.to_string(), + "lmstudio" => LMSTUDIO_BASE_URL.to_string(), + other => { + warn!("Unknown embedding provider '{other}', using OpenAI-compatible format"); + format!("https://{other}/v1") + } } }; @@ -351,8 +356,46 @@ mod tests { #[test] fn test_create_embedding_driver_ollama() { // Should succeed even without API key (ollama is local) - let driver = create_embedding_driver("ollama", "all-MiniLM-L6-v2", ""); + let driver = create_embedding_driver("ollama", "all-MiniLM-L6-v2", "", None); + assert!(driver.is_ok()); + assert_eq!(driver.unwrap().dimensions(), 384); + } + + #[test] + fn test_create_embedding_driver_with_url_override() { + let driver = create_embedding_driver( + "ollama", + "nomic-embed-text", + "", + Some("http://192.168.1.100:11434/v1"), + ); + assert!(driver.is_ok()); + + let driver = create_embedding_driver( + "ollama", + "all-MiniLM-L6-v2", + "", + Some("https://ollama.remote.com/v1"), + ); assert!(driver.is_ok()); assert_eq!(driver.unwrap().dimensions(), 384); } + + #[test] + fn test_create_embedding_driver_openai_with_url_override() { + // Test that URL override works for OpenAI provider + let driver = create_embedding_driver( + "openai", + "text-embedding-3-small", + "OPENAI_API_KEY", + Some("https://custom-openai-compatible.example.com/v1"), + ); + assert!(driver.is_ok()); + } + + #[test] + fn test_create_embedding_driver_fallback_to_default_url() { + let driver = create_embedding_driver("ollama", "nomic-embed-text", "", None); + assert!(driver.is_ok()); + } }