Spaces:
Running
Running
File size: 1,948 Bytes
0e38162 | 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 | """
Vector Memory: in-memory store for embedding-based similarity search.
Used for short-term retrieval within a single job run.
"""
from typing import Any, Dict, List, Optional
from memory.base_memory import BaseMemory
from services.vector_store_service import VectorStoreService
from utils.logger import get_logger
logger = get_logger("vector_memory")
class VectorMemory(BaseMemory):
"""In-process vector store backed by sentence-transformers."""
def __init__(self, embedder: VectorStoreService):
self.embedder = embedder
self._store: Dict[str, Dict] = {} # key -> {value, embedding, metadata}
async def store(self, key: str, value: Any, metadata: Optional[Dict] = None) -> None:
text = str(value)
embedding = self.embedder.embed(text)
self._store[key] = {
"value": value,
"text": text,
"embedding": embedding,
"metadata": metadata or {},
}
async def retrieve(self, key: str) -> Optional[Any]:
entry = self._store.get(key)
return entry["value"] if entry else None
async def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
if not self._store:
return []
q_vec = self.embedder.embed(query)
scored = []
for key, entry in self._store.items():
sim = self.embedder.cosine_similarity(q_vec, entry["embedding"])
scored.append({"key": key, "value": entry["value"], "score": sim})
scored.sort(key=lambda x: x["score"], reverse=True)
return scored[:k]
async def delete(self, key: str) -> None:
self._store.pop(key, None)
async def clear(self, scope: Optional[str] = None) -> None:
if scope:
keys = [k for k, v in self._store.items() if v["metadata"].get("job_id") == scope]
for k in keys:
del self._store[k]
else:
self._store.clear()
|