- 1
use std::collections::HashMap; - 2
use std::sync::{Arc, PoisonError, RwLock}; - 3
- 4
use tokio_util::sync::CancellationToken; - 5
- 6
use crate::Provider; - 7
use crate::anthropic::{AnthropicConfig, AnthropicProvider, DEFAULT_BASE_URL}; - 8
use crate::error::LlmError; - 9
use crate::stream::EventStream; - 10
use crate::types::ChatRequest; - 11
- 12
#[derive(Debug, Clone, Default)] - 13
pub struct ProviderAuth { - 14
pub api_key: String, - 15
pub base_url: Option<String>, - 16
/// Non-secret identity used to freeze a credential choice in a route. - 17
/// Providers continue to authenticate with `api_key`; this label is - 18
/// never sent over the wire. - 19
pub credential_id: Option<String>, - 20
/// Generic per-provider tuning knobs threaded from `vak-config` - 21
/// (docs/design/68-context-engine.md §8), e.g. Ollama's `keep_alive`/ - 22
/// `num_ctx`. Keys and meaning are entirely provider-defined; the - 23
/// registry and this struct stay agnostic to their contents. - 24
pub options: std::collections::BTreeMap<String, String>, - 25
} - 26
- 27
type Factory = Arc<dyn Fn(&ProviderAuth) -> Result<Arc<dyn Provider>, LlmError> + Send + Sync>; - 28
- 29
/// Cache identity: same provider name but a different key or base URL is - 30
/// a DIFFERENT provider — returning the cached one would silently send - 31
/// requests with stale credentials to the wrong endpoint. - 32
fn cache_key(name: &str, auth: &ProviderAuth) -> String { - 33
use std::hash::{Hash, Hasher}; - 34
let mut h = std::collections::hash_map::DefaultHasher::new(); - 35
auth.api_key.hash(&mut h); - 36
format!( - 37
"{name}\u{0}{}\u{0}{:016x}", - 38
auth.base_url.as_deref().unwrap_or(""), - 39
h.finish() - 40
) - 41
} - 42
- 43
#[derive(Default)] - 44
pub struct ProviderRegistry { - 45
factories: RwLock<HashMap<String, Factory>>, - 46
cache: RwLock<HashMap<String, Arc<dyn Provider>>>, - 47
} - 48
- 49
impl ProviderRegistry { - 50
pub fn new() -> Self { - 51
Self::default() - 52
} - 53
- 54
pub fn register( - 55
&self, - 56
name: impl Into<String>, - 57
factory: impl Fn(&ProviderAuth) -> Result<Arc<dyn Provider>, LlmError> + Send + Sync + 'static, - 58
) { - 59
self.factories - 60
.write() - 61
.unwrap_or_else(PoisonError::into_inner) - 62
.insert(name.into(), Arc::new(factory)); - 63
} - 64
- 65
pub fn get(&self, name: &str, auth: &ProviderAuth) -> Result<Arc<dyn Provider>, LlmError> { - 66
let key = cache_key(name, auth); - 67
if let Some(cached) = self - 68
.cache - 69
.read() - 70
.unwrap_or_else(PoisonError::into_inner) - 71
.get(&key) - 72
{ - 73
return Ok(cached.clone()); - 74
} - 75
let factory = self - 76
.factories - 77
.read() - 78
.unwrap_or_else(PoisonError::into_inner) - 79
.get(name) - 80
.cloned() - 81
.ok_or_else(|| LlmError::InvalidRequest(format!("unknown provider: {name}")))?; - 82
let provider = factory(auth)?; - 83
self.cache - 84
.write() - 85
.unwrap_or_else(PoisonError::into_inner) - 86
.insert(key, provider.clone()); - 87
Ok(provider) - 88
} - 89
- 90
pub fn names(&self) -> Vec<String> { - 91
let mut names: Vec<String> = self - 92
.factories - 93
.read() - 94
.unwrap_or_else(PoisonError::into_inner) - 95
.keys() - 96
.cloned() - 97
.collect(); - 98
names.sort(); - 99
names - 100
} - 101
} - 102
- 103
pub fn default_registry() -> ProviderRegistry { - 104
let registry = ProviderRegistry::new(); - 105
registry.register("anthropic", |auth| { - 106
Ok(Arc::new(AnthropicProvider::new( - 107
anthropic_config_from_auth(auth), - 108
)?)) - 109
}); - 110
- 111
use crate::openai::{OPENAI_DEFAULT_BASE_URL, OpenAiCompletionsProvider, OpenAiConfig}; - 112
let openai_compat = |default_base: &'static str, cache_key: bool, openrouter: bool| { - 113
move |auth: &ProviderAuth| { - 114
Ok(Arc::new(OpenAiCompletionsProvider::new(OpenAiConfig { - 115
api_key: auth.api_key.clone(), - 116
base_url: auth.base_url.clone().unwrap_or_else(|| default_base.into()), - 117
cache_key, - 118
openrouter, - 119
})?) as Arc<dyn Provider>) - 120
} - 121
}; - 122
registry.register( - 123
"openai", - 124
openai_compat(OPENAI_DEFAULT_BASE_URL, true, false), - 125
); - 126
- 127
use crate::openai_responses::{ - 128
OPENAI_RESPONSES_DEFAULT_BASE_URL, OpenAiResponsesConfig, OpenAiResponsesProvider, - 129
}; - 130
let responses = |default_base: &'static str, cache_key: bool, openrouter: bool| { - 131
move |auth: &ProviderAuth| { - 132
Ok( - 133
Arc::new(OpenAiResponsesProvider::new(OpenAiResponsesConfig { - 134
api_key: auth.api_key.clone(), - 135
base_url: auth.base_url.clone().unwrap_or_else(|| default_base.into()), - 136
cache_key, - 137
openrouter, - 138
})?) as Arc<dyn Provider>, - 139
) - 140
} - 141
}; - 142
registry.register( - 143
"openai-responses", - 144
responses(OPENAI_RESPONSES_DEFAULT_BASE_URL, true, false), - 145
); - 146
- 147
use crate::google::{GOOGLE_DEFAULT_BASE_URL, GoogleConfig, GoogleProvider}; - 148
registry.register("google", |auth| { - 149
Ok(Arc::new(GoogleProvider::new(GoogleConfig { - 150
api_key: auth.api_key.clone(), - 151
base_url: auth - 152
.base_url - 153
.clone() - 154
.unwrap_or_else(|| GOOGLE_DEFAULT_BASE_URL.into()), - 155
})?) as Arc<dyn Provider>) - 156
}); - 157
registry.register( - 158
"openrouter", - 159
openai_compat("https://openrouter.ai/api/v1", true, true), - 160
); - 161
registry.register( - 162
"openrouter-responses", - 163
responses("https://openrouter.ai/api/v1", true, true), - 164
); - 165
registry.register( - 166
"opencode-zen", - 167
openai_compat("https://opencode.ai/zen/v1", true, false), - 168
); - 169
- 170
use crate::ollama::OllamaProvider; - 171
registry.register("ollama", |auth| { - 172
Ok(Arc::new(OllamaProvider::new(ollama_config_from_auth(auth))?) as Arc<dyn Provider>) - 173
}); - 174
- 175
// Amazon Bedrock Mantle exposes an OpenAI-compatible API. The region is - 176
// part of the endpoint; callers may override it with VAK_BEDROCK_BASE_URL. - 177
registry.register( - 178
"bedrock", - 179
openai_compat("https://bedrock-mantle.us-east-1.api.aws/v1", true, false), - 180
); - 181
registry - 182
} - 183
- 184
/// Builds an `AnthropicConfig` from generic `ProviderAuth` fields - 185
/// (docs/design/68-context-engine.md §11 "Anthropic" row). Standalone so it - 186
/// can be unit tested without constructing a live `AnthropicProvider`. - 187
fn anthropic_config_from_auth(auth: &ProviderAuth) -> AnthropicConfig { - 188
AnthropicConfig { - 189
api_key: auth.api_key.clone(), - 190
base_url: auth - 191
.base_url - 192
.clone() - 193
.unwrap_or_else(|| DEFAULT_BASE_URL.into()), - 194
model: String::new(), - 195
fast_mode: auth - 196
.options - 197
.get("fast_mode") - 198
.and_then(|v| v.parse::<bool>().ok()) - 199
.unwrap_or(false), - 200
} - 201
} - 202
- 203
/// Builds an `OllamaConfig` from generic `ProviderAuth` fields - 204
/// (docs/design/68-context-engine.md §8). Standalone so it can be unit - 205
/// tested without constructing a live `OllamaProvider`. - 206
fn ollama_config_from_auth(auth: &ProviderAuth) -> crate::ollama::OllamaConfig { - 207
let default_config = crate::ollama::OllamaConfig::default(); - 208
// `auth.base_url` may carry the `/v1` OpenAI-compat suffix used for - 209
// discovery (models.rs::default_base_url); the native `/api/chat` wire - 210
// always wants the bare root. - 211
let base_url = auth - 212
.base_url - 213
.clone() - 214
.unwrap_or_else(|| crate::ollama::OLLAMA_DEFAULT_BASE_URL.into()); - 215
let base_url = base_url - 216
.trim_end_matches('/') - 217
.trim_end_matches("/v1") - 218
.to_string(); - 219
let keep_alive = auth - 220
.options - 221
.get("keep_alive") - 222
.cloned() - 223
.unwrap_or(default_config.keep_alive); - 224
let num_ctx = auth - 225
.options - 226
.get("num_ctx") - 227
.and_then(|v| v.parse::<u64>().ok()) - 228
.or(default_config.num_ctx); - 229
crate::ollama::OllamaConfig { - 230
base_url, - 231
api_key: auth.api_key.clone(), - 232
keep_alive, - 233
num_ctx, - 234
} - 235
} - 236
- 237
pub async fn stream_via( - 238
provider: &dyn Provider, - 239
request: ChatRequest, - 240
cancel: CancellationToken, - 241
) -> Result<EventStream, LlmError> { - 242
provider.stream(request, cancel).await - 243
} - 244
- 245
#[cfg(test)] - 246
mod tests { - 247
use super::*; - 248
- 249
#[test] - 250
fn default_registry_names_include_ollama() { - 251
assert!(default_registry().names().contains(&"ollama".to_string())); - 252
} - 253
- 254
#[test] - 255
fn anthropic_config_defaults_fast_mode_off() { - 256
let auth = ProviderAuth { - 257
api_key: "k".into(), - 258
base_url: None, - 259
credential_id: None, - 260
options: Default::default(), - 261
}; - 262
let config = anthropic_config_from_auth(&auth); - 263
assert_eq!(config.base_url, DEFAULT_BASE_URL); - 264
assert!(!config.fast_mode); - 265
} - 266
- 267
#[test] - 268
fn anthropic_config_threads_fast_mode_from_options() { - 269
let mut options = std::collections::BTreeMap::new(); - 270
options.insert("fast_mode".to_string(), "true".to_string()); - 271
let auth = ProviderAuth { - 272
api_key: "k".into(), - 273
base_url: None, - 274
credential_id: None, - 275
options, - 276
}; - 277
let config = anthropic_config_from_auth(&auth); - 278
assert!(config.fast_mode); - 279
} - 280
- 281
#[test] - 282
fn anthropic_config_ignores_unparseable_fast_mode() { - 283
let mut options = std::collections::BTreeMap::new(); - 284
options.insert("fast_mode".to_string(), "yes-please".to_string()); - 285
let auth = ProviderAuth { - 286
api_key: "k".into(), - 287
base_url: None, - 288
credential_id: None, - 289
options, - 290
}; - 291
let config = anthropic_config_from_auth(&auth); - 292
assert!(!config.fast_mode); - 293
} - 294
- 295
#[test] - 296
fn ollama_config_defaults_when_no_options_are_set() { - 297
let auth = ProviderAuth { - 298
api_key: "ollama".into(), - 299
base_url: None, - 300
credential_id: None, - 301
options: Default::default(), - 302
}; - 303
let config = ollama_config_from_auth(&auth); - 304
assert_eq!(config.base_url, crate::ollama::OLLAMA_DEFAULT_BASE_URL); - 305
assert_eq!(config.keep_alive, "30m"); - 306
assert_eq!(config.num_ctx, None); - 307
} - 308
- 309
#[test] - 310
fn ollama_config_threads_keep_alive_and_num_ctx_from_options() { - 311
let mut options = std::collections::BTreeMap::new(); - 312
options.insert("keep_alive".to_string(), "10m".to_string()); - 313
options.insert("num_ctx".to_string(), "8192".to_string()); - 314
let auth = ProviderAuth { - 315
api_key: "ollama".into(), - 316
base_url: None, - 317
credential_id: None, - 318
options, - 319
}; - 320
let config = ollama_config_from_auth(&auth); - 321
assert_eq!(config.keep_alive, "10m"); - 322
assert_eq!(config.num_ctx, Some(8192)); - 323
} - 324
- 325
#[test] - 326
fn ollama_config_strips_v1_compat_suffix_from_base_url() { - 327
let auth = ProviderAuth { - 328
api_key: "ollama".into(), - 329
base_url: Some("http://localhost:11434/v1".into()), - 330
credential_id: None, - 331
options: Default::default(), - 332
}; - 333
let config = ollama_config_from_auth(&auth); - 334
assert_eq!(config.base_url, "http://localhost:11434"); - 335
} - 336
- 337
#[test] - 338
fn ollama_config_ignores_unparseable_num_ctx() { - 339
let mut options = std::collections::BTreeMap::new(); - 340
options.insert("num_ctx".to_string(), "not-a-number".to_string()); - 341
let auth = ProviderAuth { - 342
api_key: "ollama".into(), - 343
base_url: None, - 344
credential_id: None, - 345
options, - 346
}; - 347
let config = ollama_config_from_auth(&auth); - 348
assert_eq!(config.num_ctx, None); - 349
} - 350
} - 351
Indexing the workspace…
Vakyartha documentation is discovering safe artifacts, anchors, and source references.