feat: first class Ollama provider via OLLAMA_HOST

This commit is contained in:
Aneesh 2026-06-02 16:03:20 -04:00
parent 4d3dc5b873
commit 943e7b1762
4 changed files with 61 additions and 10 deletions

View File

@ -32,16 +32,23 @@ impl ProviderClient {
OpenAiCompatConfig::xai(),
)?)),
ProviderKind::OpenAi => {
// DashScope models (qwen-*) also return ProviderKind::OpenAi because they
// speak the OpenAI wire format, but they need the DashScope config which
// reads DASHSCOPE_API_KEY and points at dashscope.aliyuncs.com.
let config = match providers::metadata_for_model(&resolved_model) {
Some(meta) if meta.auth_env == "DASHSCOPE_API_KEY" => {
OpenAiCompatConfig::dashscope()
}
_ => OpenAiCompatConfig::openai(),
};
Ok(Self::OpenAi(OpenAiCompatClient::from_env(config)?))
if std::env::var_os("OLLAMA_HOST").is_some() {
// unwrap is safe: from_ollama_env always returns Some
Ok(Self::OpenAi(
openai_compat::OpenAiCompatClient::from_ollama_env().unwrap(),
))
} else {
// DashScope models (qwen-*) also return ProviderKind::OpenAi because they
// speak the OpenAI wire format, but they need the DashScope config which
// reads DASHSCOPE_API_KEY and points at dashscope.aliyuncs.com.
let config = match providers::metadata_for_model(&resolved_model) {
Some(meta) if meta.auth_env == "DASHSCOPE_API_KEY" => {
OpenAiCompatConfig::dashscope()
}
_ => OpenAiCompatConfig::openai(),
};
Ok(Self::OpenAi(OpenAiCompatClient::from_env(config)?))
}
}
}
}

View File

@ -339,6 +339,10 @@ pub fn provider_diagnostics_for_model(model: &str) -> ProviderDiagnostics {
#[must_use]
pub fn detect_provider_kind(model: &str) -> ProviderKind {
if std::env::var_os("OLLAMA_HOST").is_some() {
return ProviderKind::OpenAi;
}
if let Some(metadata) = metadata_for_model(model) {
return metadata.provider;
}

View File

@ -48,6 +48,14 @@ const XAI_MAX_REQUEST_BODY_BYTES: usize = 52_428_800; // 50MB
const OPENAI_MAX_REQUEST_BODY_BYTES: usize = 104_857_600; // 100MB
const DASHSCOPE_MAX_REQUEST_BODY_BYTES: usize = 6_291_456; // 6MB (observed limit in dogfood)
pub const OLLAMA_CONFIG: OpenAiCompatConfig = OpenAiCompatConfig {
provider_name: "Ollama",
api_key_env: "OLLAMA_HOST",
base_url_env: "OLLAMA_HOST",
default_base_url: "http://127.0.0.1:11434/v1",
max_request_body_bytes: 104_857_600,
};
impl OpenAiCompatConfig {
#[must_use]
pub const fn xai() -> Self {
@ -140,6 +148,21 @@ impl OpenAiCompatClient {
Ok(Self::new(api_key, config))
}
pub fn from_ollama_env() -> Option<Self> {
let host =
std::env::var("OLLAMA_HOST").unwrap_or_else(|_| "http://127.0.0.1:11434".to_string());
let base_url = format!("{}/v1", host.trim_end_matches('/'));
Some(Self {
http: build_http_client_or_default(),
api_key: "ollama".to_string(),
config: OLLAMA_CONFIG,
base_url,
max_retries: DEFAULT_MAX_RETRIES,
initial_backoff: DEFAULT_INITIAL_BACKOFF,
max_backoff: DEFAULT_MAX_BACKOFF,
})
}
#[must_use]
pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
self.base_url = base_url.into();

View File

@ -1991,6 +1991,12 @@ fn resolve_model_alias_with_config(model: &str) -> String {
/// Rejects: empty, whitespace-only, strings with spaces, or invalid chars.
fn validate_model_syntax(model: &str) -> Result<(), String> {
let trimmed = model.trim();
if std::env::var_os("OLLAMA_HOST").is_some() {
if trimmed.is_empty() {
return Err("invalid model syntax: model string cannot be empty.\nUsage: --model <provider/model> e.g. --model anthropic/claude-opus-4-7".to_string());
}
return Ok(());
}
if trimmed.is_empty() {
return Err("invalid model syntax: model string cannot be empty.\nUsage: --model <provider/model> e.g. --model anthropic/claude-opus-4-7".to_string());
}
@ -16659,4 +16665,15 @@ mod alias_resolution_tests {
assert_eq!(resolve_model_alias_with_config(model), model);
assert!(validate_model_syntax(model).is_ok());
}
#[test]
fn test_ollama_host_bypasses_provider_model_validation() {
// Safety: test sets and clears env var within the test.
// May be flaky if tests run in parallel with conflicting OLLAMA_HOST usage,
// but acceptable for a unit-level smoke test.
std::env::set_var("OLLAMA_HOST", "http://127.0.0.1:11434");
let result = validate_model_syntax("qwen2.5-coder:7b");
std::env::remove_var("OLLAMA_HOST");
assert!(result.is_ok());
}
}