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    pub fn add_assistant_message(&self, content: &str) {
55        self.add_message("assistant", content, None);
56    }
57
58    pub fn add_system_message(&self, content: &str) {
59        self.add_message("system", content, None);
60    }
61
62    pub fn get_messages(&self) -> Vec<Message> {
63        let messages = self.messages.read();
64        let now = Local::now();
65        
66        messages
67            .iter()
68            .filter(|entry| {
69                let elapsed = now.signed_duration_since(entry.timestamp);
70                elapsed.num_hours() < self.expiry_hours
71            })
72            .map(|entry| entry.message.clone())
73            .collect()
74    }
75
76    pub fn get_messages_for_player(&self, player: &str) -> Vec<Message> {
77        let messages = self.messages.read();
78        let now = Local::now();
79        
80        messages
81            .iter()
82            .filter(|entry| {
83                let elapsed = now.signed_duration_since(entry.timestamp);
84                elapsed.num_hours() < self.expiry_hours
85            })
86            .filter(|entry| {
87                entry.player.as_deref() == Some(player) || entry.player.is_none()
88            })
89            .map(|entry| entry.message.clone())
90            .collect()
91    }
92
93    #[allow(dead_code)]
94    pub fn clear(&self) {
95        let mut messages = self.messages.write();
96        messages.clear();
97    }
98
99    pub fn len(&self) -> usize {
100        self.messages.read().len()
101    }
102
103    #[allow(dead_code)]
104    pub fn is_empty(&self) -> bool {
105        self.messages.read().is_empty()
106    }
107
108    #[allow(dead_code)]
109    pub fn get_recent_players(&self, count: usize) -> Vec<String> {
110        let messages = self.messages.read();
111        let mut players = Vec::new();
112        
113        for entry in messages.iter().rev() {
114            if let Some(ref player) = entry.player {
115                if !players.contains(player) {
116                    players.push(player.clone());
117                    if players.len() >= count {
118                        break;
119                    }
120                }
121            }
122        }
123        
124        players
125    }
126}
127
128impl Default for ContextManager {
129    fn default() -> Self {
130        Self::new()
131    }
132}