gitdataai/libs/agent/model/pricing.rs

144 lines
4.6 KiB
Rust

//! Model pricing management — CRUD.
//!
//! All functions take `&DatabaseConnection` instead of `&AppService`.
use chrono::Utc;
use db::database::AppDatabase;
use models::agents::PricingCurrency;
use models::agents::model_pricing;
use sea_orm::*;
use uuid::Uuid;
use crate::error::AgentError;
#[derive(Debug, Clone, serde::Deserialize, utoipa::ToSchema)]
pub struct CreateModelPricingRequest {
pub model_version_id: Uuid,
pub input_price_per_1k_tokens: String,
pub output_price_per_1k_tokens: String,
pub currency: String,
pub effective_from: chrono::DateTime<Utc>,
}
#[derive(Debug, Clone, serde::Deserialize, utoipa::ToSchema)]
pub struct UpdateModelPricingRequest {
pub input_price_per_1k_tokens: Option<String>,
pub output_price_per_1k_tokens: Option<String>,
pub currency: Option<String>,
pub effective_from: Option<chrono::DateTime<Utc>>,
}
#[derive(Debug, Clone, serde::Serialize, utoipa::ToSchema)]
pub struct ModelPricingResponse {
pub id: i64,
pub model_version_id: Uuid,
pub input_price_per_1k_tokens: String,
pub output_price_per_1k_tokens: String,
pub currency: String,
pub effective_from: chrono::DateTime<Utc>,
}
impl From<model_pricing::Model> for ModelPricingResponse {
fn from(p: model_pricing::Model) -> Self {
Self {
id: p.id,
model_version_id: p.model_version_id,
input_price_per_1k_tokens: p.input_price_per_1k_tokens,
output_price_per_1k_tokens: p.output_price_per_1k_tokens,
currency: p.currency,
effective_from: p.effective_from,
}
}
}
/// List pricing records for a model version.
pub async fn list_pricing(
db: &AppDatabase,
model_version_id: Uuid,
) -> Result<Vec<ModelPricingResponse>, AgentError> {
let records = model_pricing::Entity::find()
.filter(model_pricing::Column::ModelVersionId.eq(model_version_id))
.order_by_desc(model_pricing::Column::EffectiveFrom)
.all(db)
.await?;
Ok(records
.into_iter()
.map(ModelPricingResponse::from)
.collect())
}
/// Get a single pricing record by ID.
pub async fn get_pricing(db: &AppDatabase, id: i64) -> Result<ModelPricingResponse, AgentError> {
let record = model_pricing::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| AgentError::NotFound(format!("Pricing record not found: {}", id)))?;
Ok(ModelPricingResponse::from(record))
}
/// Create a new pricing record.
pub async fn create_pricing(
db: &AppDatabase,
request: CreateModelPricingRequest,
) -> Result<ModelPricingResponse, AgentError> {
let _ = request
.currency
.parse::<PricingCurrency>()
.map_err(|_| AgentError::InvalidInput {
field: "currency".into(),
reason: "Invalid pricing currency".into(),
})?;
let active = model_pricing::ActiveModel {
model_version_id: Set(request.model_version_id),
input_price_per_1k_tokens: Set(request.input_price_per_1k_tokens),
output_price_per_1k_tokens: Set(request.output_price_per_1k_tokens),
currency: Set(request.currency),
effective_from: Set(request.effective_from),
..Default::default()
};
let record = active.insert(db).await?;
Ok(ModelPricingResponse::from(record))
}
/// Update a pricing record.
pub async fn update_pricing(
db: &AppDatabase,
id: i64,
request: UpdateModelPricingRequest,
) -> Result<ModelPricingResponse, AgentError> {
let record = model_pricing::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| AgentError::NotFound(format!("Pricing record not found: {}", id)))?;
let mut active: model_pricing::ActiveModel = record.into();
if let Some(v) = request.input_price_per_1k_tokens {
active.input_price_per_1k_tokens = Set(v);
}
if let Some(v) = request.output_price_per_1k_tokens {
active.output_price_per_1k_tokens = Set(v);
}
if let Some(v) = request.currency {
let _ = v
.parse::<PricingCurrency>()
.map_err(|_| AgentError::InvalidInput {
field: "currency".into(),
reason: "Invalid pricing currency".into(),
})?;
active.currency = Set(v);
}
if let Some(v) = request.effective_from {
active.effective_from = Set(v);
}
let record = active.update(db).await?;
Ok(ModelPricingResponse::from(record))
}
/// Delete a pricing record.
pub async fn delete_pricing(db: &AppDatabase, id: i64) -> Result<(), AgentError> {
model_pricing::Entity::delete_by_id(id).exec(db).await?;
Ok(())
}