96 lines
4.0 KiB
Python
96 lines
4.0 KiB
Python
"""
|
|
Phase 2 #2B: MemPalace Hierarchical Memory
|
|
Wing (User) → Room (Conversation) → Drawer (Topic)
|
|
Three-level memory index with priority retrieval.
|
|
"""
|
|
from typing import List, Dict, Optional
|
|
from datetime import datetime
|
|
|
|
|
|
class MemPalace:
|
|
"""Hierarchical memory system with Wing/Room/Drawer indexing."""
|
|
|
|
LEVELS = {"wing": 0, "room": 1, "drawer": 2}
|
|
|
|
def __init__(self, user_id: str, conv_id: Optional[str] = None):
|
|
self.user_id = user_id
|
|
self.conv_id = conv_id
|
|
|
|
async def remember(self, content: str, level: str = "room",
|
|
drawer_topic: str = "", ttl_days: int = 30) -> dict:
|
|
"""Store a memory at the specified level."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
import uuid
|
|
db = SessionLocal()
|
|
mem_id = str(uuid.uuid4())
|
|
db.execute(text(
|
|
"INSERT INTO memories (id, user_id, conv_id, content, level, drawer_topic, "
|
|
"ttl_days, created_at) VALUES (:id, :uid, :cid, :c, :l, :dt, :ttl, NOW())"
|
|
), {"id": mem_id, "uid": self.user_id, "cid": self.conv_id,
|
|
"c": content[:2000], "l": level, "dt": drawer_topic, "ttl": ttl_days})
|
|
db.commit()
|
|
db.close()
|
|
return {"id": mem_id, "status": "stored", "level": level}
|
|
except Exception as e:
|
|
return {"status": "error", "error": str(e)}
|
|
|
|
async def recall(self, query: str, max_results: int = 5,
|
|
level_filter: str = None) -> List[dict]:
|
|
"""Retrieve memories by relevance, respecting hierarchy."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
db = SessionLocal()
|
|
# Priority: wing(user) > room(conv) > drawer(topic)
|
|
params = {"uid": self.user_id, "cid": self.conv_id or ""}
|
|
where = "user_id = :uid AND (conv_id = :cid OR conv_id IS NULL)"
|
|
if level_filter:
|
|
where += f" AND level = :lv"
|
|
params["lv"] = level_filter
|
|
# Simulated relevance via LIKE (in production, use pgvector)
|
|
search = f"%{query}%"
|
|
rows = db.execute(text(
|
|
f"SELECT id, content, level, drawer_topic, created_at FROM memories "
|
|
f"WHERE {where} AND (content ILIKE :q OR drawer_topic ILIKE :q) "
|
|
f"ORDER BY CASE level WHEN 'drawer' THEN 0 WHEN 'room' THEN 1 WHEN 'wing' THEN 2 END, "
|
|
f"created_at DESC LIMIT :lim"
|
|
), {**params, "q": search, "lim": max_results}).fetchall()
|
|
db.close()
|
|
return [{"id": r[0], "content": r[1], "level": r[2],
|
|
"topic": r[3], "created": str(r[4])} for r in rows]
|
|
except Exception:
|
|
return []
|
|
|
|
async def forget(self, level: str = None, older_than_days: int = None) -> int:
|
|
"""Purge old or level-specific memories."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
db = SessionLocal()
|
|
where = "user_id = :uid"
|
|
params = {"uid": self.user_id}
|
|
if level:
|
|
where += " AND level = :lv"
|
|
params["lv"] = level
|
|
if older_than_days:
|
|
where += " AND created_at < NOW() - make_interval(days => :days)"
|
|
params["days"] = older_than_days
|
|
result = db.execute(text(f"DELETE FROM memories WHERE {where}"), params)
|
|
db.commit()
|
|
db.close()
|
|
return result.rowcount
|
|
except Exception:
|
|
return 0
|
|
|
|
async def get_context(self, max_tokens: int = 2000) -> str:
|
|
"""Get compact memory context for LLM prompt injection."""
|
|
memories = await self.recall("", max_results=10)
|
|
if not memories:
|
|
return ""
|
|
lines = ["[Memory Context]"]
|
|
for m in memories:
|
|
lines.append(f"[{m['level']}] {m['content'][:200]}")
|
|
return "\n".join(lines)[:max_tokens]
|