mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-15 15:46:47 +00:00
25f2e25be1
## Description Contains 2 updates to for integration tests to run on langchain's CI. Addendum to #25057 to get release github action to succeed.
84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
"""Search index commands are only supported on Atlas Clusters >=M10"""
|
|
|
|
import os
|
|
from typing import Generator, List, Optional
|
|
|
|
import pytest
|
|
from pymongo import MongoClient
|
|
from pymongo.collection import Collection
|
|
|
|
from langchain_mongodb import index
|
|
|
|
DB_NAME = "langchain_test_index_db"
|
|
COLLECTION_NAME = "test_index"
|
|
VECTOR_INDEX_NAME = "vector_index"
|
|
|
|
TIMEOUT = 120
|
|
DIMENSIONS = 10
|
|
|
|
|
|
@pytest.fixture
|
|
def collection() -> Generator:
|
|
"""Depending on uri, this could point to any type of cluster."""
|
|
uri = os.environ.get("MONGODB_ATLAS_URI")
|
|
client: MongoClient = MongoClient(uri)
|
|
clxn = client[DB_NAME][COLLECTION_NAME]
|
|
clxn.insert_one({"foo": "bar"})
|
|
yield clxn
|
|
clxn.drop()
|
|
|
|
|
|
def test_search_index_commands(collection: Collection) -> None:
|
|
index_name = VECTOR_INDEX_NAME
|
|
dimensions = DIMENSIONS
|
|
path = "embedding"
|
|
similarity = "cosine"
|
|
filters: Optional[List[str]] = None
|
|
wait_until_complete = TIMEOUT
|
|
|
|
for index_info in collection.list_search_indexes():
|
|
index.drop_vector_search_index(
|
|
collection, index_info["name"], wait_until_complete=wait_until_complete
|
|
)
|
|
|
|
assert len(list(collection.list_search_indexes())) == 0
|
|
|
|
index.create_vector_search_index(
|
|
collection=collection,
|
|
index_name=index_name,
|
|
dimensions=dimensions,
|
|
path=path,
|
|
similarity=similarity,
|
|
filters=filters,
|
|
wait_until_complete=wait_until_complete,
|
|
)
|
|
|
|
assert index._is_index_ready(collection, index_name)
|
|
indexes = list(collection.list_search_indexes())
|
|
assert len(indexes) == 1
|
|
assert indexes[0]["name"] == index_name
|
|
|
|
new_similarity = "euclidean"
|
|
index.update_vector_search_index(
|
|
collection,
|
|
index_name,
|
|
DIMENSIONS,
|
|
"embedding",
|
|
new_similarity,
|
|
filters=[],
|
|
wait_until_complete=wait_until_complete,
|
|
)
|
|
|
|
assert index._is_index_ready(collection, index_name)
|
|
indexes = list(collection.list_search_indexes())
|
|
assert len(indexes) == 1
|
|
assert indexes[0]["name"] == index_name
|
|
assert indexes[0]["latestDefinition"]["fields"][0]["similarity"] == new_similarity
|
|
|
|
index.drop_vector_search_index(
|
|
collection, index_name, wait_until_complete=wait_until_complete
|
|
)
|
|
|
|
indexes = list(collection.list_search_indexes())
|
|
assert len(indexes) == 0
|