Skip to main content

vv_agent/memory/manager/
session_context.rs

1use super::MemoryManager;
2use crate::memory::session::SessionMemory;
3use crate::types::{Message, MessageRole};
4
5impl MemoryManager {
6    pub fn session_memory(&self) -> Option<&SessionMemory> {
7        self.session_memory.as_ref()
8    }
9
10    pub fn session_memory_mut(&mut self) -> Option<&mut SessionMemory> {
11        self.session_memory.as_mut()
12    }
13
14    pub fn apply_session_memory_context(&self, messages: &[Message]) -> Vec<Message> {
15        let Some(session_context) = self
16            .session_memory
17            .as_ref()
18            .map(SessionMemory::render_as_system_context)
19            .filter(|context| !context.is_empty())
20        else {
21            return messages.to_vec();
22        };
23        let mut updated = messages.to_vec();
24        if let Some(system_message) = updated
25            .iter_mut()
26            .find(|message| message.role == MessageRole::System)
27        {
28            if !system_message.content.contains("<Session Memory>") {
29                system_message.content.push_str("\n\n");
30                system_message.content.push_str(&session_context);
31            }
32            return updated;
33        }
34        let mut system_message = Message::system(session_context);
35        system_message.name = Some("session_memory".to_string());
36        updated.insert(0, system_message);
37        updated
38    }
39
40    pub fn strip_session_memory_context(&self, messages: &[Message]) -> Vec<Message> {
41        let mut updated = messages.to_vec();
42        let Some(system_message) = updated
43            .iter_mut()
44            .find(|message| message.role == MessageRole::System)
45        else {
46            return updated;
47        };
48        let Some(marker_index) = system_message.content.find("<Session Memory>") else {
49            return updated;
50        };
51        system_message.content = system_message.content[..marker_index]
52            .trim_end()
53            .to_string();
54        updated
55    }
56}