Create memory.py

utils/memory.py

@@ -0,0 +1,41 @@
+import streamlit as st
+from typing import List, Dict, Union
+from PIL import Image
+
+class ChatMemory:
+    def __init__(self, max_context_length: int = 2048):
+        self.max_context_length = max_context_length
+        if "memory" not in st.session_state:
+            st.session_state.memory = {
+                'history': [],
+                'context': []
+            }
+    
+    def update(self, user_input: Union[str, Image.Image], response: str):
+        """Store interaction with automatic context pruning"""
+        # Store user input
+        if isinstance(user_input, Image.Image):
+            st.session_state.memory['history'].append(('user', 'image', user_input))
+        else:
+            st.session_state.memory['history'].append(('user', 'text', user_input))
+        
+        # Store assistant response
+        st.session_state.memory['history'].append(('assistant', 'text', response))
+        
+        # Maintain context window
+        current_length = sum(len(item[2]) for item in st.session_state.memory['history'] if item[1] == 'text')
+        while current_length > self.max_context_length and len(st.session_state.memory['history']) > 2:
+            removed = st.session_state.memory['history'].pop(0)
+            if removed[1] == 'text':
+                current_length -= len(removed[2])
+    
+    def get_context(self) -> str:
+        """Generate conversation context string"""
+        return "\n".join(
+            f"{role}: {content}" 
+            for role, type_, content in st.session_state.memory['history']
+            if type_ == 'text'
+        )
+    
+    def clear(self):
+        st.session_state.memory = {'history': [], 'context': []}