INTPYTHON-655 Infer dimensions from embedding if not provided (#182)

Author: blink1073

libs/langchain-mongodb/langchain_mongodb/vectorstores.py

@@ -208,7 +208,7 @@ def __init__(
         embedding_key: str = "embedding",
         relevance_score_fn: str = "cosine",
         dimensions: int = -1,
-        auto_create_index: bool = True,
+        auto_create_index: bool | None = None,
         auto_index_timeout: int = 15,
         **kwargs: Any,
     ):
@@ -222,8 +222,9 @@ def __init__(
             embedding_key: Field that will contain the embedding for each document
             relevance_score_fn: The similarity score used for the index
                 Currently supported: 'euclidean', 'cosine', and 'dotProduct'
-            dimensions: Number of dimensions in embedding.  If the value is set and
-                the index does not exist, an index will be created.
+            auto_create_index: Whether to automatically create an index if it does not exist.
+            dimensions: Number of dimensions in embedding.  If the value is not provided, and `auto_create_index`
+                is `true`, the value will be inferred.
             auto_index_timeout: Timeout in seconds to wait for an auto-created index
                to be ready.
         """
@@ -234,18 +235,21 @@ def __init__(
         self._embedding_key = embedding_key
         self._relevance_score_fn = relevance_score_fn
 
-        if not auto_create_index or dimensions == -1:
+        if auto_create_index is False:
             return
+        if auto_create_index is None and dimensions == -1:
+            return
+        if dimensions == -1:
+            dimensions = len(embedding.embed_query("foo"))
+
         coll = self._collection
-        if not any(
-            [ix["name"] == self._index_name for ix in coll.list_search_indexes()]
-        ):
+        if not any([ix["name"] == index_name for ix in coll.list_search_indexes()]):
             create_vector_search_index(
                 collection=coll,
-                index_name=self._index_name,
+                index_name=index_name,
                 dimensions=dimensions,
-                path=self._embedding_key,
-                similarity=self._relevance_score_fn,
+                path=embedding_key,
+                similarity=relevance_score_fn,
                 wait_until_complete=auto_index_timeout,
             )
 

libs/langchain-mongodb/tests/unit_tests/test_vectorstores.py

@@ -98,6 +98,7 @@ def test_from_documents(
             collection=collection,
             index_name=INDEX_NAME,
         )
+        # TODO: test how DIMS is handled here.
         self._validate_search(
             vectorstore, collection, metadata=documents[2].metadata["c"]
         )
@@ -191,3 +192,35 @@ def test_mmr(
         assert len(output) == len(texts)
         assert output[0].page_content == "foo"
         assert output[1].page_content != "foo"
+
+    def test_auto_create_index(
+        self, embedding_openai: Embeddings, collection: MockCollection
+    ) -> None:
+        # Explicit auto_create_index
+        assert len(collection._search_indexes) == 0
+        _ = MongoDBAtlasVectorSearch(
+            embedding=embedding_openai,
+            collection=collection,
+            index_name=INDEX_NAME,
+            auto_create_index=True,
+        )
+        assert len(collection._search_indexes) == 1
+
+        # Explicit dimensions
+        collection._search_indexes = []
+        _ = MongoDBAtlasVectorSearch(
+            embedding=embedding_openai,
+            collection=collection,
+            index_name=INDEX_NAME,
+            dimensions=10,
+        )
+        assert len(collection._search_indexes) == 1
+
+        # Does not auto-create
+        collection._search_indexes = []
+        _ = MongoDBAtlasVectorSearch(
+            embedding=embedding_openai,
+            collection=collection,
+            index_name=INDEX_NAME,
+        )
+        assert len(collection._search_indexes) == 0

libs/langchain-mongodb/tests/utils.py

@@ -23,6 +23,7 @@
 from pydantic import model_validator
 from pymongo import MongoClient
 from pymongo.collection import Collection
+from pymongo.operations import SearchIndexModel
 from pymongo.results import BulkWriteResult, DeleteResult, InsertManyResult
 
 from langchain_mongodb import MongoDBAtlasVectorSearch
@@ -250,9 +251,11 @@ class MockCollection(Collection):
 
     def __init__(self, database: MockDatabase | None = None) -> None:
         self._data = []
+        self._name = "test"
         self.is_closed = False
         self._aggregate_result = []
         self._insert_result = None
+        self._search_indexes = []
         self._simulate_cache_aggregation_query = False
         self._database = database or MockDatabase()  # type:ignore[assignment]
 
@@ -263,6 +266,17 @@ def database(self):
     def close(self):
         self.is_closed = True
 
+    def list_search_indexes(self, name=None, session=None, comment=None, **kwargs):
+        return [
+            dict(name=idx.document["name"], status="READY")
+            for idx in self._search_indexes
+        ]
+
+    def create_search_index(self, model, session=None, comment=None, **kwargs):
+        if not isinstance(model, SearchIndexModel):
+            model = SearchIndexModel(model, name=f"test{len(self._search_indexes)}")
+        self._search_indexes.append(model)
+
     def delete_many(self, *args, **kwargs) -> DeleteResult:  # type: ignore
         old_len = len(self._data)
         self._data = []