//! ## 数据库查询(models) //! //! 所有数据库操作集中在此模块,包括: //! //! ### 表操作 //! * `app_state` — 认证状态的持久化(key-value 存储) //! * `chat_records` — 聊天记录(inbound/outbound) //! * `llm_usage` — Token 用量记录和统计 //! * `session_summaries` — 会话摘要存档 //! * `user_memories` — 用户长期记忆 //! * `pending_approvals` — 待审批记录 use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use crate::state::AuthState; // ─── 认证状态 ─── /// 从数据库加载认证状态 pub async fn load_auth(pool: &PgPool) -> Option { let row = sqlx::query_as::<_, (serde_json::Value,)>( "SELECT value FROM app_state WHERE key = 'auth'", ) .fetch_optional(pool) .await .ok()??; serde_json::from_value(row.0).ok() } /// 保存认证状态到数据库 pub async fn save_auth(pool: &PgPool, auth: &AuthState) -> Result<(), String> { let value = serde_json::to_value(auth).map_err(|e| format!("序列化 auth 失败: {}", e))?; sqlx::query( r#" INSERT INTO app_state (key, value, updated_at) VALUES ('auth', $1, NOW()) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = NOW() "#, ) .bind(&value) .execute(pool) .await .map_err(|e| format!("保存 auth 到数据库失败: {}", e))?; Ok(()) } // ─── 聊天记录 ─── #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct ChatRecord { pub id: i64, pub created_at: DateTime, pub direction: String, pub user_id: String, pub account_id: String, pub text: String, pub source: String, pub context_token: String, pub message_id: String, } /// 插入一条聊天记录 pub async fn insert_chat_record( pool: &PgPool, direction: &str, user_id: &str, account_id: &str, text: &str, source: &str, context_token: Option<&str>, message_id: &str, ) -> Result { let ctx = context_token.unwrap_or(""); let row: Option<(i64,)> = sqlx::query_as( r#" INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (message_id) WHERE message_id NOT IN ('', '0') DO NOTHING RETURNING id "#, ) .bind(direction) .bind(user_id) .bind(account_id) .bind(text) .bind(source) .bind(ctx) .bind(message_id) .fetch_optional(pool) .await .map_err(|e| format!("插入聊天记录失败: {}", e))?; Ok(row.map(|r| r.0).unwrap_or(0)) } /// 查询最近的聊天记录(用户维度) pub async fn list_recent_chat_records( pool: &PgPool, user_id: &str, limit: i64, ) -> Result, String> { let records = sqlx::query_as::<_, ChatRecord>( r#" SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id FROM chat_records WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2 "#, ) .bind(user_id) .bind(limit) .fetch_all(pool) .await .map_err(|e| format!("查询聊天记录失败: {}", e))?; Ok(records) } // ─── LLM 用量 ─── /// 插入一条 LLM 用量记录 /// kind/entry 为 iPet 遗留兼容字段,写入默认值 ("llm_call" / null) pub async fn insert_llm_usage( pool: &PgPool, user_id: &str, model: &str, provider: &str, prompt_tokens: u32, completion_tokens: u32, cache_hit_tokens: u32, cache_miss_tokens: u32, ) -> Result<(), String> { let total = prompt_tokens + completion_tokens; sqlx::query( r#" INSERT INTO llm_usage (user_id, model, provider, prompt_tokens, completion_tokens, total_tokens, cache_hit_tokens, cache_miss_tokens, kind, entry) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) "#, ) .bind(user_id) .bind(model) .bind(provider) .bind(prompt_tokens as i32) .bind(completion_tokens as i32) .bind(total as i32) .bind(cache_hit_tokens as i32) .bind(cache_miss_tokens as i32) .bind("llm_call") .bind(serde_json::Value::Null) .execute(pool) .await .map_err(|e| format!("插入 LLM 用量失败: {}", e))?; Ok(()) } /// LLM 用量聚合统计 #[derive(Debug, Clone, sqlx::FromRow)] pub struct LlmUsageStats { pub total_calls: i64, pub total_prompt_tokens: i64, pub total_completion_tokens: i64, pub total_tokens: i64, pub total_cache_hit: i64, pub total_cache_miss: i64, } /// 查询 LLM 用量统计(按时间范围过滤) pub async fn query_llm_usage_stats( pool: &PgPool, since: Option>, until: Option>, model_filter: Option<&str>, ) -> Result { let since = since.unwrap_or_else(|| { chrono::Utc::now() - chrono::Duration::days(7) }); let until = until.unwrap_or_else(chrono::Utc::now); let row = sqlx::query_as::<_, LlmUsageStats>( r#" SELECT COUNT(*)::bigint as total_calls, COALESCE(SUM(prompt_tokens), 0)::bigint as total_prompt_tokens, COALESCE(SUM(completion_tokens), 0)::bigint as total_completion_tokens, COALESCE(SUM(total_tokens), 0)::bigint as total_tokens, COALESCE(SUM(cache_hit_tokens), 0)::bigint as total_cache_hit, COALESCE(SUM(cache_miss_tokens), 0)::bigint as total_cache_miss FROM llm_usage WHERE created_at >= $1 AND created_at <= $2 AND ($3::text IS NULL OR model = $3) "#, ) .bind(since) .bind(until) .bind(model_filter) .fetch_one(pool) .await .map_err(|e| format!("查询 LLM 用量失败: {}", e))?; Ok(row) } // ─── 会话摘要 ─── /// 保存一条摘要到数据库 pub async fn save_summary( pool: &PgPool, user_id: &str, text: &str, reason: &str, message_count: i32, ) -> Result<(), String> { sqlx::query( "INSERT INTO session_summaries (user_id, summary_text, reason, message_count) VALUES ($1, $2, $3, $4)", ) .bind(user_id) .bind(text) .bind(reason) .bind(message_count) .execute(pool) .await .map_err(|e| format!("保存摘要失败: {}", e))?; Ok(()) } /// 加载最近的摘要(用于 read_summaries 工具),按 user_id 过滤 pub async fn load_summaries(pool: &PgPool, user_id: &str, limit: i64) -> Result, String> { let rows = sqlx::query_as::<_, (String,)>( "SELECT summary_text FROM session_summaries WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2", ) .bind(user_id) .bind(limit) .fetch_all(pool) .await .map_err(|e| format!("查询摘要失败: {}", e))?; Ok(rows.into_iter().map(|(t,)| t).collect()) }