1use std::collections::VecDeque;
2use std::sync::Arc;
3use parking_lot::RwLock;
4use chrono::{DateTime, Local};
5use crate::ai::Message;
6
7const MAX_CONTEXT_MESSAGES: usize = 20;
8const CONTEXT_EXPIRY_HOURS: i64 = 2;
9
10pub struct ContextManager {
11 messages: Arc<RwLock<VecDeque<ContextEntry>>>,
12 max_messages: usize,
13 expiry_hours: i64,
14}
15
16#[derive(Debug, Clone)]
17pub struct ContextEntry {
18 pub message: Message,
19 pub timestamp: DateTime<Local>,
20 pub player: Option<String>,
21}
22
23impl ContextManager {
24 pub fn new() -> Self {
25 Self {
26 messages: Arc::new(RwLock::new(VecDeque::with_capacity(MAX_CONTEXT_MESSAGES))),
27 max_messages: MAX_CONTEXT_MESSAGES,
28 expiry_hours: CONTEXT_EXPIRY_HOURS,
29 }
30 }
31
32 pub fn add_message(&self, role: &str, content: &str, player: Option<&str>) {
33 let entry = ContextEntry {
34 message: Message {
35 role: role.to_string(),
36 content: content.to_string(),
37 },
38 timestamp: Local::now(),
39 player: player.map(|s| s.to_string()),
40 };
41
42 let mut messages = self.messages.write();
43 messages.push_back(entry);
44
45 while messages.len() > self.max_messages {
46 messages.pop_front();
47 }
48 }
49
50 pub fn add_user_message(&self, content: &str, player: &str) {
51 self.add_message("user", content, Some(player));
52 }
53
54 #[allow(dead_code)]
55 pub fn add_assistant_message(&self, content: &str) {
56 self.add_message("assistant", content, None);
57 }
58
59 pub fn add_assistant_message_for_player(&self, content: &str, player: &str) {
60 self.add_message("assistant", content, Some(player));
61 }
62
63 pub fn add_system_message(&self, content: &str) {
64 self.add_message("system", content, None);
65 }
66
67 pub fn get_messages(&self) -> Vec<Message> {
68 let messages = self.messages.read();
69 let now = Local::now();
70
71 messages
72 .iter()
73 .filter(|entry| {
74 let elapsed = now.signed_duration_since(entry.timestamp);
75 elapsed.num_hours() < self.expiry_hours
76 })
77 .map(|entry| entry.message.clone())
78 .collect()
79 }
80
81 pub fn get_messages_for_player(&self, player: &str) -> Vec<Message> {
82 let messages = self.messages.read();
83 let now = Local::now();
84
85 messages
86 .iter()
87 .filter(|entry| {
88 let elapsed = now.signed_duration_since(entry.timestamp);
89 elapsed.num_hours() < self.expiry_hours
90 })
91 .filter(|entry| {
92 entry.player.as_deref() == Some(player) || entry.player.is_none()
93 })
94 .map(|entry| entry.message.clone())
95 .collect()
96 }
97
98 #[allow(dead_code)]
99 pub fn clear(&self) {
100 let mut messages = self.messages.write();
101 messages.clear();
102 }
103
104 pub fn len(&self) -> usize {
105 self.messages.read().len()
106 }
107
108 #[allow(dead_code)]
109 pub fn is_empty(&self) -> bool {
110 self.messages.read().is_empty()
111 }
112
113 #[allow(dead_code)]
114 pub fn get_recent_players(&self, count: usize) -> Vec<String> {
115 let messages = self.messages.read();
116 let mut players = Vec::new();
117
118 for entry in messages.iter().rev() {
119 if let Some(ref player) = entry.player {
120 if !players.contains(player) {
121 players.push(player.clone());
122 if players.len() >= count {
123 break;
124 }
125 }
126 }
127 }
128
129 players
130 }
131}
132
133impl Default for ContextManager {
134 fn default() -> Self {
135 Self::new()
136 }
137}