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

helios_engine/
chat.rs
