File size: 2,025 Bytes
945aae0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import time
from typing import Any
class MemoryRecord:
def __init__(self, role: str, content: str):
self.role = role
self.content = content
self.timestamp = time.time()
class MemorySession:
def __init__(self):
self.records: list[MemoryRecord] = []
def write(self, role: str, content: str) -> None:
self.records.append(MemoryRecord(role, content))
def read(self, limit: int = 10) -> list[dict]:
return [
{"role": r.role, "content": r.content}
for r in self.records[-limit:]
]
@property
def size(self) -> int:
return len(self.records)
class MemoryEngine:
def __init__(self):
self._users: dict[str, dict[str, MemorySession]] = {}
def _ensure_session(self, user_id: str, session_id: str) -> MemorySession:
if user_id not in self._users:
self._users[user_id] = {}
if session_id not in self._users[user_id]:
self._users[user_id][session_id] = MemorySession()
return self._users[user_id][session_id]
def write(self, user_id: str, session_id: str, role: str, content: str) -> None:
session = self._ensure_session(user_id, session_id)
session.write(role, content)
def read(self, user_id: str, session_id: str, limit: int = 10) -> list[dict]:
session = self._ensure_session(user_id, session_id)
return session.read(limit)
def get_user_ids(self) -> list[str]:
return list(self._users.keys())
def get_session_ids(self, user_id: str) -> list[str]:
return list(self._users.get(user_id, {}).keys())
def get_stats(self) -> dict:
total_users = len(self._users)
total_sessions = sum(len(s) for s in self._users.values())
total_records = sum(
s.size for u in self._users.values() for s in u.values()
)
return {
"users": total_users,
"sessions": total_sessions,
"records": total_records,
}
|