| 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, |
| } |
|
|
