use crate::validation::Validate;
use acorn_core::prelude::alloc::{String, ToString, Vec};
use acorn_core::util::MimeType;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelArtifactKind {
#[default]
Annotated,
ModelKit,
ModelPack,
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelExtraction {
#[default]
None,
Tar,
TarGzip,
TarZstd,
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelInventoryStatus {
Complete,
Declared,
#[default]
Deferred,
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum ModelLayerRole {
#[default]
Model,
ModelWeight,
WeightConfig,
Documentation,
Code,
Dataset,
McpBundle,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "camelCase")]
pub struct ModelDetails {
pub registry: String,
pub repository: String,
pub requested_reference: String,
pub resolved_digest: String,
pub artifact_type: Option<String>,
pub package_format: ModelArtifactKind,
pub inventory_status: ModelInventoryStatus,
pub total_size: u64,
#[serde(default)]
pub transfer_size: u64,
pub installed_size: Option<u64>,
#[serde(default)]
#[validate(nested)]
pub layers: Vec<ModelLayer>,
#[serde(default)]
#[validate(nested)]
pub ignored_layers: Vec<ModelLayer>,
#[serde(default)]
#[validate(nested)]
pub materialized_files: Vec<ModelFile>,
pub files: Vec<String>,
pub harbor_tags: Option<Vec<String>>,
pub harbor_labels: Option<Vec<String>>,
pub harbor_push_time: Option<String>,
}
#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "camelCase")]
pub struct ModelFile {
pub path: String,
pub role: ModelLayerRole,
#[validate(digest)]
pub layer_digest: String,
pub installed_size: Option<u64>,
}
#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, Validate)]
#[serde(rename_all = "camelCase")]
pub struct ModelLayer {
#[serde(deserialize_with = "deserialize_oci_media_type", serialize_with = "serialize_oci_media_type")]
#[schemars(with = "String")]
pub media_type: MimeType,
#[validate(digest)]
pub digest: String,
#[validate(range(min = 0))]
pub transfer_size: u64,
pub role: ModelLayerRole,
pub path: Option<String>,
pub extraction: ModelExtraction,
pub inventory_deferred: bool,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelPackCapabilities {
#[serde(default)]
pub input_types: Option<Vec<String>>,
#[serde(default)]
pub output_types: Option<Vec<String>>,
#[serde(default)]
pub knowledge_cutoff: Option<String>,
#[serde(default)]
pub reasoning: Option<bool>,
#[serde(default)]
pub tool_usage: Option<bool>,
#[serde(default)]
pub reward: Option<bool>,
#[serde(default)]
pub languages: Option<Vec<String>>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelPackConfig {
pub descriptor: ModelPackDescriptor,
pub config: ModelPackModelConfig,
pub modelfs: ModelPackFileSystem,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelPackDescriptor {
#[serde(default)]
pub created_at: Option<String>,
#[serde(default)]
pub authors: Option<Vec<String>>,
#[serde(default)]
pub family: Option<String>,
#[serde(default)]
pub name: Option<String>,
#[serde(default, rename = "docURL")]
pub doc_url: Option<String>,
#[serde(default, rename = "sourceURL")]
pub source_url: Option<String>,
#[serde(default, rename = "datasetsURL")]
pub datasets_url: Option<Vec<String>>,
#[serde(default)]
pub version: Option<String>,
#[serde(default)]
pub revision: Option<String>,
#[serde(default)]
pub vendor: Option<String>,
#[serde(default)]
pub licenses: Option<Vec<String>>,
#[serde(default)]
pub title: Option<String>,
#[serde(default)]
pub description: Option<String>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelPackFileMetadata {
pub name: String,
pub mode: u32,
pub uid: u32,
pub gid: u32,
pub size: i64,
pub mtime: String,
pub typeflag: u8,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelPackFileSystem {
#[serde(rename = "type")]
pub kind: String,
pub diff_ids: Vec<String>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelPackModelConfig {
#[serde(default)]
pub architecture: Option<String>,
#[serde(default)]
pub format: Option<String>,
#[serde(default)]
pub param_size: Option<String>,
#[serde(default)]
pub precision: Option<String>,
#[serde(default)]
pub quantization: Option<String>,
#[serde(default)]
pub capabilities: Option<ModelPackCapabilities>,
}
impl From<&str> for ModelExtraction {
fn from(value: &str) -> Self {
match value {
| "tar" => Self::Tar,
| "tar_gzip" => Self::TarGzip,
| "tar_zstd" => Self::TarZstd,
| _ => Self::None,
}
}
}
fn deserialize_oci_media_type<'de, D>(deserializer: D) -> Result<MimeType, D::Error>
where
D: serde::Deserializer<'de>,
{
String::deserialize(deserializer).map(MimeType::OciLayer)
}
fn serialize_oci_media_type<S>(value: &MimeType, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&value.to_string())
}