| import json |
| import os |
| from collections import deque |
| from pathlib import Path |
| from typing import Any, Dict, List, Optional |
| from .config import WORKING_MEMORY_SIZE, MEMORY_DIR |
|
|
| class MemorySystem: |
| def __init__(self): |
| self.working = deque(maxlen=WORKING_MEMORY_SIZE) |
| self._long_term: List[Dict[str, Any]] = [] |
| self._load_long_term() |
|
|
| def _long_term_path(self) -> Path: |
| return MEMORY_DIR / "long_term.jsonl" |
|
|
| def _load_long_term(self): |
| path = self._long_term_path() |
| if path.exists(): |
| with open(path) as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| self._long_term.append(json.loads(line)) |
|
|
| def _save_long_term(self): |
| path = self._long_term_path() |
| with open(path, "w") as f: |
| for entry in self._long_term[-1000:]: |
| f.write(json.dumps(entry) + "\n") |
|
|
| def add_working(self, entry: Dict[str, Any]): |
| self.working.append(entry) |
|
|
| def get_working(self) -> List[Dict[str, Any]]: |
| return list(self.working) |
|
|
| def add_long_term(self, entry: Dict[str, Any]): |
| self._long_term.append(entry) |
| self._save_long_term() |
|
|
| def query_long_term(self, key: str, top_k: int = 5) -> List[Dict[str, Any]]: |
| matches = [] |
| for entry in reversed(self._long_term): |
| if key.lower() in str(entry).lower(): |
| matches.append(entry) |
| if len(matches) >= top_k: |
| break |
| return matches |
|
|
| def get_context(self, query: str) -> str: |
| working_summary = "\n".join( |
| f"{e.get('role', 'system')}: {e.get('content', '')[:200]}" |
| for e in self.working[-5:] |
| ) |
| lt_matches = self.query_long_term(query, top_k=3) |
| lt_summary = "\n".join( |
| f"[memory] {m.get('summary', str(m)[:200])}" for m in lt_matches |
| ) |
| parts = [] |
| if working_summary: |
| parts.append(f"[working memory]\n{working_summary}") |
| if lt_summary: |
| parts.append(f"[long-term memory]\n{lt_summary}") |
| return "\n\n".join(parts) |
|
|
| def summarize_to_long_term(self, role: str, content: str): |
| summary = content[:300] if len(content) > 300 else content |
| self.add_long_term({"role": role, "summary": summary, "content": content[:1000]}) |
|
|
