funera_core/chat/
message.rs

1use std::{fmt, sync::Arc};
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::{Value as JsonValue, json};
6use uuid::Uuid;
7
8#[cfg(feature = "tool")]
9use crate::re_act::tool::ToolType;
10
11pub trait Message {
12    fn to_prompt_content(&self) -> String;
13}
14
15#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
16pub enum Role {
17    User,
18    Assistant,
19    System,
20    Tool,
21}
22impl fmt::Display for Role {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        match self {
25            Role::User => write!(f, "user"),
26            Role::Assistant => write!(f, "assistant"),
27            Role::System => write!(f, "system"),
28            Role::Tool => write!(f, "tool"),
29        }
30    }
31}
32
33#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
34pub enum MsgVariant {
35    Text(TextMessage),
36    #[cfg(feature = "tool")]
37    ToolRequest(ToolRequestMessage),
38    #[cfg(feature = "tool")]
39    ToolResponse(ToolResponseMessage),
40}
41
42#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
43pub struct FuneraMessage {
44    id: Uuid,
45    role: Role,
46    timestamp: DateTime<Utc>,
47    msg_variant: MsgVariant,
48}
49
50#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
51pub struct TextMessage {
52    pub text: Arc<str>,
53    pub reasoning_content: Option<Arc<str>>,
54}
55
56#[cfg(feature = "tool")]
57#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
58pub struct ToolRequestMessage {
59    pub tool_call_id: Arc<str>,
60    pub tool_type: ToolType,
61    pub function_name: Arc<str>,
62    pub function_args: JsonValue,
63    pub reasoning_content: Option<Arc<str>>,
64}
65
66#[cfg(feature = "tool")]
67#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
68pub struct ToolResponseMessage {
69    pub tool_call_id: Arc<str>,
70    pub result: Arc<str>,
71}
72
73impl Message for TextMessage {
74    fn to_prompt_content(&self) -> String {
75        self.text.to_string()
76    }
77}
78
79#[cfg(feature = "tool")]
80impl Message for ToolRequestMessage {
81    fn to_prompt_content(&self) -> String {
82        format!(
83            "Tool Call -> id: {}, tool_type: {}, function_name: {}, function_args: {}",
84            self.tool_call_id, self.tool_type, self.function_name, self.function_args
85        )
86    }
87}
88
89#[cfg(feature = "tool")]
90impl Message for ToolResponseMessage {
91    fn to_prompt_content(&self) -> String {
92        format!(
93            "Tool Response -> tool_call_id: {}, result: {}",
94            self.tool_call_id, self.result
95        )
96    }
97}
98
99impl Message for FuneraMessage {
100    fn to_prompt_content(&self) -> String {
101        match &self.msg_variant {
102            MsgVariant::Text(text_msg) => text_msg.to_prompt_content(),
103            #[cfg(feature = "tool")]
104            MsgVariant::ToolRequest(tool_requet_msg) => tool_requet_msg.to_prompt_content(),
105            #[cfg(feature = "tool")]
106            MsgVariant::ToolResponse(tool_response_msg) => tool_response_msg.to_prompt_content(),
107        }
108    }
109}
110
111impl FuneraMessage {
112    pub fn new(role: Role, msg_variant: MsgVariant) -> Self {
113        Self {
114            id: Uuid::new_v4(),
115            role,
116            timestamp: Utc::now(),
117            msg_variant,
118        }
119    }
120
121    pub fn id(&self) -> Uuid {
122        self.id
123    }
124
125    pub fn role(&self) -> &Role {
126        &self.role
127    }
128
129    pub fn timestamp(&self) -> &DateTime<Utc> {
130        &self.timestamp
131    }
132
133    pub fn msg_variant(&self) -> &MsgVariant {
134        &self.msg_variant
135    }
136
137    pub fn format_json(&self) -> JsonValue {
138        match &self.msg_variant {
139            MsgVariant::Text(text_msg) => {
140                let mut obj = json!({
141                    "role": self.role.to_string(),
142                    "content": text_msg.to_prompt_content(),
143                });
144                if let Some(ref rc) = text_msg.reasoning_content {
145                    obj.as_object_mut()
146                        .unwrap()
147                        .insert("reasoning_content".into(), json!(rc.as_ref()));
148                }
149                obj
150            }
151            #[cfg(feature = "tool")]
152            MsgVariant::ToolRequest(tool_request_msg) => {
153                let mut obj = json!({
154                    "role": self.role.to_string(),
155                    "tool_calls": [
156                        {
157                            "id": tool_request_msg.tool_call_id.as_ref(),
158                            "type": tool_request_msg.tool_type.to_string(),
159                            "function": {
160                                "name": tool_request_msg.function_name.as_ref(),
161                                "arguments": serde_json::to_string(&tool_request_msg.function_args).unwrap_or_default(),
162                            }
163                        }
164                    ]
165                });
166                if let Some(ref rc) = tool_request_msg.reasoning_content {
167                    obj.as_object_mut()
168                        .unwrap()
169                        .insert("reasoning_content".into(), json!(rc.as_ref()));
170                }
171                obj
172            }
173            #[cfg(feature = "tool")]
174            MsgVariant::ToolResponse(tool_response_msg) => {
175                json!({
176                    "role": self.role.to_string(),
177                    "tool_call_id": tool_response_msg.tool_call_id.as_ref(),
178                    "content": tool_response_msg.to_prompt_content(),
179                })
180            }
181        }
182    }
183}
funera_core/chat/message.rs

funera_core/chat/
message.rs
