Spaces:
Sleeping
Sleeping
| # memory.py — lightweight conversation memory | |
| from __future__ import annotations | |
| from dataclasses import dataclass, asdict | |
| from typing import List, Dict, Optional | |
| import json, time, os | |
| 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() |