Skip to main content

mc_minder/context/
mod.rs

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}