Skip to main content

mc_minder/
context.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3use parking_lot::RwLock;
4use chrono::{DateTime, Local};
5use crate::ai_client::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    pub fn clear(&self) {
94        let mut messages = self.messages.write();
95        messages.clear();
96    }
97
98    pub fn len(&self) -> usize {
99        self.messages.read().len()
100    }
101
102    pub fn is_empty(&self) -> bool {
103        self.messages.read().is_empty()
104    }
105
106    pub fn get_recent_players(&self, count: usize) -> Vec<String> {
107        let messages = self.messages.read();
108        let mut players = Vec::new();
109        
110        for entry in messages.iter().rev() {
111            if let Some(ref player) = entry.player {
112                if !players.contains(player) {
113                    players.push(player.clone());
114                    if players.len() >= count {
115                        break;
116                    }
117                }
118            }
119        }
120        
121        players
122    }
123}
124
125impl Default for ContextManager {
126    fn default() -> Self {
127        Self::new()
128    }
129}