diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index d52367c..45750c2 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -21,6 +21,9 @@ //! //! [`OpenAiEmbeddingModel`] adds a hosted provider backed by the OpenAI //! embeddings endpoint (always compiled). +//! [`OllamaEmbeddingModel::embed_discovering_dimensions`] handles +//! provider-managed local models whose vector width is unknown without exposing +//! an adapter whose dimensional identity can change after construction. //! //! The design mirrors LangChain's separation of concerns: chat models generate //! messages, embedding models generate vectors, vector stores search vectors, diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs index 32151b1..1d291f8 100644 --- a/src/harness/embeddings/ollama.rs +++ b/src/harness/embeddings/ollama.rs @@ -2,6 +2,10 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; use super::EmbeddingModel; use crate::error::{Result, TinyAgentsError}; @@ -17,7 +21,7 @@ pub struct OllamaEmbeddingModel { client: reqwest::Client, base_url: String, model: String, - dimensions: usize, + dimensions: Arc, options: Option, } @@ -27,15 +31,51 @@ impl OllamaEmbeddingModel { client: reqwest::Client::new(), base_url: normalize_base_url(base_url)?, model: normalize_model(model)?, - dimensions: if dimensions == 0 { + dimensions: Arc::new(AtomicUsize::new(if dimensions == 0 { DEFAULT_OLLAMA_DIMENSIONS } else { dimensions - }, + })), options: None, }) } + pub(super) fn try_new_unresolved(base_url: &str, model: &str) -> Result { + Ok(Self { + client: reqwest::Client::new(), + base_url: normalize_base_url(base_url)?, + model: normalize_model(model)?, + dimensions: Arc::new(AtomicUsize::new(0)), + options: None, + }) + } + + /// Embeds text for a model whose vector width is not known in advance. + /// + /// The temporary adapter learns and validates the response width internally; + /// callers receive only the resolved width and vectors, so no model with an + /// unstable `dimensions()` or signature escapes this operation. + pub async fn embed_discovering_dimensions( + base_url: &str, + model: &str, + client: reqwest::Client, + texts: &[String], + num_ctx: u32, + num_batch: u32, + ) -> Result<(usize, Vec>)> { + if !texts.iter().any(|text| !text.trim().is_empty()) { + return Err(TinyAgentsError::Validation( + "dynamic embedding dimension discovery requires at least one nonblank input" + .to_string(), + )); + } + let adapter = Self::try_new_unresolved(base_url, model)? + .with_client(client) + .with_context_options(num_ctx, num_batch); + let vectors = adapter.embed(texts).await?; + Ok((adapter.dimensions(), vectors)) + } + pub fn new(base_url: &str, model: &str, dimensions: usize) -> Self { Self::try_new(base_url, model, dimensions).expect("invalid Ollama embedding configuration") } @@ -134,11 +174,25 @@ impl OllamaEmbeddingModel { Ok(output) } - fn validate_dimensions(&self, index: usize, vector: &[f32]) -> Result<()> { - if vector.len() != self.dimensions { + pub(super) fn validate_dimensions(&self, index: usize, vector: &[f32]) -> Result<()> { + if vector.is_empty() { + return Err(TinyAgentsError::Embedding(format!( + "ollama embed returned an empty vector at index {index}" + ))); + } + let expected = match self.dimensions.compare_exchange( + 0, + vector.len(), + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => vector.len(), + Err(expected) => expected, + }; + if vector.len() != expected { return Err(TinyAgentsError::Embedding(format!( "ollama embed dimension mismatch at index {index}: expected {}, got {}", - self.dimensions, + expected, vector.len() ))); } @@ -284,7 +338,7 @@ impl EmbeddingModel for OllamaEmbeddingModel { } fn dimensions(&self) -> usize { - self.dimensions + self.dimensions.load(Ordering::Acquire) } async fn embed(&self, texts: &[String]) -> Result>> { diff --git a/src/harness/embeddings/test.rs b/src/harness/embeddings/test.rs index 42559ec..2583c40 100644 --- a/src/harness/embeddings/test.rs +++ b/src/harness/embeddings/test.rs @@ -69,6 +69,30 @@ async fn mock_model_batches_in_order() { } } +#[test] +fn unresolved_ollama_dimensions_are_learned_and_enforced_internally() { + let model = OllamaEmbeddingModel::try_new_unresolved("http://host:11434", "custom").unwrap(); + assert_eq!(model.dimensions(), 0); + model.validate_dimensions(0, &[0.0; 7]).unwrap(); + assert_eq!(model.dimensions(), 7); + assert!(model.validate_dimensions(1, &[0.0; 8]).is_err()); +} + +#[tokio::test] +async fn dynamic_ollama_discovery_rejects_blank_only_batches() { + let error = OllamaEmbeddingModel::embed_discovering_dimensions( + "http://host:11434", + "custom", + reqwest::Client::new(), + &[" ".to_string()], + 1, + 1, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("nonblank")); +} + #[tokio::test] async fn mock_model_empty_input_returns_empty() { let model = MockEmbeddingModel::new(8); diff --git a/src/harness/providers/openai/README.md b/src/harness/providers/openai/README.md index 408def0..58e7b66 100644 --- a/src/harness/providers/openai/README.md +++ b/src/harness/providers/openai/README.md @@ -37,12 +37,21 @@ of the rest of the harness. - `OpenAiModel::from_spec(spec, api_key)` / `from_spec_env(spec)` — build from a `providers::ProviderSpec` (base URL, default model, provider id already resolved). -- **Compatibility presets** — thin wrappers over `new` + `with_base_url` + - `with_model` for endpoints that speak the same Chat Completions wire format: +- **Compatibility presets** for hosted endpoints that speak the same Chat + Completions wire format: `compatible(base_url, model)` / `compatible_provider(..)` (arbitrary endpoint), `deepseek`, `anthropic` (compat endpoint, not the native - Anthropic API), `groq`, `xai`, `openrouter`, `together`, `mistral`, `ollama`. + Anthropic API), `groq`, `xai`, `openrouter`, `together`, and `mistral`. Override the preset's default model with `.with_model(..)`. +- **Local-runtime presets** — `ollama()`, fallible + `ollama_at(base_url, model)`, and fallible + `lm_studio(base_url, api_key, model)`. These normalize a server or + `/v1/models` URL to its `/v1` API base and use conservative local defaults: + no native or parallel tool calls, no streamed tool chunks, and no image + input. `ollama()` and `ollama_at()` send no authorization header; LM Studio + sends a bearer token only when its API key is non-empty. + `ProviderSpec::Ollama` uses the same defaults unless `requires_api_key` is + enabled, in which case it sends a bearer token. Accessors: `.model()`, `.provider()`, `.base_url()`. diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index e19d18e..f842485 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -639,29 +639,46 @@ fn compatible_presets_set_base_url_and_default_model() { #[test] fn local_runtime_presets_normalize_endpoint_and_model() { - let ollama = OpenAiModel::ollama_at("127.0.0.1:11434/", "qwen3:8b"); + let ollama = OpenAiModel::ollama_at("127.0.0.1:11434/", "qwen3:8b").unwrap(); assert_eq!(ollama.provider(), "ollama"); assert_eq!(ollama.base_url(), "http://127.0.0.1:11434/v1"); assert_eq!(ollama.model(), "qwen3:8b"); - let lm_studio = OpenAiModel::lm_studio("http://127.0.0.1:1234/v1/models", "", "local-model"); + let lm_studio = + OpenAiModel::lm_studio("http://127.0.0.1:1234/v1/models", "", "local-model").unwrap(); assert_eq!(lm_studio.provider(), "lm_studio"); assert_eq!(lm_studio.base_url(), "http://127.0.0.1:1234/v1"); assert_eq!(lm_studio.model(), "local-model"); assert_eq!( - OpenAiModel::ollama_at("http://models", "qwen3").base_url(), + OpenAiModel::ollama_at("http://models", "qwen3") + .unwrap() + .base_url(), "http://models/v1" ); assert_eq!( - OpenAiModel::ollama_at("http://v1", "qwen3").base_url(), + OpenAiModel::ollama_at("http://v1", "qwen3") + .unwrap() + .base_url(), "http://v1/v1" ); let overridden = OpenAiModel::ollama().with_model("qwen3:8b"); let profile = >::profile(&overridden).unwrap(); assert!(!profile.tool_calling); + assert!(!profile.parallel_tool_calls); + assert!(!profile.streaming_tool_chunks); assert!(!profile.modalities.image_in); + + assert!(OpenAiModel::ollama_at("http://[::1", "qwen3").is_err()); + assert!(OpenAiModel::ollama_at("ftp://host", "qwen3").is_err()); + + let caller_client = OpenAiModel::ollama().with_client(reqwest::Client::new()); + assert_eq!(caller_client.effective_request_timeout(None, false), None); + assert_eq!( + caller_client.effective_request_timeout(Some(25), false), + Some(std::time::Duration::from_millis(25)) + ); } #[test] @@ -679,6 +696,19 @@ fn provider_spec_builds_compatible_model() { .as_deref(), Some("ollama") ); + let profile = >::profile(&model).unwrap(); + assert!(!profile.tool_calling); + assert!(!profile.parallel_tool_calls); + assert!(!profile.streaming_tool_chunks); + assert!(!profile.modalities.image_in); + + let mut authenticated = ProviderSpec::for_kind(ProviderKind::Ollama); + authenticated.requires_api_key = true; + let authenticated = OpenAiModel::from_spec(authenticated, "proxy-secret").unwrap(); + assert_eq!( + authenticated.auth_config(), + ("proxy-secret", &AuthStyle::Bearer) + ); } #[test] diff --git a/src/harness/providers/openai/transport.rs b/src/harness/providers/openai/transport.rs index 1272ef3..282eb7f 100644 --- a/src/harness/providers/openai/transport.rs +++ b/src/harness/providers/openai/transport.rs @@ -39,6 +39,8 @@ pub enum AuthStyle { pub struct OpenAiModel { /// Shared HTTP client. client: reqwest::Client, + /// Whether the client was supplied by the caller and owns default deadlines. + caller_owned_client: bool, /// API credential; how it is sent is governed by [`Self::auth`]. api_key: String, /// How `api_key` is attached to each request (default [`AuthStyle::Bearer`]). @@ -317,6 +319,7 @@ impl OpenAiModel { .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS)) .build() .expect("default reqwest client builds"), + caller_owned_client: false, api_key: api_key.into(), auth: AuthStyle::Bearer, extra_headers: Vec::new(), @@ -393,6 +396,14 @@ impl OpenAiModel { self } + /// Reuses a caller-configured HTTP client, including its connection pool, + /// proxy settings, and request deadlines. + pub fn with_client(mut self, client: reqwest::Client) -> Self { + self.client = client; + self.caller_owned_client = true; + self + } + /// Overrides how the API credential is sent (default [`AuthStyle::Bearer`]). /// /// Use this for OpenAI-compatible endpoints that authenticate with @@ -545,12 +556,19 @@ impl OpenAiModel { /// Overrides the default model id. pub fn with_model(mut self, model: impl Into) -> Self { self.model = model.into(); - let local_capabilities = self - .local_capabilities_locked - .then_some((self.profile.tool_calling, self.profile.modalities.image_in)); + let local_capabilities = self.local_capabilities_locked.then_some(( + self.profile.tool_calling, + self.profile.parallel_tool_calls, + self.profile.streaming_tool_chunks, + self.profile.modalities.image_in, + )); self.profile = derive_profile(&self.provider, &self.model); - if let Some((tool_calling, image_in)) = local_capabilities { + if let Some((tool_calling, parallel_tool_calls, streaming_tool_chunks, image_in)) = + local_capabilities + { self.profile.tool_calling = tool_calling; + self.profile.parallel_tool_calls = parallel_tool_calls; + self.profile.streaming_tool_chunks = streaming_tool_chunks; self.profile.modalities.image_in = image_in; } self @@ -559,12 +577,19 @@ impl OpenAiModel { /// Overrides the provider family id used in profiles and normalized errors. pub fn with_provider(mut self, provider: impl Into) -> Self { self.provider = provider.into(); - let local_capabilities = self - .local_capabilities_locked - .then_some((self.profile.tool_calling, self.profile.modalities.image_in)); + let local_capabilities = self.local_capabilities_locked.then_some(( + self.profile.tool_calling, + self.profile.parallel_tool_calls, + self.profile.streaming_tool_chunks, + self.profile.modalities.image_in, + )); self.profile = derive_profile(&self.provider, &self.model); - if let Some((tool_calling, image_in)) = local_capabilities { + if let Some((tool_calling, parallel_tool_calls, streaming_tool_chunks, image_in)) = + local_capabilities + { self.profile.tool_calling = tool_calling; + self.profile.parallel_tool_calls = parallel_tool_calls; + self.profile.streaming_tool_chunks = streaming_tool_chunks; self.profile.modalities.image_in = image_in; } self @@ -614,6 +639,7 @@ impl OpenAiModel { /// Builds an OpenAI-compatible model from a provider spec and explicit API /// key. pub fn from_spec(spec: ProviderSpec, api_key: impl Into) -> Result { + let api_key = api_key.into(); if spec.model.trim().is_empty() { return Err(TinyAgentsError::Validation( "provider spec model must not be empty".to_string(), @@ -624,6 +650,20 @@ impl OpenAiModel { "provider spec base_url must not be empty".to_string(), )); } + if spec.kind == crate::harness::providers::ProviderKind::Ollama { + let auth = if spec.requires_api_key { + AuthStyle::Bearer + } else { + AuthStyle::None + }; + return Ok(Self::local_runtime( + &spec.provider, + normalize_local_v1_base_url(spec.base_url, "http://localhost:11434")?, + api_key, + spec.model, + ) + .with_auth_style(auth)); + } Ok(Self::compatible_provider( spec.provider, api_key, @@ -797,16 +837,17 @@ impl OpenAiModel { /// `llama3.2`. pub fn ollama() -> Self { Self::ollama_at("http://localhost:11434", "llama3.2") + .expect("the built-in Ollama URL is valid") } /// An Ollama server exposed through its OpenAI-compatible HTTP API. - pub fn ollama_at(base_url: impl Into, model: impl Into) -> Self { - Self::local_runtime( + pub fn ollama_at(base_url: impl Into, model: impl Into) -> Result { + Ok(Self::local_runtime( "ollama", - normalize_local_v1_base_url(base_url.into(), "http://localhost:11434"), + normalize_local_v1_base_url(base_url.into(), "http://localhost:11434")?, "", model, - ) + )) } /// An LM Studio server exposed through its OpenAI-compatible HTTP API. @@ -817,20 +858,20 @@ impl OpenAiModel { base_url: impl Into, api_key: impl Into, model: impl Into, - ) -> Self { + ) -> Result { let api_key = api_key.into(); let auth = if api_key.trim().is_empty() { AuthStyle::None } else { AuthStyle::Bearer }; - Self::local_runtime( + Ok(Self::local_runtime( "lm_studio", - normalize_local_v1_base_url(base_url.into(), "http://localhost:1234"), + normalize_local_v1_base_url(base_url.into(), "http://localhost:1234")?, api_key, model, ) - .with_auth_style(auth) + .with_auth_style(auth)) } fn local_runtime( @@ -851,6 +892,11 @@ impl OpenAiModel { self } + #[cfg(test)] + pub(super) fn auth_config(&self) -> (&str, &AuthStyle) { + (&self.api_key, &self.auth) + } + /// Returns the default model id this instance will request. pub fn model(&self) -> &str { &self.model @@ -1109,7 +1155,7 @@ impl OpenAiModel { url: &str, ) -> Result { let mut builder = self.authorized(self.client.post(url)).json(body); - if let Some(timeout) = request_timeout(timeout_ms, false) { + if let Some(timeout) = self.effective_request_timeout(timeout_ms, false) { builder = builder.timeout(timeout); } self.send_checked(builder, "responses request", url).await @@ -1160,12 +1206,24 @@ impl OpenAiModel { ) -> Result { let url = format!("{}/chat/completions", self.base_url); let mut builder = self.authorized(self.client.post(&url)).json(body); - if let Some(timeout) = request_timeout(timeout_ms, streaming) { + if let Some(timeout) = self.effective_request_timeout(timeout_ms, streaming) { builder = builder.timeout(timeout); } self.send_checked(builder, what, &url).await } + pub(super) fn effective_request_timeout( + &self, + timeout_ms: Option, + streaming: bool, + ) -> Option { + if self.caller_owned_client && timeout_ms.is_none() { + None + } else { + request_timeout(timeout_ms, streaming) + } + } + /// Builds the chat-completions wire body for `request` under the given /// `degrade`, setting the streaming fields when `streaming` is `true`. fn build_chat_body( @@ -1278,7 +1336,7 @@ impl OpenAiModel { } } -fn normalize_local_v1_base_url(raw: String, default_root: &str) -> String { +fn normalize_local_v1_base_url(raw: String, default_root: &str) -> Result { let trimmed = raw.trim().trim_end_matches('/'); let root = if trimmed.is_empty() { default_root.to_owned() @@ -1287,8 +1345,15 @@ fn normalize_local_v1_base_url(raw: String, default_root: &str) -> String { } else { format!("http://{trimmed}") }; - let mut url = - reqwest::Url::parse(&root).expect("local runtime URL is normalized with a scheme"); + let mut url = reqwest::Url::parse(&root).map_err(|error| { + TinyAgentsError::Validation(format!("invalid local runtime URL `{root}`: {error}")) + })?; + if !matches!(url.scheme(), "http" | "https") { + return Err(TinyAgentsError::Validation(format!( + "local runtime URL must use http or https, got `{}`", + url.scheme() + ))); + } let mut segments: Vec<&str> = url .path() .split('/') @@ -1305,7 +1370,7 @@ fn normalize_local_v1_base_url(raw: String, default_root: &str) -> String { url.set_path(&format!("/{}", segments.join("/"))); url.set_query(None); url.set_fragment(None); - url.to_string().trim_end_matches('/').to_owned() + Ok(url.to_string().trim_end_matches('/').to_owned()) } /// Request-shape degradations to apply when building an OpenAI wire body.