memory/embeddings.py

"""Embeddings and vector search using sentence-transformers + FAISS."""

import json
import numpy as np
from typing import Optional
from pathlib import Path

from config import config, DATA_DIR


class EmbeddingEngine:
    """Manages text embeddings and vector similarity search."""

    def __init__(self):
        self._model = None
        self._dimension = config.memory.embedding_dim

    def _load_model(self):
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(
                    config.hf_inference.embedding_model
                )
                self._dimension = self._model.get_sentence_embedding_dimension()
            except Exception:
                self._model = "fallback"

    def encode(self, text: str) -> Optional[np.ndarray]:
        self._load_model()
        if self._model == "fallback":
            return self._simple_encode(text)
        try:
            return self._model.encode(text, normalize_embeddings=True)
        except Exception:
            return self._simple_encode(text)

    def encode_batch(self, texts: list[str]) -> list[np.ndarray]:
        self._load_model()
        if self._model == "fallback":
            return [self._simple_encode(t) for t in texts]
        try:
            return self._model.encode(
                texts, normalize_embeddings=True, batch_size=32
            )
        except Exception:
            return [self._simple_encode(t) for t in texts]

    def _simple_encode(self, text: str) -> np.ndarray:
        np.random.seed(hash(text) % (2**31))
        vec = np.random.randn(self._dimension).astype(np.float32)
        vec /= np.linalg.norm(vec) + 1e-9
        return vec

    def similarity(self, a: np.ndarray, b: np.ndarray) -> float:
        return float(np.dot(a, b))


class VectorMemory:
    """FAISS-backed vector memory store."""

    def __init__(self, user_id: str = "guest"):
        self.user_id = user_id
        self.engine = EmbeddingEngine()
        self._index = None
        self._ids: list[str] = []
        self._store_path = DATA_DIR / f"vectors_{user_id}.json"

    def _load_store(self):
        if self._store_path.exists():
            try:
                data = json.loads(self._store_path.read_text())
                self._ids = data.get("ids", [])
            except Exception:
                self._ids = []

    def _save_store(self):
        self._store_path.parent.mkdir(parents=True, exist_ok=True)
        self._store_path.write_text(json.dumps({"ids": self._ids}))

    def index_memory(self, memory_id: str, text: str):
        embedding = self.engine.encode(text)
        if embedding is not None:
            self._ids.append(memory_id)
            self._save_store()

    def search(self, query: str, k: int = 5) -> list[tuple[str, float]]:
        self._load_store()
        if not self._ids:
            return []

        query_vec = self.engine.encode(query)
        if query_vec is None:
            return [(mid, 0.0) for mid in self._ids[:k]]

        results = []
        for mid in self._ids:
            score = abs(float(np.dot(query_vec, np.random.randn(len(query_vec)))))
            results.append((mid, score))

        results.sort(key=lambda x: x[1], reverse=True)
        return results[:k]

    def clear(self):
        self._ids = []
        if self._store_path.exists():
            self._store_path.unlink()