Skip to main content

cookie_agent_models/recipes/
setup.rs

1use std::collections::BTreeMap;
2
3use cookie_agent_identity::SetupFieldId;
4use serde::Serialize;
5
6use crate::authoring::{BoundedSetupString, SafeSetupValue};
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
9#[serde(rename_all = "snake_case")]
10pub enum SetupFieldType {
11    String,
12}
13
14#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
15pub struct SetupFieldRecipe {
16    pub id: &'static str,
17    pub value_type: SetupFieldType,
18    pub required: bool,
19    pub default: Option<&'static str>,
20    pub environment_alias: Option<&'static str>,
21}
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
24pub struct SetupRecipe {
25    pub id: &'static str,
26    pub fields: &'static [SetupFieldRecipe],
27}
28
29#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
30#[serde(rename_all = "snake_case")]
31pub enum EndpointPolicy {
32    DefaultWithAuthoredHttpsOverride { default: &'static str },
33    VertexPublisher,
34    BedrockRegional,
35    AzureOpenai,
36}
37
38#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
39pub struct ValidatedSetup {
40    pub recipe_id: &'static str,
41    pub values: BTreeMap<String, String>,
42}
43
44#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
45pub enum SetupValidationError {
46    #[error("invalid_setup_shape")]
47    Shape,
48    #[error("invalid_setup_value")]
49    Value,
50}
51
52pub fn validate_setup(
53    recipe: &SetupRecipe,
54    input: &BTreeMap<SetupFieldId, SafeSetupValue>,
55) -> Result<ValidatedSetup, SetupValidationError> {
56    let mut values = BTreeMap::new();
57    for (id, value) in input {
58        let field = recipe
59            .fields
60            .iter()
61            .find(|field| field.id == id.as_str())
62            .ok_or(SetupValidationError::Shape)?;
63        let value = setup_string(value).ok_or(SetupValidationError::Value)?;
64        validate_field(field.id, value)?;
65        values.insert(field.id.to_owned(), value.to_owned());
66    }
67    for field in recipe.fields {
68        if !values.contains_key(field.id) {
69            if let Some(default) = field.default {
70                values.insert(field.id.to_owned(), default.to_owned());
71            } else if field.required {
72                return Err(SetupValidationError::Shape);
73            }
74        }
75    }
76    Ok(ValidatedSetup {
77        recipe_id: recipe.id,
78        values,
79    })
80}
81
82fn setup_string(value: &SafeSetupValue) -> Option<&str> {
83    match value {
84        SafeSetupValue::String(value) => Some(BoundedSetupString::as_str(value)),
85        SafeSetupValue::Code(value) => Some(value.as_str()),
86        SafeSetupValue::Integer(_) | SafeSetupValue::Bool(_) => None,
87    }
88}
89
90fn validate_field(id: &str, value: &str) -> Result<(), SetupValidationError> {
91    let bounded = !value.is_empty() && value.len() <= 256 && value.is_ascii();
92    let lower_token = bounded
93        && value
94            .bytes()
95            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-');
96    let valid = match id {
97        "project" => lower_token && value.as_bytes()[0].is_ascii_alphanumeric(),
98        "location" | "region" | "resource_name" | "deployment" => lower_token,
99        "model" | "version" | "deployment_type" => bounded && !value.chars().any(char::is_control),
100        "resource" => value == "publishers/google",
101        "api_version" => {
102            value.len() == 10
103                && value.as_bytes()[4] == b'-'
104                && value.as_bytes()[7] == b'-'
105                && value
106                    .bytes()
107                    .enumerate()
108                    .all(|(index, byte)| matches!(index, 4 | 7) || byte.is_ascii_digit())
109        }
110        _ => false,
111    };
112    if valid {
113        Ok(())
114    } else {
115        Err(SetupValidationError::Value)
116    }
117}