Skip to content

Commit

Permalink
INTPYTHON-447 Ensure index is ready for retriever tests (#25)
Browse files Browse the repository at this point in the history
  • Loading branch information
blink1073 authored Dec 6, 2024
1 parent 04e667b commit 2f5f1fc
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 2 deletions.
2 changes: 1 addition & 1 deletion libs/mongodb/langchain_mongodb/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ def _is_index_ready(collection: Collection, index_name: str) -> bool:
bool : True if the index is present and READY false otherwise
"""
for index in collection.list_search_indexes(index_name):
if index["type"] == "vectorSearch" and index["status"] == "READY":
if index["status"] == "READY":
return True
return False

Expand Down
15 changes: 14 additions & 1 deletion libs/mongodb/tests/integration_tests/test_retrievers.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
from typing import Generator, List
from time import sleep, time

import pytest
from langchain_core.documents import Document
Expand All @@ -10,7 +11,7 @@
from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_mongodb.index import (
create_fulltext_search_index,
create_vector_search_index,
create_vector_search_index
)
from langchain_mongodb.retrievers import (
MongoDBAtlasFullTextSearchRetriever,
Expand Down Expand Up @@ -160,6 +161,18 @@ def test_fulltext_retriever(
search_field=PAGE_CONTENT_FIELD,
)

# Wait for the search index to complete.
search_content = dict(index=SEARCH_INDEX_NAME, wildcard=dict(query="*", path=PAGE_CONTENT_FIELD, allowAnalyzedField=True))
n_docs = collection.count_documents({})
t0 = time()
while True:
if (time() - t0) > TIMEOUT:
raise TimeoutError(f'Search index {SEARCH_INDEX_NAME} did not complete in {TIMEOUT}')
results = collection.aggregate([{ "$search": search_content }])
if len(list(results)) == n_docs:
break
sleep(INTERVAL)

query = "When was the last time I visited new orleans?"
results = retriever.invoke(query)
assert "New Orleans" in results[0].page_content
Expand Down

0 comments on commit 2f5f1fc

Please sign in to comment.