aiagents4pharma 1.19.1__py3-none-any.whl → 1.20.1__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_integration.py

@@ -0,0 +1,237 @@
+"""
+Integration tests for talk2scholars system.
+
+These tests ensure that:
+1. The main agent and sub-agent work together.
+2. The agents correctly interact with tools (search, recommendations).
+3. The full pipeline processes queries and updates state correctly.
+"""
+
+# pylint: disable=redefined-outer-name
+from unittest.mock import patch, Mock
+import pytest
+from langchain_core.messages import HumanMessage
+from ..agents.main_agent import get_app as get_main_app
+from ..agents.s2_agent import get_app as get_s2_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+
+@pytest.fixture(autouse=True)
+def mock_hydra():
+    """Mock Hydra configuration to prevent external dependencies."""
+    with patch("hydra.initialize"), patch("hydra.compose") as mock_compose:
+        cfg_mock = Mock()
+        cfg_mock.agents.talk2scholars.main_agent.temperature = 0
+        cfg_mock.agents.talk2scholars.main_agent.main_agent = "Test main agent prompt"
+        cfg_mock.agents.talk2scholars.s2_agent.temperature = 0
+        cfg_mock.agents.talk2scholars.s2_agent.s2_agent = "Test s2 agent prompt"
+        mock_compose.return_value = cfg_mock
+        yield mock_compose
+
+
+@pytest.fixture(autouse=True)
+def mock_tools():
+    """Mock tools to prevent execution of real API calls."""
+    with (
+        patch(
+            "aiagents4pharma.talk2scholars.tools.s2.search.search_tool"
+        ) as mock_s2_search,
+        patch(
+            "aiagents4pharma.talk2scholars.tools.s2.display_results.display_results"
+        ) as mock_s2_display,
+        patch(
+            "aiagents4pharma.talk2scholars.tools.s2.single_paper_rec."
+            "get_single_paper_recommendations"
+        ) as mock_s2_single_rec,
+        patch(
+            "aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec."
+            "get_multi_paper_recommendations"
+        ) as mock_s2_multi_rec,
+    ):
+
+        mock_s2_search.return_value = {"papers": {"id123": "Mock Paper"}}
+        mock_s2_display.return_value = "Displaying Mock Results"
+        mock_s2_single_rec.return_value = {"recommendations": ["Paper A", "Paper B"]}
+        mock_s2_multi_rec.return_value = {
+            "multi_recommendations": ["Paper X", "Paper Y"]
+        }
+
+        yield {
+            "search_tool": mock_s2_search,
+            "display_results": mock_s2_display,
+            "single_paper_rec": mock_s2_single_rec,
+            "multi_paper_rec": mock_s2_multi_rec,
+        }
+
+
+def test_full_workflow():
+    """Test the full workflow from main agent to S2 agent."""
+    thread_id = "test_thread"
+    main_app = get_main_app(thread_id)
+
+    # Define expected mock response with the actual structure
+    expected_paper = {
+        "530a059cb48477ad1e3d4f8f4b153274c8997332": {
+            "Title": "Explainable Artificial Intelligence",
+            "Abstract": None,
+            "Citation Count": 5544,
+            "Year": "2024",
+            "URL": "https://example.com/paper",
+        }
+    }
+
+    # Mock the search tool instead of the app
+    with patch(
+        "aiagents4pharma.talk2scholars.tools.s2.search.search_tool",
+        return_value={"papers": expected_paper},
+    ):
+        state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
+        result = main_app.invoke(
+            state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+    # Check values
+    assert "papers" in result
+    assert "messages" in result
+    assert len(result["papers"]) > 0
+
+
+def test_s2_agent_execution():
+    """Test if the S2 agent processes requests correctly and updates state."""
+    thread_id = "test_thread"
+    s2_app = get_s2_app(thread_id)
+
+    state = Talk2Scholars(messages=[HumanMessage(content="Get recommendations")])
+
+    result = s2_app.invoke(
+        state,
+        config={
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": "test_ns",
+                "checkpoint_id": "test_checkpoint",
+            }
+        },
+    )
+
+    assert "messages" in result
+    assert "multi_papers" in result
+    assert result["multi_papers"] is not None
+
+
+def test_tool_integration(mock_tools):
+    """Test if the tools interact correctly with the workflow."""
+    thread_id = "test_thread"
+    s2_app = get_s2_app(thread_id)
+
+    state = Talk2Scholars(
+        messages=[HumanMessage(content="Search for AI ethics papers")]
+    )
+
+    mock_paper_id = "11159bdb213aaa243916f42f576396d483ba474b"
+    mock_response = {
+        "papers": {
+            mock_paper_id: {
+                "Title": "Mock AI Ethics Paper",
+                "Abstract": "A study on AI ethics",
+                "Citation Count": 100,
+                "URL": "https://example.com/mock-paper",
+            }
+        }
+    }
+
+    # Update both the fixture mock and patch the actual tool
+    mock_tools["search_tool"].return_value = {"papers": mock_response["papers"]}
+
+    with patch(
+        "aiagents4pharma.talk2scholars.tools.s2.search.search_tool",
+        return_value={"papers": mock_response["papers"]},
+    ):
+        result = s2_app.invoke(
+            state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+    assert "papers" in result
+    assert len(result["papers"]) > 0  # Verify we have papers
+    assert isinstance(result["papers"], dict)  # Verify it's a dictionary
+
+
+def test_empty_query():
+    """Test how the system handles an empty query."""
+    thread_id = "test_thread"
+    main_app = get_main_app(thread_id)
+
+    state = Talk2Scholars(messages=[HumanMessage(content="")])
+
+    # Mock the s2_agent app
+    mock_s2_app = get_s2_app(thread_id)
+
+    with patch(
+        "aiagents4pharma.talk2scholars.agents.s2_agent.get_app",
+        return_value=mock_s2_app,
+    ):
+        result = main_app.invoke(
+            state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+    assert "messages" in result
+    last_message = result["messages"][-1].content.lower()
+    assert any(
+        phrase in last_message
+        for phrase in ["no valid input", "how can i assist", "please provide a query"]
+    )
+
+
+def test_api_failure_handling():
+    """Test if the system gracefully handles an API failure."""
+    thread_id = "test_thread"
+    s2_app = get_s2_app(thread_id)
+
+    expected_error = "API Timeout: Connection failed"
+    with patch("requests.get", side_effect=Exception(expected_error)):
+        state = Talk2Scholars(messages=[HumanMessage(content="Find latest NLP papers")])
+
+        result = s2_app.invoke(
+            state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+        assert "messages" in result
+        last_message = result["messages"][-1].content.lower()
+
+        # Update assertions to match actual error message
+        assert any(
+            [
+                "unable to retrieve" in last_message,
+                "connection issue" in last_message,
+                "please try again later" in last_message,
+            ]
+        )
+        assert "nlp papers" in last_message  # Verify context is maintained

aiagents4pharma/talk2scholars/tests/test_main_agent.py

@@ -0,0 +1,180 @@
+"""
+Unit tests for main agent functionality.
+Tests the supervisor agent's routing logic and state management.
+"""
+
+# pylint: disable=redefined-outer-name
+from unittest.mock import Mock, patch, MagicMock
+import pytest
+from langchain_core.messages import HumanMessage, AIMessage
+from ..agents.main_agent import make_supervisor_node, get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+
+@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.main_agent = "Test 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"
+    llm_model = "gpt-4o-mini"
+
+    # Mock the LLM
+    mock_llm = Mock()
+
+    with patch(
+        "aiagents4pharma.talk2scholars.agents.main_agent.ChatOpenAI",
+        return_value=mock_llm,
+    ):
+
+        app = get_app(thread_id, llm_model)
+        assert app is not None
+        assert "supervisor" in app.nodes
+        assert "s2_agent" in app.nodes  # Ensure nodes exist
+
+
+def test_supervisor_node_execution():
+    """Test that the supervisor node processes messages and makes a decision."""
+    mock_llm = Mock()
+    thread_id = "test_thread"
+
+    # Mock the supervisor agent's response
+    mock_supervisor = Mock()
+    mock_supervisor.invoke.return_value = {"messages": [AIMessage(content="s2_agent")]}
+
+    with patch(
+        "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent",
+        return_value=mock_supervisor,
+    ):
+        supervisor_node = make_supervisor_node(mock_llm, thread_id)
+
+        # Create a mock state
+        mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
+
+        # Execute
+        result = supervisor_node(mock_state)
+
+        # Validate
+        assert result.goto == "s2_agent"
+        mock_supervisor.invoke.assert_called_once_with(
+            mock_state, {"configurable": {"thread_id": thread_id}}
+        )  # Ensure invoke was called correctly
+
+
+def test_call_s2_agent():
+    """Test the call to S2 agent and its integration with the state."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
+
+    with patch("aiagents4pharma.talk2scholars.agents.s2_agent.get_app") as mock_s2_app:
+        mock_s2_app.return_value.invoke.return_value = {
+            "messages": [AIMessage(content="Here are the papers")],
+            "papers": {"id123": "Sample Paper"},
+        }
+
+        app = get_app(thread_id)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+        assert "messages" in result
+        assert "papers" in result
+        assert result["papers"]["id123"] == "Sample Paper"
+
+        mock_s2_app.return_value.invoke.assert_called_once()
+
+
+def test_hydra_failure():
+    """Test exception handling when Hydra fails to load config."""
+    thread_id = "test_thread"
+    with patch("hydra.initialize", side_effect=Exception("Hydra error")):
+        with pytest.raises(Exception) as exc_info:
+            get_app(thread_id)
+
+        assert "Hydra error" in str(exc_info.value)
+
+
+class TestMainAgent:
+    """Basic tests for the main agent initialization and configuration"""
+
+    def test_supervisor_node_creation(self, mock_hydra):
+        """Test that supervisor node can be created with correct config"""
+        mock_llm = Mock()
+        thread_id = "test_thread"
+
+        with patch(
+            "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent"
+        ) as mock_create:
+            mock_create.return_value = Mock()
+            supervisor = make_supervisor_node(mock_llm, thread_id)
+
+            assert supervisor is not None
+            assert mock_create.called
+            # Verify Hydra was called with correct parameters
+            assert mock_hydra.call_count == 1  # Updated assertion
+
+    def test_supervisor_config_loading(self, mock_hydra):
+        """Test that supervisor loads configuration correctly"""
+        mock_llm = Mock()
+        thread_id = "test_thread"
+
+        with patch(
+            "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent"
+        ):
+            make_supervisor_node(mock_llm, thread_id)
+
+            # Verify Hydra initialization
+            assert mock_hydra.call_count == 1
+            assert "agents/talk2scholars/main_agent=default" in str(
+                mock_hydra.call_args_list[0]
+            )
+
+    def test_react_agent_params(self):
+        """Test that react agent is created with correct parameters"""
+        mock_llm = Mock()
+        thread_id = "test_thread"
+
+        with patch(
+            "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent"
+        ) as mock_create:
+            mock_create.return_value = Mock()
+            make_supervisor_node(mock_llm, thread_id)
+
+            # Verify create_react_agent was called
+            assert mock_create.called
+
+            # Verify the parameters
+            args, kwargs = mock_create.call_args
+            assert args[0] == mock_llm  # First argument should be the LLM
+            assert "state_schema" in kwargs  # Should have state_schema
+            assert hasattr(
+                mock_create.return_value, "invoke"
+            )  # Should have invoke method
+
+    def test_supervisor_custom_config(self, mock_hydra):
+        """Test supervisor with custom configuration"""
+        mock_llm = Mock()
+        thread_id = "test_thread"
+
+        with patch(
+            "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent"
+        ):
+            make_supervisor_node(mock_llm, thread_id)
+
+            # Verify Hydra was called
+            mock_hydra.assert_called_once()

aiagents4pharma/talk2scholars/tests/test_s2_agent.py

@@ -0,0 +1,138 @@
+"""
+Unit tests for the S2 agent (Semantic Scholar sub-agent).
+"""
+
+from unittest import mock
+import pytest
+from langchain_core.messages import HumanMessage, AIMessage
+from ..agents.s2_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+
+@pytest.fixture(autouse=True)
+def mock_hydra_fixture():
+    """Mock Hydra configuration to prevent external dependencies."""
+    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.s2_agent.s2_agent = "Test prompt"
+        mock_compose.return_value = cfg_mock
+        yield mock_compose
+
+
+@pytest.fixture
+def mock_tools_fixture():
+    """Mock tools to prevent execution of real API calls."""
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.search.search_tool"
+        ) as mock_s2_search,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.display_results.display_results"
+        ) as mock_s2_display,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.single_paper_rec."
+            "get_single_paper_recommendations"
+        ) as mock_s2_single_rec,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec."
+            "get_multi_paper_recommendations"
+        ) as mock_s2_multi_rec,
+    ):
+
+        mock_s2_search.return_value = mock.Mock()
+        mock_s2_display.return_value = mock.Mock()
+        mock_s2_single_rec.return_value = mock.Mock()
+        mock_s2_multi_rec.return_value = mock.Mock()
+
+        yield [mock_s2_search, mock_s2_display, mock_s2_single_rec, mock_s2_multi_rec]
+
+
+@pytest.mark.usefixtures("mock_hydra_fixture")
+def test_s2_agent_initialization():
+    """Test that S2 agent initializes correctly with mock configuration."""
+    thread_id = "test_thread"
+
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.s2_agent.create_react_agent"
+    ) as mock_create:
+        mock_create.return_value = mock.Mock()
+
+        app = get_app(thread_id)
+
+        assert app is not None
+        assert mock_create.called
+
+
+def test_s2_agent_invocation():
+    """Test that the S2 agent processes user input and returns a valid response."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
+
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.s2_agent.create_react_agent"
+    ) as mock_create:
+        mock_agent = mock.Mock()
+        mock_create.return_value = mock_agent
+        mock_agent.invoke.return_value = {
+            "messages": [AIMessage(content="Here are some AI papers")],
+            "papers": {"id123": "AI Research Paper"},
+        }
+
+        app = get_app(thread_id)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+
+        assert "messages" in result
+        assert "papers" in result
+        assert result["papers"]["id123"] == "AI Research Paper"
+
+
+def test_s2_agent_tools_assignment(request):
+    """Ensure that the correct tools are assigned to the agent."""
+    thread_id = "test_thread"
+
+    # Dynamically retrieve the fixture to avoid redefining it in the function signature
+    mock_tools = request.getfixturevalue("mock_tools_fixture")
+
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.s2_agent.create_react_agent"
+        ) as mock_create,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.s2_agent.ToolNode"
+        ) as mock_toolnode,
+    ):
+
+        mock_agent = mock.Mock()
+        mock_create.return_value = mock_agent
+
+        # Mock ToolNode behavior
+        mock_tool_instance = mock.Mock()
+        mock_tool_instance.tools = mock_tools  # Use the dynamically retrieved fixture
+        mock_toolnode.return_value = mock_tool_instance
+
+        get_app(thread_id)
+
+        # Ensure the agent was created with the mocked ToolNode
+        assert mock_toolnode.called
+        assert len(mock_tool_instance.tools) == 4
+
+
+def test_s2_agent_hydra_failure():
+    """Test exception handling when Hydra fails to load config."""
+    thread_id = "test_thread"
+
+    with mock.patch("hydra.initialize", side_effect=Exception("Hydra error")):
+        with pytest.raises(Exception) as exc_info:
+            get_app(thread_id)
+
+        assert "Hydra error" in str(exc_info.value)