mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
b051bba1a9
- **Description:** finishes adding the you.com functionality including:
- add async functions to utility and retriever
- add the You.com Tool
- add async testing for utility, retriever, and tool
- add a tool integration notebook page
- **Dependencies:** any dependencies required for this change
- **Twitter handle:** @scottnath
112 lines
4.0 KiB
Python
112 lines
4.0 KiB
Python
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
import responses
|
|
|
|
from langchain_community.retrievers.you import YouRetriever
|
|
|
|
from ..utilities.test_you import (
|
|
LIMITED_PARSED_OUTPUT,
|
|
MOCK_PARSED_OUTPUT,
|
|
MOCK_RESPONSE_RAW,
|
|
NEWS_RESPONSE_PARSED,
|
|
NEWS_RESPONSE_RAW,
|
|
TEST_ENDPOINT,
|
|
)
|
|
|
|
|
|
class TestYouRetriever:
|
|
@responses.activate
|
|
def test_get_relevant_documents(self) -> None:
|
|
responses.add(
|
|
responses.GET, f"{TEST_ENDPOINT}/search", json=MOCK_RESPONSE_RAW, status=200
|
|
)
|
|
query = "Test query text"
|
|
you_wrapper = YouRetriever(ydc_api_key="test")
|
|
results = you_wrapper.get_relevant_documents(query)
|
|
expected_result = MOCK_PARSED_OUTPUT
|
|
assert results == expected_result
|
|
|
|
@responses.activate
|
|
def test_invoke(self) -> None:
|
|
responses.add(
|
|
responses.GET, f"{TEST_ENDPOINT}/search", json=MOCK_RESPONSE_RAW, status=200
|
|
)
|
|
query = "Test query text"
|
|
you_wrapper = YouRetriever(ydc_api_key="test")
|
|
results = you_wrapper.invoke(query)
|
|
expected_result = MOCK_PARSED_OUTPUT
|
|
assert results == expected_result
|
|
|
|
@responses.activate
|
|
def test_invoke_max_docs(self) -> None:
|
|
responses.add(
|
|
responses.GET, f"{TEST_ENDPOINT}/search", json=MOCK_RESPONSE_RAW, status=200
|
|
)
|
|
query = "Test query text"
|
|
you_wrapper = YouRetriever(k=2, ydc_api_key="test")
|
|
results = you_wrapper.invoke(query)
|
|
expected_result = [MOCK_PARSED_OUTPUT[0], MOCK_PARSED_OUTPUT[1]]
|
|
assert results == expected_result
|
|
|
|
@responses.activate
|
|
def test_invoke_limit_snippets(self) -> None:
|
|
responses.add(
|
|
responses.GET, f"{TEST_ENDPOINT}/search", json=MOCK_RESPONSE_RAW, status=200
|
|
)
|
|
|
|
query = "Test query text"
|
|
you_wrapper = YouRetriever(n_snippets_per_hit=1, ydc_api_key="test")
|
|
results = you_wrapper.results(query)
|
|
expected_result = LIMITED_PARSED_OUTPUT
|
|
assert results == expected_result
|
|
|
|
@responses.activate
|
|
def test_invoke_news(self) -> None:
|
|
responses.add(
|
|
responses.GET, f"{TEST_ENDPOINT}/news", json=NEWS_RESPONSE_RAW, status=200
|
|
)
|
|
|
|
query = "Test news text"
|
|
# ensure limit on number of docs returned
|
|
you_wrapper = YouRetriever(endpoint_type="news", ydc_api_key="test")
|
|
results = you_wrapper.results(query)
|
|
expected_result = NEWS_RESPONSE_PARSED
|
|
assert results == expected_result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aget_relevant_documents(self) -> None:
|
|
instance = YouRetriever(ydc_api_key="test_api_key")
|
|
|
|
# Mock response object to simulate aiohttp response
|
|
mock_response = AsyncMock()
|
|
mock_response.__aenter__.return_value = (
|
|
mock_response # Make the context manager return itself
|
|
)
|
|
mock_response.__aexit__.return_value = None # No value needed for exit
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(return_value=MOCK_RESPONSE_RAW)
|
|
|
|
# Patch the aiohttp.ClientSession object
|
|
with patch("aiohttp.ClientSession.get", return_value=mock_response):
|
|
results = await instance.aget_relevant_documents("test query")
|
|
assert results == MOCK_PARSED_OUTPUT
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ainvoke(self) -> None:
|
|
instance = YouRetriever(ydc_api_key="test_api_key")
|
|
|
|
# Mock response object to simulate aiohttp response
|
|
mock_response = AsyncMock()
|
|
mock_response.__aenter__.return_value = (
|
|
mock_response # Make the context manager return itself
|
|
)
|
|
mock_response.__aexit__.return_value = None # No value needed for exit
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(return_value=MOCK_RESPONSE_RAW)
|
|
|
|
# Patch the aiohttp.ClientSession object
|
|
with patch("aiohttp.ClientSession.get", return_value=mock_response):
|
|
results = await instance.ainvoke("test query")
|
|
assert results == MOCK_PARSED_OUTPUT
|