1use std::sync::OnceLock;
2
3use cookie_agent_identity::RecipeRegistryRevision;
4use serde::Serialize;
5use sha2::{Digest as _, Sha256};
6
7use crate::{
8 adapters::OvenAdapterFamily,
9 authoring::ManagedModelShape,
10 catalog::{CatalogModelRecord, CatalogProviderRecord},
11};
12
13pub const COMPILER_VERSION: &str = "family-registry-compiler-v1";
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
16#[serde(rename_all = "kebab-case")]
17pub enum FamilyKind {
18 OpenAiCompatibleChat,
19 Anthropic,
20 OpenAi,
21 Google,
22 Vertex,
23 VertexAnthropic,
24 Bedrock,
25 Azure,
26 Cohere,
27}
28
29impl FamilyKind {
30 #[must_use]
31 pub const fn id(self) -> &'static str {
32 match self {
33 Self::OpenAiCompatibleChat => "openai-compatible-chat",
34 Self::Anthropic => "anthropic",
35 Self::OpenAi => "openai",
36 Self::Google => "google",
37 Self::Vertex => "vertex",
38 Self::VertexAnthropic => "vertex-anthropic",
39 Self::Bedrock => "bedrock",
40 Self::Azure => "azure",
41 Self::Cohere => "cohere",
42 }
43 }
44}
45
46#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
47pub struct FamilyRecipe {
48 pub npm: &'static str,
49 pub family: FamilyKind,
50 pub default_endpoint: Option<&'static str>,
51 pub default_auth_method: &'static str,
52 pub allowed_auth_methods: &'static [&'static str],
53}
54
55const RECIPES: &[FamilyRecipe] = &[
56 compatible("@ai-sdk/openai-compatible", None),
57 compatible("@ai-sdk/groq", Some("https://api.groq.com/openai/v1")),
58 compatible("@ai-sdk/mistral", Some("https://api.mistral.ai/v1")),
59 compatible("@ai-sdk/xai", Some("https://api.x.ai/v1")),
60 compatible("@ai-sdk/cerebras", Some("https://api.cerebras.ai/v1")),
61 compatible("@ai-sdk/togetherai", Some("https://api.together.xyz/v1")),
62 compatible(
63 "@ai-sdk/deepinfra",
64 Some("https://api.deepinfra.com/v1/openai"),
65 ),
66 compatible("@ai-sdk/perplexity", Some("https://api.perplexity.ai")),
67 compatible(
68 "venice-ai-sdk-provider",
69 Some("https://api.venice.ai/api/v1"),
70 ),
71 compatible(
72 "@openrouter/ai-sdk-provider",
73 Some("https://openrouter.ai/api/v1"),
74 ),
75 compatible("@qvac/ai-sdk-provider", None),
76 family(
77 "@ai-sdk/anthropic",
78 FamilyKind::Anthropic,
79 Some("https://api.anthropic.com/v1"),
80 "anthropic-api-key-v1",
81 &["anthropic-api-key-v1", "bearer-api-key-v1"],
82 ),
83 family(
84 "@ai-sdk/openai",
85 FamilyKind::OpenAi,
86 Some("https://api.openai.com/v1"),
87 "bearer-api-key-v1",
88 &["bearer-api-key-v1"],
89 ),
90 family(
91 "@ai-sdk/google",
92 FamilyKind::Google,
93 Some("https://generativelanguage.googleapis.com/v1beta"),
94 "google-api-key-header-v1",
95 &["google-api-key-header-v1"],
96 ),
97 family(
98 "@ai-sdk/google-vertex",
99 FamilyKind::Vertex,
100 None,
101 "oauth-access-token-v1",
102 &["oauth-access-token-v1"],
103 ),
104 family(
105 "@ai-sdk/google-vertex/anthropic",
106 FamilyKind::VertexAnthropic,
107 None,
108 "oauth-access-token-v1",
109 &["oauth-access-token-v1"],
110 ),
111 family(
112 "@ai-sdk/amazon-bedrock",
113 FamilyKind::Bedrock,
114 Some("https://bedrock-runtime.${AWS_REGION}.amazonaws.com"),
115 "aws-sigv4-credentials-v1",
116 &["aws-sigv4-credentials-v1", "bearer-api-key-v1"],
117 ),
118 family(
119 "@ai-sdk/amazon-bedrock/mantle",
120 FamilyKind::Bedrock,
121 None,
122 "bearer-api-key-v1",
123 &["bearer-api-key-v1"],
124 ),
125 family(
126 "@ai-sdk/azure",
127 FamilyKind::Azure,
128 Some("https://${AZURE_RESOURCE_NAME}.openai.azure.com"),
129 "azure-api-key-v1",
130 &["azure-api-key-v1", "bearer-api-key-v1"],
131 ),
132 family(
133 "@ai-sdk/cohere",
134 FamilyKind::Cohere,
135 Some("https://api.cohere.com/v2/chat"),
136 "bearer-api-key-v1",
137 &["bearer-api-key-v1"],
138 ),
139];
140
141const fn compatible(npm: &'static str, endpoint: Option<&'static str>) -> FamilyRecipe {
142 family(
143 npm,
144 FamilyKind::OpenAiCompatibleChat,
145 endpoint,
146 "bearer-api-key-v1",
147 &["bearer-api-key-v1"],
148 )
149}
150
151const fn family(
152 npm: &'static str,
153 kind: FamilyKind,
154 endpoint: Option<&'static str>,
155 default_auth_method: &'static str,
156 allowed_auth_methods: &'static [&'static str],
157) -> FamilyRecipe {
158 FamilyRecipe {
159 npm,
160 family: kind,
161 default_endpoint: endpoint,
162 default_auth_method,
163 allowed_auth_methods,
164 }
165}
166
167#[derive(Clone, Copy, Debug, Default)]
168pub struct FamilyRecipeRegistry;
169
170#[must_use]
171pub const fn family_registry() -> FamilyRecipeRegistry {
172 FamilyRecipeRegistry
173}
174
175impl FamilyRecipeRegistry {
176 pub fn recipes(self) -> impl ExactSizeIterator<Item = &'static FamilyRecipe> {
177 RECIPES.iter()
178 }
179
180 #[must_use]
181 pub fn by_npm(self, npm: &str) -> Option<&'static FamilyRecipe> {
182 RECIPES.iter().find(|recipe| recipe.npm == npm)
183 }
184
185 #[must_use]
186 pub fn classify(self, provider: &CatalogProviderRecord) -> Option<&'static FamilyRecipe> {
187 self.by_npm(&provider.npm)
188 }
189
190 #[must_use]
191 pub fn revision(self) -> RecipeRegistryRevision {
192 static REVISION: OnceLock<RecipeRegistryRevision> = OnceLock::new();
193 REVISION
194 .get_or_init(|| {
195 let bytes = serde_json::to_vec(&(
196 RECIPES,
197 crate::recipes::auth_methods().collect::<Vec<_>>(),
198 ))
199 .expect("static family registry serializes");
200 let mut digest = Sha256::new();
201 digest.update(b"cookie-agent/family-recipe-registry/v1\0");
202 digest.update(bytes);
203 RecipeRegistryRevision::new(format!("sha256:{:x}", digest.finalize()))
204 .expect("registry digest is valid")
205 })
206 .clone()
207 }
208}
209
210#[must_use]
215pub fn compatible_auth_method(
216 source_method: &str,
217 effective: &FamilyRecipe,
218) -> Option<&'static str> {
219 if let Some(method) = effective
220 .allowed_auth_methods
221 .iter()
222 .copied()
223 .find(|method| *method == source_method)
224 {
225 return Some(method);
226 }
227 let source_is_api_key = matches!(
228 source_method,
229 "anthropic-api-key-v1"
230 | "azure-api-key-v1"
231 | "google-api-key-header-v1"
232 | "bearer-api-key-v1"
233 );
234 let source_is_access_token = source_method == "oauth-access-token-v1";
235 if !source_is_api_key && !source_is_access_token {
236 return None;
237 }
238 let preferred = match effective.family {
239 FamilyKind::Anthropic => {
240 if source_method == "bearer-api-key-v1" || source_is_access_token {
243 "bearer-api-key-v1"
244 } else {
245 "anthropic-api-key-v1"
246 }
247 }
248 FamilyKind::OpenAi | FamilyKind::OpenAiCompatibleChat | FamilyKind::Cohere => {
249 "bearer-api-key-v1"
250 }
251 FamilyKind::Google => "google-api-key-header-v1",
252 FamilyKind::Azure => {
253 if source_method == "bearer-api-key-v1" || source_is_access_token {
254 "bearer-api-key-v1"
255 } else {
256 "azure-api-key-v1"
257 }
258 }
259 FamilyKind::Vertex | FamilyKind::VertexAnthropic => return None,
260 FamilyKind::Bedrock => return None,
261 };
262 effective
263 .allowed_auth_methods
264 .contains(&preferred)
265 .then_some(preferred)
266}
267
268#[must_use]
269pub fn compatible_credential_field(
270 source_method: &str,
271 target_field: &str,
272) -> Option<&'static str> {
273 match target_field {
274 "api_key" if source_method == "oauth-access-token-v1" => Some("access_token"),
275 "api_key" => Some("api_key"),
276 "access_token" if source_method == "bearer-api-key-v1" => Some("api_key"),
277 "access_token" => Some("access_token"),
278 "access_key_id" if source_method == "aws-sigv4-credentials-v1" => Some("access_key_id"),
279 "secret_access_key" if source_method == "aws-sigv4-credentials-v1" => {
280 Some("secret_access_key")
281 }
282 "session_token" if source_method == "aws-sigv4-credentials-v1" => Some("session_token"),
283 _ => None,
284 }
285}
286
287#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
288#[serde(rename_all = "lowercase")]
289pub enum ResolvedShape {
290 Chat,
291 Responses,
292}
293
294#[derive(Clone, Debug, Eq, PartialEq)]
295pub struct ResolvedFamilyModel {
296 pub recipe: &'static FamilyRecipe,
297 pub npm: String,
298 pub endpoint_template: Option<String>,
299 pub shape: ResolvedShape,
300 pub adapter: OvenAdapterFamily,
301}
302
303#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
304pub enum FamilyResolutionError {
305 #[error("no_known_protocol_family")]
306 UnknownFamily,
307 #[error("unsupported_model_shape")]
308 UnsupportedShape,
309}
310
311pub fn resolve_model(
312 provider: &CatalogProviderRecord,
313 model: &CatalogModelRecord,
314 provider_shape: Option<ManagedModelShape>,
315 model_shape: Option<ManagedModelShape>,
316) -> Result<ResolvedFamilyModel, FamilyResolutionError> {
317 let override_provider = model.provider.as_ref();
318 let npm = override_provider
319 .and_then(|value| value.npm.as_deref())
320 .unwrap_or(&provider.npm);
321 let recipe = family_registry()
322 .by_npm(npm)
323 .ok_or(FamilyResolutionError::UnknownFamily)?;
324 let endpoint_template = override_provider
325 .and_then(|value| value.api.clone())
326 .or_else(|| provider.api.clone())
327 .or_else(|| recipe.default_endpoint.map(str::to_owned));
328 let catalog_shape = override_provider
329 .and_then(|value| value.shape.as_deref())
330 .or(model.shape.as_deref())
331 .or(provider.shape.as_deref());
332 let catalog_shape = catalog_shape
333 .map(|shape| match shape {
334 "responses" => Ok(ResolvedShape::Responses),
335 "completions" | "chat" => Ok(ResolvedShape::Chat),
336 _ => Err(FamilyResolutionError::UnsupportedShape),
337 })
338 .transpose()?;
339 let shape = model_shape
340 .or(provider_shape)
341 .map(|shape| match shape {
342 ManagedModelShape::Chat => ResolvedShape::Chat,
343 ManagedModelShape::Responses => ResolvedShape::Responses,
344 })
345 .or(catalog_shape)
346 .unwrap_or_else(|| {
347 if recipe.family == FamilyKind::OpenAi {
348 ResolvedShape::Responses
349 } else {
350 ResolvedShape::Chat
351 }
352 });
353 let adapter = adapter(recipe.family, shape)?;
354 Ok(ResolvedFamilyModel {
355 recipe,
356 npm: npm.to_owned(),
357 endpoint_template,
358 shape,
359 adapter,
360 })
361}
362
363fn adapter(
364 kind: FamilyKind,
365 shape: ResolvedShape,
366) -> Result<OvenAdapterFamily, FamilyResolutionError> {
367 Ok(match kind {
368 FamilyKind::OpenAiCompatibleChat => match shape {
369 ResolvedShape::Chat => OvenAdapterFamily::OpenaiCompatible,
370 ResolvedShape::Responses => OvenAdapterFamily::OpenaiResponses,
371 },
372 FamilyKind::Anthropic => OvenAdapterFamily::AnthropicCompatible,
373 FamilyKind::OpenAi => match shape {
374 ResolvedShape::Chat => OvenAdapterFamily::OpenaiChat,
375 ResolvedShape::Responses => OvenAdapterFamily::OpenaiResponses,
376 },
377 FamilyKind::Google => OvenAdapterFamily::GoogleGemini,
378 FamilyKind::Vertex | FamilyKind::VertexAnthropic => OvenAdapterFamily::GoogleVertexGemini,
379 FamilyKind::Bedrock => match shape {
380 ResolvedShape::Chat => OvenAdapterFamily::AwsBedrockConverse,
381 ResolvedShape::Responses => OvenAdapterFamily::OpenaiResponses,
382 },
383 FamilyKind::Azure => match shape {
384 ResolvedShape::Chat => OvenAdapterFamily::AzureOpenaiChat,
385 ResolvedShape::Responses => OvenAdapterFamily::AzureOpenaiResponses,
386 },
387 FamilyKind::Cohere => OvenAdapterFamily::CohereV2Chat,
388 })
389}
390
391#[must_use]
392pub fn environment_aliases<'a>(
393 provider: &'a CatalogProviderRecord,
394 credential: &str,
395) -> Vec<&'a str> {
396 match credential {
397 "api_key" | "access_token" => provider
398 .environment
399 .first()
400 .map(String::as_str)
401 .into_iter()
402 .collect(),
403 "access_key_id" => provider
404 .environment
405 .iter()
406 .find(|value| value.as_str() == "AWS_ACCESS_KEY_ID")
407 .map(String::as_str)
408 .into_iter()
409 .collect(),
410 "secret_access_key" => provider
411 .environment
412 .iter()
413 .find(|value| value.as_str() == "AWS_SECRET_ACCESS_KEY")
414 .map(String::as_str)
415 .into_iter()
416 .collect(),
417 "session_token" => provider
418 .environment
419 .iter()
420 .find(|value| value.as_str() == "AWS_SESSION_TOKEN")
421 .map(String::as_str)
422 .into_iter()
423 .collect(),
424 _ => Vec::new(),
425 }
426}
427
428#[must_use]
429pub fn placeholders(template: &str) -> Vec<String> {
430 let mut result = Vec::new();
431 let mut rest = template;
432 while let Some(start) = rest.find("${") {
433 rest = &rest[start + 2..];
434 let Some(end) = rest.find('}') else { break };
435 let name = &rest[..end];
436 if !name.is_empty()
437 && name
438 .bytes()
439 .all(|byte| byte.is_ascii_uppercase() || byte.is_ascii_digit() || byte == b'_')
440 && !result.iter().any(|value| value == name)
441 {
442 result.push(name.to_owned());
443 }
444 rest = &rest[end + 1..];
445 }
446 result.sort();
447 result
448}
449
450pub fn substitute_placeholders(
451 template: &str,
452 values: &std::collections::BTreeMap<String, String>,
453) -> Option<String> {
454 let mut endpoint = template.to_owned();
455 for name in placeholders(template) {
456 let field = setup_field_name(&name);
457 let value = values.get(&name).or_else(|| values.get(&field))?;
458 endpoint = endpoint.replace(&format!("${{{name}}}"), value);
459 }
460 Some(endpoint)
461}
462
463#[must_use]
464pub fn setup_field_name(name: &str) -> String {
465 match name {
466 "AWS_REGION" => "region".to_owned(),
467 "AZURE_RESOURCE_NAME" => "resource_name".to_owned(),
468 "GOOGLE_VERTEX_PROJECT" => "project".to_owned(),
469 "GOOGLE_VERTEX_LOCATION" => "location".to_owned(),
470 "GOOGLE_VERTEX_ENDPOINT" => "endpoint".to_owned(),
471 _ => name.to_ascii_lowercase(),
472 }
473}