gitdataai/libs/service/agent/model_pricing.rs

use crate::AppService;
use crate::error::AppError;
use chrono::Utc;
use models::agents::PricingCurrency;
use models::agents::model_pricing;
use sea_orm::*;
use serde::{Deserialize, Serialize};
use session::Session;
use utoipa::ToSchema;
use uuid::Uuid;

use super::provider::require_system_caller;

#[derive(Debug, Clone, Deserialize, 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, Deserialize, 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, Serialize, 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,
        }
    }
}

impl AppService {
    pub async fn agent_model_pricing_list(
        &self,
        model_version_id: Uuid,
        _ctx: &Session,
    ) -> Result<Vec<ModelPricingResponse>, AppError> {
        let records = model_pricing::Entity::find()
            .filter(model_pricing::Column::ModelVersionId.eq(model_version_id))
            .order_by_desc(model_pricing::Column::EffectiveFrom)
            .all(&self.db)
            .await?;
        Ok(records
            .into_iter()
            .map(ModelPricingResponse::from)
            .collect())
    }

    pub async fn agent_model_pricing_get(
        &self,
        id: i64,
        _ctx: &Session,
    ) -> Result<ModelPricingResponse, AppError> {
        let record = model_pricing::Entity::find_by_id(id)
            .one(&self.db)
            .await?
            .ok_or(AppError::NotFound("Pricing record not found".to_string()))?;
        Ok(ModelPricingResponse::from(record))
    }

    pub async fn agent_model_pricing_create(
        &self,
        request: CreateModelPricingRequest,
        ctx: &Session,
    ) -> Result<ModelPricingResponse, AppError> {
        require_system_caller(ctx)?;

        let _ = request
            .currency
            .parse::<PricingCurrency>()
            .map_err(|_| AppError::BadRequest("Invalid pricing currency".to_string()))?;

        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(&self.db).await?;
        Ok(ModelPricingResponse::from(record))
    }

    pub async fn agent_model_pricing_update(
        &self,
        id: i64,
        request: UpdateModelPricingRequest,
        ctx: &Session,
    ) -> Result<ModelPricingResponse, AppError> {
        require_system_caller(ctx)?;

        let record = model_pricing::Entity::find_by_id(id)
            .one(&self.db)
            .await?
            .ok_or(AppError::NotFound("Pricing record not found".to_string()))?;

        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(|_| AppError::BadRequest("Invalid pricing currency".to_string()))?;
            active.currency = Set(v);
        }
        if let Some(v) = request.effective_from {
            active.effective_from = Set(v);
        }

        let record = active.update(&self.db).await?;
        Ok(ModelPricingResponse::from(record))
    }

    pub async fn agent_model_pricing_delete(&self, id: i64, ctx: &Session) -> Result<(), AppError> {
        require_system_caller(ctx)?;
        model_pricing::Entity::delete_by_id(id)
            .exec(&self.db)
            .await?;
        Ok(())
    }
}