use std::collections::HashSet;
use std::hash::BuildHasher;
use chrono::NaiveDate;
use systemprompt_identifiers::ModelId;
use systemprompt_manifest::services::catalog::{CatalogDocument, DeploymentTarget};
use systemprompt_manifest::services::{DiscoveryReport, VertexRateCard, VertexRateCardEntry};
use systemprompt_models::providers::ClientFamily;
use systemprompt_wire::WireProtocol;
use super::source::DiscoveryEndpoint;
fn deployed_on(document: &CatalogDocument, model: &ModelId, endpoint: &DiscoveryEndpoint) -> bool {
document
.deployments
.get(model.as_str())
.is_some_and(|targets| targets.iter().any(|t| *t.endpoint() == endpoint.endpoint))
}
fn declared_as(
document: &CatalogDocument,
endpoint: &DiscoveryEndpoint,
entry: &VertexRateCardEntry,
) -> Option<ModelId> {
if document.models.get(entry.id.as_str()).is_some() {
return Some(entry.id.clone());
}
document
.models
.iter()
.find(|(id, model)| {
model.upstream_name(id) == entry.upstream_model && deployed_on(document, id, endpoint)
})
.map(|(id, _)| id.clone())
}
const fn wire_family(wire: WireProtocol) -> ClientFamily {
match wire {
WireProtocol::Anthropic => ClientFamily::Claude,
WireProtocol::OpenAiChat | WireProtocol::OpenAiResponses => ClientFamily::OpenAi,
WireProtocol::Gemini => ClientFamily::Gemini,
}
}
fn endpoint_clients(document: &CatalogDocument, endpoint: &DiscoveryEndpoint) -> Vec<ClientFamily> {
let mut clients: Vec<ClientFamily> = Vec::new();
for (id, model) in document.models.iter() {
if deployed_on(document, id, endpoint) {
for family in &model.clients {
if !clients.contains(family) {
clients.push(*family);
}
}
}
}
if clients.is_empty() {
clients.push(wire_family(endpoint.wire));
}
clients
}
pub fn publish(
document: &mut CatalogDocument,
endpoint: &DiscoveryEndpoint,
entry: &VertexRateCardEntry,
today: NaiveDate,
report: &mut DiscoveryReport,
) {
let id = entry.id.as_str().to_owned();
let declared = declared_as(document, endpoint, entry);
if !entry.is_supported(today) {
match declared {
Some(model) if entry.is_retired(today) => {
tracing::error!(
model = %model,
retires_on = ?entry.retires_on,
docs = %entry.docs,
"declared Vertex model is past its retirement date and no longer exists \
upstream; it is no longer advertised. Remove it from the catalog"
);
if let Some(declaration) = document.models.models.get_mut(&model) {
declaration.clients.clear();
}
},
Some(model) => {
tracing::warn!(
model = %model,
retires_on = ?entry.retires_on,
docs = %entry.docs,
"explicitly declared Vertex model is retiring or unsupported; it stays served \
because the catalog declares it, but it should be removed"
);
},
None => {},
}
report.retiring.push(id);
return;
}
if declared.is_some() {
report.explicit_wins.push(id);
return;
}
let clients = endpoint_clients(document, endpoint);
document
.models
.models
.insert(entry.id.clone(), entry.to_model_entry(clients));
document.deployments.deployments.insert(
entry.id.clone(),
vec![DeploymentTarget::Endpoint(endpoint.endpoint.clone())],
);
report.discovered_priced.push(id);
}
pub fn record_unpriced(upstream: String, report: &mut DiscoveryReport) {
if !report.discovered_unpriced.contains(&upstream) {
report.discovered_unpriced.push(upstream);
}
}
pub fn record_unseen<S: BuildHasher>(
card: &VertexRateCard,
wire: WireProtocol,
seen: &HashSet<String, S>,
report: &mut DiscoveryReport,
) {
for entry in card.entries_for(wire) {
if !seen.contains(&entry.upstream) {
report
.priced_not_published
.push(entry.id.as_str().to_owned());
}
}
}