INTPYTHON-752 Integrate pymongo-vectorsearch-utils

Author: aclark4life

libs/langchain-mongodb/langchain_mongodb/index.py

@@ -2,14 +2,17 @@
 
 import logging
 from time import monotonic, sleep
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional
 
 from pymongo.collection import Collection
-from pymongo.operations import SearchIndexModel
 
 logger = logging.getLogger(__file__)
 
 
+# Don't break imports for modules that expect these functions
+# to be in this module.
+
+
 def _vector_search_index_definition(
     dimensions: int,
     path: str,

libs/langchain-mongodb/langchain_mongodb/utils.py

@@ -24,7 +24,6 @@
 from typing import Any, Dict, List, Union
 
 import numpy as np
-from pymongo import MongoClient
 from pymongo.driver_info import DriverInfo
 
 logger = logging.getLogger(__name__)
@@ -33,11 +32,8 @@
 
 DRIVER_METADATA = DriverInfo(name="Langchain", version=version("langchain-mongodb"))
 
-
-def _append_client_metadata(client: MongoClient) -> None:
-    # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
-    if callable(client.append_metadata):
-        client.append_metadata(DRIVER_METADATA)
+# Don't break imports for modules that expect this function
+# to be in this module.
 
 
 def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:

libs/langchain-mongodb/langchain_mongodb/vectorstores.py

@@ -22,9 +22,10 @@
 from langchain_core.embeddings import Embeddings
 from langchain_core.runnables.config import run_in_executor
 from langchain_core.vectorstores import VectorStore
-from pymongo import MongoClient, ReplaceOne
+from pymongo import MongoClient
 from pymongo.collection import Collection
 from pymongo.errors import CollectionInvalid
+from pymongo_vectorsearch_utils import bulk_embed_and_insert_texts
 
 from langchain_mongodb.index import (
     create_vector_search_index,
@@ -362,11 +363,11 @@ def add_texts(
                 metadatas_batch.append(metadata)
                 if (j + 1) % batch_size == 0 or size >= 47_000_000:
                     if ids:
-                        batch_res = self.bulk_embed_and_insert_texts(
+                        batch_res = bulk_embed_and_insert_texts(
                             texts_batch, metadatas_batch, ids[i : j + 1]
                         )
                     else:
-                        batch_res = self.bulk_embed_and_insert_texts(
+                        batch_res = bulk_embed_and_insert_texts(
                             texts_batch, metadatas_batch
                         )
                     result_ids.extend(batch_res)
@@ -376,13 +377,11 @@ def add_texts(
                     i = j + 1
         if texts_batch:
             if ids:
-                batch_res = self.bulk_embed_and_insert_texts(
+                batch_res = bulk_embed_and_insert_texts(
                     texts_batch, metadatas_batch, ids[i : j + 1]
                 )
             else:
-                batch_res = self.bulk_embed_and_insert_texts(
-                    texts_batch, metadatas_batch
-                )
+                batch_res = bulk_embed_and_insert_texts(texts_batch, metadatas_batch)
             result_ids.extend(batch_res)
         return result_ids
 
@@ -419,37 +418,6 @@ def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:
             docs.append(Document(page_content=text, id=oid_to_str(_id), metadata=doc))
         return docs
 
-    def bulk_embed_and_insert_texts(
-        self,
-        texts: Union[List[str], Iterable[str]],
-        metadatas: Union[List[dict], Generator[dict, Any, Any]],
-        ids: Optional[List[str]] = None,
-    ) -> List[str]:
-        """Bulk insert single batch of texts, embeddings, and optionally ids.
-
-        See add_texts for additional details.
-        """
-        if not texts:
-            return []
-        # Compute embedding vectors
-        embeddings = self._embedding.embed_documents(list(texts))
-        if not ids:
-            ids = [str(ObjectId()) for _ in range(len(list(texts)))]
-        docs = [
-            {
-                "_id": str_to_oid(i),
-                self._text_key: t,
-                self._embedding_key: embedding,
-                **m,
-            }
-            for i, t, m, embedding in zip(ids, texts, metadatas, embeddings)
-        ]
-        operations = [ReplaceOne({"_id": doc["_id"]}, doc, upsert=True) for doc in docs]
-        # insert the documents in MongoDB Atlas
-        result = self._collection.bulk_write(operations)
-        assert result.upserted_ids is not None
-        return [oid_to_str(_id) for _id in result.upserted_ids.values()]
-
     def add_documents(
         self,
         documents: List[Document],
@@ -481,7 +449,7 @@ def add_documents(
                 *[(doc.page_content, doc.metadata) for doc in documents[start:end]]
             )
             result_ids.extend(
-                self.bulk_embed_and_insert_texts(
+                bulk_embed_and_insert_texts(
                     texts=texts, metadatas=metadatas, ids=ids[start:end]
                 )
             )

libs/langchain-mongodb/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "langchain-text-splitters>=0.3",
     "numpy>=1.26",
     "lark<2.0.0,>=1.1.9",
+    # "pymongo-vectorsearch-utils",
 ]
 
 [dependency-groups]

libs/langchain-mongodb/tests/utils.py

@@ -26,6 +26,7 @@
 from pymongo.driver_info import DriverInfo
 from pymongo.operations import SearchIndexModel
 from pymongo.results import BulkWriteResult, DeleteResult, InsertManyResult
+from pymongo_vectorsearch_utils import bulk_embed_and_insert_texts
 
 from langchain_mongodb import MongoDBAtlasVectorSearch
 from langchain_mongodb.agent_toolkit.database import MongoDBDatabase
@@ -63,7 +64,7 @@ def bulk_embed_and_insert_texts(
         ids: Optional[List[str]] = None,
     ) -> List:
         """Patched insert_texts that waits for data to be indexed before returning"""
-        ids_inserted = super().bulk_embed_and_insert_texts(texts, metadatas, ids)
+        ids_inserted = bulk_embed_and_insert_texts(texts, metadatas, ids)
         n_docs = self.collection.count_documents({})
         start = monotonic()
         while monotonic() - start <= TIMEOUT: