memory/vector_memory.py

"""
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()