tibba_model_token/
price.rs

1// Copyright 2026 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::{
16    Error, JsonSnafu, ModelListParams, SERVICE_API, SERVICE_LLM, SERVICE_STORAGE, Schema,
17    SchemaAllowCreate, SchemaAllowEdit, SchemaType, SchemaView, SqlxSnafu, Status, format_datetime,
18    new_schema_options,
19};
20use serde::{Deserialize, Serialize};
21use snafu::ResultExt;
22use sqlx::FromRow;
23use sqlx::{Pool, Postgres, QueryBuilder};
24use std::collections::HashMap;
25use tibba_model::Model;
26use time::PrimitiveDateTime;
27
28type Result<T> = std::result::Result<T, Error>;
29
30#[derive(FromRow)]
31struct TokenPriceSchema {
32    id: i64,
33    service: String,
34    model: String,
35    input_price: i64,
36    output_price: i64,
37    fixed_price: i64,
38    unit_size: i32,
39    status: i16,
40    remark: String,
41    created: PrimitiveDateTime,
42    modified: PrimitiveDateTime,
43}
44
45#[derive(Debug, Clone, Deserialize, Serialize)]
46pub struct TokenPrice {
47    pub id: i64,
48    pub service: String,
49    pub model: String,
50    /// 每 unit_size 个输入 token 扣除的积分数
51    pub input_price: i64,
52    /// 每 unit_size 个输出 token 扣除的积分数
53    pub output_price: i64,
54    /// 每次调用固定扣除积分数
55    pub fixed_price: i64,
56    /// 计费基数，默认 1000（per 1K tokens）
57    pub unit_size: i32,
58    pub status: i16,
59    pub remark: String,
60    pub created: String,
61    pub modified: String,
62}
63
64impl From<TokenPriceSchema> for TokenPrice {
65    fn from(s: TokenPriceSchema) -> Self {
66        Self {
67            id: s.id,
68            service: s.service,
69            model: s.model,
70            input_price: s.input_price,
71            output_price: s.output_price,
72            fixed_price: s.fixed_price,
73            unit_size: s.unit_size,
74            status: s.status,
75            remark: s.remark,
76            created: format_datetime(s.created),
77            modified: format_datetime(s.modified),
78        }
79    }
80}
81
82#[derive(Debug, Clone, Deserialize)]
83pub struct TokenPriceInsertParams {
84    pub service: String,
85    pub model: Option<String>,
86    pub input_price: i64,
87    pub output_price: i64,
88    pub fixed_price: Option<i64>,
89    pub unit_size: Option<i32>,
90    pub status: Option<i16>,
91    pub remark: Option<String>,
92}
93
94#[derive(Debug, Clone, Deserialize, Default)]
95pub struct TokenPriceUpdateParams {
96    pub input_price: Option<i64>,
97    pub output_price: Option<i64>,
98    pub fixed_price: Option<i64>,
99    pub unit_size: Option<i32>,
100    pub status: Option<i16>,
101    pub remark: Option<String>,
102}
103
104#[derive(Default)]
105pub struct TokenPriceModel {}
106
107impl TokenPriceModel {
108    /// 按服务类型和模型名查询定价配置。
109    /// 先精确匹配 (service, model)，找不到时退回匹配 (service, "default")。
110    pub async fn get_by_service_model(
111        &self,
112        pool: &Pool<Postgres>,
113        service: &str,
114        model: &str,
115    ) -> Result<Option<TokenPrice>> {
116        // 精确匹配
117        let result = sqlx::query_as::<_, TokenPriceSchema>(
118            r#"SELECT * FROM token_prices
119               WHERE service = $1 AND model = $2 AND status = 1 AND deleted_at IS NULL
120               LIMIT 1"#,
121        )
122        .bind(service)
123        .bind(model)
124        .fetch_optional(pool)
125        .await
126        .context(SqlxSnafu)?;
127
128        if result.is_some() {
129            return Ok(result.map(Into::into));
130        }
131
132        // 回退：匹配该服务的 "default" 定价（避免 model 已是 default 时重复查询）
133        if model != "default" {
134            let fallback = sqlx::query_as::<_, TokenPriceSchema>(
135                r#"SELECT * FROM token_prices
136                   WHERE service = $1 AND model = 'default' AND status = 1 AND deleted_at IS NULL
137                   LIMIT 1"#,
138            )
139            .bind(service)
140            .fetch_optional(pool)
141            .await
142            .context(SqlxSnafu)?;
143            return Ok(fallback.map(Into::into));
144        }
145
146        Ok(None)
147    }
148
149    /// 根据定价配置和 token 用量计算本次消耗积分。
150    /// 使用整数向上取整，避免浮点误差。
151    pub fn calculate_cost(price: &TokenPrice, input_tokens: i32, output_tokens: i32) -> i64 {
152        let unit = price.unit_size.max(1) as i64;
153        // 向上取整：(n * p + unit - 1) / unit
154        let input_cost = (input_tokens as i64 * price.input_price + unit - 1) / unit;
155        let output_cost = (output_tokens as i64 * price.output_price + unit - 1) / unit;
156        price.fixed_price + input_cost + output_cost
157    }
158}
159
160impl Model for TokenPriceModel {
161    type Output = TokenPrice;
162    fn new() -> Self {
163        Self::default()
164    }
165
166    async fn schema_view(&self, _pool: &Pool<Postgres>) -> SchemaView {
167        SchemaView {
168            schemas: vec![
169                Schema::new_id(),
170                Schema {
171                    name: "service".to_string(),
172                    category: SchemaType::String,
173                    required: true,
174                    fixed: true,
175                    filterable: true,
176                    options: Some(new_schema_options(&[
177                        SERVICE_LLM,
178                        SERVICE_API,
179                        SERVICE_STORAGE,
180                    ])),
181                    ..Default::default()
182                },
183                Schema {
184                    name: "model".to_string(),
185                    category: SchemaType::String,
186                    fixed: true,
187                    filterable: true,
188                    options: Some(new_schema_options(&["default", "mimo-v2.5-pro"])),
189                    ..Default::default()
190                },
191                Schema {
192                    name: "input_price".to_string(),
193                    category: SchemaType::Number,
194                    required: true,
195                    ..Default::default()
196                },
197                Schema {
198                    name: "output_price".to_string(),
199                    category: SchemaType::Number,
200                    required: true,
201                    ..Default::default()
202                },
203                Schema {
204                    name: "fixed_price".to_string(),
205                    category: SchemaType::Number,
206                    ..Default::default()
207                },
208                Schema {
209                    name: "unit_size".to_string(),
210                    category: SchemaType::Number,
211                    default_value: Some(serde_json::json!(1000)),
212                    ..Default::default()
213                },
214                Schema::new_status(),
215                Schema::new_remark(),
216                Schema::new_created(),
217                Schema::new_modified(),
218            ],
219            allow_edit: SchemaAllowEdit {
220                roles: vec!["su".to_string()],
221                ..Default::default()
222            },
223            allow_create: SchemaAllowCreate {
224                roles: vec!["su".to_string()],
225                ..Default::default()
226            },
227        }
228    }
229
230    async fn insert(&self, pool: &Pool<Postgres>, data: serde_json::Value) -> Result<u64> {
231        let p: TokenPriceInsertParams = serde_json::from_value(data).context(JsonSnafu)?;
232        let row: (i64,) = sqlx::query_as(
233            r#"INSERT INTO token_prices
234               (service, model, input_price, output_price, fixed_price, unit_size, status, remark)
235               VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
236               RETURNING id"#,
237        )
238        .bind(&p.service)
239        .bind(p.model.unwrap_or_default())
240        .bind(p.input_price)
241        .bind(p.output_price)
242        .bind(p.fixed_price.unwrap_or(0))
243        .bind(p.unit_size.unwrap_or(1000))
244        .bind(p.status.unwrap_or(Status::Enabled as i16))
245        .bind(p.remark.unwrap_or_default())
246        .fetch_one(pool)
247        .await
248        .context(SqlxSnafu)?;
249        Ok(row.0 as u64)
250    }
251
252    async fn get_by_id(&self, pool: &Pool<Postgres>, id: u64) -> Result<Option<Self::Output>> {
253        let result = sqlx::query_as::<_, TokenPriceSchema>(
254            r#"SELECT * FROM token_prices WHERE id = $1 AND deleted_at IS NULL"#,
255        )
256        .bind(id as i64)
257        .fetch_optional(pool)
258        .await
259        .context(SqlxSnafu)?;
260        Ok(result.map(Into::into))
261    }
262
263    async fn update_by_id(
264        &self,
265        pool: &Pool<Postgres>,
266        id: u64,
267        data: serde_json::Value,
268    ) -> Result<()> {
269        let p: TokenPriceUpdateParams = serde_json::from_value(data).context(JsonSnafu)?;
270        let mut qb: QueryBuilder<Postgres> =
271            QueryBuilder::new("UPDATE token_prices SET modified = NOW()");
272        if let Some(v) = p.input_price {
273            qb.push(", input_price = ").push_bind(v);
274        }
275        if let Some(v) = p.output_price {
276            qb.push(", output_price = ").push_bind(v);
277        }
278        if let Some(v) = p.fixed_price {
279            qb.push(", fixed_price = ").push_bind(v);
280        }
281        if let Some(v) = p.unit_size {
282            qb.push(", unit_size = ").push_bind(v);
283        }
284        if let Some(v) = p.status {
285            qb.push(", status = ").push_bind(v);
286        }
287        if let Some(v) = p.remark {
288            qb.push(", remark = ").push_bind(v);
289        }
290        qb.push(" WHERE id = ")
291            .push_bind(id as i64)
292            .push(" AND deleted_at IS NULL");
293        qb.build().execute(pool).await.context(SqlxSnafu)?;
294        Ok(())
295    }
296
297    async fn delete_by_id(&self, pool: &Pool<Postgres>, id: u64) -> Result<()> {
298        sqlx::query(
299            r#"UPDATE token_prices SET deleted_at = NOW(), modified = NOW() WHERE id = $1 AND deleted_at IS NULL"#,
300        )
301        .bind(id as i64)
302        .execute(pool)
303        .await
304        .context(SqlxSnafu)?;
305        Ok(())
306    }
307
308    async fn count(&self, pool: &Pool<Postgres>, params: &ModelListParams) -> Result<i64> {
309        let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("SELECT COUNT(*) FROM token_prices");
310        self.push_conditions(&mut qb, params)?;
311        let row: (i64,) = qb
312            .build_query_as()
313            .fetch_one(pool)
314            .await
315            .context(SqlxSnafu)?;
316        Ok(row.0)
317    }
318
319    async fn list(
320        &self,
321        pool: &Pool<Postgres>,
322        params: &ModelListParams,
323    ) -> Result<Vec<Self::Output>> {
324        let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("SELECT * FROM token_prices");
325        self.push_conditions(&mut qb, params)?;
326        params.push_pagination(&mut qb);
327        let rows = qb
328            .build_query_as::<TokenPriceSchema>()
329            .fetch_all(pool)
330            .await
331            .context(SqlxSnafu)?;
332        Ok(rows.into_iter().map(Into::into).collect())
333    }
334
335    fn push_filter_conditions<'args>(
336        &self,
337        qb: &mut QueryBuilder<'args, Postgres>,
338        filters: &HashMap<String, String>,
339    ) -> Result<()> {
340        if let Some(service) = filters.get("service") {
341            qb.push(" AND service = ").push_bind(service.clone());
342        }
343        if let Some(status) = filters.get("status") {
344            if let Ok(v) = status.parse::<i16>() {
345                qb.push(" AND status = ").push_bind(v);
346            }
347        }
348        Ok(())
349    }
350}
tibba_model_token/price.rs

tibba_model_token/
price.rs
