systemprompt_models/services/providers/
mod.rs1mod discovery_report;
24mod error;
25mod protocol;
26mod rate_card;
27mod surface;
28
29use std::collections::{HashMap, HashSet};
30
31use serde::{Deserialize, Serialize};
32use systemprompt_identifiers::{ModelId, ProviderId, SecretName};
33
34use crate::services::ai::{ModelCapabilities, ModelGovernance, ModelLimits, ModelPricing};
35
36pub use discovery_report::DiscoveryReport;
37pub use error::{ProviderRegistryError, ProviderRegistryResult};
38pub use protocol::WireProtocol;
39pub use rate_card::{
40 DocumentedLaunchStage, RETIREMENT_NOTICE_DAYS, VertexRateCard, VertexRateCardEntry,
41};
42pub use surface::ApiSurface;
43
44const DEFAULT_CATALOG_YAML: &str = include_str!("default_catalog.yaml");
45
46#[derive(Deserialize)]
47struct DefaultCatalogFile {
48 providers: Vec<ProviderEntry>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
52#[serde(deny_unknown_fields)]
53pub struct ProviderModel {
54 pub id: ModelId,
55
56 #[serde(default, skip_serializing_if = "Vec::is_empty")]
57 pub aliases: Vec<ModelId>,
58
59 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub upstream_model: Option<String>,
61
62 #[serde(default)]
63 pub pricing: ModelPricing,
64
65 #[serde(default)]
66 pub capabilities: ModelCapabilities,
67
68 #[serde(default)]
69 pub limits: ModelLimits,
70
71 #[serde(default, skip_serializing_if = "Option::is_none")]
72 pub governance: Option<ModelGovernance>,
73}
74
75impl ProviderModel {
76 #[must_use]
77 pub fn matches(&self, requested: &str) -> bool {
78 self.id.as_str() == requested || self.aliases.iter().any(|a| a.as_str() == requested)
79 }
80
81 #[must_use]
82 pub fn effective_upstream_model<'a>(&'a self, requested: &'a str) -> &'a str {
83 self.upstream_model.as_deref().unwrap_or(requested)
84 }
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
88#[serde(deny_unknown_fields)]
89pub struct ProviderEntry {
90 pub name: ProviderId,
91
92 pub wire: WireProtocol,
93
94 pub surface: ApiSurface,
95
96 pub endpoint: String,
97
98 pub api_key_secret: SecretName,
99
100 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
101 pub extra_headers: HashMap<String, String>,
102
103 #[serde(default, skip_serializing_if = "Vec::is_empty")]
104 pub models: Vec<ProviderModel>,
105
106 #[serde(default)]
107 pub governance: ModelGovernance,
108}
109
110impl ProviderEntry {
111 #[must_use]
112 pub fn find_model(&self, requested: &str) -> Option<&ProviderModel> {
113 self.models.iter().find(|m| m.matches(requested))
114 }
115
116 #[must_use]
117 pub fn upstream_model_for<'a>(
118 &'a self,
119 route_override: Option<&'a str>,
120 requested: &'a str,
121 ) -> &'a str {
122 route_override.unwrap_or_else(|| {
123 self.find_model(requested)
124 .map_or(requested, |model| model.effective_upstream_model(requested))
125 })
126 }
127
128 #[must_use]
129 pub fn effective_governance(&self, requested: &str) -> ModelGovernance {
130 self.find_model(requested)
131 .and_then(|m| m.governance)
132 .unwrap_or(self.governance)
133 }
134}
135
136#[derive(Debug, Clone, Default, Serialize, Deserialize, schemars::JsonSchema)]
137#[serde(transparent)]
138pub struct ProviderRegistry {
139 pub providers: Vec<ProviderEntry>,
140}
141
142impl ProviderRegistry {
143 pub fn default_seed() -> ProviderRegistryResult<Self> {
144 let file: DefaultCatalogFile = serde_yaml::from_str(DEFAULT_CATALOG_YAML)
145 .map_err(|e| ProviderRegistryError::InvalidDefaultCatalog(e.to_string()))?;
146 Ok(Self {
147 providers: file.providers,
148 })
149 }
150
151 #[must_use]
152 pub fn find_provider(&self, name: &str) -> Option<&ProviderEntry> {
153 self.providers.iter().find(|p| p.name.as_str() == name)
154 }
155
156 #[must_use]
157 pub fn contains_model(&self, requested: &str) -> bool {
158 self.providers
159 .iter()
160 .any(|p| p.find_model(requested).is_some())
161 }
162
163 pub fn advertised_providers(&self) -> impl Iterator<Item = &ProviderEntry> {
164 self.providers
165 .iter()
166 .filter(|entry| entry.surface.is_advertised())
167 }
168
169 #[must_use]
170 pub fn advertised_model_ids(&self, surfaces: &[ApiSurface]) -> Vec<String> {
171 self.advertised_providers()
172 .filter(|entry| surfaces.is_empty() || surfaces.contains(&entry.surface))
173 .flat_map(|entry| {
174 entry.models.iter().flat_map(|m| {
175 std::iter::once(m.id.as_str().to_owned())
176 .chain(m.aliases.iter().map(|a| a.as_str().to_owned()))
177 })
178 })
179 .collect()
180 }
181
182 pub fn validate(&self) -> ProviderRegistryResult<()> {
183 let trusted = crate::net::trusted_http_hosts_from_env();
184 let mut seen_providers: HashSet<&str> = HashSet::with_capacity(self.providers.len());
185 let mut seen_models: HashSet<&str> = HashSet::new();
186
187 for provider in &self.providers {
188 if !seen_providers.insert(provider.name.as_str()) {
189 return Err(ProviderRegistryError::DuplicateProvider {
190 name: provider.name.as_str().to_owned(),
191 });
192 }
193 if provider.endpoint.is_empty() {
194 return Err(ProviderRegistryError::EmptyEndpoint {
195 name: provider.name.as_str().to_owned(),
196 });
197 }
198 crate::net::validate_outbound_url_with_trust(&provider.endpoint, &trusted).map_err(
199 |e| ProviderRegistryError::BlockedEndpoint {
200 provider: provider.name.as_str().to_owned(),
201 endpoint: provider.endpoint.clone(),
202 reason: e.to_string(),
203 },
204 )?;
205 if names_a_project_literally(&provider.endpoint) {
206 return Err(ProviderRegistryError::LiteralProjectInEndpoint {
207 provider: provider.name.as_str().to_owned(),
208 endpoint: provider.endpoint.clone(),
209 });
210 }
211
212 for model in &provider.models {
213 if model.id.as_str().is_empty() {
214 return Err(ProviderRegistryError::EmptyModelId {
215 id: provider.name.as_str().to_owned(),
216 });
217 }
218 if !seen_models.insert(model.id.as_str()) {
219 return Err(ProviderRegistryError::DuplicateModel {
220 id: model.id.as_str().to_owned(),
221 });
222 }
223 for alias in &model.aliases {
224 if !seen_models.insert(alias.as_str()) {
225 return Err(ProviderRegistryError::DuplicateModel {
226 id: alias.as_str().to_owned(),
227 });
228 }
229 }
230 }
231 }
232 Ok(())
233 }
234}
235
236pub const PROJECT_PLACEHOLDER: &str = "{project}";
237
238pub const REGION_PLACEHOLDER: &str = "{region}";
245
246#[must_use]
253pub fn names_a_project_literally(endpoint: &str) -> bool {
254 let Ok(url) = url::Url::parse(endpoint) else {
255 return false;
256 };
257 let on_vertex = url.host_str().is_some_and(|host| {
258 host.eq_ignore_ascii_case("aiplatform.googleapis.com")
259 || host
260 .to_ascii_lowercase()
261 .ends_with("-aiplatform.googleapis.com")
262 });
263 if !on_vertex {
264 return false;
265 }
266 let mut segments = url.path_segments().into_iter().flatten();
267 while let Some(segment) = segments.next() {
268 if segment == "projects" {
269 return segments
270 .next()
271 .is_some_and(|id| id != "%7Bproject%7D" && id != PROJECT_PLACEHOLDER);
272 }
273 }
274 false
275}