SHARKAGENTDEMO / memory.py
OSOCONSULT's picture
Upload 14 files
0cdae6c verified
import os
import json
class FileMemory:
def __init__(self, base_path="memory_data"):
self.base_path = base_path
os.makedirs(self.base_path, exist_ok=True)
def _get_file_path(self, user_id):
return os.path.join(self.base_path, f"{user_id}.json")
def get_memory(self, user_id):
path = self._get_file_path(user_id)
if os.path.exists(path):
with open(path, "r") as f:
return json.load(f)
return []
def store_user_message(self, user_id, message):
memory = self.get_memory(user_id)
memory.append({"role": "user", "content": message})
self._save_memory(user_id, memory)
def store_agent_message(self, user_id, message):
memory = self.get_memory(user_id)
memory.append({"role": "assistant", "content": message})
self._save_memory(user_id, memory)
def clear_memory(self, user_id):
path = self._get_file_path(user_id)
if os.path.exists(path):
os.remove(path)
def _save_memory(self, user_id, memory):
path = self._get_file_path(user_id)
with open(path, "w") as f:
json.dump(memory, f, indent=2)
def get_context(self, user_id, max_turns=10):
memory = self.get_memory(user_id)
return "\n".join([f"{m['role'].upper()}: {m['content']}" for m in memory[-max_turns:]])
memory = FileMemory()