codei_session/
model.rs

1use std::path::PathBuf;
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use uuid::Uuid;
6
7use codei_llm::{Message, Role as LlmRole, ToolCall};
8
9pub type SessionId = String;
10pub type MessageId = String;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Session {
14    pub id: SessionId,
15    pub title: Option<String>,
16    pub cwd: PathBuf,
17    pub created_at: DateTime<Utc>,
18    pub updated_at: DateTime<Utc>,
19    pub messages: Vec<StoredMessage>,
20}
21
22#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
23#[serde(rename_all = "lowercase")]
24pub enum Role {
25    System,
26    User,
27    Assistant,
28    Tool,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct StoredMessage {
33    pub id: MessageId,
34    pub role: Role,
35    pub content: MessageContent,
36    pub tool_calls: Option<Vec<ToolCallRecord>>,
37    pub tool_call_id: Option<String>,
38    pub created_at: DateTime<Utc>,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub enum MessageContent {
43    Text(String),
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ToolCallRecord {
48    pub id: String,
49    pub name: String,
50    pub arguments: String,
51}
52
53impl Session {
54    pub fn new(cwd: PathBuf) -> Self {
55        let now = Utc::now();
56        Self {
57            id: Uuid::new_v4().to_string(),
58            title: None,
59            cwd,
60            created_at: now,
61            updated_at: now,
62            messages: Vec::new(),
63        }
64    }
65
66    pub fn touch(&mut self) {
67        self.updated_at = Utc::now();
68    }
69
70    pub fn push_user(&mut self, content: impl Into<String>) -> &StoredMessage {
71        self.push_message(Role::User, MessageContent::Text(content.into()), None, None)
72    }
73
74    pub fn push_assistant(
75        &mut self,
76        content: String,
77        tool_calls: Option<Vec<ToolCallRecord>>,
78    ) -> &StoredMessage {
79        self.push_message(
80            Role::Assistant,
81            MessageContent::Text(content),
82            tool_calls,
83            None,
84        )
85    }
86
87    pub fn push_tool(
88        &mut self,
89        tool_call_id: impl Into<String>,
90        content: impl Into<String>,
91    ) -> &StoredMessage {
92        self.push_message(
93            Role::Tool,
94            MessageContent::Text(content.into()),
95            None,
96            Some(tool_call_id.into()),
97        )
98    }
99
100    pub fn clear_messages(&mut self) {
101        self.messages.clear();
102        self.touch();
103    }
104
105    /// Drop older messages, keeping the most recent `keep_recent` entries.
106    pub fn compact(&mut self, keep_recent: usize) {
107        if self.messages.len() <= keep_recent {
108            return;
109        }
110        let remove = self.messages.len() - keep_recent;
111        self.messages.drain(0..remove);
112        self.touch();
113    }
114
115    fn push_message(
116        &mut self,
117        role: Role,
118        content: MessageContent,
119        tool_calls: Option<Vec<ToolCallRecord>>,
120        tool_call_id: Option<String>,
121    ) -> &StoredMessage {
122        let msg = StoredMessage {
123            id: Uuid::new_v4().to_string(),
124            role,
125            content,
126            tool_calls,
127            tool_call_id,
128            created_at: Utc::now(),
129        };
130        self.messages.push(msg);
131        self.touch();
132        self.messages.last().expect("message just pushed")
133    }
134}
135
136impl From<ToolCall> for ToolCallRecord {
137    fn from(value: ToolCall) -> Self {
138        Self {
139            id: value.id,
140            name: value.name,
141            arguments: value.arguments,
142        }
143    }
144}
145
146impl StoredMessage {
147    pub fn text(&self) -> Option<&str> {
148        match &self.content {
149            MessageContent::Text(s) => Some(s),
150        }
151    }
152}
153
154pub fn to_llm_messages(session: &Session, system_prompt: &str) -> Vec<Message> {
155    let mut messages = vec![Message::system(system_prompt)];
156    for msg in &session.messages {
157        match msg.role {
158            Role::User => {
159                if let Some(text) = msg.text() {
160                    messages.push(Message::user(text));
161                }
162            }
163            Role::Assistant => {
164                let tool_calls: Option<Vec<ToolCall>> = msg.tool_calls.as_ref().map(|calls| {
165                    calls
166                        .iter()
167                        .map(|c| ToolCall {
168                            id: c.id.clone(),
169                            name: c.name.clone(),
170                            arguments: c.arguments.clone(),
171                        })
172                        .collect()
173                });
174                let text = msg.text().unwrap_or("").to_string();
175                if let Some(calls) = &tool_calls {
176                    if !calls.is_empty() {
177                        messages.push(Message {
178                            role: LlmRole::Assistant,
179                            content: if text.is_empty() { None } else { Some(text) },
180                            tool_calls,
181                            tool_call_id: None,
182                        });
183                    } else {
184                        messages.push(Message::assistant(text, None));
185                    }
186                } else {
187                    messages.push(Message::assistant(text, None));
188                }
189            }
190            Role::Tool => {
191                if let (Some(id), Some(text)) = (&msg.tool_call_id, msg.text()) {
192                    messages.push(Message::tool(id.clone(), text));
193                }
194            }
195            Role::System => {}
196        }
197    }
198    messages
199}
codei_session/model.rs

codei_session/
model.rs
