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,
@@ -34,133 +37,6 @@ def _vector_search_index_definition(
     return definition
 
 
-def create_vector_search_index(
-    collection: Collection,
-    index_name: str,
-    dimensions: int,
-    path: str,
-    similarity: str,
-    filters: Optional[List[str]] = None,
-    *,
-    wait_until_complete: Optional[float] = None,
-    **kwargs: Any,
-) -> None:
-    """Experimental Utility function to create a vector search index
-
-    Args:
-        collection (Collection): MongoDB Collection
-        index_name (str): Name of Index
-        dimensions (int): Number of dimensions in embedding
-        path (str): field with vector embedding
-        similarity (str): The similarity score used for the index
-        filters (List[str]): Fields/paths to index to allow filtering in $vectorSearch
-        wait_until_complete (Optional[float]): If provided, number of seconds to wait
-            until search index is ready.
-        kwargs: Keyword arguments supplying any additional options to SearchIndexModel.
-    """
-    logger.info("Creating Search Index %s on %s", index_name, collection.name)
-
-    if collection.name not in collection.database.list_collection_names():
-        collection.database.create_collection(collection.name)
-
-    result = collection.create_search_index(
-        SearchIndexModel(
-            definition=_vector_search_index_definition(
-                dimensions=dimensions,
-                path=path,
-                similarity=similarity,
-                filters=filters,
-                **kwargs,
-            ),
-            name=index_name,
-            type="vectorSearch",
-        )
-    )
-
-    if wait_until_complete:
-        _wait_for_predicate(
-            predicate=lambda: _is_index_ready(collection, index_name),
-            err=f"{index_name=} did not complete in {wait_until_complete}!",
-            timeout=wait_until_complete,
-        )
-    logger.info(result)
-
-
-def drop_vector_search_index(
-    collection: Collection,
-    index_name: str,
-    *,
-    wait_until_complete: Optional[float] = None,
-) -> None:
-    """Drop a created vector search index
-
-    Args:
-        collection (Collection): MongoDB Collection with index to be dropped
-        index_name (str): Name of the MongoDB index
-        wait_until_complete (Optional[float]): If provided, number of seconds to wait
-            until search index is ready.
-    """
-    logger.info(
-        "Dropping Search Index %s from Collection: %s", index_name, collection.name
-    )
-    collection.drop_search_index(index_name)
-    if wait_until_complete:
-        _wait_for_predicate(
-            predicate=lambda: len(list(collection.list_search_indexes())) == 0,
-            err=f"Index {index_name} did not drop in {wait_until_complete}!",
-            timeout=wait_until_complete,
-        )
-    logger.info("Vector Search index %s.%s dropped", collection.name, index_name)
-
-
-def update_vector_search_index(
-    collection: Collection,
-    index_name: str,
-    dimensions: int,
-    path: str,
-    similarity: str,
-    filters: Optional[List[str]] = None,
-    *,
-    wait_until_complete: Optional[float] = None,
-    **kwargs: Any,
-) -> None:
-    """Update a search index.
-
-    Replace the existing index definition with the provided definition.
-
-    Args:
-        collection (Collection): MongoDB Collection
-        index_name (str): Name of Index
-        dimensions (int): Number of dimensions in embedding
-        path (str): field with vector embedding
-        similarity (str): The similarity score used for the index.
-        filters (List[str]): Fields/paths to index to allow filtering in $vectorSearch
-        wait_until_complete (Optional[float]): If provided, number of seconds to wait
-            until search index is ready.
-        kwargs: Keyword arguments supplying any additional options to SearchIndexModel.
-    """
-    logger.info(
-        "Updating Search Index %s from Collection: %s", index_name, collection.name
-    )
-    collection.update_search_index(
-        name=index_name,
-        definition=_vector_search_index_definition(
-            dimensions=dimensions,
-            path=path,
-            similarity=similarity,
-            filters=filters,
-            **kwargs,
-        ),
-    )
-    if wait_until_complete:
-        _wait_for_predicate(
-            predicate=lambda: _is_index_ready(collection, index_name),
-            err=f"Index {index_name} update did not complete in {wait_until_complete}!",
-            timeout=wait_until_complete,
-        )
-    logger.info("Update succeeded")
-
-
 def _is_index_ready(collection: Collection, index_name: str) -> bool:
     """Check for the index name in the list of available search indexes to see if the
     specified index is of status READY
@@ -197,48 +73,3 @@ def _wait_for_predicate(
         if monotonic() - start > timeout:
             raise TimeoutError(err)
         sleep(interval)
-
-
-def create_fulltext_search_index(
-    collection: Collection,
-    index_name: str,
-    field: Union[str, List[str]],
-    *,
-    wait_until_complete: Optional[float] = None,
-    **kwargs: Any,
-) -> None:
-    """Experimental Utility function to create an Atlas Search index
-
-    Args:
-        collection (Collection): MongoDB Collection
-        index_name (str): Name of Index
-        field (str): Field to index
-        wait_until_complete (Optional[float]): If provided, number of seconds to wait
-            until search index is ready
-        kwargs: Keyword arguments supplying any additional options to SearchIndexModel.
-    """
-    logger.info("Creating Search Index %s on %s", index_name, collection.name)
-
-    if collection.name not in collection.database.list_collection_names():
-        collection.database.create_collection(collection.name)
-
-    if isinstance(field, str):
-        fields_definition = {field: [{"type": "string"}]}
-    else:
-        fields_definition = {f: [{"type": "string"}] for f in field}
-    definition = {"mappings": {"dynamic": False, "fields": fields_definition}}
-    result = collection.create_search_index(
-        SearchIndexModel(
-            definition=definition,
-            name=index_name,
-            type="search",
-            **kwargs,
-        )
-    )
-    if wait_until_complete:
-        _wait_for_predicate(
-            predicate=lambda: _is_index_ready(collection, index_name),
-            err=f"{index_name=} did not complete in {wait_until_complete}!",
-            timeout=wait_until_complete,
-        )
-    logger.info(result)

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

@@ -21,9 +21,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,
@@ -360,11 +361,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)
@@ -374,13 +375,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
 
@@ -417,37 +416,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],
@@ -479,7 +447,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: