mofa_foundation/persistence/
entities.rs

1//! 持久化实体定义
2//!
3//! 对应数据库表结构的实体类型
4
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use uuid::Uuid;
8
9/// 消息角色
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum MessageRole {
13    /// 系统消息
14    System,
15    /// 用户消息
16    User,
17    /// 助手消息
18    Assistant,
19    /// 工具消息
20    Tool,
21}
22
23impl std::fmt::Display for MessageRole {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        match self {
26            MessageRole::System => write!(f, "system"),
27            MessageRole::User => write!(f, "user"),
28            MessageRole::Assistant => write!(f, "assistant"),
29            MessageRole::Tool => write!(f, "tool"),
30        }
31    }
32}
33
34impl std::str::FromStr for MessageRole {
35    type Err = String;
36
37    fn from_str(s: &str) -> Result<Self, Self::Err> {
38        match s.to_lowercase().as_str() {
39            "system" => Ok(MessageRole::System),
40            "user" => Ok(MessageRole::User),
41            "assistant" => Ok(MessageRole::Assistant),
42            "tool" => Ok(MessageRole::Tool),
43            _ => Err(format!("Unknown role: {}", s)),
44        }
45    }
46}
47
48/// API 调用状态
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
50#[serde(rename_all = "lowercase")]
51#[derive(Default)]
52pub enum ApiCallStatus {
53    /// 成功
54    #[default]
55    Success,
56    /// 失败
57    Failed,
58    /// 超时
59    Timeout,
60    /// 限流
61    RateLimited,
62    /// 取消
63    Cancelled,
64}
65
66impl std::fmt::Display for ApiCallStatus {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        match self {
69            ApiCallStatus::Success => write!(f, "success"),
70            ApiCallStatus::Failed => write!(f, "failed"),
71            ApiCallStatus::Timeout => write!(f, "timeout"),
72            ApiCallStatus::RateLimited => write!(f, "rate_limited"),
73            ApiCallStatus::Cancelled => write!(f, "cancelled"),
74        }
75    }
76}
77
78/// 消息内容
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct MessageContent {
81    /// 文本内容
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub text: Option<String>,
84    /// 工具调用
85    #[serde(skip_serializing_if = "Option::is_none")]
86    pub tool_calls: Option<Vec<ToolCallContent>>,
87    /// 工具结果
88    #[serde(skip_serializing_if = "Option::is_none")]
89    pub tool_result: Option<ToolResultContent>,
90    /// 附加数据
91    #[serde(flatten)]
92    pub extra: HashMap<String, serde_json::Value>,
93}
94
95impl MessageContent {
96    /// 创建文本消息内容
97    pub fn text(content: impl Into<String>) -> Self {
98        Self {
99            text: Some(content.into()),
100            tool_calls: None,
101            tool_result: None,
102            extra: HashMap::new(),
103        }
104    }
105
106    /// 创建工具调用消息内容
107    pub fn tool_calls(calls: Vec<ToolCallContent>) -> Self {
108        Self {
109            text: None,
110            tool_calls: Some(calls),
111            tool_result: None,
112            extra: HashMap::new(),
113        }
114    }
115
116    /// 创建工具结果消息内容
117    pub fn tool_result(result: ToolResultContent) -> Self {
118        Self {
119            text: None,
120            tool_calls: None,
121            tool_result: Some(result),
122            extra: HashMap::new(),
123        }
124    }
125}
126
127/// 工具调用内容
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct ToolCallContent {
130    /// 工具调用 ID
131    pub id: String,
132    /// 工具名称
133    pub name: String,
134    /// 工具参数
135    pub arguments: serde_json::Value,
136}
137
138/// 工具结果内容
139#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct ToolResultContent {
141    /// 工具调用 ID
142    pub tool_call_id: String,
143    /// 结果内容
144    pub content: String,
145    /// 是否错误
146    #[serde(default)]
147    pub is_error: bool,
148}
149
150/// LLM 消息实体
151///
152/// 对应 `entity_llm_message` 表
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct LLMMessage {
155    /// 消息 ID
156    pub id: Uuid,
157    /// 父消息 ID (用于构建对话树)
158    #[serde(skip_serializing_if = "Option::is_none")]
159    pub parent_message_id: Option<Uuid>,
160    /// 会话 ID
161    pub chat_session_id: Uuid,
162    /// Agent ID
163    pub agent_id: Uuid,
164    /// 消息内容
165    pub content: MessageContent,
166    /// 消息角色
167    pub role: MessageRole,
168    /// 用户 ID
169    pub user_id: Uuid,
170    /// 租户 ID
171    pub tenant_id: Uuid,
172    /// 创建时间
173    pub create_time: chrono::DateTime<chrono::Utc>,
174    /// 更新时间
175    pub update_time: chrono::DateTime<chrono::Utc>,
176}
177
178impl LLMMessage {
179    /// 创建新消息
180    pub fn new(
181        chat_session_id: Uuid,
182        agent_id: Uuid,
183        user_id: Uuid,
184        tenant_id: Uuid,
185        role: MessageRole,
186        content: MessageContent,
187    ) -> Self {
188        let now = chrono::Utc::now();
189        Self {
190            id: Uuid::now_v7(),
191            parent_message_id: None,
192            chat_session_id,
193            agent_id,
194            content,
195            role,
196            user_id,
197            tenant_id,
198            create_time: now,
199            update_time: now,
200        }
201    }
202
203    /// 设置父消息 ID
204    pub fn with_parent(mut self, parent_id: Uuid) -> Self {
205        self.parent_message_id = Some(parent_id);
206        self
207    }
208
209    /// 设置租户 ID
210    pub fn with_tenant(mut self, tenant_id: Uuid) -> Self {
211        self.tenant_id = tenant_id;
212        self
213    }
214}
215
216/// Token 使用详情
217#[derive(Debug, Clone, Default, Serialize, Deserialize)]
218pub struct TokenDetails {
219    /// 缓存 tokens
220    #[serde(skip_serializing_if = "Option::is_none")]
221    pub cached_tokens: Option<i32>,
222    /// 推理 tokens
223    #[serde(skip_serializing_if = "Option::is_none")]
224    pub reasoning_tokens: Option<i32>,
225    /// 附加信息
226    #[serde(flatten)]
227    pub extra: HashMap<String, serde_json::Value>,
228}
229
230/// 价格详情
231#[derive(Debug, Clone, Default, Serialize, Deserialize)]
232pub struct PriceDetails {
233    /// 输入价格
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub input_price: Option<f64>,
236    /// 输出价格
237    #[serde(skip_serializing_if = "Option::is_none")]
238    pub output_price: Option<f64>,
239    /// 货币单位
240    #[serde(default = "default_currency")]
241    pub currency: String,
242    /// 附加信息
243    #[serde(flatten)]
244    pub extra: HashMap<String, serde_json::Value>,
245}
246
247fn default_currency() -> String {
248    "USD".to_string()
249}
250
251/// LLM API 调用记录实体
252///
253/// 对应 `entity_llm_api_call` 表
254#[derive(Debug, Clone, Serialize, Deserialize)]
255pub struct LLMApiCall {
256    /// 记录 ID
257    pub id: Uuid,
258    /// 会话 ID
259    pub chat_session_id: Uuid,
260    /// Agent ID
261    pub agent_id: Uuid,
262    /// 用户 ID
263    pub user_id: Uuid,
264    // 租户 ID
265    pub tenant_id: Uuid,
266    /// 请求消息 ID
267    pub request_message_id: Uuid,
268    /// 响应消息 ID
269    pub response_message_id: Uuid,
270    /// 模型名称
271    pub model_name: String,
272    /// 提示词 tokens
273    pub prompt_tokens: i32,
274    /// 提示词 tokens 详情
275    #[serde(skip_serializing_if = "Option::is_none")]
276    pub prompt_tokens_details: Option<TokenDetails>,
277    /// 补全 tokens
278    pub completion_tokens: i32,
279    /// 补全 tokens 详情
280    #[serde(skip_serializing_if = "Option::is_none")]
281    pub completion_tokens_details: Option<TokenDetails>,
282    /// 总 tokens
283    pub total_tokens: i32,
284    /// 总价格
285    #[serde(skip_serializing_if = "Option::is_none")]
286    pub total_price: Option<f64>,
287    /// 价格详情
288    #[serde(skip_serializing_if = "Option::is_none")]
289    pub price_details: Option<PriceDetails>,
290    /// 延迟 (毫秒)
291    #[serde(skip_serializing_if = "Option::is_none")]
292    pub latency_ms: Option<i32>,
293    /// 首 token 时间 (毫秒)
294    #[serde(skip_serializing_if = "Option::is_none")]
295    pub time_to_first_token_ms: Option<i32>,
296    /// tokens/秒
297    #[serde(skip_serializing_if = "Option::is_none")]
298    pub tokens_per_second: Option<f64>,
299    /// API 响应 ID
300    #[serde(skip_serializing_if = "Option::is_none")]
301    pub api_response_id: Option<String>,
302    /// 调用状态
303    pub status: ApiCallStatus,
304    /// 错误消息
305    #[serde(skip_serializing_if = "Option::is_none")]
306    pub error_message: Option<String>,
307    /// 错误代码
308    #[serde(skip_serializing_if = "Option::is_none")]
309    pub error_code: Option<String>,
310    /// 创建时间
311    pub create_time: chrono::DateTime<chrono::Utc>,
312    /// 更新时间
313    pub update_time: chrono::DateTime<chrono::Utc>,
314}
315
316impl LLMApiCall {
317    /// 创建成功的 API 调用记录
318    pub fn success(
319        chat_session_id: Uuid,
320        agent_id: Uuid,
321        user_id: Uuid,
322        tenant_id: Uuid,
323        request_message_id: Uuid,
324        response_message_id: Uuid,
325        model_name: impl Into<String>,
326        prompt_tokens: i32,
327        completion_tokens: i32,
328        request_time: chrono::DateTime<chrono::Utc>,
329        response_time: chrono::DateTime<chrono::Utc>,
330    ) -> Self {
331        let latency_ms = (response_time - request_time).num_milliseconds() as i32;
332        let tokens_per_second = if latency_ms > 0 {
333            Some(completion_tokens as f64 / (latency_ms as f64 / 1000.0))
334        } else {
335            None
336        };
337
338        Self {
339            id: Uuid::now_v7(),
340            chat_session_id,
341            agent_id,
342            user_id,
343            tenant_id,
344            request_message_id,
345            response_message_id,
346            model_name: model_name.into(),
347            prompt_tokens,
348            prompt_tokens_details: None,
349            completion_tokens,
350            completion_tokens_details: None,
351            total_tokens: prompt_tokens + completion_tokens,
352            total_price: None,
353            price_details: None,
354            latency_ms: Some(latency_ms),
355            time_to_first_token_ms: None,
356            tokens_per_second,
357            api_response_id: None,
358            status: ApiCallStatus::Success,
359            error_message: None,
360            error_code: None,
361            create_time: request_time,
362            update_time: response_time,
363        }
364    }
365
366    /// 创建失败的 API 调用记录
367    pub fn failed(
368        chat_session_id: Uuid,
369        agent_id: Uuid,
370        user_id: Uuid,
371        tenant_id: Uuid,
372        request_message_id: Uuid,
373        model_name: impl Into<String>,
374        error_message: impl Into<String>,
375        error_code: Option<String>,
376        request_time: chrono::DateTime<chrono::Utc>,
377    ) -> Self {
378        let now = chrono::Utc::now();
379        Self {
380            id: Uuid::now_v7(),
381            chat_session_id,
382            agent_id,
383            user_id,
384            tenant_id,
385            request_message_id,
386            response_message_id: Uuid::nil(),
387            model_name: model_name.into(),
388            prompt_tokens: 0,
389            prompt_tokens_details: None,
390            completion_tokens: 0,
391            completion_tokens_details: None,
392            total_tokens: 0,
393            total_price: None,
394            price_details: None,
395            latency_ms: Some((now - request_time).num_milliseconds() as i32),
396            time_to_first_token_ms: None,
397            tokens_per_second: None,
398            api_response_id: None,
399            status: ApiCallStatus::Failed,
400            error_message: Some(error_message.into()),
401            error_code,
402            create_time: request_time,
403            update_time: now,
404        }
405    }
406
407    /// 设置 API 响应 ID
408    pub fn with_api_response_id(mut self, id: impl Into<String>) -> Self {
409        self.api_response_id = Some(id.into());
410        self
411    }
412
413    /// 设置价格信息
414    pub fn with_price(mut self, total_price: f64, details: Option<PriceDetails>) -> Self {
415        self.total_price = Some(total_price);
416        self.price_details = details;
417        self
418    }
419
420    /// 设置首 token 时间
421    pub fn with_time_to_first_token(mut self, ttft_ms: i32) -> Self {
422        self.time_to_first_token_ms = Some(ttft_ms);
423        self
424    }
425
426    /// 设置 token 详情
427    pub fn with_token_details(
428        mut self,
429        prompt_details: Option<TokenDetails>,
430        completion_details: Option<TokenDetails>,
431    ) -> Self {
432        self.prompt_tokens_details = prompt_details;
433        self.completion_tokens_details = completion_details;
434        self
435    }
436}
437
438/// 会话实体
439#[derive(Debug, Clone, Serialize, Deserialize)]
440pub struct ChatSession {
441    /// 会话 ID
442    pub id: Uuid,
443    /// 用户 ID
444    pub user_id: Uuid,
445    /// Agent ID
446    pub agent_id: Uuid,
447    /// 租户 ID
448    pub tenant_id: Uuid,
449    /// 会话标题
450    pub title: Option<String>,
451    /// 会话元数据
452    pub metadata: HashMap<String, serde_json::Value>,
453    /// 创建时间
454    pub create_time: chrono::DateTime<chrono::Utc>,
455    /// 更新时间
456    pub update_time: chrono::DateTime<chrono::Utc>,
457}
458
459impl ChatSession {
460    /// 创建新会话
461    pub fn new(user_id: Uuid, agent_id: Uuid) -> Self {
462        let now = chrono::Utc::now();
463        Self {
464            id: Uuid::now_v7(),
465            user_id,
466            agent_id,
467            tenant_id: Uuid::nil(), // 默认为 nil UUID，可以通过 with_tenant_id 设置
468            title: None,
469            metadata: HashMap::new(),
470            create_time: now,
471            update_time: now,
472        }
473    }
474
475    /// 设置标题
476    pub fn with_title(mut self, title: impl Into<String>) -> Self {
477        self.title = Some(title.into());
478        self
479    }
480
481    /// 设置 ID
482    pub fn with_id(mut self, id: Uuid) -> Self {
483        self.id = id;
484        self
485    }
486
487    /// 设置租户 ID
488    pub fn with_tenant_id(mut self, tenant_id: Uuid) -> Self {
489        self.tenant_id = tenant_id;
490        self
491    }
492
493    /// 添加元数据
494    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
495        self.metadata.insert(key.into(), value);
496        self
497    }
498}
499
500/// 查询过滤器
501#[derive(Debug, Clone, Default)]
502pub struct QueryFilter {
503    /// 用户 ID
504    pub user_id: Option<Uuid>,
505    /// 会话 ID
506    pub session_id: Option<Uuid>,
507    /// Agent ID
508    pub agent_id: Option<Uuid>,
509    /// 开始时间
510    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
511    /// 结束时间
512    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
513    /// 状态过滤
514    pub status: Option<ApiCallStatus>,
515    /// 模型名称
516    pub model_name: Option<String>,
517    /// 分页: 偏移量
518    pub offset: Option<i64>,
519    /// 分页: 限制数量
520    pub limit: Option<i64>,
521}
522
523impl QueryFilter {
524    pub fn new() -> Self {
525        Self::default()
526    }
527
528    pub fn user(mut self, user_id: Uuid) -> Self {
529        self.user_id = Some(user_id);
530        self
531    }
532
533    pub fn session(mut self, session_id: Uuid) -> Self {
534        self.session_id = Some(session_id);
535        self
536    }
537
538    pub fn agent(mut self, agent_id: Uuid) -> Self {
539        self.agent_id = Some(agent_id);
540        self
541    }
542
543    pub fn time_range(
544        mut self,
545        start: chrono::DateTime<chrono::Utc>,
546        end: chrono::DateTime<chrono::Utc>,
547    ) -> Self {
548        self.start_time = Some(start);
549        self.end_time = Some(end);
550        self
551    }
552
553    pub fn with_status(mut self, status: ApiCallStatus) -> Self {
554        self.status = Some(status);
555        self
556    }
557
558    pub fn model(mut self, model_name: impl Into<String>) -> Self {
559        self.model_name = Some(model_name.into());
560        self
561    }
562
563    pub fn paginate(mut self, offset: i64, limit: i64) -> Self {
564        self.offset = Some(offset);
565        self.limit = Some(limit);
566        self
567    }
568}
569
570/// 统计摘要
571#[derive(Debug, Clone, Default, Serialize, Deserialize)]
572pub struct UsageStatistics {
573    /// 总调用次数
574    pub total_calls: i64,
575    /// 成功次数
576    pub success_count: i64,
577    /// 失败次数
578    pub failed_count: i64,
579    /// 总 tokens
580    pub total_tokens: i64,
581    /// 总提示词 tokens
582    pub total_prompt_tokens: i64,
583    /// 总补全 tokens
584    pub total_completion_tokens: i64,
585    /// 总费用
586    pub total_cost: Option<f64>,
587    /// 平均延迟 (毫秒)
588    pub avg_latency_ms: Option<f64>,
589    /// 平均 tokens/秒
590    pub avg_tokens_per_second: Option<f64>,
591}
592
593/// Provider Entity - maps to entity_provider table
594#[derive(Debug, Clone, Serialize, Deserialize)]
595pub struct Provider {
596    pub id: Uuid,
597    pub tenant_id: Uuid,
598    pub provider_name: String,
599    pub provider_type: String,
600    pub api_base: String,
601    pub api_key: String,
602    pub enabled: bool,
603    pub create_time: chrono::DateTime<chrono::Utc>,
604    pub update_time: chrono::DateTime<chrono::Utc>,
605}
606
607/// Agent Entity - maps to entity_agent table
608#[derive(Debug, Clone, Serialize, Deserialize)]
609pub struct Agent {
610    pub id: Uuid,
611    pub tenant_id: Uuid,
612    pub agent_code: String,
613    pub agent_name: String,
614    pub agent_order: i32,
615    pub agent_status: bool,
616    pub context_limit: Option<i32>,
617    pub custom_params: Option<serde_json::Value>,
618    pub max_completion_tokens: Option<i32>,
619    pub model_name: String,
620    pub provider_id: Uuid,
621    pub response_format: Option<String>,
622    pub system_prompt: String,
623    pub temperature: Option<f32>,
624    pub stream: Option<bool>,
625    pub thinking: Option<serde_json::Value>,
626    pub create_time: chrono::DateTime<chrono::Utc>,
627    pub update_time: chrono::DateTime<chrono::Utc>,
628}
629
630/// Agent Configuration with Provider
631#[derive(Debug, Clone)]
632pub struct AgentConfig {
633    pub provider: Provider,
634    pub agent: Agent,
635}
mofa_foundation/persistence/entities.rs

mofa_foundation/persistence/
entities.rs
