mirror of
https://github.com/hwchase17/langchain
synced 2024-10-29 17:07:25 +00:00
82f3e32d8d
Multiple people have asked in #5081 for a way to limit the documents returned from an AzureCognitiveSearchRetriever. This PR adds the `top_n` parameter to allow that. Twitter handle: [@UmerHAdil](twitter.com/umerHAdil)
29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
"""Test Azure Cognitive Search wrapper."""
|
|
import pytest
|
|
|
|
from langchain.retrievers.azure_cognitive_search import AzureCognitiveSearchRetriever
|
|
from langchain.schema import Document
|
|
|
|
|
|
def test_azure_cognitive_search_get_relevant_documents() -> None:
|
|
"""Test valid call to Azure Cognitive Search."""
|
|
retriever = AzureCognitiveSearchRetriever()
|
|
documents = retriever.get_relevant_documents("what is langchain")
|
|
for doc in documents:
|
|
assert isinstance(doc, Document)
|
|
assert doc.page_content
|
|
|
|
retriever = AzureCognitiveSearchRetriever(top_k=1)
|
|
documents = retriever.get_relevant_documents("what is langchain")
|
|
assert len(documents) <= 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_cognitive_search_aget_relevant_documents() -> None:
|
|
"""Test valid async call to Azure Cognitive Search."""
|
|
retriever = AzureCognitiveSearchRetriever()
|
|
documents = await retriever.aget_relevant_documents("what is langchain")
|
|
for doc in documents:
|
|
assert isinstance(doc, Document)
|
|
assert doc.page_content
|