refactor: use abtract VecStoreFromDocIndex for in memory and hnswlib implementation

Signed-off-by: anna-charlotte <[email protected]>

Author: anna-charlotte

langchain/vectorstores/hnsw_lib.py

@@ -1,52 +1,44 @@
 """Wrapper around in-memory DocArray store."""
 from __future__ import annotations
 
-from operator import itemgetter
 from typing import List, Optional, Any, Tuple, Iterable, Type, Callable, Sequence, TYPE_CHECKING
+from docarray.typing import NdArray
 
 from langchain.embeddings.base import Embeddings
-from langchain.schema import Document
-from langchain.vectorstores import VectorStore
 from langchain.vectorstores.base import VST
-from langchain.vectorstores.utils import maximal_marginal_relevance
-
-from docarray import BaseDoc
-from docarray.typing import NdArray
+from langchain.vectorstores.vector_store_from_doc_index import VecStoreFromDocIndex, _check_docarray_import
 
 
-class HnswLib(VectorStore):
+class HnswLib(VecStoreFromDocIndex):
     """Wrapper around HnswLib storage.
 
-    To use it, you should have the ``docarray`` package with version >=0.30.0 installed.
+    To use it, you should have the ``docarray`` package with version >=0.31.0 installed.
     """
     def __init__(
         self,
-        work_dir: str,
-        n_dim: int,
         texts: List[str],
         embedding: Embeddings,
+        work_dir: str,
+        n_dim: int,
         metadatas: Optional[List[dict]],
-        sim_metric: str = 'cosine',
-        kwargs: dict = None
+        dist_metric: str = 'cosine',
+        **kwargs,
     ) -> None:
-        """Initialize HnswLib store."""
-        try:
-            import docarray
-            da_version = docarray.__version__.split('.')
-            if int(da_version[0]) == 0 and int(da_version[1]) <= 21:
-                raise ValueError(
-                    f'To use the HnswLib VectorStore the docarray version >=0.30.0 is expected, '
-                    f'received: {docarray.__version__}.'
-                    f'To upgrade, please run: `pip install -U docarray`.'
-                )
-            else:
-                from docarray import DocList
-                from docarray.index import HnswDocumentIndex
-        except ImportError:
-            raise ImportError(
-                "Could not import docarray python package. "
-                "Please install it with `pip install -U docarray`."
-            )
+        """Initialize HnswLib store.
+
+        Args:
+            texts (List[str]): Text data.
+            embedding (Embeddings): Embedding function.
+            metadatas (Optional[List[dict]]): Metadata for each text if it exists.
+                Defaults to None.
+            work_dir (str): path to the location where all the data will be stored.
+            n_dim (int): dimension of an embedding.
+            dist_metric (str): Distance metric for HnswLib can be one of: 'cosine',
+                'ip', and 'l2'. Defaults to 'cosine'.
+        """
+        _check_docarray_import()
+        from docarray.index import HnswDocumentIndex
+
         try:
             import google.protobuf
         except ImportError:
@@ -55,27 +47,13 @@ def __init__(
                 "Please install it with `pip install -U protobuf`."
             )
 
-        if metadatas is None:
-            metadatas = [{} for _ in range(len(texts))]
-
-        self.embedding = embedding
-
-        self.doc_cls = self._get_doc_cls(n_dim, sim_metric)
-        self.doc_index = HnswDocumentIndex[self.doc_cls](work_dir=work_dir)
-        embeddings = self.embedding.embed_documents(texts)
-        docs = DocList[self.doc_cls](
-            [
-                self.doc_cls(
-                    text=t,
-                    embedding=e,
-                    metadata=m,
-                ) for t, m, e in zip(texts, metadatas, embeddings)
-            ]
-        )
-        self.doc_index.index(docs)
+        doc_cls = self._get_doc_cls(n_dim, dist_metric)
+        doc_index = HnswDocumentIndex[doc_cls](work_dir=work_dir)
+        super().__init__(doc_index, texts, embedding, metadatas)
 
     @staticmethod
     def _get_doc_cls(n_dim: int, sim_metric: str):
+        from docarray import BaseDoc
         from pydantic import Field
 
         class DocArrayDoc(BaseDoc):
@@ -93,6 +71,7 @@ def from_texts(
         metadatas: Optional[List[dict]] = None,
         work_dir: str = None,
         n_dim: int = None,
+        dist_metric: str = 'cosine',
         **kwargs: Any
     ) -> HnswLib:
 
@@ -107,129 +86,6 @@ def from_texts(
             texts=texts,
             embedding=embedding,
             metadatas=metadatas,
-            kwargs=kwargs
+            dist_metric=dist_metric,
+            kwargs=kwargs,
         )
-
-    def add_texts(
-        self,
-        texts: Iterable[str],
-        metadatas: Optional[List[dict]] = None,
-        **kwargs: Any
-    ) -> List[str]:
-        """Run more texts through the embeddings and add to the vectorstore.
-
-        Args:
-            texts: Iterable of strings to add to the vectorstore.
-            metadatas: Optional list of metadatas associated with the texts.
-
-        Returns:
-            List of ids from adding the texts into the vectorstore.
-        """
-        if metadatas is None:
-            metadatas = [{} for _ in range(len(list(texts)))]
-
-        ids = []
-        embeddings = self.embedding.embed_documents(texts)
-        for t, m, e in zip(texts, metadatas, embeddings):
-            doc = self.doc_cls(
-                text=t,
-                embedding=e,
-                metadata=m
-            )
-            self.doc_index.index(doc)
-            ids.append(doc.id)  # TODO return index of self.docs ?
-
-        return ids
-
-    def similarity_search_with_score(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> List[Tuple[Document, float]]:
-        """Return docs most similar to query.
-
-        Args:
-            query: Text to look up documents similar to.
-            k: Number of Documents to return. Defaults to 4.
-
-        Returns:
-            List of Documents most similar to the query and score for each.
-        """
-        query_embedding = self.embedding.embed_query(query)
-        query_embedding = [1., 1., 1., 1., 1., 1., 1., 1., 1., 0.]
-        print(f"query_embedding = {query_embedding}")
-        query_doc = self.doc_cls(embedding=query_embedding)
-        docs, scores = self.doc_index.find(query_doc, search_field='embedding', limit=k)
-
-        result = [(Document(page_content=doc.text), score) for doc, score in zip(docs, scores)]
-        return result
-
-    def similarity_search(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> List[Document]:
-        """Return docs most similar to query.
-
-        Args:
-            query: Text to look up documents similar to.
-            k: Number of Documents to return. Defaults to 4.
-
-        Returns:
-            List of Documents most similar to the query.
-        """
-        results = self.similarity_search_with_score(query, k)
-        return list(map(itemgetter(0), results))
-
-    def _similarity_search_with_relevance_scores(
-        self,
-        query: str,
-        k: int = 4,
-        **kwargs: Any,
-    ) -> List[Tuple[Document, float]]:
-        """Return docs and relevance scores, normalized on a scale from 0 to 1.
-
-        0 is dissimilar, 1 is most similar.
-        """
-        raise NotImplementedError
-
-    def similarity_search_by_vector(self, embedding: List[float], k: int = 4, **kwargs: Any) -> List[Document]:
-        """Return docs most similar to embedding vector.
-
-        Args:
-            embedding: Embedding to look up documents similar to.
-            k: Number of Documents to return. Defaults to 4.
-
-        Returns:
-            List of Documents most similar to the query vector.
-        """
-
-        query_doc = self.doc_cls(embedding=embedding)
-        docs = self.doc_index.find(query_doc, search_field='embedding', limit=k).documents
-
-        result = [Document(page_content=doc.text) for doc in docs]
-        return result
-
-    def max_marginal_relevance_search(
-        self, query: str, k: int = 4, fetch_k: int = 20, **kwargs: Any
-    ) -> List[Document]:
-        """Return docs selected using the maximal marginal relevance.
-
-        Maximal marginal relevance optimizes for similarity to query AND diversity
-        among selected documents.
-
-        Args:
-            query: Text to look up documents similar to.
-            k: Number of Documents to return. Defaults to 4.
-            fetch_k: Number of Documents to fetch to pass to MMR algorithm.
-
-        Returns:
-            List of Documents selected by maximal marginal relevance.
-        """
-        query_embedding = self.embedding.embed_query(query)
-        query_doc = self.doc_cls(embedding=query_embedding)
-
-        docs, scores = self.doc_index.find(query_doc, search_field='embedding', limit=fetch_k)
-
-        embeddings = [emb for emb in docs.emb]
-
-        mmr_selected = maximal_marginal_relevance(query_embedding, embeddings, k=k)
-        results = [Document(page_content=self.doc_index[idx].text) for idx in mmr_selected]
-        return results
-