systemprompt_models/profile/providers/
mod.rs1mod error;
20mod protocol;
21mod surface;
22
23use std::collections::{HashMap, HashSet};
24
25use serde::{Deserialize, Serialize};
26use systemprompt_identifiers::{ModelId, ProviderId, SecretName};
27
28use crate::services::ai::{ModelCapabilities, ModelLimits, ModelPricing};
29
30pub use error::{ProviderRegistryError, ProviderRegistryResult};
31pub use protocol::WireProtocol;
32pub use surface::ApiSurface;
33
34const DEFAULT_CATALOG_YAML: &str = include_str!("default_catalog.yaml");
35
36#[derive(Deserialize)]
37struct DefaultCatalogFile {
38 providers: Vec<ProviderEntry>,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
42#[serde(deny_unknown_fields)]
43pub struct ProviderModel {
44 pub id: ModelId,
45
46 #[serde(default, skip_serializing_if = "Vec::is_empty")]
47 pub aliases: Vec<ModelId>,
48
49 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub upstream_model: Option<String>,
54
55 #[serde(default)]
56 pub pricing: ModelPricing,
57
58 #[serde(default)]
59 pub capabilities: ModelCapabilities,
60
61 #[serde(default)]
62 pub limits: ModelLimits,
63}
64
65impl ProviderModel {
66 #[must_use]
67 pub fn matches(&self, requested: &str) -> bool {
68 self.id.as_str() == requested || self.aliases.iter().any(|a| a.as_str() == requested)
69 }
70
71 #[must_use]
72 pub fn effective_upstream_model<'a>(&'a self, requested: &'a str) -> &'a str {
73 self.upstream_model.as_deref().unwrap_or(requested)
74 }
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
78#[serde(deny_unknown_fields)]
79pub struct ProviderEntry {
80 pub name: ProviderId,
81
82 pub wire: WireProtocol,
85
86 pub surface: ApiSurface,
90
91 pub endpoint: String,
92
93 pub api_key_secret: SecretName,
94
95 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
96 pub extra_headers: HashMap<String, String>,
97
98 #[serde(default, skip_serializing_if = "Vec::is_empty")]
99 pub models: Vec<ProviderModel>,
100}
101
102impl ProviderEntry {
103 #[must_use]
104 pub fn find_model(&self, requested: &str) -> Option<&ProviderModel> {
105 self.models.iter().find(|m| m.matches(requested))
106 }
107}
108
109#[derive(Debug, Clone, Default, Serialize, Deserialize, schemars::JsonSchema)]
110#[serde(transparent)]
111pub struct ProviderRegistry {
112 pub providers: Vec<ProviderEntry>,
113}
114
115impl ProviderRegistry {
116 pub fn default_seed() -> ProviderRegistryResult<Self> {
117 let file: DefaultCatalogFile = serde_yaml::from_str(DEFAULT_CATALOG_YAML)
118 .map_err(|e| ProviderRegistryError::InvalidDefaultCatalog(e.to_string()))?;
119 Ok(Self {
120 providers: file.providers,
121 })
122 }
123
124 #[must_use]
125 pub fn find_provider(&self, name: &str) -> Option<&ProviderEntry> {
126 self.providers.iter().find(|p| p.name.as_str() == name)
127 }
128
129 #[must_use]
130 pub fn contains_model(&self, requested: &str) -> bool {
131 self.providers
132 .iter()
133 .any(|p| p.find_model(requested).is_some())
134 }
135
136 pub fn advertised_providers(&self) -> impl Iterator<Item = &ProviderEntry> {
142 self.providers
143 .iter()
144 .filter(|entry| entry.surface.is_advertised())
145 }
146
147 #[must_use]
154 pub fn advertised_model_ids(&self, surfaces: &[ApiSurface]) -> Vec<String> {
155 self.advertised_providers()
156 .filter(|entry| surfaces.is_empty() || surfaces.contains(&entry.surface))
157 .flat_map(|entry| {
158 entry.models.iter().flat_map(|m| {
159 std::iter::once(m.id.as_str().to_owned())
160 .chain(m.aliases.iter().map(|a| a.as_str().to_owned()))
161 })
162 })
163 .collect()
164 }
165
166 pub fn validate(&self) -> ProviderRegistryResult<()> {
167 let trusted = crate::net::trusted_http_hosts_from_env();
168 let mut seen_providers: HashSet<&str> = HashSet::with_capacity(self.providers.len());
169 let mut seen_models: HashSet<&str> = HashSet::new();
170
171 for provider in &self.providers {
172 if !seen_providers.insert(provider.name.as_str()) {
173 return Err(ProviderRegistryError::DuplicateProvider {
174 name: provider.name.as_str().to_owned(),
175 });
176 }
177 if provider.endpoint.is_empty() {
178 return Err(ProviderRegistryError::EmptyEndpoint {
179 name: provider.name.as_str().to_owned(),
180 });
181 }
182 crate::net::validate_outbound_url_with_trust(&provider.endpoint, &trusted).map_err(
183 |e| ProviderRegistryError::BlockedEndpoint {
184 provider: provider.name.as_str().to_owned(),
185 endpoint: provider.endpoint.clone(),
186 reason: e.to_string(),
187 },
188 )?;
189
190 for model in &provider.models {
191 if model.id.as_str().is_empty() {
192 return Err(ProviderRegistryError::EmptyModelId {
193 id: provider.name.as_str().to_owned(),
194 });
195 }
196 if !seen_models.insert(model.id.as_str()) {
197 return Err(ProviderRegistryError::DuplicateModel {
198 id: model.id.as_str().to_owned(),
199 });
200 }
201 for alias in &model.aliases {
202 if !seen_models.insert(alias.as_str()) {
203 return Err(ProviderRegistryError::DuplicateModel {
204 id: alias.as_str().to_owned(),
205 });
206 }
207 }
208 }
209 }
210 Ok(())
211 }
212}