systemprompt_loader/vertex_discovery/

classify.rs

//! Deciding which Vertex listing entries are models we can actually serve.
//!
//! A Model Garden listing mixes three populations under one JSON shape:
//! serverless models Google hosts and bills per token; *deployable
//! checkpoints*, which are weights plus a serving container and are callable
//! only after you stand up an endpoint yourself; and, under
//! `publishers/google`, every non-chat modality Google sells — embeddings,
//! speech, image, video, robotics. There is no field that separates them by
//! modality.
//!
//! So the rule is in two halves. Shape rules out what cannot be called at all
//! (a checkpoint with a `deploy` action; a partner entry that is not `MaaS`).
//! The rate card rules in what we are willing to serve — it is the only place
//! that knows `gemini-2.5-flash` is chat and `gemini-embedding-001` is not.
//!
//! Shape (`is_serverless`): Google's own publisher is served serverlessly
//! across the board, so shape says nothing there and everything is a
//! candidate. A partner model qualifies only as `MaaS` — the `-maas` suffix
//! Vertex gives every serverless partner model, the third-party OSS category,
//! and no actions of its own. Anything with a `deploy` action is a checkpoint.
//!
//! Pricing (`classify_discovered`) keys the card by the listing endpoint's
//! wire and knows nothing about who listed the model;
//! every [`CatalogSource`](super::source::CatalogSource) is judged by exactly
//! this rule, on the provider-agnostic [`DiscoveredModel`] shape.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use serde::Deserialize;
use systemprompt_manifest::services::{VertexRateCard, VertexRateCardEntry};
use systemprompt_wire::WireProtocol;

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";

/// One entry of `publisherModels`, named
/// `publishers/{publisher}/models/{model}`; the rate card names an upstream as
/// `{publisher}/{model}`.
///
/// Unknown fields are ignored on purpose: the listing carries presentation
/// data (notebook links, container specs, regional availability) that grows
/// without notice, and a boot-time reader that fails on a new field would turn
/// a Google release note into an outage.
#[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)]
    // JSON: Vertex AI `supportedActions` — undocumented action map, only probed for keys.
    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 declares_actions(&self) -> bool {
        match self.supported_actions.as_ref() {
            None | Some(serde_json::Value::Null) => false,
            Some(serde_json::Value::Object(actions)) => !actions.is_empty(),
            Some(_) => true,
        }
    }
}

#[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.declares_actions()
}

/// What discovery decided about one listing entry.
///
/// Not serverless (needs a deploy, or not a `MaaS` partner model); callable
/// but unpriced by the rate card, so never served; priced but not GA with a
/// card entry that does not opt into previews; or priced and publishable.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Classification {
    NotServerless,
    Unpriced,
    PreviewWithheld,
    Publish,
}

#[must_use]
pub fn classify<'a>(
    model: &PublisherModel,
    card: &'a VertexRateCard,
    wire: WireProtocol,
) -> (Classification, Option<&'a VertexRateCardEntry>) {
    classify_discovered(&model.discovered(), card, wire)
}

#[must_use]
pub fn classify_discovered<'a>(
    model: &DiscoveredModel,
    card: &'a VertexRateCard,
    wire: WireProtocol,
) -> (Classification, Option<&'a VertexRateCardEntry>) {
    if !model.serverless {
        return (Classification::NotServerless, None);
    }
    let Some(entry) = card.lookup(&model.upstream).filter(|e| e.wire == wire) else {
        return (Classification::Unpriced, None);
    };
    if !model.launch_stage.is_generally_available() && !entry.allow_preview {
        return (Classification::PreviewWithheld, Some(entry));
    }
    (Classification::Publish, Some(entry))
}