from https://github.com/ArneBinder/pie-document-level/pull/233

use QdrantVectorStore again (since it works with download/upload now)

vector_store.py

@@ -12,6 +12,16 @@ 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."""
@@ -22,6 +32,11 @@ class VectorStore(Generic[T, E], abc.ABC):
         """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:
@@ -67,16 +82,41 @@ class VectorStore(Generic[T, E], abc.ABC):
     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."""
-        raise NotImplementedError
+        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.
         """
-        raise NotImplementedError
+        if replace:
+            self.clear()
+        self._add_from_directory(directory)
 
 
 def vector_norm(vector: List[float]) -> float:
@@ -88,10 +128,7 @@ def cosine_similarity(a: List[float], b: List[float]) -> float:
 
 
 class SimpleVectorStore(VectorStore[T, List[float]]):
-
-    INDEX_FILE = "vectors_index.json"
-    EMBEDDINGS_FILE = "vectors_data.npy"
-    PAYLOADS_FILE = "vectors_payloads.json"
+    """Simple in-memory vector store using a dictionary."""
 
     def __init__(self):
         self.vectors: dict[str, List[float]] = {}
@@ -148,8 +185,7 @@ class SimpleVectorStore(VectorStore[T, List[float]]):
 
         return [(emb_id, self.payloads[emb_id], sim) for emb_id, sim in similar_entries]
 
-    def save_to_directory(self, directory: str) -> None:
-        os.makedirs(directory, exist_ok=True)
+    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)
@@ -159,20 +195,15 @@ class SimpleVectorStore(VectorStore[T, List[float]]):
         with open(os.path.join(directory, self.PAYLOADS_FILE), "w") as f:
             json.dump(payloads, f)
 
-    def load_from_directory(self, directory: str, replace: bool = False) -> None:
-        if replace:
-            self.clear()
-        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.vectors[emb_id] = emb.tolist()
-            self.payloads[emb_id] = payload
+    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
@@ -184,8 +215,8 @@ class QdrantVectorStore(VectorStore[T, List[float]]):
         distance: Distance = Distance.COSINE,
     ):
         self.client = QdrantClient(location=location)
-        self.id2idx = {}
-        self.idx2id = {}
+        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),
@@ -196,21 +227,26 @@ class QdrantVectorStore(VectorStore[T, List[float]]):
 
     def _add(self, emb_id: str, payload: T, embedding: List[float]) -> None:
 
-        # we use the length of the id2idx 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
-        _id = len(self.id2idx)
+        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=_id, vector=embedding, payload=payload)],
+            points=[PointStruct(id=point_id, vector=embedding, payload=payload)],
         )
-        self.id2idx[emb_id] = _id
-        self.idx2id[_id] = emb_id
 
     def _get(self, emb_id: str) -> Optional[List[float]]:
         points = self.client.retrieve(
             collection_name=self.COLLECTION_NAME,
-            ids=[self.id2idx[emb_id]],
+            ids=[self.emb_id2point_id[emb_id]],
             with_vectors=True,
         )
         if len(points) == 0:
@@ -225,11 +261,14 @@ class QdrantVectorStore(VectorStore[T, List[float]]):
     ) -> List[Tuple[str, T, float]]:
         similar_entries = self.client.recommend(
             collection_name=self.COLLECTION_NAME,
-            positive=[self.id2idx[ref_id]],
+            positive=[self.emb_id2point_id[ref_id]],
             limit=top_k or self.MAX_LIMIT,
             score_threshold=min_similarity,
         )
-        return [(self.idx2id[entry.id], entry.payload, entry.score) for entry in similar_entries]
+        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(
@@ -240,5 +279,15 @@ class QdrantVectorStore(VectorStore[T, List[float]]):
             collection_name=self.COLLECTION_NAME,
             vectors_config=vectors_config,
         )
-        self.id2idx.clear()
-        self.idx2id.clear()
+        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