Spaces:
Sleeping
Sleeping
| """ | |
| memory.py — Conversation memory using ChromaDB + sentence-transformers. | |
| Stores past conversation turns as vector embeddings. | |
| Retrieves semantically relevant past context for each new query. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| import os | |
| from typing import List | |
| _EMBED_MODEL_NAME = "all-MiniLM-L6-v2" # 22 MB, fast CPU embed | |
| _embed = None | |
| _client = None | |
| _col = None | |
| def _model(): | |
| global _embed | |
| if _embed is None: | |
| from sentence_transformers import SentenceTransformer | |
| print("[memory] Loading embedding model ...", flush=True) | |
| _embed = SentenceTransformer(_EMBED_MODEL_NAME) | |
| print("[memory] Embedding model ready", flush=True) | |
| return _embed | |
| def _collection(): | |
| global _client, _col | |
| if _col is None: | |
| import chromadb | |
| _client = chromadb.PersistentClient(path="./memory_db") | |
| _col = _client.get_or_create_collection( | |
| name="conversations", | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| return _col | |
| def store(user_msg: str, ai_response: str, session: str = "default") -> None: | |
| """Store a conversation turn as an embedding.""" | |
| try: | |
| text = f"User: {user_msg}\nAssistant: {ai_response}" | |
| emb = _model().encode(text).tolist() | |
| doc_id = f"{session}_{int(time.time() * 1000)}" | |
| _collection().add( | |
| documents=[text], | |
| embeddings=[emb], | |
| ids=[doc_id], | |
| metadatas=[{ | |
| "session": session, | |
| "timestamp": time.time(), | |
| "user_msg": user_msg[:200], | |
| }], | |
| ) | |
| except Exception as e: | |
| print(f"[memory] store error: {e}", flush=True) | |
| def retrieve(query: str, session: str = "default", top_k: int = 3) -> List[str]: | |
| """Return top_k semantically similar past exchanges.""" | |
| try: | |
| col = _collection() | |
| if col.count() == 0: | |
| return [] | |
| emb = _model().encode(query).tolist() | |
| n = min(top_k, col.count()) | |
| results = col.query( | |
| query_embeddings=[emb], | |
| n_results=n, | |
| where={"session": session}, | |
| ) | |
| return results["documents"][0] if results["documents"] else [] | |
| except Exception as e: | |
| print(f"[memory] retrieve error: {e}", flush=True) | |
| return [] | |
| def clear(session: str = "default") -> int: | |
| """Delete all memory entries for a session.""" | |
| try: | |
| col = _collection() | |
| ids = col.get(where={"session": session})["ids"] | |
| if ids: | |
| col.delete(ids=ids) | |
| return len(ids) | |
| except Exception as e: | |
| print(f"[memory] clear error: {e}", flush=True) | |
| return 0 | |
| def list_recent(session: str = "default", limit: int = 10) -> List[dict]: | |
| """Return recent conversation entries for a session.""" | |
| try: | |
| col = _collection() | |
| res = col.get(where={"session": session}, limit=limit, | |
| include=["documents", "metadatas"]) | |
| out = [] | |
| for doc, meta in zip(res["documents"], res["metadatas"]): | |
| out.append({"text": doc, "timestamp": meta.get("timestamp"), | |
| "user_msg": meta.get("user_msg", "")}) | |
| out.sort(key=lambda x: x["timestamp"] or 0, reverse=True) | |
| return out[:limit] | |
| except Exception as e: | |
| print(f"[memory] list error: {e}", flush=True) | |
| return [] | |
| def stats() -> dict: | |
| """Global memory statistics.""" | |
| try: | |
| col = _collection() | |
| return { | |
| "total_entries": col.count(), | |
| "embed_model": _EMBED_MODEL_NAME, | |
| "db_path": "./memory_db", | |
| } | |
| except Exception: | |
| return {"total_entries": 0, "embed_model": _EMBED_MODEL_NAME} | |