138 lines
4.0 KiB
Rust
138 lines
4.0 KiB
Rust
use crate::AppService;
|
|
use crate::error::AppError;
|
|
use chrono::Utc;
|
|
use models::agents::CapabilityType;
|
|
use models::agents::model_capability;
|
|
use sea_orm::*;
|
|
use serde::{Deserialize, Serialize};
|
|
use session::Session;
|
|
use utoipa::ToSchema;
|
|
|
|
use super::provider::require_system_caller;
|
|
|
|
#[derive(Debug, Clone, Deserialize, ToSchema)]
|
|
pub struct CreateModelCapabilityRequest {
|
|
pub model_version_id: i64,
|
|
pub capability: String,
|
|
#[serde(default)]
|
|
pub is_supported: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, ToSchema)]
|
|
pub struct UpdateModelCapabilityRequest {
|
|
pub is_supported: Option<bool>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, ToSchema)]
|
|
pub struct ModelCapabilityResponse {
|
|
pub id: i64,
|
|
pub model_version_id: i64,
|
|
pub capability: String,
|
|
pub is_supported: bool,
|
|
pub created_at: chrono::DateTime<Utc>,
|
|
}
|
|
|
|
impl From<model_capability::Model> for ModelCapabilityResponse {
|
|
fn from(mc: model_capability::Model) -> Self {
|
|
Self {
|
|
id: mc.id,
|
|
model_version_id: mc.model_version_id,
|
|
capability: mc.capability,
|
|
is_supported: mc.is_supported,
|
|
created_at: mc.created_at,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl AppService {
|
|
pub async fn agent_model_capability_list(
|
|
&self,
|
|
model_version_id: i64,
|
|
_ctx: &Session,
|
|
) -> Result<Vec<ModelCapabilityResponse>, AppError> {
|
|
let caps = model_capability::Entity::find()
|
|
.filter(model_capability::Column::ModelVersionId.eq(model_version_id))
|
|
.order_by_asc(model_capability::Column::Capability)
|
|
.all(&self.db)
|
|
.await?;
|
|
Ok(caps
|
|
.into_iter()
|
|
.map(ModelCapabilityResponse::from)
|
|
.collect())
|
|
}
|
|
|
|
pub async fn agent_model_capability_get(
|
|
&self,
|
|
id: i64,
|
|
_ctx: &Session,
|
|
) -> Result<ModelCapabilityResponse, AppError> {
|
|
let cap = model_capability::Entity::find_by_id(id)
|
|
.one(&self.db)
|
|
.await?
|
|
.ok_or(AppError::NotFound(
|
|
"Capability record not found".to_string(),
|
|
))?;
|
|
Ok(ModelCapabilityResponse::from(cap))
|
|
}
|
|
|
|
pub async fn agent_model_capability_create(
|
|
&self,
|
|
request: CreateModelCapabilityRequest,
|
|
ctx: &Session,
|
|
) -> Result<ModelCapabilityResponse, AppError> {
|
|
require_system_caller(ctx)?;
|
|
|
|
let _ = request
|
|
.capability
|
|
.parse::<CapabilityType>()
|
|
.map_err(|_| AppError::BadRequest("Invalid capability type".to_string()))?;
|
|
|
|
let now = Utc::now();
|
|
let active = model_capability::ActiveModel {
|
|
model_version_id: Set(request.model_version_id),
|
|
capability: Set(request.capability),
|
|
is_supported: Set(request.is_supported),
|
|
created_at: Set(now),
|
|
..Default::default()
|
|
};
|
|
let cap = active.insert(&self.db).await?;
|
|
Ok(ModelCapabilityResponse::from(cap))
|
|
}
|
|
|
|
pub async fn agent_model_capability_update(
|
|
&self,
|
|
id: i64,
|
|
request: UpdateModelCapabilityRequest,
|
|
ctx: &Session,
|
|
) -> Result<ModelCapabilityResponse, AppError> {
|
|
require_system_caller(ctx)?;
|
|
|
|
let cap = model_capability::Entity::find_by_id(id)
|
|
.one(&self.db)
|
|
.await?
|
|
.ok_or(AppError::NotFound(
|
|
"Capability record not found".to_string(),
|
|
))?;
|
|
|
|
let mut active: model_capability::ActiveModel = cap.into();
|
|
if let Some(is_supported) = request.is_supported {
|
|
active.is_supported = Set(is_supported);
|
|
}
|
|
|
|
let cap = active.update(&self.db).await?;
|
|
Ok(ModelCapabilityResponse::from(cap))
|
|
}
|
|
|
|
pub async fn agent_model_capability_delete(
|
|
&self,
|
|
id: i64,
|
|
ctx: &Session,
|
|
) -> Result<(), AppError> {
|
|
require_system_caller(ctx)?;
|
|
model_capability::Entity::delete_by_id(id)
|
|
.exec(&self.db)
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
}
|