Skip to main content

cookie_agent_models/recipes/
families.rs

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/// Maps one provider-level credential method onto the wire auth expected by an
211/// effective model family. Only equivalent credential semantics are mapped:
212/// single API keys may change header/bearer encoding, and an explicit access
213/// token may become bearer auth. AWS static credentials never cross families.
214#[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            // `@ai-sdk/anthropic` sends configured API keys as `x-api-key`, and
241            // Microsoft Foundry accepts `x-api-key` on its Anthropic endpoint.
242            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}