vector_store.py

import abc
import json
import os
from typing import Any, Generic, List, Optional, Tuple, TypeVar

import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

T = TypeVar("T", bound=dict[str, Any])
E = TypeVar("E")


class VectorStore(Generic[T, E], abc.ABC):
    """Abstract base class for a vector store.

    A vector store is a key-value store that maps an ID to a vector embedding and a payload. The
    payload can be any JSON-serializable object, e.g. a dictionary.
    """

    INDEX_FILE = "vectors_index.json"
    EMBEDDINGS_FILE = "vectors_data.npy"
    PAYLOADS_FILE = "vectors_payloads.json"

    @abc.abstractmethod
    def _add(self, embedding: E, payload: T, emb_id: str) -> None:
        """Save an embedding with payload for a given ID."""
        pass

    @abc.abstractmethod
    def _get(self, emb_id: str) -> Optional[E]:
        """Get the embedding for a given ID."""
        pass

    @abc.abstractmethod
    def clear(self) -> None:
        """Clear the store."""
        pass

    def _get_emb_id(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> str:
        if emb_id is None:
            if payload is None:
                raise ValueError("Either emb_id or payload must be provided.")
            emb_id = json.dumps(payload, sort_keys=True)
        return emb_id

    def add(self, embedding: E, payload: T, emb_id: Optional[str] = None) -> None:
        if emb_id is None:
            emb_id = json.dumps(payload, sort_keys=True)
        self._add(embedding=embedding, payload=payload, emb_id=emb_id)

    def get(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> Optional[E]:
        return self._get(emb_id=self._get_emb_id(emb_id=emb_id, payload=payload))

    def has(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> bool:
        return self.get(emb_id=emb_id, payload=payload) is not None

    @abc.abstractmethod
    def _retrieve_similar(
        self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
    ) -> List[Tuple[T, float]]:
        """Retrieve IDs, payloads and the respective similarity scores with respect to the
        reference entry. In the case that the reference entry is not in the store itself, an empty
        list will be returned.

        Args:
            ref_id: The ID of the reference entry.
            top_k: If provided, only the top-k most similar entries will be returned.
            min_similarity: If provided, only entries with a similarity score greater or equal to
                this value will be returned.

        Returns:
            A list of tuples consisting of the ID and the similarity score, sorted by similarity
            score in descending order.
        """
        pass

    def retrieve_similar(
        self, ref_id: Optional[str] = None, ref_payload: Optional[T] = None, **kwargs
    ) -> List[Tuple[T, float]]:
        if not self.has(emb_id=ref_id, payload=ref_payload):
            return []
        return self._retrieve_similar(
            ref_id=self._get_emb_id(emb_id=ref_id, payload=ref_payload), **kwargs
        )

    @abc.abstractmethod
    def __len__(self):
        pass

    def _add_from_directory(self, directory: str) -> None:
        with open(os.path.join(directory, self.INDEX_FILE), "r") as f:
            index = json.load(f)
        embeddings_np = np.load(os.path.join(directory, self.EMBEDDINGS_FILE))
        with open(os.path.join(directory, self.PAYLOADS_FILE), "r") as f:
            payloads = json.load(f)
        for emb_id, emb, payload in zip(index, embeddings_np, payloads):
            self._add(emb_id=emb_id, payload=payload, embedding=emb.tolist())

    @abc.abstractmethod
    def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
        """Return a tuple of indices, vectors and payloads."""
        pass

    def _save_to_directory(self, directory: str) -> None:
        indices, vectors, payloads = self.as_indices_vectors_payloads()
        np.save(os.path.join(directory, self.EMBEDDINGS_FILE), vectors)
        with open(os.path.join(directory, self.PAYLOADS_FILE), "w") as f:
            json.dump(payloads, f)
        with open(os.path.join(directory, self.INDEX_FILE), "w") as f:
            json.dump(indices, f)

    def save_to_directory(self, directory: str) -> None:
        """Save the vector store to a directory."""
        os.makedirs(directory, exist_ok=True)
        self._save_to_directory(directory)

    def load_from_directory(self, directory: str, replace: bool = False) -> None:
        """Load the vector store from a directory.

        If `replace` is True, the current content of the store will be replaced.
        """
        if replace:
            self.clear()
        self._add_from_directory(directory)


def vector_norm(vector: List[float]) -> float:
    return sum(x**2 for x in vector) ** 0.5


def cosine_similarity(a: List[float], b: List[float]) -> float:
    return sum(a * b for a, b in zip(a, b)) / (vector_norm(a) * vector_norm(b))


class SimpleVectorStore(VectorStore[T, List[float]]):
    """Simple in-memory vector store using a dictionary."""

    def __init__(self):
        self.vectors: dict[str, List[float]] = {}
        self.payloads: dict[str, T] = {}
        self._cache = {}
        self._sim = cosine_similarity

    def _add(self, embedding: E, payload: T, emb_id: str) -> None:
        self.vectors[emb_id] = embedding
        self.payloads[emb_id] = payload

    def _get(self, emb_id: str) -> Optional[E]:
        return self.vectors.get(emb_id)

    def delete(self, emb_id: str) -> None:
        if emb_id in self.vectors:
            del self.vectors[emb_id]
            del self.payloads[emb_id]
            # remove from cache
            self._cache = {k: v for k, v in self._cache.items() if emb_id not in k}

    def clear(self) -> None:
        self.vectors.clear()
        self._cache.clear()
        self.payloads.clear()

    def __len__(self):
        return len(self.vectors)

    def _retrieve_similar(
        self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
    ) -> List[Tuple[str, T, float]]:
        ref_embedding = self.get(emb_id=ref_id)
        if ref_embedding is None:
            raise ValueError(f"Reference embedding '{ref_id}' not found.")

        # calculate similarity to all embeddings
        similarities = {}
        for emb_id, embedding in self.vectors.items():
            if (emb_id, ref_id) not in self._cache:
                # use cosine similarity
                self._cache[(emb_id, ref_id)] = self._sim(ref_embedding, embedding)
            similarities[emb_id] = self._cache[(emb_id, ref_id)]

        # sort by similarity
        similar_entries = sorted(similarities.items(), key=lambda x: x[1], reverse=True)

        if min_similarity is not None:
            similar_entries = [
                (emb_id, sim) for emb_id, sim in similar_entries if sim >= min_similarity
            ]
        if top_k is not None:
            similar_entries = similar_entries[:top_k]

        return [(emb_id, self.payloads[emb_id], sim) for emb_id, sim in similar_entries]

    def _save_to_directory(self, directory: str) -> None:
        indices = list(self.vectors.keys())
        with open(os.path.join(directory, self.INDEX_FILE), "w") as f:
            json.dump(indices, f)
        embeddings_np = np.array(list(self.vectors.values()))
        np.save(os.path.join(directory, self.EMBEDDINGS_FILE), embeddings_np)
        payloads = [self.payloads[idx] for idx in indices]
        with open(os.path.join(directory, self.PAYLOADS_FILE), "w") as f:
            json.dump(payloads, f)

    def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
        indices = list(self.vectors.keys())
        embeddings_np = np.array(list(self.vectors.values()))
        payloads = [self.payloads[idx] for idx in indices]
        return indices, embeddings_np, payloads


class QdrantVectorStore(VectorStore[T, List[float]]):
    """Vector store using Qdrant as a backend."""

    COLLECTION_NAME = "ADUs"
    MAX_LIMIT = 100

    def __init__(
        self,
        location: str = ":memory:",
        vector_size: int = 768,
        distance: Distance = Distance.COSINE,
    ):
        self.client = QdrantClient(location=location)
        self.emb_id2point_id = {}
        self.point_id2emb_id = {}
        self.client.create_collection(
            collection_name=self.COLLECTION_NAME,
            vectors_config=VectorParams(size=vector_size, distance=distance),
        )

    def __len__(self):
        return self.client.get_collection(collection_name=self.COLLECTION_NAME).points_count

    def _add(self, emb_id: str, payload: T, embedding: List[float]) -> None:

        if emb_id in self.emb_id2point_id:
            # update existing entry
            point_id = self.emb_id2point_id[emb_id]
        else:
            # we use the length of the emb_id2point_id dict as the index,
            # because we assume that, even when we delete an entry from
            # the store, we do not delete it from the index
            point_id = len(self.emb_id2point_id)
            self.emb_id2point_id[emb_id] = point_id
            self.point_id2emb_id[point_id] = emb_id

        self.client.upsert(
            collection_name=self.COLLECTION_NAME,
            points=[PointStruct(id=point_id, vector=embedding, payload=payload)],
        )

    def _get(self, emb_id: str) -> Optional[List[float]]:
        if emb_id not in self.emb_id2point_id:
            return None
        points = self.client.retrieve(
            collection_name=self.COLLECTION_NAME,
            ids=[self.emb_id2point_id[emb_id]],
            with_vectors=True,
        )
        if len(points) == 0:
            return None
        elif len(points) == 1:
            return points[0].vector
        else:
            raise ValueError(f"Multiple points found for ID '{emb_id}'.")

    def _retrieve_similar(
        self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
    ) -> List[Tuple[str, T, float]]:
        similar_entries = self.client.recommend(
            collection_name=self.COLLECTION_NAME,
            positive=[self.emb_id2point_id[ref_id]],
            limit=top_k or self.MAX_LIMIT,
            score_threshold=min_similarity,
        )
        return [
            (self.point_id2emb_id[entry.id], entry.payload, entry.score)
            for entry in similar_entries
        ]

    def clear(self) -> None:
        vectors_config = self.client.get_collection(
            collection_name=self.COLLECTION_NAME
        ).vectors_config
        self.client.delete_collection(collection_name=self.COLLECTION_NAME)
        self.client.create_collection(
            collection_name=self.COLLECTION_NAME,
            vectors_config=vectors_config,
        )
        self.emb_id2point_id.clear()
        self.point_id2emb_id.clear()

    def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
        num_entries = self.client.get_collection(collection_name=self.COLLECTION_NAME).points_count
        data, point_ids = self.client.scroll(
            collection_name=self.COLLECTION_NAME, with_vectors=True, limit=num_entries
        )
        vectors_np = np.array([point.vector for point in data])
        payloads = [point.payload for point in data]
        emb_ids = [self.point_id2emb_id[point.id] for point in data]
        return emb_ids, vectors_np, payloads