langchain/libs/partners/exa/tests/integration_tests/test_retriever.py
Erick Friis 3a2eb6e12b
infra: add print rule to ruff (#16221)
Added noqa for existing prints. Can slowly remove / will prevent more
being intro'd
2024-02-09 16:13:30 -08:00

27 lines
901 B
Python

from langchain_core.documents import Document
from langchain_exa import ExaSearchRetriever
def test_exa_retriever() -> None:
retriever = ExaSearchRetriever()
res = retriever.invoke("best time to visit japan")
print(res) # noqa: T201
assert len(res) == 10 # default k
assert isinstance(res, list)
assert isinstance(res[0], Document)
def test_exa_retriever_highlights() -> None:
retriever = ExaSearchRetriever(highlights=True)
res = retriever.invoke("best time to visit japan")
print(res) # noqa: T201
assert isinstance(res, list)
assert isinstance(res[0], Document)
highlights = res[0].metadata["highlights"]
highlight_scores = res[0].metadata["highlight_scores"]
assert isinstance(highlights, list)
assert isinstance(highlight_scores, list)
assert isinstance(highlights[0], str)
assert isinstance(highlight_scores[0], float)