vector_store.py

import numpy as np
from typing import List, Dict, Any, Optional
import chromadb
from chromadb.config import Settings


class VectorStore:
    def __init__(self, persist_dir: str, embedding_dim: int = 384):
        self.client = chromadb.PersistentClient(
            path=persist_dir,
            settings=Settings(anonymized_telemetry=False, allow_reset=True),
        )
        self.embedding_dim = embedding_dim
        self.collection: Optional[Any] = None

    def create_collection(
        self,
        name: str = "documents",
        ef_construction: int = 200,
        m: int = 32,
        ef_search: int = 256,
        overwrite: bool = True,
    ):
        if overwrite:
            try:
                self.client.delete_collection(name)
            except Exception:
                pass
        self.collection = self.client.create_collection(
            name=name,
            metadata={
                "hnsw:space": "cosine",
                "hnsw:construction_ef": ef_construction,
                "hnsw:M": m,
                "hnsw:search_ef": ef_search,
            },
        )
        return self.collection

    def add(
        self,
        ids: List[str],
        embeddings: np.ndarray,
        texts: List[str],
        metadatas: List[Dict[str, Any]],
    ):
        if self.collection is None:
            self.create_collection()
        if len(ids) == 0:
            return
        self.collection.add(
            ids=ids,
            embeddings=embeddings.tolist(),
            documents=texts,
            metadatas=metadatas,
        )

    def search(
        self, query_embedding: np.ndarray, n_results: int = 10
    ) -> Dict[str, Any]:
        if self.collection is None:
            return {"ids": [[]], "distances": [[]], "documents": [[]], "metadatas": [[]]}
        return self.collection.query(
            query_embeddings=query_embedding.tolist(),
            n_results=n_results,
            include=["documents", "metadatas", "distances"],
        )

    def count(self) -> int:
        if self.collection is None:
            return 0
        return self.collection.count()

    def get_all_ids(self) -> List[str]:
        if self.collection is None:
            return []
        result = self.collection.get(include=[])
        return result.get("ids", [])

    def reset(self):
        self.client.reset()
        self.collection = None