Integrate Bedrock model catalog metadata

2026-08-22 12:51:31 -05:00
parent f291cfe803
commit f17642fc62
24 changed files with 27662 additions and 26 deletions
@@ -0,0 +1,65 @@
+//! A versioned, offline snapshot of the public Amazon Bedrock model catalog.
+//!
+//! The JSON is kept verbatim so consumers can adopt new upstream fields without
+//! first changing this crate. Run `./script/update_bedrock_model_catalog` from
+//! the workspace root to refresh the snapshot.
+
+mod metadata;
+
+pub use metadata::{ModelMetadata, model_metadata, models};
+
+/// Repository from which the bundled snapshot was downloaded.
+pub const UPSTREAM_REPOSITORY: &str =
+    "https://github.com/amazonbedrockmodels/amazonbedrockmodels.github.io";
+
+/// The catalog's public API manifest.
+pub const API_MANIFEST_JSON: &str = include_str!("../data/api.json");
+
+/// Foundation models discovered through the AWS-native Bedrock APIs.
+pub const FOUNDATION_MODELS_JSON: &str = include_str!("../data/models.json");
+
+/// System-defined inference profiles discovered across AWS regions.
+pub const INFERENCE_PROFILES_JSON: &str = include_str!("../data/profiles.json");
+
+/// Models discovered before they appear in the main AWS documentation.
+pub const BETA_MODELS_JSON: &str = include_str!("../data/beta_models.json");
+
+/// Models available through the Bedrock Mantle endpoint.
+pub const MANTLE_MODELS_JSON: &str = include_str!("../data/mantle_models.json");
+
+/// Enriched model-card metadata maintained by the upstream catalog.
+pub const MODEL_CARDS_JSON: &str = include_str!("../data/model_cards.json");
+
+/// Returns the exact upstream commit used for the bundled snapshot.
+pub fn upstream_commit() -> &'static str {
+    include_str!("../data/UPSTREAM_COMMIT").trim()
+}
+
+/// All catalog documents embedded in this crate.
+#[derive(Debug, Clone, Copy)]
+pub struct CatalogSnapshot {
+    pub api_manifest: &'static str,
+    pub foundation_models: &'static str,
+    pub inference_profiles: &'static str,
+    pub beta_models: &'static str,
+    pub mantle_models: &'static str,
+    pub model_cards: &'static str,
+}
+
+/// The bundled catalog snapshot.
+pub const SNAPSHOT: CatalogSnapshot = CatalogSnapshot {
+    api_manifest: API_MANIFEST_JSON,
+    foundation_models: FOUNDATION_MODELS_JSON,
+    inference_profiles: INFERENCE_PROFILES_JSON,
+    beta_models: BETA_MODELS_JSON,
+    mantle_models: MANTLE_MODELS_JSON,
+    model_cards: MODEL_CARDS_JSON,
+};
+
+#[cfg(test)]
+#[path = "lib_tests.rs"]
+mod tests;
+
+#[cfg(test)]
+#[path = "metadata_tests.rs"]
+mod metadata_tests;
@@ -0,0 +1,42 @@
+use serde_json::Value;
+
+use super::*;
+
+#[test]
+fn bundled_catalog_documents_are_valid_json() {
+    for document in [
+        API_MANIFEST_JSON,
+        FOUNDATION_MODELS_JSON,
+        INFERENCE_PROFILES_JSON,
+        BETA_MODELS_JSON,
+        MANTLE_MODELS_JSON,
+        MODEL_CARDS_JSON,
+    ] {
+        serde_json::from_str::<Value>(document).expect("catalog document should be valid JSON");
+    }
+}
+
+#[test]
+fn bundled_catalog_has_models_and_provenance() {
+    let foundation_models: Value =
+        serde_json::from_str(FOUNDATION_MODELS_JSON).expect("foundation models should parse");
+    let mantle_models: Value =
+        serde_json::from_str(MANTLE_MODELS_JSON).expect("Mantle models should parse");
+
+    assert!(
+        foundation_models
+            .as_array()
+            .is_some_and(|models| !models.is_empty())
+    );
+    assert!(
+        mantle_models
+            .as_object()
+            .is_some_and(|models| !models.is_empty())
+    );
+    assert_eq!(upstream_commit().len(), 40);
+    assert!(
+        upstream_commit()
+            .bytes()
+            .all(|byte| byte.is_ascii_hexdigit())
+    );
+}
@@ -0,0 +1,228 @@
+use std::collections::HashMap;
+use std::sync::OnceLock;
+
+use serde::Deserialize;
+
+use crate::FOUNDATION_MODELS_JSON;
+
+const CROSS_REGION_PREFIXES: &[&str] = &["global.", "apac.", "us.", "eu.", "jp.", "au."];
+
+/// Provider-neutral metadata for one Bedrock foundation model.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ModelMetadata {
+    pub model_id: String,
+    pub model_name: String,
+    pub provider_name: String,
+    pub input_modalities: Vec<String>,
+    pub output_modalities: Vec<String>,
+    pub response_streaming_supported: Option<bool>,
+    pub regions: Vec<String>,
+    pub context_window_tokens: Option<u32>,
+    pub max_output_tokens: Option<u32>,
+    pub converse_supported: Option<bool>,
+    pub invoke_supported: Option<bool>,
+    pub bedrock_runtime_supported: Option<bool>,
+    pub bedrock_mantle_supported: Option<bool>,
+    pub mantle_regions: Vec<String>,
+    pub model_card_url: Option<String>,
+}
+
+impl ModelMetadata {
+    pub fn supports_vision_input(&self) -> bool {
+        self.input_modalities
+            .iter()
+            .any(|modality| modality == "IMAGE")
+    }
+
+    /// Reports whether the catalog has enough evidence to classify this as an
+    /// Agent Mode model. `None` deliberately leaves new or incomplete catalog
+    /// entries usable until their metadata catches up.
+    pub fn supports_agent_runtime(&self) -> Option<bool> {
+        if !self.output_modalities.is_empty()
+            && !self
+                .output_modalities
+                .iter()
+                .any(|modality| modality == "TEXT")
+        {
+            return Some(false);
+        }
+
+        for support in [
+            self.response_streaming_supported,
+            self.converse_supported,
+            self.bedrock_runtime_supported,
+        ] {
+            if support == Some(false) {
+                return Some(false);
+            }
+        }
+
+        if self.output_modalities.is_empty()
+            || self.response_streaming_supported.is_none()
+            || self.converse_supported.is_none()
+            || self.bedrock_runtime_supported.is_none()
+        {
+            None
+        } else {
+            Some(true)
+        }
+    }
+}
+
+/// Returns all parsed foundation-model metadata in the bundled snapshot.
+pub fn models() -> &'static [ModelMetadata] {
+    &catalog_index().models
+}
+
+/// Finds metadata for a foundation model, accepting Bedrock ARNs,
+/// cross-region inference prefixes, and Galaxy's optional `[1m]` suffix.
+pub fn model_metadata(model_id: &str) -> Option<&'static ModelMetadata> {
+    let index = catalog_index();
+    let normalized = normalize_model_id(model_id);
+    index
+        .by_id
+        .get(normalized)
+        .and_then(|position| index.models.get(*position))
+}
+
+fn catalog_index() -> &'static CatalogIndex {
+    static INDEX: OnceLock<CatalogIndex> = OnceLock::new();
+    INDEX.get_or_init(CatalogIndex::load)
+}
+
+struct CatalogIndex {
+    models: Vec<ModelMetadata>,
+    by_id: HashMap<String, usize>,
+}
+
+impl CatalogIndex {
+    fn load() -> Self {
+        let raw_models =
+            serde_json::from_str::<Vec<RawModel>>(FOUNDATION_MODELS_JSON).unwrap_or_default();
+        let models = raw_models
+            .into_iter()
+            .map(ModelMetadata::from)
+            .collect::<Vec<_>>();
+        let by_id = models
+            .iter()
+            .enumerate()
+            .map(|(position, model)| (model.model_id.clone(), position))
+            .collect();
+        Self { models, by_id }
+    }
+}
+
+impl From<RawModel> for ModelMetadata {
+    fn from(raw: RawModel) -> Self {
+        let card = raw.model_card.unwrap_or_default();
+        Self {
+            model_id: raw.model_id,
+            model_name: raw.model_name,
+            provider_name: raw.provider_name,
+            input_modalities: raw.input_modalities,
+            output_modalities: raw.output_modalities,
+            response_streaming_supported: raw.response_streaming_supported,
+            regions: raw.regions,
+            context_window_tokens: card.context_window.as_deref().and_then(parse_token_count),
+            max_output_tokens: card
+                .max_output_tokens
+                .as_deref()
+                .and_then(parse_token_count),
+            converse_supported: card.apis_supported.as_ref().and_then(|apis| apis.converse),
+            invoke_supported: card.apis_supported.as_ref().and_then(|apis| apis.invoke),
+            bedrock_runtime_supported: card
+                .endpoints_supported
+                .as_ref()
+                .and_then(|endpoints| endpoints.bedrock_runtime),
+            bedrock_mantle_supported: card
+                .endpoints_supported
+                .as_ref()
+                .and_then(|endpoints| endpoints.bedrock_mantle),
+            mantle_regions: card.mantle_regions,
+            model_card_url: card.model_card_url,
+        }
+    }
+}
+
+fn normalize_model_id(model_id: &str) -> &str {
+    let model_id = model_id.rsplit('/').next().unwrap_or(model_id);
+    let model_id = model_id
+        .strip_suffix("[1m]")
+        .or_else(|| model_id.strip_suffix("[1M]"))
+        .unwrap_or(model_id);
+    CROSS_REGION_PREFIXES
+        .iter()
+        .find_map(|prefix| model_id.strip_prefix(prefix))
+        .unwrap_or(model_id)
+}
+
+fn parse_token_count(value: &str) -> Option<u32> {
+    let value = value
+        .trim()
+        .strip_suffix("tokens")
+        .or_else(|| value.trim().strip_suffix("token"))
+        .unwrap_or(value)
+        .trim();
+    let suffix_start = value
+        .find(|character: char| !character.is_ascii_digit() && character != ',' && character != '.')
+        .unwrap_or(value.len());
+    let (number, suffix) = value.split_at(suffix_start);
+    let multiplier = match suffix.trim().to_ascii_uppercase().as_str() {
+        "" => 1.0,
+        "K" => 1_000.0,
+        "M" => 1_000_000.0,
+        "B" => 1_000_000_000.0,
+        _ => return None,
+    };
+    let number = number.replace(',', "").parse::<f64>().ok()?;
+    let tokens = number * multiplier;
+    if !tokens.is_finite() || tokens < 0.0 || tokens > u32::MAX as f64 {
+        return None;
+    }
+    Some(tokens.round() as u32)
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RawModel {
+    model_id: String,
+    #[serde(default)]
+    model_name: String,
+    #[serde(default)]
+    provider_name: String,
+    #[serde(default)]
+    input_modalities: Vec<String>,
+    #[serde(default)]
+    output_modalities: Vec<String>,
+    #[serde(default)]
+    response_streaming_supported: Option<bool>,
+    #[serde(default)]
+    regions: Vec<String>,
+    #[serde(default)]
+    model_card: Option<RawModelCard>,
+}
+
+#[derive(Default, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RawModelCard {
+    context_window: Option<String>,
+    max_output_tokens: Option<String>,
+    apis_supported: Option<RawApisSupported>,
+    endpoints_supported: Option<RawEndpointsSupported>,
+    #[serde(default)]
+    mantle_regions: Vec<String>,
+    model_card_url: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct RawApisSupported {
+    converse: Option<bool>,
+    invoke: Option<bool>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RawEndpointsSupported {
+    bedrock_runtime: Option<bool>,
+    bedrock_mantle: Option<bool>,
+}
@@ -0,0 +1,50 @@
+use super::*;
+
+#[test]
+fn reads_context_and_capabilities_for_an_agent_model() {
+    let metadata = model_metadata("anthropic.claude-opus-4-8").unwrap();
+
+    assert_eq!(metadata.model_name, "Claude Opus 4.8");
+    assert_eq!(metadata.provider_name, "Anthropic");
+    assert_eq!(metadata.context_window_tokens, Some(1_000_000));
+    assert_eq!(metadata.max_output_tokens, Some(128_000));
+    assert!(metadata.supports_vision_input());
+    assert_eq!(metadata.supports_agent_runtime(), Some(true));
+}
+
+#[test]
+fn normalizes_inference_profile_ids_and_galaxy_suffixes() {
+    let metadata = model_metadata(
+        "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-opus-4-8[1M]",
+    )
+    .unwrap();
+
+    assert_eq!(metadata.model_id, "anthropic.claude-opus-4-8");
+}
+
+#[test]
+fn preserves_unknown_and_invalid_metadata_as_unknown() {
+    assert!(model_metadata("example.future-model-v1").is_none());
+    assert_eq!(
+        model_metadata("cohere.embed-english-v3")
+            .unwrap()
+            .context_window_tokens,
+        None
+    );
+}
+
+#[test]
+fn rejects_known_non_agent_models() {
+    assert_eq!(
+        model_metadata("cohere.embed-english-v3")
+            .unwrap()
+            .supports_agent_runtime(),
+        Some(false)
+    );
+    assert_eq!(
+        model_metadata("amazon.nova-2-sonic-v1:0")
+            .unwrap()
+            .supports_agent_runtime(),
+        Some(false)
+    );
+}