File size: 4,238 Bytes
bb10045 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """
STEW — Vector Memory Engine (ChromaDB)
Long-term semantic memory: stores facts, documents, and project context
that can be retrieved by semantic similarity.
"""
import hashlib
from typing import Optional
from config import Config
class VectorMemory:
"""
Persistent semantic memory using ChromaDB.
Falls back to in-memory store if ChromaDB is not available.
"""
def __init__(self, persist_dir: str = "/tmp/stew_vector_db"):
self._client = None
self._collection = None
self._persist_dir = persist_dir
self._fallback: list[dict] = []
def _init_chroma(self):
if self._collection is not None:
return True
try:
import chromadb
self._client = chromadb.PersistentClient(path=self._persist_dir)
self._collection = self._client.get_or_create_collection(
name="stew_memory",
metadata={"hnsw:space": "cosine"},
)
return True
except Exception:
return False
# ── Document storage ─────────────────────────────────────────────────────
def store(self, text: str, metadata: dict | None = None, doc_id: str | None = None) -> str:
"""Store a text chunk in vector memory."""
if not doc_id:
doc_id = hashlib.sha256(text.encode()).hexdigest()[:16]
meta = metadata or {}
if self._init_chroma():
try:
self._collection.upsert(
ids=[doc_id],
documents=[text],
metadatas=[meta],
)
return f"✅ Stored in vector memory (id: {doc_id})"
except Exception as e:
return f"Vector store error: {e}"
else:
# Fallback: in-memory list
self._fallback.append({"id": doc_id, "text": text, "meta": meta})
return f"✅ Stored in fallback memory (id: {doc_id})"
def search(self, query: str, n_results: int = 5) -> str:
"""Semantic search over stored memory."""
if self._init_chroma():
try:
results = self._collection.query(
query_texts=[query],
n_results=min(n_results, self._collection.count() or 1),
)
docs = results.get("documents", [[]])[0]
if not docs:
return "No relevant memories found."
lines = [f"[{i+1}] {d[:300]}" for i, d in enumerate(docs)]
return "🧠 Relevant memories:\n" + "\n\n".join(lines)
except Exception as e:
return f"Vector search error: {e}"
else:
# Simple substring fallback
matches = [f["text"] for f in self._fallback if query.lower() in f["text"].lower()]
return ("🧠 Matches:\n" + "\n".join(m[:300] for m in matches[:5])) if matches else "No matches."
def store_project(self, project_name: str, description: str) -> str:
"""Store a project description."""
return self.store(
f"Project: {project_name}\n{description}",
metadata={"type": "project", "name": project_name},
doc_id=f"project_{project_name[:20]}",
)
def count(self) -> int:
if self._init_chroma():
try:
return self._collection.count()
except Exception:
pass
return len(self._fallback)
def clear(self) -> str:
if self._init_chroma():
try:
self._client.delete_collection("stew_memory")
self._collection = None
return "🗑️ Vector memory cleared."
except Exception as e:
return f"Clear error: {e}"
self._fallback.clear()
return "🗑️ Fallback memory cleared."
# Module-level singleton
_vector_memory: Optional[VectorMemory] = None
def get_vector_memory() -> VectorMemory:
global _vector_memory
if _vector_memory is None:
_vector_memory = VectorMemory()
return _vector_memory
|