Replace vector search index funcs

Remove vector search index funcs from langchain-mongodb and import from
pymongo-search-utils.

Author: aclark4life

libs/langchain-mongodb/langchain_mongodb/index.py

@@ -6,6 +6,11 @@
 
 from pymongo.collection import Collection
 from pymongo.operations import SearchIndexModel
+from pymongo_search_utils import (
+    create_vector_search_index,  # noqa: F401
+    drop_vector_search_index,  # noqa: F401
+    update_vector_search_index,  # noqa: F401
+)
 
 logger = logging.getLogger(__file__)
 
@@ -34,135 +39,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(
-        authorizedCollections=True
-    ):
-        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

libs/langchain-mongodb/pyproject.toml

@@ -65,6 +65,7 @@ filterwarnings = [
 
 [tool.mypy]
 disallow_untyped_defs = true
+disable_error_code = ["import-untyped"]
 
 [[tool.mypy.overrides]]
 module = ["tests.*"]