feat: API endpoint for fetching an SBOM's AI models

Fixes #2254

Author: jcrossley3

entity/src/qualified_purl.rs

@@ -70,6 +70,12 @@ pub enum Relation {
         to = "super::sbom_package_purl_ref::Column::QualifiedPurlId"
     )]
     SbomPackage,
+    #[sea_orm(
+        belongs_to = "super::sbom_ai::Entity",
+        from = "Column::Id",
+        to = "super::sbom_ai::Column::QualifiedPurlId"
+    )]
+    AI,
 }
 
 impl Related<super::versioned_purl::Entity> for Entity {
@@ -84,6 +90,12 @@ impl Related<super::sbom_package_purl_ref::Entity> for Entity {
     }
 }
 
+impl Related<super::sbom_ai::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::AI.def()
+    }
+}
+
 impl Related<super::base_purl::Entity> for Entity {
     fn to() -> RelationDef {
         super::versioned_purl::Relation::BasePurl.def()

entity/src/sbom_ai.rs

@@ -8,6 +8,7 @@ pub struct Model {
     #[sea_orm(primary_key)]
     pub node_id: String,
     pub properties: serde_json::Value,
+    pub qualified_purl_id: Uuid,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
@@ -20,11 +21,7 @@ pub enum Relation {
         to = "super::sbom::Column::SbomId"
     )]
     Sbom,
-    #[sea_orm(
-        belongs_to = "super::sbom_package_purl_ref::Entity",
-        from = "(Column::SbomId, Column::NodeId)",
-        to = "(super::sbom_package_purl_ref::Column::SbomId, super::sbom_package_purl_ref::Column::NodeId)"
-    )]
+    #[sea_orm(has_one = "super::qualified_purl::Entity")]
     Purl,
 }
 
@@ -40,7 +37,7 @@ impl Related<super::sbom::Entity> for Entity {
     }
 }
 
-impl Related<super::sbom_package_purl_ref::Entity> for Entity {
+impl Related<super::qualified_purl::Entity> for Entity {
     fn to() -> RelationDef {
         Relation::Purl.def()
     }

migration/src/lib.rs

@@ -50,6 +50,7 @@ mod m0002080_add_cargo_version_scheme;
 mod m0002090_vulnerability_id_sort_index;
 mod m0002100_analysis_perf_indexes;
 mod m0002110_license_query_performance;
+mod m0002120_add_ai_model_purl;
 
 pub trait MigratorExt: Send {
     fn build_migrations() -> Migrations;
@@ -116,6 +117,7 @@ impl MigratorExt for Migrator {
             .normal(m0002090_vulnerability_id_sort_index::Migration)
             .normal(m0002100_analysis_perf_indexes::Migration)
             .normal(m0002110_license_query_performance::Migration)
+            .normal(m0002120_add_ai_model_purl::Migration)
     }
 }
 

migration/src/m0002120_add_ai_model_purl.rs

@@ -0,0 +1,44 @@
+use sea_orm_migration::prelude::*;
+
+#[derive(DeriveMigrationName)]
+pub struct Migration;
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(SbomAi::Table)
+                    .add_column(
+                        ColumnDef::new(SbomAi::QualifiedPurlId)
+                            .uuid()
+                            .not_null()
+                            .to_owned(),
+                    )
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(SbomAi::Table)
+                    .drop_column(SbomAi::QualifiedPurlId)
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+}
+
+#[derive(DeriveIden)]
+enum SbomAi {
+    Table,
+    QualifiedPurlId,
+}

modules/fundamental/src/sbom/endpoints/mod.rs

@@ -5,6 +5,7 @@ mod query;
 mod test;
 
 pub use query::*;
+use uuid::Uuid;
 
 use crate::{
     Error,
@@ -16,8 +17,8 @@ use crate::{
     },
     sbom::{
         model::{
-            SbomExternalPackageReference, SbomNodeReference, SbomPackage, SbomPackageRelation,
-            SbomSummary, Which, details::SbomAdvisory,
+            SbomExternalPackageReference, SbomModel, SbomNodeReference, SbomPackage,
+            SbomPackageRelation, SbomSummary, Which, details::SbomAdvisory,
         },
         service::{SbomService, sbom::FetchOptions},
     },
@@ -65,6 +66,7 @@ pub fn configure(
         .service(get_sbom_advisories)
         .service(delete)
         .service(packages)
+        .service(models)
         .service(related)
         .service(upload)
         .service(download)
@@ -397,6 +399,36 @@ pub async fn packages(
     Ok(HttpResponse::Ok().json(result))
 }
 
+/// Search for AI models associated with an SBOM
+#[utoipa::path(
+    tag = "sbom",
+    operation_id = "listModels",
+    params(
+        ("id", Path, description = "ID of the SBOM to get models for"),
+        Query,
+        Paginated,
+    ),
+    responses(
+        (status = 200, description = "AI Models", body = PaginatedResults<SbomModel>),
+    ),
+)]
+#[get("/v2/sbom/{id}/models")]
+pub async fn models(
+    fetch: web::Data<SbomService>,
+    db: web::Data<Database>,
+    id: web::Path<Uuid>,
+    web::Query(search): web::Query<Query>,
+    web::Query(paginated): web::Query<Paginated>,
+    _: Require<ReadSbom>,
+) -> actix_web::Result<impl Responder> {
+    let tx = db.begin_read().await?;
+    let result = fetch
+        .fetch_sbom_models(id.into_inner(), search, paginated, &tx)
+        .await?;
+
+    Ok(HttpResponse::Ok().json(result))
+}
+
 #[derive(Clone, Debug, serde::Deserialize, utoipa::IntoParams)]
 struct RelatedQuery {
     /// The Package to use as reference

modules/fundamental/src/sbom/endpoints/test.rs

@@ -1679,7 +1679,7 @@ async fn get_cbom(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
 
 #[test_context(TrustifyContext)]
 #[test(actix_web::test)]
-async fn get_aibom(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
+async fn get_aibom_packages(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
     let app = caller(ctx).await?;
 
     // First upload it via the normal SBOM endpoint
@@ -1759,6 +1759,82 @@ async fn get_aibom(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
     Ok(())
 }
 
+#[test_context(TrustifyContext)]
+#[test(actix_web::test)]
+async fn get_aibom_models(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
+    let app = caller(ctx).await?;
+
+    let id = ctx
+        .ingest_document("cyclonedx/ai/ibm-granite_granite-docling-258M_aibom.json")
+        .await?
+        .id
+        .to_string();
+
+    let uri = format!("/api/v2/sbom/{id}/models");
+    let req = TestRequest::get().uri(&uri).to_request();
+    let response: Value = app.call_and_read_body_json(req).await;
+    log::debug!("response:\n{:#}", json!(response));
+    let expected = json!({
+        "items": [
+            {
+                "id": "pkg:huggingface/ibm-granite/granite-docling-258M@1.0",
+                "name": "granite-docling-258M",
+                "purl": "pkg:huggingface/ibm-granite/granite-docling-258M@1.0",
+                "properties": {
+                    "version": "1.0.0",
+                    "licenses": "apache-2.0",
+                    "bomFormat": "CycloneDX",
+                    "suppliedBy": "ibm-granite",
+                    "specVersion": "1.6",
+                    "typeOfModel": "idefics3",
+                    "serialNumber": "urn:uuid:ibm-granite-granite-docling-258M",
+                    "primaryPurpose": "image-text-to-text",
+                    "downloadLocation": "https://huggingface.co/ibm-granite/granite-docling-258M/tree/main",
+                    "external_references": "[{\"type\": \"website\", \"url\": \"https://huggingface.co/ibm-granite/granite-docling-258M\", \"comment\": \"Model repository\"}, {\"type\": \"distribution\", \"url\": \"https://huggingface.co/ibm-granite/granite-docling-258M/tree/main\", \"comment\": \"Model files\"}]",
+                    "safetyRiskAssessment": "and fairness, misinformation, and autonomous decision-making, and ethical considerations, including but not limited to: bias and fairness, misinformation, and autonomous decision-making, considerations, the model may in some cases produce inaccurate, biased, offensive or unwanted responses to user prompts, in prompts and responses across key dimensions outlined in the IBM AI Risk Atlas, of triggering unwanted output"
+                },
+            },
+        ],
+        "total": 1
+    });
+    assert_eq!(expected, response);
+
+    Ok(())
+}
+
+#[test_context(TrustifyContext)]
+#[rstest]
+#[case("hugging")]
+#[case("granite")]
+#[case("pkg:huggingface/ibm-granite")]
+#[case("pkg:huggingface/ibm-granite/granite-docling-258M")]
+#[case("pkg:huggingface/ibm-granite/granite-docling-258M@1.0")]
+#[case("purl=pkg:huggingface/ibm-granite/granite-docling-258M@1.0")]
+#[case("purl~granite")]
+#[case("purl:namespace=ibm-granite&purl:version=1.0&purl:type=huggingface")]
+#[case("name~granite")]
+#[case("name=granite-docling-258M")]
+#[case("properties:typeOfModel=idefics3")]
+#[case("properties:typeOfModel=idefics3&properties:primaryPurpose=image-text-to-text")]
+#[test_log::test(actix_web::test)]
+async fn query_aibom_models(ctx: &TrustifyContext, #[case] q: &str) -> Result<(), anyhow::Error> {
+    let app = caller(ctx).await?;
+
+    let id = ctx
+        .ingest_document("cyclonedx/ai/ibm-granite_granite-docling-258M_aibom.json")
+        .await?
+        .id
+        .to_string();
+
+    let uri = format!("/api/v2/sbom/{id}/models?q={}", encode(q));
+    let req = TestRequest::get().uri(&uri).to_request();
+    let response: Value = app.call_and_read_body_json(req).await;
+
+    assert_eq!(response["total"].as_i64(), Some(1), "q: {q}");
+
+    Ok(())
+}
+
 #[test_context(TrustifyContext)]
 #[rstest]
 #[case::no_filter([], 3)]

modules/fundamental/src/sbom/model/mod.rs

@@ -8,13 +8,14 @@ use crate::{
     purl::model::summary::purl::PurlSummary,
     source_document::model::SourceDocument,
 };
-use sea_orm::{ConnectionTrait, ModelTrait, PaginatorTrait, prelude::Uuid};
+use sea_orm::{ConnectionTrait, FromQueryResult, ModelTrait, PaginatorTrait, prelude::Uuid};
 use serde::{Deserialize, Serialize};
 use time::OffsetDateTime;
 use tracing::instrument;
 use trustify_common::{cpe::Cpe, purl::Purl};
 use trustify_entity::{
-    labels::Labels, relationship::Relationship, sbom, sbom_node, sbom_package, source_document,
+    labels::Labels, qualified_purl::CanonicalPurl, relationship::Relationship, sbom, sbom_node,
+    sbom_package, source_document,
 };
 use utoipa::ToSchema;
 
@@ -109,6 +110,34 @@ impl SbomSummary {
     }
 }
 
+#[derive(FromQueryResult, Serialize, Deserialize, Debug, Clone, PartialEq, Eq, ToSchema)]
+pub struct SbomModel {
+    /// The internal ID of a model
+    pub id: String,
+    /// The name of the model in the SBOM
+    pub name: String,
+    /// The model's PURL
+    pub purl: serde_json::Value,
+    /// The properties associated with the model
+    pub properties: serde_json::Value,
+}
+
+impl SbomModel {
+    pub fn stringify_purl(self) -> SbomModel {
+        if self.purl.is_object() {
+            let mut result = self.clone();
+            let purl = match serde_json::from_value::<CanonicalPurl>(self.purl) {
+                Ok(cp) => serde_json::Value::String(Purl::from(cp).to_string()),
+                Err(_) => serde_json::Value::Null,
+            };
+            result.purl = purl;
+            result
+        } else {
+            self
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, ToSchema, Default)]
 pub struct SbomPackage {
     /// The SBOM internal ID of a package