vector_db.py

"""
vector_db.py — A minimal vector database backed by a JSON file.

Flow:
  1. Embed text using a sentence-transformer model (runs locally).
  2. Store the embedding + metadata as a record in data.json.
  3. At query time, embed the query and rank all records by cosine similarity.
"""

import json
import math
import os
from typing import Any


class VectorDB:
    def __init__(self, db_path: str = "data.json", model_name: str = "all-MiniLM-L6-v2"):
        self.db_path = db_path
        self._model_name = model_name
        self._model = None          # loaded lazily so imports stay fast
        self._records: list[dict] = []
        if os.path.exists(db_path):
            self.load()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _get_model(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            print(f"[VectorDB] Loading model '{self._model_name}' (first run downloads ~80 MB)...")
            self._model = SentenceTransformer(self._model_name)
            print("[VectorDB] Model ready.\n")
        return self._model

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        mag_a = math.sqrt(sum(x * x for x in a))
        mag_b = math.sqrt(sum(x * x for x in b))
        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def add(self, text: str, metadata: dict[str, Any] | None = None) -> None:
        """Embed `text` and append a record to the in-memory database."""
        model = self._get_model()
        embedding: list[float] = model.encode(text).tolist()
        record = {
            "id": len(self._records),
            "text": text,
            "metadata": metadata or {},
            "embedding": embedding,
        }
        self._records.append(record)

    def save(self) -> None:
        """Persist all records (including raw embeddings) to JSON."""
        with open(self.db_path, "w", encoding="utf-8") as f:
            json.dump(self._records, f, indent=2)
        print(f"[VectorDB] Saved {len(self._records)} records -> {self.db_path}")

    def load(self) -> None:
        """Load records from the JSON file (embeddings included)."""
        with open(self.db_path, encoding="utf-8") as f:
            self._records = json.load(f)
        print(f"[VectorDB] Loaded {len(self._records)} records <- {self.db_path}")

    def search(self, query: str, top_k: int = 5) -> list[dict]:
        """
        Embed the query and return the top_k most similar records,
        ordered by cosine similarity (highest first).
        """
        if not self._records:
            return []

        model = self._get_model()
        query_vec: list[float] = model.encode(query).tolist()

        scored = [
            {
                "score": self._cosine_similarity(query_vec, rec["embedding"]),
                "id": rec["id"],
                "text": rec["text"],
                "metadata": rec["metadata"],
            }
            for rec in self._records
        ]

        scored.sort(key=lambda r: r["score"], reverse=True)
        return scored[:top_k]