Spaces:
Sleeping
Sleeping
File size: 3,453 Bytes
19de729 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | # memory.py — lightweight conversation memory
from __future__ import annotations
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional
import json, time, os
@dataclass
class Turn:
role: str # "user" | "assistant" | "tool"
content: str
meta: Dict = None
ts: float = 0.0
def to_dict(self):
d = asdict(self)
d["ts"] = self.ts or time.time()
d["meta"] = self.meta or {}
return d
class ConversationMemory:
"""
Simple memory with:
- buffer: last N turns (short-term)
- summary: rolling abstractive summary (long-term)
- entities: key entities/ids spotted so far
"""
def __init__(self, path: str, buffer_size: int = 12):
self.path = path
self.buffer_size = buffer_size
self.buffer: List[Turn] = []
self.summary: str = ""
self.entities: Dict[str, List[str]] = {} # e.g., {"competition_id": ["2313", ...]}
self._load()
# ------------ persistence ------------
def _load(self):
if not os.path.exists(self.path): return
with open(self.path, "r") as f:
data = json.load(f)
self.buffer = [Turn(**t) for t in data.get("buffer", [])]
self.summary = data.get("summary", "")
self.entities = data.get("entities", {})
def _save(self):
os.makedirs(os.path.dirname(self.path), exist_ok=True)
with open(self.path, "w") as f:
json.dump({
"buffer": [t.to_dict() for t in self.buffer],
"summary": self.summary,
"entities": self.entities,
}, f, ensure_ascii=False, indent=2)
# ------------ public API ------------
def add_turn(self, role: str, content: str, meta: Optional[Dict]=None):
self.buffer.append(Turn(role=role, content=content, meta=meta or {}, ts=time.time()))
if len(self.buffer) > self.buffer_size:
self.buffer = self.buffer[-self.buffer_size:]
self._save()
def get_context(self) -> Dict:
"""What to feed into prompts/tools."""
return {
"summary": self.summary,
"recent": [{"role": t.role, "content": t.content} for t in self.buffer[-self.buffer_size:]],
"entities": self.entities,
}
def update_summary(self, llm_summarize_fn):
"""
Call with a function that maps (summary, recent) -> new_summary.
Only do this occasionally (e.g., every 6–10 user turns).
"""
if not self.buffer:
return
recent_text = "\n".join(
f"{t.role.upper()}: {t.content}" for t in self.buffer[-self.buffer_size:]
)
prompt = (
"You are a diligent note-taker. Update the long-term summary of this conversation.\n"
"Keep it under 150 words. Capture tasks, preferences, important grant IDs/themes, and open questions.\n\n"
f"EXISTING SUMMARY:\n{self.summary or '(none)'}\n\n"
f"RECENT TURNS:\n{recent_text}\n\n"
"Return ONLY the updated summary text."
)
new_sum = llm_summarize_fn(prompt).strip()
if new_sum:
self.summary = new_sum
self._save()
def add_entity(self, kind: str, value: str):
if not value: return
arr = self.entities.setdefault(kind, [])
if value not in arr:
arr.append(value)
self._save() |