tibba_model_token/
llm.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, Schema, SchemaAllowCreate, SchemaAllowEdit, SchemaType,
17    SchemaView, SqlxSnafu, Status, format_datetime, new_schema_options,
18};
19use serde::{Deserialize, Serialize};
20use snafu::ResultExt;
21use sqlx::FromRow;
22use sqlx::{Pool, Postgres, QueryBuilder};
23use std::collections::HashMap;
24use tibba_model::Model;
25use time::PrimitiveDateTime;
26
27type Result<T> = std::result::Result<T, Error>;
28
29/// 后端协议：openai（默认）或 anthropic
30pub const LLM_PROVIDER_OPENAI: &str = "openai";
31pub const LLM_PROVIDER_ANTHROPIC: &str = "anthropic";
32
33#[derive(FromRow)]
34struct TokenLlmSchema {
35    id: i64,
36    name: String,
37    url: String,
38    model: String,
39    api_key: String,
40    provider: String,
41    status: i16,
42    remark: String,
43    created: PrimitiveDateTime,
44    modified: PrimitiveDateTime,
45}
46
47#[derive(Debug, Clone, Deserialize, Serialize)]
48pub struct TokenLlm {
49    pub id: i64,
50    pub name: String,
51    pub url: String,
52    pub model: String,
53    pub api_key: String,
54    /// 后端协议：openai 或 anthropic，留空时按 openai 处理
55    pub provider: String,
56    pub status: i16,
57    pub remark: String,
58    pub created: String,
59    pub modified: String,
60}
61
62impl From<TokenLlmSchema> for TokenLlm {
63    fn from(s: TokenLlmSchema) -> Self {
64        Self {
65            id: s.id,
66            name: s.name,
67            url: s.url,
68            model: s.model,
69            api_key: s.api_key,
70            provider: s.provider,
71            status: s.status,
72            remark: s.remark,
73            created: format_datetime(s.created),
74            modified: format_datetime(s.modified),
75        }
76    }
77}
78
79#[derive(Debug, Clone, Deserialize)]
80pub struct TokenLlmInsertParams {
81    pub name: String,
82    pub url: String,
83    pub model: String,
84    pub api_key: String,
85    pub provider: Option<String>,
86    pub status: Option<i16>,
87    pub remark: Option<String>,
88}
89
90#[derive(Debug, Clone, Deserialize, Default)]
91pub struct TokenLlmUpdateParams {
92    pub url: Option<String>,
93    pub model: Option<String>,
94    pub api_key: Option<String>,
95    pub provider: Option<String>,
96    pub status: Option<i16>,
97    pub remark: Option<String>,
98}
99
100#[derive(Default)]
101pub struct TokenLlmModel {}
102
103impl TokenLlmModel {
104    /// 按 name 查询启用状态的 LLM 配置；未命中时回退到 name = "default"。
105    pub async fn get_by_name(&self, pool: &Pool<Postgres>, name: &str) -> Result<Option<TokenLlm>> {
106        let result = sqlx::query_as::<_, TokenLlmSchema>(
107            r#"SELECT * FROM token_llms
108               WHERE name = $1 AND status = 1 AND deleted_at IS NULL
109               LIMIT 1"#,
110        )
111        .bind(name)
112        .fetch_optional(pool)
113        .await
114        .context(SqlxSnafu)?;
115
116        if result.is_some() {
117            return Ok(result.map(Into::into));
118        }
119
120        // 回退：匹配 name = "default"（避免 name 已是 default 时重复查询）
121        if name != "default" {
122            let fallback = sqlx::query_as::<_, TokenLlmSchema>(
123                r#"SELECT * FROM token_llms
124                   WHERE name = 'default' AND status = 1 AND deleted_at IS NULL
125                   LIMIT 1"#,
126            )
127            .fetch_optional(pool)
128            .await
129            .context(SqlxSnafu)?;
130            return Ok(fallback.map(Into::into));
131        }
132
133        Ok(None)
134    }
135}
136
137impl Model for TokenLlmModel {
138    type Output = TokenLlm;
139    fn new() -> Self {
140        Self::default()
141    }
142
143    async fn schema_view(&self, _pool: &Pool<Postgres>) -> SchemaView {
144        SchemaView {
145            schemas: vec![
146                Schema::new_id(),
147                Schema {
148                    name: "name".to_string(),
149                    category: SchemaType::String,
150                    required: true,
151                    fixed: true,
152                    filterable: true,
153                    ..Default::default()
154                },
155                Schema {
156                    name: "url".to_string(),
157                    category: SchemaType::String,
158                    required: true,
159                    ..Default::default()
160                },
161                Schema {
162                    name: "model".to_string(),
163                    category: SchemaType::String,
164                    required: true,
165                    filterable: true,
166                    options: Some(new_schema_options(&["mimo-v2.5-pro"])),
167                    ..Default::default()
168                },
169                Schema {
170                    name: "api_key".to_string(),
171                    category: SchemaType::String,
172                    required: true,
173                    ..Default::default()
174                },
175                Schema {
176                    name: "provider".to_string(),
177                    category: SchemaType::String,
178                    filterable: true,
179                    options: Some(new_schema_options(&[
180                        LLM_PROVIDER_OPENAI,
181                        LLM_PROVIDER_ANTHROPIC,
182                    ])),
183                    default_value: Some(serde_json::json!(LLM_PROVIDER_OPENAI)),
184                    ..Default::default()
185                },
186                Schema::new_status(),
187                Schema::new_remark(),
188                Schema::new_created(),
189                Schema::new_modified(),
190            ],
191            allow_edit: SchemaAllowEdit {
192                roles: vec!["su".to_string()],
193                ..Default::default()
194            },
195            allow_create: SchemaAllowCreate {
196                roles: vec!["su".to_string()],
197                ..Default::default()
198            },
199        }
200    }
201
202    async fn insert(&self, pool: &Pool<Postgres>, data: serde_json::Value) -> Result<u64> {
203        let p: TokenLlmInsertParams = serde_json::from_value(data).context(JsonSnafu)?;
204        let provider = p
205            .provider
206            .filter(|s| !s.is_empty())
207            .unwrap_or_else(|| LLM_PROVIDER_OPENAI.to_string());
208        let row: (i64,) = sqlx::query_as(
209            r#"INSERT INTO token_llms
210               (name, url, model, api_key, provider, status, remark)
211               VALUES ($1, $2, $3, $4, $5, $6, $7)
212               RETURNING id"#,
213        )
214        .bind(&p.name)
215        .bind(&p.url)
216        .bind(&p.model)
217        .bind(&p.api_key)
218        .bind(provider)
219        .bind(p.status.unwrap_or(Status::Enabled as i16))
220        .bind(p.remark.unwrap_or_default())
221        .fetch_one(pool)
222        .await
223        .context(SqlxSnafu)?;
224        Ok(row.0 as u64)
225    }
226
227    async fn get_by_id(&self, pool: &Pool<Postgres>, id: u64) -> Result<Option<Self::Output>> {
228        let result = sqlx::query_as::<_, TokenLlmSchema>(
229            r#"SELECT * FROM token_llms WHERE id = $1 AND deleted_at IS NULL"#,
230        )
231        .bind(id as i64)
232        .fetch_optional(pool)
233        .await
234        .context(SqlxSnafu)?;
235        Ok(result.map(Into::into))
236    }
237
238    async fn update_by_id(
239        &self,
240        pool: &Pool<Postgres>,
241        id: u64,
242        data: serde_json::Value,
243    ) -> Result<()> {
244        let p: TokenLlmUpdateParams = serde_json::from_value(data).context(JsonSnafu)?;
245        let mut qb: QueryBuilder<Postgres> =
246            QueryBuilder::new("UPDATE token_llms SET modified = NOW()");
247        if let Some(v) = p.url {
248            qb.push(", url = ").push_bind(v);
249        }
250        if let Some(v) = p.model {
251            qb.push(", model = ").push_bind(v);
252        }
253        if let Some(v) = p.api_key {
254            qb.push(", api_key = ").push_bind(v);
255        }
256        if let Some(v) = p.provider {
257            qb.push(", provider = ").push_bind(v);
258        }
259        if let Some(v) = p.status {
260            qb.push(", status = ").push_bind(v);
261        }
262        if let Some(v) = p.remark {
263            qb.push(", remark = ").push_bind(v);
264        }
265        qb.push(" WHERE id = ")
266            .push_bind(id as i64)
267            .push(" AND deleted_at IS NULL");
268        qb.build().execute(pool).await.context(SqlxSnafu)?;
269        Ok(())
270    }
271
272    async fn delete_by_id(&self, pool: &Pool<Postgres>, id: u64) -> Result<()> {
273        sqlx::query(
274            r#"UPDATE token_llms SET deleted_at = NOW(), modified = NOW() WHERE id = $1 AND deleted_at IS NULL"#,
275        )
276        .bind(id as i64)
277        .execute(pool)
278        .await
279        .context(SqlxSnafu)?;
280        Ok(())
281    }
282
283    async fn count(&self, pool: &Pool<Postgres>, params: &ModelListParams) -> Result<i64> {
284        let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("SELECT COUNT(*) FROM token_llms");
285        self.push_conditions(&mut qb, params)?;
286        let row: (i64,) = qb
287            .build_query_as()
288            .fetch_one(pool)
289            .await
290            .context(SqlxSnafu)?;
291        Ok(row.0)
292    }
293
294    async fn list(
295        &self,
296        pool: &Pool<Postgres>,
297        params: &ModelListParams,
298    ) -> Result<Vec<Self::Output>> {
299        let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("SELECT * FROM token_llms");
300        self.push_conditions(&mut qb, params)?;
301        params.push_pagination(&mut qb);
302        let rows = qb
303            .build_query_as::<TokenLlmSchema>()
304            .fetch_all(pool)
305            .await
306            .context(SqlxSnafu)?;
307        Ok(rows.into_iter().map(Into::into).collect())
308    }
309
310    fn push_filter_conditions<'args>(
311        &self,
312        qb: &mut QueryBuilder<'args, Postgres>,
313        filters: &HashMap<String, String>,
314    ) -> Result<()> {
315        if let Some(name) = filters.get("name") {
316            qb.push(" AND name = ").push_bind(name.clone());
317        }
318        if let Some(model) = filters.get("model") {
319            qb.push(" AND model = ").push_bind(model.clone());
320        }
321        if let Some(provider) = filters.get("provider") {
322            qb.push(" AND provider = ").push_bind(provider.clone());
323        }
324        if let Some(status) = filters.get("status") {
325            if let Ok(v) = status.parse::<i16>() {
326                qb.push(" AND status = ").push_bind(v);
327            }
328        }
329        Ok(())
330    }
331}
tibba_model_token/llm.rs

tibba_model_token/
llm.rs
