aiagents4pharma 1.28.0__py3-none-any.whl → 1.30.0__py3-none-any.whl

aiagents4pharma/talk2scholars/configs/tools/zotero_read/default.yaml

@@ -12,4 +12,45 @@ search_params:
 # Item Types and Limit
 zotero:
   max_limit: 100
-  filter_item_types: ["journalArticle", "conferencePaper", "preprint"]
+  filter_item_types:
+    [
+      "Artwork",
+      "Audio Recording",
+      "Bill",
+      "Blog Post",
+      "Book",
+      "Book Section",
+      "Case",
+      "Conference Paper",
+      "Dataset",
+      "Dictionary Entry",
+      "Document",
+      "E-mail",
+      "Encyclopedia Article",
+      "Film",
+      "Forum Post",
+      "Hearing",
+      "Instant Message",
+      "Interview",
+      "Journal Article",
+      "Letter",
+      "Magazine Article",
+      "Manuscript",
+      "Map",
+      "Newspaper Article",
+      "Patent",
+      "Podcast",
+      "Preprint",
+      "Presentation",
+      "Radio Broadcast",
+      "Report",
+      "Software",
+      "Standard",
+      "Statute",
+      "Thesis",
+      "TV Broadcast",
+      "Video Recording",
+      "Web Page",
+    ]
+
+  filter_excluded_types: ["attachment", "note", "annotation"]

aiagents4pharma/talk2scholars/configs/tools/zotero_write/__inti__.py

@@ -0,0 +1,3 @@
+"""
+Import all the modules in the package
+"""

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -63,3 +63,4 @@ class Talk2Scholars(AgentState):
     pdf_data: Annotated[Dict[str, Any], replace_dict]
     zotero_read: Annotated[Dict[str, Any], replace_dict]
     llm_model: BaseChatModel
+    pdf_data: Annotated[Dict[str, Any], replace_dict]

aiagents4pharma/talk2scholars/tests/test_main_agent.py

@@ -3,119 +3,194 @@ Unit tests for main agent functionality.
 Tests the supervisor agent's routing logic and state management.
 """
 
-# pylint: disable=redefined-outer-name
 # pylint: disable=redefined-outer-name,too-few-public-methods
-import random
-from unittest.mock import Mock, patch, MagicMock
+
+from types import SimpleNamespace
 import pytest
-from langchain_core.messages import HumanMessage, AIMessage
-from langgraph.graph import END
-from ..agents.main_agent import make_supervisor_node, get_hydra_config, get_app
-from ..state.state_talk2scholars import Talk2Scholars
+import hydra
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_openai import ChatOpenAI
+from pydantic import Field
+from aiagents4pharma.talk2scholars.agents.main_agent import get_app
+
+# --- Dummy LLM Implementation ---
+
+
+class DummyLLM(BaseChatModel):
+    """A dummy language model implementation for testing purposes."""
+
+    model_name: str = Field(...)
+
+    def _generate(self, prompt, stop=None):
+        """Generate a response given a prompt."""
+        DummyLLM.called_prompt = prompt
+        return "dummy output"
+
+    @property
+    def _llm_type(self):
+        """Return the type of the language model."""
+        return "dummy"
+
+
+# --- Dummy Workflow and Sub-agent Functions ---
+
+
+class DummyWorkflow:
+    """A dummy workflow class that records arguments for verification."""
+
+    def __init__(self, supervisor_args=None):
+        """Initialize the workflow with the given supervisor arguments."""
+        self.supervisor_args = supervisor_args or {}
+        self.checkpointer = None
+        self.name = None
+
+    def compile(self, checkpointer, name):
+        """Compile the workflow with the given checkpointer and name."""
+        self.checkpointer = checkpointer
+        self.name = name
+        return self
+
+
+def dummy_get_app_s2(uniq_id, llm_model):
+    """Return a DummyWorkflow for the S2 agent."""
+    dummy_get_app_s2.called_uniq_id = uniq_id
+    dummy_get_app_s2.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "s2", "uniq_id": uniq_id})
+
+
+def dummy_get_app_zotero(uniq_id, llm_model):
+    """Return a DummyWorkflow for the Zotero agent."""
+    dummy_get_app_zotero.called_uniq_id = uniq_id
+    dummy_get_app_zotero.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "zotero", "uniq_id": uniq_id})
+
+
+def dummy_create_supervisor(apps, model, state_schema, **kwargs):
+    """Return a DummyWorkflow for the supervisor."""
+    dummy_create_supervisor.called_kwargs = kwargs
+    return DummyWorkflow(
+        supervisor_args={
+            "apps": apps,
+            "model": model,
+            "state_schema": state_schema,
+            **kwargs,
+        }
+    )
+
+
+# --- Dummy Hydra Configuration Setup ---
+
+
+class DummyHydraContext:
+    """A dummy context manager for mocking Hydra's initialize and compose functions."""
+
+    def __enter__(self):
+        """Return None when entering the context."""
+        return None
+
+    def __exit__(self, exc_type, exc_val, traceback):
+        """Exit function that does nothing."""
+        return None
+
+
+def dict_to_namespace(d):
+    """Convert a dictionary to a SimpleNamespace object."""
+    return SimpleNamespace(
+        **{
+            key: dict_to_namespace(val) if isinstance(val, dict) else val
+            for key, val in d.items()
+        }
+    )
+
+
+dummy_config = {
+    "agents": {
+        "talk2scholars": {"main_agent": {"system_prompt": "Dummy system prompt"}}
+    }
+}
+
+
+class DummyHydraCompose:
+    """A dummy class that returns a namespace from a dummy config dictionary."""
+
+    def __init__(self, config):
+        """Constructor that stores the dummy config."""
+        self.config = config
+
+    def __getattr__(self, item):
+        """Return a namespace from the dummy config."""
+        return dict_to_namespace(self.config.get(item, {}))
+
+
+# --- Pytest Fixtures to Patch Dependencies ---
+
+
+@pytest.fixture(autouse=True)
+def patch_hydra(monkeypatch):
+    """Patch the hydra.initialize and hydra.compose functions to return dummy objects."""
+    monkeypatch.setattr(
+        hydra, "initialize", lambda version_base, config_path: DummyHydraContext()
+    )
+    monkeypatch.setattr(
+        hydra, "compose", lambda config_name, overrides: DummyHydraCompose(dummy_config)
+    )
 
 
 @pytest.fixture(autouse=True)
-def mock_hydra():
-    """Mock Hydra configuration."""
-    with patch("hydra.initialize"), patch("hydra.compose") as mock_compose:
-        cfg_mock = MagicMock()
-        cfg_mock.agents.talk2scholars.main_agent.temperature = 0
-        cfg_mock.agents.talk2scholars.main_agent.system_prompt = "System prompt"
-        cfg_mock.agents.talk2scholars.main_agent.router_prompt = "Router prompt"
-        mock_compose.return_value = cfg_mock
-        yield mock_compose
-
-
-def test_get_app():
-    """Test the full initialization of the LangGraph application."""
-    thread_id = "test_thread"
-    mock_llm = Mock()
-    app = get_app(thread_id, mock_llm)
-    assert app is not None
-    assert "supervisor" in app.nodes
-    assert "s2_agent" in app.nodes  # Ensure nodes exist
-    assert "zotero_agent" in app.nodes
-
-
-def test_get_app_with_default_llm():
-    """Test app initialization with default LLM parameters."""
-    thread_id = "test_thread"
-    llm_mock = Mock()
-
-    # We need to explicitly pass the mock instead of patching, since the function uses
-    # ChatOpenAI as a default argument value which is evaluated at function definition time
-    app = get_app(thread_id, llm_mock)
-    assert app is not None
-    # We can only verify the app was created successfully
-
-
-def test_get_hydra_config():
-    """Test that Hydra configuration loads correctly."""
-    with patch("hydra.initialize"), patch("hydra.compose") as mock_compose:
-        cfg_mock = MagicMock()
-        cfg_mock.agents.talk2scholars.main_agent.temperature = 0
-        mock_compose.return_value = cfg_mock
-        cfg = get_hydra_config()
-        assert cfg is not None
-        assert cfg.temperature == 0
-
-
-def test_hydra_failure():
-    """Test exception handling when Hydra fails to load config."""
-    thread_id = "test_thread"
-    llm_mock = Mock()
-    with patch("hydra.initialize", side_effect=Exception("Hydra error")):
-        with pytest.raises(Exception) as exc_info:
-            get_app(thread_id, llm_model=llm_mock)
-        assert "Hydra error" in str(exc_info.value)
-
-
-def test_supervisor_node_execution():
-    """Test that the supervisor node routes correctly."""
-    mock_llm = Mock()
-    thread_id = "test_thread"
-
-    class MockRouter:
-        """Mock router class."""
-
-        next = random.choice(["s2_agent", "zotero_agent"])
-
-    with (
-        patch.object(mock_llm, "with_structured_output", return_value=mock_llm),
-        patch.object(mock_llm, "invoke", return_value=MockRouter()),
-    ):
-        supervisor_node = make_supervisor_node(mock_llm, thread_id)
-        mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
-        result = supervisor_node(mock_state)
-
-        # Accept either "s2_agent" or "zotero_agent"
-        assert result.goto in ["s2_agent", "zotero_agent"]
-
-
-def test_supervisor_node_finish():
-    """Test that supervisor node correctly handles FINISH case."""
-    mock_llm = Mock()
-    thread_id = "test_thread"
-
-    class MockRouter:
-        """Mock router class."""
-
-        next = "FINISH"
-
-    class MockAIResponse:
-        """Mock AI response class."""
-
-        def __init__(self):
-            self.content = "Final AI Response"
-
-    with (
-        patch.object(mock_llm, "with_structured_output", return_value=mock_llm),
-        patch.object(mock_llm, "invoke", side_effect=[MockRouter(), MockAIResponse()]),
-    ):
-        supervisor_node = make_supervisor_node(mock_llm, thread_id)
-        mock_state = Talk2Scholars(messages=[HumanMessage(content="End conversation")])
-        result = supervisor_node(mock_state)
-        assert result.goto == END
-        assert "messages" in result.update
-        assert isinstance(result.update["messages"], AIMessage)
-        assert result.update["messages"].content == "Final AI Response"
+def patch_sub_agents_and_supervisor(monkeypatch):
+    """Patch the sub-agents and supervisor creation functions."""
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_s2", dummy_get_app_s2
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_zotero",
+        dummy_get_app_zotero,
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.create_supervisor",
+        dummy_create_supervisor,
+    )
+
+
+# --- Test Cases ---
+
+
+def test_dummy_llm_generate():
+    """Test the dummy LLM's generate function."""
+    dummy = DummyLLM(model_name="test-model")
+    output = getattr(dummy, "_generate")("any prompt")
+    assert output == "dummy output"
+
+
+def test_dummy_llm_llm_type():
+    """Test the dummy LLM's _llm_type property."""
+    dummy = DummyLLM(model_name="test-model")
+    assert getattr(dummy, "_llm_type") == "dummy"
+
+
+def test_get_app_with_gpt4o_mini():
+    """
+    Test that get_app replaces a 'gpt-4o-mini' LLM with a new ChatOpenAI instance.
+    """
+    uniq_id = "test_thread"
+    dummy_llm = DummyLLM(model_name="gpt-4o-mini")
+    app = get_app(uniq_id, dummy_llm)
+
+    supervisor_args = getattr(app, "supervisor_args", {})
+    assert isinstance(supervisor_args.get("model"), ChatOpenAI)
+    assert supervisor_args.get("prompt") == "Dummy system prompt"
+    assert getattr(app, "name", "") == "Talk2Scholars_MainAgent"
+
+
+def test_get_app_with_other_model():
+    """
+    Test that get_app does not replace the LLM if its model_name is not 'gpt-4o-mini'.
+    """
+    uniq_id = "test_thread_2"
+    dummy_llm = DummyLLM(model_name="other-model")
+    app = get_app(uniq_id, dummy_llm)
+
+    supervisor_args = getattr(app, "supervisor_args", {})
+    assert supervisor_args.get("model") is dummy_llm
+    assert supervisor_args.get("prompt") == "Dummy system prompt"
+    assert getattr(app, "name", "") == "Talk2Scholars_MainAgent"

aiagents4pharma/talk2scholars/tests/test_paper_download_agent.py

@@ -0,0 +1,142 @@
+"""Unit tests for the paper download agent in Talk2Scholars."""
+
+from unittest import mock
+import pytest
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.language_models.chat_models import BaseChatModel
+from ..agents.paper_download_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+
+@pytest.fixture(autouse=True)
+def mock_hydra_fixture():
+    """Mocks Hydra configuration for tests."""
+    with mock.patch("hydra.initialize"), mock.patch("hydra.compose") as mock_compose:
+        cfg_mock = mock.MagicMock()
+        cfg_mock.agents.talk2scholars.s2_agent.temperature = 0
+        cfg_mock.agents.talk2scholars.paper_download_agent.prompt = "Test prompt"
+        mock_compose.return_value = cfg_mock
+        yield mock_compose
+
+
+@pytest.fixture
+def mock_tools_fixture():
+    """Mocks paper download tools to prevent real HTTP calls."""
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.paper_download."
+            "download_arxiv_input.download_arxiv_paper"
+        ) as mock_download_arxiv_paper,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.query_results.query_results"
+        ) as mock_query_results,
+    ):
+        mock_download_arxiv_paper.return_value = {
+            "pdf_data": {"dummy_key": "dummy_value"}
+        }
+        mock_query_results.return_value = {
+        "result": "Mocked Query Result"
+        }
+        yield [mock_download_arxiv_paper, mock_query_results]
+
+@pytest.mark.usefixtures("mock_hydra_fixture")
+def test_paper_download_agent_initialization():
+    """Ensures the paper download agent initializes properly with a prompt."""
+    thread_id = "test_thread_paper_dl"
+    llm_mock = mock.Mock(spec=BaseChatModel)  # Mock LLM
+
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.paper_download_agent.create_react_agent"
+    ) as mock_create_agent:
+        mock_create_agent.return_value = mock.Mock()
+
+        app = get_app(thread_id, llm_mock)
+        assert app is not None, "The agent app should be successfully created."
+        assert mock_create_agent.called
+
+def test_paper_download_agent_invocation():
+    """Verifies agent processes queries and updates state correctly."""
+    _ = mock_tools_fixture  # Prevents unused-argument warning
+    thread_id = "test_thread_paper_dl"
+    mock_state = Talk2Scholars(
+        messages=[HumanMessage(content="Download paper 1234.5678")]
+    )
+    llm_mock = mock.Mock(spec=BaseChatModel)
+
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.paper_download_agent.create_react_agent"
+    ) as mock_create_agent:
+        mock_agent = mock.Mock()
+        mock_create_agent.return_value = mock_agent
+        mock_agent.invoke.return_value = {
+            "messages": [AIMessage(content="Here is the paper")],
+            "pdf_data": {"file_bytes": b"FAKE_PDF_CONTENTS"},
+        }
+
+
+        app = get_app(thread_id, llm_mock)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+        assert "messages" in result
+        assert "pdf_data" in result
+
+
+def test_paper_download_agent_tools_assignment(request):  # Keep fixture name
+    """Checks correct tool assignment (download_arxiv_paper, query_results)."""
+    thread_id = "test_thread_paper_dl"
+    mock_tools = request.getfixturevalue("mock_tools_fixture")
+    llm_mock = mock.Mock(spec=BaseChatModel)
+
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.paper_download_agent.create_react_agent"
+        ) as mock_create_agent,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.paper_download_agent.ToolNode"
+        ) as mock_toolnode,
+    ):
+        mock_agent = mock.Mock()
+        mock_create_agent.return_value = mock_agent
+        mock_tool_instance = mock.Mock()
+        mock_tool_instance.tools = mock_tools
+        mock_toolnode.return_value= mock_tool_instance
+
+        get_app(thread_id, llm_mock)
+        assert mock_toolnode.called
+        assert len(mock_tool_instance.tools) == 2
+
+
+def test_paper_download_agent_hydra_failure():
+    """Confirms the agent gracefully handles exceptions if Hydra fails."""
+    thread_id = "test_thread_paper_dl"
+    llm_mock = mock.Mock(spec=BaseChatModel)
+
+    with mock.patch("hydra.initialize", side_effect=Exception("Mock Hydra failure")):
+        with pytest.raises(Exception) as exc_info:
+            get_app(thread_id, llm_mock)
+        assert "Mock Hydra failure" in str(exc_info.value)
+
+
+def test_paper_download_agent_model_failure():
+    """Ensures agent handles model-related failures gracefully."""
+    thread_id = "test_thread_paper_dl"
+    llm_mock = mock.Mock(spec=BaseChatModel)
+
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.paper_download_agent.create_react_agent",
+        side_effect=Exception("Mock model failure"),
+    ):
+        with pytest.raises(Exception) as exc_info:
+            get_app(thread_id, llm_mock)
+        assert "Mock model failure" in str(exc_info.value), (
+            "Model initialization failure should raise an exception."
+        )

aiagents4pharma/talk2scholars/tests/test_paper_download_tools.py

@@ -0,0 +1,154 @@
+"""
+Unit tests for arXiv paper downloading functionality, including:
+- AbstractPaperDownloader (base class)
+- ArxivPaperDownloader (arXiv-specific implementation)
+- download_arxiv_paper tool function.
+"""
+
+from unittest.mock import patch, MagicMock
+import pytest
+import requests
+from requests.exceptions import HTTPError
+from langgraph.types import Command
+from langchain_core.messages import ToolMessage
+
+# Import the classes and function under test
+from aiagents4pharma.talk2scholars.tools.paper_download.abstract_downloader import (
+    AbstractPaperDownloader,
+)
+from aiagents4pharma.talk2scholars.tools.paper_download.arxiv_downloader import (
+    ArxivPaperDownloader,
+)
+from aiagents4pharma.talk2scholars.tools.paper_download.download_arxiv_input import (
+    download_arxiv_paper,
+)
+
+@pytest.mark.parametrize("class_obj", [AbstractPaperDownloader])
+
+def test_abstract_downloader_cannot_be_instantiated(class_obj):
+    """
+    Validates that AbstractPaperDownloader is indeed abstract and raises TypeError
+    if anyone attempts to instantiate it directly.
+    """
+    with pytest.raises(TypeError):
+        class_obj()
+
+
+@pytest.fixture(name="arxiv_downloader_fixture")
+@pytest.mark.usefixtures("mock_hydra_config_setup")
+def fixture_arxiv_downloader():
+    """
+    Provides an ArxivPaperDownloader instance with a mocked Hydra config.
+    """
+    return ArxivPaperDownloader()
+
+
+def test_fetch_metadata_success(arxiv_downloader_fixture,):
+    """
+    Ensures fetch_metadata retrieves XML data correctly, given a successful HTTP response.
+    """
+    mock_response = MagicMock()
+    mock_response.text = "<xml>Mock ArXiv Metadata</xml>"
+    mock_response.raise_for_status = MagicMock()
+
+    with patch.object(requests, "get", return_value=mock_response) as mock_get:
+        paper_id = "1234.5678"
+        result = arxiv_downloader_fixture.fetch_metadata(paper_id)
+        mock_get.assert_called_once_with(
+            "http://export.arxiv.org/api/query?search_query=id:1234.5678&start=0&max_results=1",
+            timeout=10,
+        )
+        assert result["xml"] == "<xml>Mock ArXiv Metadata</xml>"
+
+
+def test_fetch_metadata_http_error(arxiv_downloader_fixture):
+    """
+    Validates that fetch_metadata raises HTTPError when the response indicates a failure.
+    """
+    mock_response = MagicMock()
+    mock_response.raise_for_status.side_effect = HTTPError("Mocked HTTP failure")
+
+    with patch.object(requests, "get", return_value=mock_response):
+        with pytest.raises(HTTPError):
+            arxiv_downloader_fixture.fetch_metadata("invalid_id")
+
+
+def test_download_pdf_success(arxiv_downloader_fixture):
+    """
+    Tests that download_pdf fetches the PDF link from metadata and successfully
+    retrieves the binary content.
+    """
+    mock_metadata = {
+        "xml": """
+        <feed xmlns="http://www.w3.org/2005/Atom">
+            <entry>
+                <link title="pdf" href="http://test.arxiv.org/pdf/1234.5678v1.pdf"/>
+            </entry>
+        </feed>
+        """
+    }
+
+    mock_pdf_response = MagicMock()
+    mock_pdf_response.raise_for_status = MagicMock()
+    mock_pdf_response.iter_content = lambda chunk_size: [b"FAKE_PDF_CONTENT"]
+
+    with patch.object(arxiv_downloader_fixture, "fetch_metadata", return_value=mock_metadata):
+        with patch.object(requests, "get", return_value=mock_pdf_response) as mock_get:
+            result = arxiv_downloader_fixture.download_pdf("1234.5678")
+            assert result["pdf_object"] == b"FAKE_PDF_CONTENT"
+            assert result["pdf_url"] == "http://test.arxiv.org/pdf/1234.5678v1.pdf"
+            assert result["arxiv_id"] == "1234.5678"
+            mock_get.assert_called_once_with(
+                "http://test.arxiv.org/pdf/1234.5678v1.pdf",
+                stream=True,
+                timeout=10,
+            )
+
+
+def test_download_pdf_no_pdf_link(arxiv_downloader_fixture):
+    """
+    Ensures a RuntimeError is raised if no <link> with title="pdf" is found in the XML.
+    """
+    mock_metadata = {"xml": "<feed></feed>"}
+
+    with patch.object(arxiv_downloader_fixture, "fetch_metadata", return_value=mock_metadata):
+        with pytest.raises(RuntimeError, match="Failed to download PDF"):
+            arxiv_downloader_fixture.download_pdf("1234.5678")
+
+
+def test_download_arxiv_paper_tool_success(arxiv_downloader_fixture):
+    """
+    Validates download_arxiv_paper orchestrates the ArxivPaperDownloader correctly,
+    returning a Command with PDF data and success messages.
+    """
+    mock_metadata = {"xml": "<mockxml></mockxml>"}
+    mock_pdf_response = {
+        "pdf_object": b"FAKE_PDF_CONTENT",
+        "pdf_url": "http://test.arxiv.org/mock.pdf",
+        "arxiv_id": "9999.8888",
+    }
+
+    with patch(
+        "aiagents4pharma.talk2scholars.tools.paper_download.download_arxiv_input."
+        "ArxivPaperDownloader",
+        return_value=arxiv_downloader_fixture,
+    ):
+        with patch.object(arxiv_downloader_fixture, "fetch_metadata", return_value=mock_metadata):
+            with patch.object(
+                arxiv_downloader_fixture,
+                "download_pdf",
+                return_value=mock_pdf_response,
+            ):
+                command_result = download_arxiv_paper.invoke(
+                    {"arxiv_id": "9999.8888", "tool_call_id": "test_tool_call"}
+                )
+
+                assert isinstance(command_result, Command)
+                assert "pdf_data" in command_result.update
+                assert command_result.update["pdf_data"] == mock_pdf_response
+
+                messages = command_result.update.get("messages", [])
+                assert len(messages) == 1
+                assert isinstance(messages[0], ToolMessage)
+                assert "Successfully downloaded PDF" in messages[0].content
+                assert "9999.8888" in messages[0].content