use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
pub key: String,
pub messages: Vec<ChatMessage>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub metadata: serde_json::Value,
#[serde(default)]
pub last_consolidated: usize,
}
impl Session {
pub fn new(key: impl Into<String>) -> Self {
let now = Utc::now();
Self {
key: key.into(),
messages: Vec::new(),
created_at: now,
updated_at: now,
metadata: serde_json::Value::Object(serde_json::Map::new()),
last_consolidated: 0,
}
}
pub fn add_message(&mut self, role: impl Into<String>, content: impl Into<String>) {
self.messages.push(ChatMessage {
role: role.into(),
content: content.into(),
timestamp: Utc::now(),
tool_call_id: None,
tool_calls: None,
name: None,
reasoning_content: None,
thinking_blocks: None,
});
self.updated_at = Utc::now();
}
pub fn add_full_message(&mut self, msg: ChatMessage) {
self.messages.push(msg);
self.updated_at = Utc::now();
}
pub fn get_history(&self, max_messages: usize) -> Vec<ChatMessage> {
let consolidated = self.last_consolidated.min(self.messages.len());
let unconsolidated = &self.messages[consolidated..];
let start = unconsolidated.len().saturating_sub(max_messages);
let mut sliced: Vec<ChatMessage> = unconsolidated[start..]
.iter()
.filter(|m| matches!(m.role.as_str(), "user" | "assistant" | "tool"))
.cloned()
.collect();
if let Some(pos) = sliced.iter().position(|m| m.role == "user") {
sliced = sliced[pos..].to_vec();
}
sliced
}
pub fn clear(&mut self) {
self.messages.clear();
self.last_consolidated = 0;
self.updated_at = Utc::now();
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: String,
pub content: String,
pub timestamp: DateTime<Utc>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub tool_calls: Option<Vec<serde_json::Value>>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub reasoning_content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub thinking_blocks: Option<Vec<serde_json::Value>>,
}
impl ChatMessage {
pub fn new(role: impl Into<String>, content: impl Into<String>) -> Self {
Self {
role: role.into(),
content: content.into(),
timestamp: Utc::now(),
tool_call_id: None,
tool_calls: None,
name: None,
reasoning_content: None,
thinking_blocks: None,
}
}
pub fn with_tool_metadata(
role: impl Into<String>,
content: impl Into<String>,
tool_call_id: Option<String>,
tool_calls: Option<Vec<serde_json::Value>>,
name: Option<String>,
) -> Self {
Self {
role: role.into(),
content: content.into(),
timestamp: Utc::now(),
tool_call_id,
tool_calls,
name,
reasoning_content: None,
thinking_blocks: None,
}
}
pub fn to_llm_format(&self) -> serde_json::Value {
serde_json::json!({
"role": &self.role,
"content": &self.content,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_session_creation() {
let session = Session::new("telegram:12345");
assert_eq!(session.key, "telegram:12345");
assert!(session.messages.is_empty());
}
#[test]
fn test_add_message() {
let mut session = Session::new("test");
session.add_message("user", "Hello");
session.add_message("assistant", "Hi there!");
assert_eq!(session.messages.len(), 2);
assert_eq!(session.messages[0].role, "user");
assert_eq!(session.messages[1].role, "assistant");
}
#[test]
fn test_get_history() {
let mut session = Session::new("test");
for i in 0..60 {
session.add_message("user", format!("Message {}", i));
}
let history = session.get_history(50);
assert_eq!(history.len(), 50);
}
}