langchain/libs/partners/exa/tests/integration_tests/test_retriever.py
2024-01-24 20:57:17 -07:00

27 lines
873 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)
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)
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)