Skip to content
3 changes: 3 additions & 0 deletions src/harness/embeddings/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
70 changes: 64 additions & 6 deletions src/harness/embeddings/ollama.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -17,7 +21,7 @@ pub struct OllamaEmbeddingModel {
client: reqwest::Client,
base_url: String,
model: String,
dimensions: usize,
dimensions: Arc<AtomicUsize>,
options: Option<OllamaOptions>,
}

Expand All @@ -27,15 +31,45 @@ 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,
})
}

fn try_new_unresolved(base_url: &str, model: &str) -> Result<Self> {
Ok(Self {
client: reqwest::Client::new(),
base_url: normalize_base_url(base_url)?,
model: normalize_model(model)?,
dimensions: Arc::new(AtomicUsize::new(0)),
Comment thread
senamakel marked this conversation as resolved.
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<Vec<f32>>)> {
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))
Comment thread
senamakel marked this conversation as resolved.
}

pub fn new(base_url: &str, model: &str, dimensions: usize) -> Self {
Self::try_new(base_url, model, dimensions).expect("invalid Ollama embedding configuration")
}
Expand Down Expand Up @@ -135,10 +169,24 @@ impl OllamaEmbeddingModel {
}

fn validate_dimensions(&self, index: usize, vector: &[f32]) -> Result<()> {
if vector.len() != self.dimensions {
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()
)));
}
Expand Down Expand Up @@ -284,7 +332,7 @@ impl EmbeddingModel for OllamaEmbeddingModel {
}

fn dimensions(&self) -> usize {
self.dimensions
self.dimensions.load(Ordering::Acquire)
}

async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
Expand Down Expand Up @@ -376,6 +424,16 @@ mod tests {
assert!(OllamaEmbeddingModel::try_new("http://host:11434", "local-v1", 1).is_err());
}

#[test]
fn unresolved_dimensions_are_learned_and_enforced_internally() {
Comment thread
senamakel marked this conversation as resolved.
Outdated
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 blank_inputs_preserve_positions_without_network() {
let model = OllamaEmbeddingModel::default();
Expand Down
14 changes: 11 additions & 3 deletions src/harness/providers/openai/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,20 @@ 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 sends no authorization header; LM Studio sends a bearer token
only when its API key is non-empty. `ProviderSpec::Ollama` uses the same
defaults.
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

Accessors: `.model()`, `.provider()`, `.base_url()`.

Expand Down
38 changes: 34 additions & 4 deletions src/harness/providers/openai/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 = <OpenAiModel as ChatModel<()>>::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]
Expand All @@ -679,6 +696,19 @@ fn provider_spec_builds_compatible_model() {
.as_deref(),
Some("ollama")
);
let profile = <OpenAiModel as ChatModel<()>>::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]
Expand Down
Loading