cookie_agent_models/adapters/
endpoints.rs

1use crate::{
2    adapters::OvenAdapterFamily,
3    authoring::EndpointUrl,
4    recipes::{EndpointPolicy, ValidatedSetup},
5};
6use url::Url;
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub enum BaseUrlOverridePolicy {
10    ManagedHttps,
11    Forbidden,
12    CustomHttpsOrReviewedLoopback { path: &'static str },
13}
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
16pub enum EndpointBuildError {
17    #[error("authored_base_url_forbidden")]
18    AuthoredOverrideForbidden,
19    #[error("invalid_endpoint_setup")]
20    Setup,
21    #[error("endpoint_override_policy_violation")]
22    Policy,
23}
24
25#[must_use]
26pub const fn managed_base_url_policy(policy: EndpointPolicy) -> BaseUrlOverridePolicy {
27    match policy {
28        EndpointPolicy::DefaultWithAuthoredHttpsOverride { .. } => {
29            BaseUrlOverridePolicy::ManagedHttps
30        }
31        EndpointPolicy::VertexPublisher
32        | EndpointPolicy::BedrockRegional
33        | EndpointPolicy::AzureOpenai => BaseUrlOverridePolicy::Forbidden,
34    }
35}
36
37#[must_use]
38pub const fn custom_endpoint_policy(family: OvenAdapterFamily) -> BaseUrlOverridePolicy {
39    BaseUrlOverridePolicy::CustomHttpsOrReviewedLoopback {
40        path: match family {
41            OvenAdapterFamily::Anthropic
42            | OvenAdapterFamily::AnthropicCompatible
43            | OvenAdapterFamily::OpenaiChat
44            | OvenAdapterFamily::OpenaiResponses
45            | OvenAdapterFamily::OpenaiCompatible
46            | OvenAdapterFamily::GoogleVertexGemini => "/v1",
47            OvenAdapterFamily::GoogleGemini => "/v1beta",
48            OvenAdapterFamily::AwsBedrockConverse
49            | OvenAdapterFamily::AzureOpenaiChat
50            | OvenAdapterFamily::AzureOpenaiResponses => "/",
51            OvenAdapterFamily::CohereV2Chat => "/v2",
52        },
53    }
54}
55
56pub fn validate_managed_base_url(
57    policy: EndpointPolicy,
58    authored: Option<&EndpointUrl>,
59) -> Result<(), EndpointBuildError> {
60    let Some(authored) = authored else {
61        return Ok(());
62    };
63    match managed_base_url_policy(policy) {
64        BaseUrlOverridePolicy::ManagedHttps => {
65            let parsed = Url::parse(authored.as_str()).map_err(|_| EndpointBuildError::Policy)?;
66            if parsed.scheme() == "https" {
67                Ok(())
68            } else {
69                Err(EndpointBuildError::Policy)
70            }
71        }
72        BaseUrlOverridePolicy::Forbidden => Err(EndpointBuildError::AuthoredOverrideForbidden),
73        BaseUrlOverridePolicy::CustomHttpsOrReviewedLoopback { .. } => unreachable!(),
74    }
75}
76
77pub fn validate_custom_endpoint(
78    family: OvenAdapterFamily,
79    endpoint: &EndpointUrl,
80) -> Result<(), EndpointBuildError> {
81    let parsed = Url::parse(endpoint.as_str()).map_err(|_| EndpointBuildError::Policy)?;
82    if parsed.scheme() == "https" {
83        return Ok(());
84    }
85    let BaseUrlOverridePolicy::CustomHttpsOrReviewedLoopback { path } =
86        custom_endpoint_policy(family)
87    else {
88        unreachable!()
89    };
90    let host = parsed.host_str().ok_or(EndpointBuildError::Policy)?;
91    let exact_loopback = matches!(host, "localhost" | "127.0.0.1" | "::1");
92    if parsed.scheme() == "http"
93        && exact_loopback
94        && parsed.port().is_some()
95        && parsed.path() == path
96    {
97        Ok(())
98    } else {
99        Err(EndpointBuildError::Policy)
100    }
101}
102
103pub fn build_endpoint(
104    policy: EndpointPolicy,
105    authored: Option<&EndpointUrl>,
106    setup: &ValidatedSetup,
107) -> Result<String, EndpointBuildError> {
108    validate_managed_base_url(policy, authored)?;
109    match policy {
110        EndpointPolicy::DefaultWithAuthoredHttpsOverride { default } => Ok(authored
111            .map(EndpointUrl::as_str)
112            .unwrap_or(default)
113            .trim_end_matches('/')
114            .to_owned()),
115        EndpointPolicy::VertexPublisher => {
116            if authored.is_some() {
117                return Err(EndpointBuildError::AuthoredOverrideForbidden);
118            }
119            let project = field(setup, "project")?;
120            let location = field(setup, "location")?;
121            let resource = field(setup, "resource")?;
122            Ok(format!(
123                "https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/{resource}"
124            ))
125        }
126        EndpointPolicy::BedrockRegional => {
127            if authored.is_some() {
128                return Err(EndpointBuildError::AuthoredOverrideForbidden);
129            }
130            Ok(format!(
131                "https://bedrock-runtime.{}.amazonaws.com",
132                field(setup, "region")?
133            ))
134        }
135        EndpointPolicy::AzureOpenai => {
136            if authored.is_some() {
137                return Err(EndpointBuildError::AuthoredOverrideForbidden);
138            }
139            Ok(format!(
140                "https://{}.openai.azure.com",
141                field(setup, "resource_name")?
142            ))
143        }
144    }
145}
146
147fn field<'a>(setup: &'a ValidatedSetup, id: &str) -> Result<&'a str, EndpointBuildError> {
148    setup
149        .values
150        .get(id)
151        .map(String::as_str)
152        .ok_or(EndpointBuildError::Setup)
153}
cookie_agent_models/adapters/endpoints.rs

cookie_agent_models/adapters/
endpoints.rs
