langchain/libs/community/tests/integration_tests/utilities/test_arxiv.py

"""Integration test for Arxiv API Wrapper."""

from typing import Any, List

import pytest
from langchain_core.documents import Document
from langchain_core.tools import BaseTool

from langchain_community.tools import ArxivQueryRun
from langchain_community.utilities import ArxivAPIWrapper


@pytest.fixture
def api_client() -> ArxivAPIWrapper:
    return ArxivAPIWrapper()


def test_run_success_paper_name(api_client: ArxivAPIWrapper) -> None:
    """Test a query of paper name that returns the correct answer"""

    output = api_client.run("Heat-bath random walks with Markov bases")
    assert "Probability distributions for Markov chains based quantum walks" in output
    assert (
        "Transformations of random walks on groups via Markov stopping times" in output
    )
    assert (
        "Recurrence of Multidimensional Persistent Random Walks. Fourier and Series "
        "Criteria" in output
    )


def test_run_success_arxiv_identifier(api_client: ArxivAPIWrapper) -> None:
    """Test a query of an arxiv identifier returns the correct answer"""

    output = api_client.run("1605.08386v1")
    assert "Heat-bath random walks with Markov bases" in output


def test_run_success_multiple_arxiv_identifiers(api_client: ArxivAPIWrapper) -> None:
    """Test a query of multiple arxiv identifiers that returns the correct answer"""

    output = api_client.run("1605.08386v1 2212.00794v2 2308.07912")
    assert "Heat-bath random walks with Markov bases" in output
    assert "Scaling Language-Image Pre-training via Masking" in output
    assert (
        "Ultra-low mass PBHs in the early universe can explain the PTA signal" in output
    )


def test_run_returns_several_docs(api_client: ArxivAPIWrapper) -> None:
    """Test that returns several docs"""

    output = api_client.run("Caprice Stanley")
    assert "On Mixing Behavior of a Family of Random Walks" in output


def test_run_returns_no_result(api_client: ArxivAPIWrapper) -> None:
    """Test that gives no result."""

    output = api_client.run("1605.08386WWW")
    assert "No good Arxiv Result was found" == output


def assert_docs(docs: List[Document]) -> None:
    for doc in docs:
        assert doc.page_content
        assert doc.metadata
        assert set(doc.metadata) == {"Published", "Title", "Authors", "Summary"}


def test_load_success_paper_name(api_client: ArxivAPIWrapper) -> None:
    """Test a query of paper name that returns one document"""

    docs = api_client.load("Heat-bath random walks with Markov bases")
    assert len(docs) == 3
    assert_docs(docs)


def test_load_success_arxiv_identifier(api_client: ArxivAPIWrapper) -> None:
    """Test a query of an arxiv identifier that returns one document"""

    docs = api_client.load("1605.08386v1")
    assert len(docs) == 1
    assert_docs(docs)


def test_load_success_multiple_arxiv_identifiers(api_client: ArxivAPIWrapper) -> None:
    """Test a query of arxiv identifiers that returns the correct answer"""

    docs = api_client.load("1605.08386v1 2212.00794v2 2308.07912")
    assert len(docs) == 3
    assert_docs(docs)


def test_load_returns_no_result(api_client: ArxivAPIWrapper) -> None:
    """Test that returns no docs"""

    docs = api_client.load("1605.08386WWW")
    assert len(docs) == 0


def test_load_returns_limited_docs() -> None:
    """Test that returns several docs"""
    expected_docs = 2
    api_client = ArxivAPIWrapper(load_max_docs=expected_docs)
    docs = api_client.load("ChatGPT")
    assert len(docs) == expected_docs
    assert_docs(docs)


def test_load_returns_limited_doc_content_chars() -> None:
    """Test that returns limited doc_content_chars_max"""

    doc_content_chars_max = 100
    api_client = ArxivAPIWrapper(doc_content_chars_max=doc_content_chars_max)
    docs = api_client.load("1605.08386")
    assert len(docs[0].page_content) == doc_content_chars_max


def test_load_returns_unlimited_doc_content_chars() -> None:
    """Test that returns unlimited doc_content_chars_max"""

    doc_content_chars_max = None
    api_client = ArxivAPIWrapper(doc_content_chars_max=doc_content_chars_max)
    docs = api_client.load("1605.08386")
    assert len(docs[0].page_content) == pytest.approx(54338, rel=1e-2)


def test_load_returns_full_set_of_metadata() -> None:
    """Test that returns several docs"""
    api_client = ArxivAPIWrapper(load_max_docs=1, load_all_available_meta=True)
    docs = api_client.load("ChatGPT")
    assert len(docs) == 1
    for doc in docs:
        assert doc.page_content
        assert doc.metadata
        assert set(doc.metadata).issuperset(
            {"Published", "Title", "Authors", "Summary"}
        )
        print(doc.metadata)  # noqa: T201
        assert len(set(doc.metadata)) > 4


def _load_arxiv_from_universal_entry(**kwargs: Any) -> BaseTool:
    from langchain.agents.load_tools import load_tools

    tools = load_tools(["arxiv"], **kwargs)
    assert len(tools) == 1, "loaded more than 1 tool"
    return tools[0]


def test_load_arxiv_from_universal_entry() -> None:
    arxiv_tool = _load_arxiv_from_universal_entry()
    output = arxiv_tool("Caprice Stanley")
    assert (
        "On Mixing Behavior of a Family of Random Walks" in output
    ), "failed to fetch a valid result"


def test_load_arxiv_from_universal_entry_with_params() -> None:
    params = {
        "top_k_results": 1,
        "load_max_docs": 10,
        "load_all_available_meta": True,
    }
    arxiv_tool = _load_arxiv_from_universal_entry(**params)
    assert isinstance(arxiv_tool, ArxivQueryRun)
    wp = arxiv_tool.api_wrapper
    assert wp.top_k_results == 1, "failed to assert top_k_results"
    assert wp.load_max_docs == 10, "failed to assert load_max_docs"
    assert (
        wp.load_all_available_meta is True
    ), "failed to assert load_all_available_meta"