use async_graphql::{Context, InputObject, Object, Result};
use llm_registry_core::{
AssetId, Checksum, HashAlgorithm, StorageBackend, StorageLocation,
};
use llm_registry_service::{RegisterAssetRequest, ServiceRegistry, UpdateAssetRequest};
use semver::Version;
use std::sync::Arc;
use super::types::{
GqlAsset, GqlAssetStatus, GqlAssetType, GqlDeleteResult, GqlRegisterResult, GqlUpdateResult,
};
use crate::auth::AuthUser;
use crate::error::ApiError;
pub struct Mutation;
#[derive(InputObject)]
pub struct RegisterAssetInput {
pub asset_type: GqlAssetType,
pub name: String,
pub version: String,
pub description: Option<String>,
pub license: Option<String>,
#[graphql(default)]
pub tags: Vec<String>,
#[graphql(default)]
pub annotations: Vec<AnnotationInput>,
pub storage_path: String,
#[graphql(default = "local")]
pub storage_backend: String,
pub storage_uri: Option<String>,
pub checksum: String,
#[graphql(default = "SHA256")]
pub checksum_algorithm: String,
pub size_bytes: Option<u64>,
pub content_type: Option<String>,
}
#[derive(InputObject)]
pub struct UpdateAssetInput {
pub asset_id: String,
pub status: Option<GqlAssetStatus>,
pub description: Option<String>,
pub license: Option<String>,
#[graphql(default)]
pub add_tags: Vec<String>,
#[graphql(default)]
pub remove_tags: Vec<String>,
#[graphql(default)]
pub add_annotations: Vec<AnnotationInput>,
#[graphql(default)]
pub remove_annotations: Vec<String>,
}
#[derive(InputObject)]
pub struct AnnotationInput {
pub key: String,
pub value: String,
}
#[Object]
impl Mutation {
async fn register_asset(
&self,
ctx: &Context<'_>,
input: RegisterAssetInput,
) -> Result<GqlRegisterResult> {
let services = ctx.data::<Arc<ServiceRegistry>>()?;
let _user = ctx.data_opt::<AuthUser>();
let version = Version::parse(&input.version)
.map_err(|e| ApiError::bad_request(format!("Invalid version: {}", e)))?;
let algorithm = match input.checksum_algorithm.to_uppercase().as_str() {
"SHA256" => HashAlgorithm::SHA256,
"SHA3_256" | "SHA3-256" => HashAlgorithm::SHA3_256,
"BLAKE3" => HashAlgorithm::BLAKE3,
_ => return Err(ApiError::bad_request("Invalid checksum algorithm"))?
};
let backend = match input.storage_backend.to_lowercase().as_str() {
"filesystem" | "local" => StorageBackend::FileSystem {
base_path: "/var/lib/llm-registry".to_string(),
},
_ => StorageBackend::FileSystem {
base_path: "/var/lib/llm-registry".to_string(),
},
};
let storage = StorageLocation {
backend,
path: input.storage_path,
uri: input.storage_uri,
};
let checksum = Checksum {
algorithm,
value: input.checksum,
};
let request = RegisterAssetRequest {
asset_type: input.asset_type.to_core(),
name: input.name,
version,
description: input.description,
license: input.license,
tags: input.tags,
annotations: input
.annotations
.into_iter()
.map(|a| (a.key, a.value))
.collect(),
storage,
checksum,
provenance: None,
dependencies: vec![],
size_bytes: input.size_bytes,
content_type: input.content_type,
};
let response = services
.registration()
.register_asset(request)
.await
.map_err(|e| ApiError::from(e))?;
Ok(GqlRegisterResult {
asset: GqlAsset(response.asset),
message: "Asset registered successfully".to_string(),
})
}
async fn update_asset(
&self,
ctx: &Context<'_>,
input: UpdateAssetInput,
) -> Result<GqlUpdateResult> {
let services = ctx.data::<Arc<ServiceRegistry>>()?;
let _user = ctx.data_opt::<AuthUser>();
let asset_id = input
.asset_id
.parse::<AssetId>()
.map_err(|e| ApiError::bad_request(format!("Invalid asset ID: {}", e)))?;
let request = UpdateAssetRequest {
asset_id,
status: input.status.map(|s| s.to_core()),
description: input.description,
license: input.license,
add_tags: input.add_tags,
remove_tags: input.remove_tags,
add_annotations: input
.add_annotations
.into_iter()
.map(|a| (a.key, a.value))
.collect(),
remove_annotations: input.remove_annotations,
};
let response = services
.registration()
.update_asset(request)
.await
.map_err(|e| ApiError::from(e))?;
Ok(GqlUpdateResult {
asset: GqlAsset(response.asset),
message: "Asset updated successfully".to_string(),
})
}
async fn delete_asset(&self, ctx: &Context<'_>, id: String) -> Result<GqlDeleteResult> {
let services = ctx.data::<Arc<ServiceRegistry>>()?;
let _user = ctx.data_opt::<AuthUser>();
let asset_id = id
.parse::<AssetId>()
.map_err(|e| ApiError::bad_request(format!("Invalid asset ID: {}", e)))?;
services
.registration()
.delete_asset(&asset_id)
.await
.map_err(|e| ApiError::from(e))?;
Ok(GqlDeleteResult {
asset_id: id,
message: "Asset deleted successfully".to_string(),
})
}
}