helios_engine/
chat.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
4#[serde(rename_all = "lowercase")]
5pub enum Role {
6    System,
7    User,
8    Assistant,
9    Tool,
10}
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ChatMessage {
14    pub role: Role,
15    #[serde(default)]
16    pub content: String,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub name: Option<String>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub tool_calls: Option<Vec<ToolCall>>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub tool_call_id: Option<String>,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct ToolCall {
27    pub id: String,
28    #[serde(rename = "type")]
29    pub call_type: String,
30    pub function: FunctionCall,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct FunctionCall {
35    pub name: String,
36    pub arguments: String,
37}
38
39impl ChatMessage {
40    pub fn system(content: impl Into<String>) -> Self {
41        Self {
42            role: Role::System,
43            content: content.into(),
44            name: None,
45            tool_calls: None,
46            tool_call_id: None,
47        }
48    }
49
50    pub fn user(content: impl Into<String>) -> Self {
51        Self {
52            role: Role::User,
53            content: content.into(),
54            name: None,
55            tool_calls: None,
56            tool_call_id: None,
57        }
58    }
59
60    pub fn assistant(content: impl Into<String>) -> Self {
61        Self {
62            role: Role::Assistant,
63            content: content.into(),
64            name: None,
65            tool_calls: None,
66            tool_call_id: None,
67        }
68    }
69
70    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
71        Self {
72            role: Role::Tool,
73            content: content.into(),
74            name: None,
75            tool_calls: None,
76            tool_call_id: Some(tool_call_id.into()),
77        }
78    }
79}
80
81#[derive(Debug, Clone)]
82pub struct ChatSession {
83    pub messages: Vec<ChatMessage>,
84    pub system_prompt: Option<String>,
85}
86
87impl ChatSession {
88    pub fn new() -> Self {
89        Self {
90            messages: Vec::new(),
91            system_prompt: None,
92        }
93    }
94
95    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
96        self.system_prompt = Some(prompt.into());
97        self
98    }
99
100    pub fn add_message(&mut self, message: ChatMessage) {
101        self.messages.push(message);
102    }
103
104    pub fn add_user_message(&mut self, content: impl Into<String>) {
105        self.messages.push(ChatMessage::user(content));
106    }
107
108    pub fn add_assistant_message(&mut self, content: impl Into<String>) {
109        self.messages.push(ChatMessage::assistant(content));
110    }
111
112    pub fn get_messages(&self) -> Vec<ChatMessage> {
113        let mut messages = Vec::new();
114        
115        if let Some(ref system_prompt) = self.system_prompt {
116            messages.push(ChatMessage::system(system_prompt.clone()));
117        }
118        
119        messages.extend(self.messages.clone());
120        messages
121    }
122
123    pub fn clear(&mut self) {
124        self.messages.clear();
125    }
126}
127
128impl Default for ChatSession {
129    fn default() -> Self {
130        Self::new()
131    }
132}
helios_engine/chat.rs

helios_engine/
chat.rs
