File size: 2,401 Bytes
454a778 | 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 67 68 69 70 | 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]})
|