use serde::Deserialize;
use systemprompt_models::services::{VertexRateCard, VertexRateCardEntry};
use super::source::{DiscoveredModel, LaunchStage};
const MAAS_SUFFIX: &str = "-maas";
const THIRD_PARTY_OSS: &str = "THIRD_PARTY_OWNED_OSS";
const GOOGLE_PUBLISHER: &str = "google";
const GA: &str = "GA";
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PublisherModel {
pub name: String,
#[serde(default)]
pub version_id: String,
#[serde(default)]
pub launch_stage: String,
#[serde(default)]
pub supported_actions: Option<serde_json::Value>,
#[serde(default)]
pub open_source_category: Option<String>,
}
impl PublisherModel {
#[must_use]
pub fn publisher(&self) -> &str {
let mut parts = self.name.split('/');
if parts.next() == Some("publishers") {
parts.next().unwrap_or_default()
} else {
""
}
}
#[must_use]
pub fn model_name(&self) -> &str {
self.name.rsplit('/').next().unwrap_or(&self.name)
}
#[must_use]
pub fn upstream(&self) -> String {
format!("{}/{}", self.publisher(), self.model_name())
}
#[must_use]
pub fn is_generally_available(&self) -> bool {
self.launch_stage == GA
}
#[must_use]
pub fn discovered(&self) -> DiscoveredModel {
DiscoveredModel {
upstream: self.upstream(),
launch_stage: if self.is_generally_available() {
LaunchStage::GenerallyAvailable
} else {
LaunchStage::Preview
},
serverless: is_serverless(self),
}
}
fn is_deployable_checkpoint(&self) -> bool {
self.supported_actions
.as_ref()
.is_some_and(|actions| actions.get("deploy").is_some())
}
fn has_no_actions(&self) -> bool {
self.supported_actions.as_ref().is_none_or(|actions| {
actions.as_object().is_some_and(serde_json::Map::is_empty) || actions.is_null()
})
}
}
#[must_use]
pub fn is_serverless(model: &PublisherModel) -> bool {
if model.is_deployable_checkpoint() {
return false;
}
if model.publisher() == GOOGLE_PUBLISHER {
return true;
}
model.model_name().ends_with(MAAS_SUFFIX)
&& model.open_source_category.as_deref() == Some(THIRD_PARTY_OSS)
&& model.has_no_actions()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Classification {
NotServerless,
Unpriced,
PreviewWithheld,
Publish,
}
#[must_use]
pub fn classify<'a>(
model: &PublisherModel,
card: &'a VertexRateCard,
provider: &str,
) -> (Classification, Option<&'a VertexRateCardEntry>) {
classify_discovered(&model.discovered(), card, provider)
}
#[must_use]
pub fn classify_discovered<'a>(
model: &DiscoveredModel,
card: &'a VertexRateCard,
provider: &str,
) -> (Classification, Option<&'a VertexRateCardEntry>) {
if !model.serverless {
return (Classification::NotServerless, None);
}
let Some(entry) = card
.lookup(&model.upstream)
.filter(|e| e.provider.as_str() == provider)
else {
return (Classification::Unpriced, None);
};
if !model.launch_stage.is_generally_available() && !entry.allow_preview {
return (Classification::PreviewWithheld, Some(entry));
}
(Classification::Publish, Some(entry))
}