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

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_s2_display.py

@@ -0,0 +1,74 @@
+"""
+Unit tests for S2 tools functionality.
+"""
+
+# pylint: disable=redefined-outer-name
+import pytest
+from langgraph.types import Command
+from ..tools.s2.display_results import (
+    display_results,
+    NoPapersFoundError as raised_error,
+)
+
+
+@pytest.fixture
+def initial_state():
+    """Provides an empty initial state for tests."""
+    return {"papers": {}, "multi_papers": {}}
+
+
+# Fixed test data for deterministic results
+MOCK_SEARCH_RESPONSE = {
+    "data": [
+        {
+            "paperId": "123",
+            "title": "Machine Learning Basics",
+            "abstract": "An introduction to ML",
+            "year": 2023,
+            "citationCount": 100,
+            "url": "https://example.com/paper1",
+            "authors": [{"name": "Test Author"}],
+        }
+    ]
+}
+
+MOCK_STATE_PAPER = {
+    "123": {
+        "Title": "Machine Learning Basics",
+        "Abstract": "An introduction to ML",
+        "Year": 2023,
+        "Citation Count": 100,
+        "URL": "https://example.com/paper1",
+    }
+}
+
+
+class TestS2Tools:
+    """Unit tests for individual S2 tools"""
+
+    def test_display_results_empty_state(self, initial_state):
+        """Verifies display_results tool behavior when state is empty and raises an exception"""
+        with pytest.raises(
+            raised_error,
+            match="No papers found. A search/rec needs to be performed first.",
+        ):
+            display_results.invoke({"state": initial_state, "tool_call_id": "test123"})
+
+    def test_display_results_shows_papers(self, initial_state):
+        """Verifies display_results tool correctly returns papers from state"""
+        state = initial_state.copy()
+        state["last_displayed_papers"] = "papers"
+        state["papers"] = MOCK_STATE_PAPER
+
+        result = display_results.invoke(
+            input={"state": state, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)  # Expect a Command object
+        assert isinstance(result.update, dict)  # Ensure update is a dictionary
+        assert "messages" in result.update
+        assert len(result.update["messages"]) == 1
+        assert (
+            "1 papers found. Papers are attached as an artifact."
+            in result.update["messages"][0].content
+        )

aiagents4pharma/talk2scholars/tests/test_s2_multi.py

@@ -0,0 +1,282 @@
+"""
+Unit tests for S2 tools functionality.
+"""
+
+import json
+from types import SimpleNamespace
+import pytest
+import requests
+from langgraph.types import Command
+from langchain_core.messages import ToolMessage
+import hydra
+from aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec import (
+    get_multi_paper_recommendations,
+)
+
+# --- Dummy Hydra Config Setup ---
+
+
+class DummyHydraContext:
+    """dummy context manager for mocking Hydra's initialize and compose functions."""
+
+    def __enter__(self):
+        """enter function that returns None."""
+        return None
+
+    def __exit__(self, exc_type, exc_val, traceback):
+        """exit function that does nothing."""
+        return None
+
+
+# Create a dummy configuration that mimics the expected hydra config.
+dummy_config = SimpleNamespace(
+    tools=SimpleNamespace(
+        multi_paper_recommendation=SimpleNamespace(
+            api_endpoint="http://dummy.endpoint/multi",
+            headers={"Content-Type": "application/json"},
+            api_fields=["paperId", "title", "authors"],
+            request_timeout=10,
+        )
+    )
+)
+
+# --- Dummy Response Classes and Functions for requests.post ---
+
+
+class DummyResponse:
+    """A dummy response class for mocking HTTP responses."""
+
+    def __init__(self, json_data, status_code=200):
+        """Initialize a DummyResponse with the given JSON data and status code."""
+        self._json_data = json_data
+        self.status_code = status_code
+
+    def json(self):
+        """Return the JSON data from the response."""
+        return self._json_data
+
+    def raise_for_status(self):
+        """raise an HTTP error for status codes >= 400."""
+        if self.status_code >= 400:
+            raise requests.HTTPError("HTTP Error")
+
+
+def test_dummy_response_no_error():
+    """Test that raise_for_status does not raise an exception for a successful response."""
+    # Create a DummyResponse with a successful status code.
+    response = DummyResponse({"data": "success"}, status_code=200)
+    # Calling raise_for_status should not raise an exception and should return None.
+    assert response.raise_for_status() is None
+
+
+def test_dummy_response_raise_error():
+    """Test that raise_for_status raises an exception for a failing response."""
+    # Create a DummyResponse with a failing status code.
+    response = DummyResponse({"error": "fail"}, status_code=400)
+    # Calling raise_for_status should raise an HTTPError.
+    with pytest.raises(requests.HTTPError):
+        response.raise_for_status()
+
+
+def dummy_requests_post_success(url, headers, params, data, timeout):
+    """dummy_requests_post_success"""
+    # Record call parameters for assertions.
+    dummy_requests_post_success.called_url = url
+    dummy_requests_post_success.called_headers = headers
+    dummy_requests_post_success.called_params = params
+    dummy_requests_post_success.called_data = data
+    dummy_requests_post_success.called_timeout = timeout
+
+    # Simulate a valid API response with three recommended papers;
+    # one paper missing authors should be filtered out.
+    dummy_data = {
+        "recommendedPapers": [
+            {
+                "paperId": "paperA",
+                "title": "Multi Rec Paper A",
+                "authors": ["Author X"],
+                "year": 2019,
+                "citationCount": 12,
+                "url": "http://paperA",
+                "externalIds": {"ArXiv": "arxivA"},
+            },
+            {
+                "paperId": "paperB",
+                "title": "Multi Rec Paper B",
+                "authors": ["Author Y"],
+                "year": 2020,
+                "citationCount": 18,
+                "url": "http://paperB",
+                "externalIds": {},
+            },
+            {
+                "paperId": "paperC",
+                "title": "Multi Rec Paper C",
+                "authors": None,  # This one should be filtered out.
+                "year": 2021,
+                "citationCount": 25,
+                "url": "http://paperC",
+                "externalIds": {"ArXiv": "arxivC"},
+            },
+        ]
+    }
+    return DummyResponse(dummy_data)
+
+
+def dummy_requests_post_unexpected(url, headers, params, data, timeout):
+    """dummy_requests_post_unexpected"""
+    dummy_requests_post_unexpected.called_url = url
+    dummy_requests_post_unexpected.called_headers = headers
+    dummy_requests_post_unexpected.called_params = params
+    dummy_requests_post_unexpected.called_data = data
+    dummy_requests_post_unexpected.called_timeout = timeout
+    # Simulate a response missing the 'recommendedPapers' key.
+    return DummyResponse({"error": "Invalid format"})
+
+
+def dummy_requests_post_no_recs(url, headers, params, data, timeout):
+    """dummy_requests_post_no_recs"""
+    dummy_requests_post_no_recs.called_url = url
+    dummy_requests_post_no_recs.called_headers = headers
+    dummy_requests_post_no_recs.called_params = params
+    dummy_requests_post_no_recs.called_data = data
+    dummy_requests_post_no_recs.called_timeout = timeout
+    # Simulate a response with an empty recommendations list.
+    return DummyResponse({"recommendedPapers": []})
+
+
+def dummy_requests_post_exception(url, headers, params, data, timeout):
+    """dummy_requests_post_exception"""
+    dummy_requests_post_exception.called_url = url
+    dummy_requests_post_exception.called_headers = headers
+    dummy_requests_post_exception.called_params = params
+    dummy_requests_post_exception.called_data = data
+    dummy_requests_post_exception.called_timeout = timeout
+    # Simulate a network exception.
+    raise requests.exceptions.RequestException("Connection error")
+
+
+# --- Pytest Fixture to Patch Hydra ---
+@pytest.fixture(autouse=True)
+def patch_hydra(monkeypatch):
+    """Patch Hydra's initialize and compose functions to return dummy objects."""
+    # Patch hydra.initialize to return our dummy context manager.
+    monkeypatch.setattr(
+        hydra, "initialize", lambda version_base, config_path: DummyHydraContext()
+    )
+    # Patch hydra.compose to return our dummy config.
+    monkeypatch.setattr(hydra, "compose", lambda config_name, overrides: dummy_config)
+
+
+# --- Test Cases ---
+
+
+def test_multi_paper_rec_success(monkeypatch):
+    """
+    Test that get_multi_paper_recommendations returns a valid Command object
+    when the API response is successful. Also, ensure that recommendations missing
+    required fields (like authors) are filtered out.
+    """
+    monkeypatch.setattr(requests, "post", dummy_requests_post_success)
+
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_ids": ["p1", "p2"],
+        "tool_call_id": tool_call_id,
+        "limit": 2,
+        "year": "2020",
+    }
+    # Call the tool using .run() with a dictionary input.
+    result = get_multi_paper_recommendations.run(input_data)
+
+    # Validate that the result is a Command with the expected update structure.
+    assert isinstance(result, Command)
+    update = result.update
+    assert "multi_papers" in update
+
+    papers = update["multi_papers"]
+    # Papers with valid 'title' and 'authors' should be included.
+    assert "paperA" in papers
+    assert "paperB" in papers
+    # Paper "paperC" is missing authors and should be filtered out.
+    assert "paperC" not in papers
+
+    # Check that a ToolMessage is included in the messages.
+    messages = update.get("messages", [])
+    assert len(messages) == 1
+    msg = messages[0]
+    assert isinstance(msg, ToolMessage)
+    assert "Recommendations based on multiple papers were successful" in msg.content
+
+    # Verify that the correct parameters were sent to requests.post.
+    called_params = dummy_requests_post_success.called_params
+    assert called_params["limit"] == 2  # Should be min(limit, 500)
+    assert called_params["fields"] == "paperId,title,authors"
+    # The year parameter should be present.
+    assert called_params["year"] == "2020"
+
+    # Also check the payload sent in the data.
+    sent_payload = json.loads(dummy_requests_post_success.called_data)
+    assert sent_payload["positivePaperIds"] == ["p1", "p2"]
+    assert sent_payload["negativePaperIds"] == []
+
+
+def test_multi_paper_rec_unexpected_format(monkeypatch):
+    """
+    Test that get_multi_paper_recommendations raises a RuntimeError when the API
+    response does not include the expected 'recommendedPapers' key.
+    """
+    monkeypatch.setattr(requests, "post", dummy_requests_post_unexpected)
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_ids": ["p1", "p2"],
+        "tool_call_id": tool_call_id,
+    }
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "Unexpected response from Semantic Scholar API. The results could not be "
+            "retrieved due to an unexpected format. "
+            "Please modify your search query and try again."
+        ),
+    ):
+        get_multi_paper_recommendations.run(input_data)
+
+
+def test_multi_paper_rec_no_recommendations(monkeypatch):
+    """
+    Test that get_multi_paper_recommendations raises a RuntimeError when the API
+    returns no recommendations.
+    """
+    monkeypatch.setattr(requests, "post", dummy_requests_post_no_recs)
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_ids": ["p1", "p2"],
+        "tool_call_id": tool_call_id,
+    }
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "No recommendations were found for your query. Consider refining your search "
+            "by using more specific keywords or different terms."
+        ),
+    ):
+        get_multi_paper_recommendations.run(input_data)
+
+
+def test_multi_paper_rec_requests_exception(monkeypatch):
+    """
+    Test that get_multi_paper_recommendations raises a RuntimeError when requests.post
+    throws an exception.
+    """
+    monkeypatch.setattr(requests, "post", dummy_requests_post_exception)
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_ids": ["p1", "p2"],
+        "tool_call_id": tool_call_id,
+    }
+    with pytest.raises(
+        RuntimeError,
+        match="Failed to connect to Semantic Scholar API. Please retry the same query.",
+    ):
+        get_multi_paper_recommendations.run(input_data)