aiagents4pharma 1.20.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_llm_main_integration.py

@@ -0,0 +1,58 @@
+"""
+Integration tests for talk2scholars system with OpenAI.
+"""
+
+import os
+import pytest
+import hydra
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from ..agents.main_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+# pylint: disable=redefined-outer-name
+
+
+@pytest.mark.skipif(
+    not os.getenv("OPENAI_API_KEY"), reason="Requires OpenAI API key to run"
+)
+def test_main_agent_real_llm():
+    """
+    Test that the main agent invokes S2 agent correctly
+    and updates the state with real LLM execution.
+    """
+
+    # Load Hydra Configuration EXACTLY like in main_agent.py
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
+        )
+    hydra_cfg = cfg.agents.talk2scholars.main_agent
+
+    assert hydra_cfg is not None, "Hydra config failed to load"
+
+    # Use the real OpenAI API (ensure env variable is set)
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=hydra_cfg.temperature)
+
+    # Initialize main agent workflow (WITH real Hydra config)
+    thread_id = "test_thread"
+    app = get_app(thread_id, llm)
+
+    # Provide an actual user query
+    initial_state = Talk2Scholars(
+        messages=[HumanMessage(content="Find AI papers on transformers")]
+    )
+
+    # Invoke the agent (triggers supervisor → s2_agent)
+    result = app.invoke(
+        initial_state,
+        {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
+    )
+
+    # Assert that the supervisor routed correctly
+    assert "messages" in result, "Expected messages in response"
+
+    # Fix: Accept AIMessage as a valid response type
+    assert isinstance(
+        result["messages"][-1], (HumanMessage, AIMessage, str)
+    ), "Last message should be a valid response"

aiagents4pharma/talk2scholars/tests/test_main_agent.py

@@ -4,10 +4,12 @@ Tests the supervisor agent's routing logic and state management.
 """
 
 # pylint: disable=redefined-outer-name
+# pylint: disable=redefined-outer-name,too-few-public-methods
 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 langgraph.graph import END
+from ..agents.main_agent import make_supervisor_node, get_hydra_config, get_app
 from ..state.state_talk2scholars import Talk2Scholars
 
 
@@ -17,7 +19,8 @@ def mock_hydra():
     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"
+        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
 
@@ -25,156 +28,129 @@ def mock_hydra():
 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()
+    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
 
-    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_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"
+    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)
 
 
 def test_supervisor_node_execution():
-    """Test that the supervisor node processes messages and makes a decision."""
+    """Test that the supervisor node routes correctly."""
     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")]}
+    class MockRouter:
+        """Mock router class."""
 
-    with patch(
-        "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent",
-        return_value=mock_supervisor,
+        next = "s2_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)
-
-        # 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."""
+def test_supervisor_node_finish():
+    """Test that supervisor node correctly handles FINISH case."""
+    mock_llm = Mock()
     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"},
-        }
+    class MockRouter:
+        """Mock router class."""
 
-        app = get_app(thread_id)
-        result = app.invoke(
-            mock_state,
-            config={
-                "configurable": {
-                    "thread_id": thread_id,
-                    "checkpoint_ns": "test_ns",
-                    "checkpoint_id": "test_checkpoint",
-                }
-            },
-        )
+        next = "FINISH"
 
-        assert "messages" in result
-        assert "papers" in result
-        assert result["papers"]["id123"] == "Sample Paper"
+    class MockAIResponse:
+        """Mock AI response class."""
 
-        mock_s2_app.return_value.invoke.assert_called_once()
+        def __init__(self):
+            self.content = "Final AI Response"
 
-
-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)
+    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"
 
 
-class TestMainAgent:
-    """Basic tests for the main agent initialization and configuration"""
+def test_call_s2_agent_failure_in_get_app():
+    """Test handling failure when calling s2_agent.get_app()."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
 
-    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.s2_agent.get_app",
+        side_effect=Exception("S2 Agent Failure"),
+    ):
+        with pytest.raises(Exception) as exc_info:
+            app = get_app(thread_id)  # Get the compiled workflow
+            app.invoke(
+                mock_state,
+                {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
+            )
 
-        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 "S2 Agent Failure" in str(exc_info.value)
 
-            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"
+def test_call_s2_agent_failure_in_invoke():
+    """Test handling failure when invoking s2_agent app."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
 
-        with patch(
-            "aiagents4pharma.talk2scholars.agents.main_agent.create_react_agent"
-        ):
-            make_supervisor_node(mock_llm, thread_id)
+    mock_app = Mock()
+    mock_app.invoke.side_effect = Exception("S2 Agent Invoke Failure")
 
-            # Verify Hydra initialization
-            assert mock_hydra.call_count == 1
-            assert "agents/talk2scholars/main_agent=default" in str(
-                mock_hydra.call_args_list[0]
+    with patch(
+        "aiagents4pharma.talk2scholars.agents.s2_agent.get_app", return_value=mock_app
+    ):
+        with pytest.raises(Exception) as exc_info:
+            app = get_app(thread_id)  # Get the compiled workflow
+            app.invoke(
+                mock_state,
+                {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
             )
 
-    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()
+        assert "S2 Agent Invoke Failure" in str(exc_info.value)

aiagents4pharma/talk2scholars/tests/test_s2_agent.py

@@ -1,7 +1,8 @@
 """
-Unit tests for the S2 agent (Semantic Scholar sub-agent).
+Updated Unit Tests for the S2 agent (Semantic Scholar sub-agent).
 """
 
+# pylint: disable=redefined-outer-name
 from unittest import mock
 import pytest
 from langchain_core.messages import HumanMessage, AIMessage
@@ -35,31 +36,42 @@ def mock_tools_fixture():
             "get_single_paper_recommendations"
         ) as mock_s2_single_rec,
         mock.patch(
-            "aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec."
-            "get_multi_paper_recommendations"
+            "aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec.get_multi_paper_recommendations"
         ) as mock_s2_multi_rec,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.query_results.query_results"
+        ) as mock_s2_query_results,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.retrieve_semantic_scholar_paper_id."
+            "retrieve_semantic_scholar_paper_id"
+        ) as mock_s2_retrieve_id,
     ):
-
-        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]
+        mock_s2_search.return_value = {"result": "Mock Search Result"}
+        mock_s2_display.return_value = {"result": "Mock Display Result"}
+        mock_s2_single_rec.return_value = {"result": "Mock Single Rec"}
+        mock_s2_multi_rec.return_value = {"result": "Mock Multi Rec"}
+        mock_s2_query_results.return_value = {"result": "Mock Query Result"}
+        mock_s2_retrieve_id.return_value = {"paper_id": "MockPaper123"}
+
+        yield [
+            mock_s2_search,
+            mock_s2_display,
+            mock_s2_single_rec,
+            mock_s2_multi_rec,
+            mock_s2_query_results,
+            mock_s2_retrieve_id,
+        ]
 
 
 @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
 
@@ -68,7 +80,6 @@ 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:
@@ -78,7 +89,6 @@ def test_s2_agent_invocation():
             "messages": [AIMessage(content="Here are some AI papers")],
             "papers": {"id123": "AI Research Paper"},
         }
-
         app = get_app(thread_id)
         result = app.invoke(
             mock_state,
@@ -90,7 +100,6 @@ def test_s2_agent_invocation():
                 }
             },
         )
-
         assert "messages" in result
         assert "papers" in result
         assert result["papers"]["id123"] == "AI Research Paper"
@@ -99,10 +108,7 @@ def test_s2_agent_invocation():
 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"
@@ -111,28 +117,88 @@ def test_s2_agent_tools_assignment(request):
             "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_tool_instance.tools = mock_tools
         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
+        assert len(mock_tool_instance.tools) == 6
+
+
+def test_s2_query_results_tool():
+    """Test if the query_results tool is correctly utilized by the agent."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(
+        messages=[HumanMessage(content="Query results for 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": [HumanMessage(content="Query results for AI papers")],
+            "last_displayed_papers": {},
+            "papers": {
+                "query_results": "Mock Query Result"
+            },  # Ensure the expected key is inside 'papers'
+            "multi_papers": {},
+        }
+        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 "query_results" in result["papers"]
+        assert mock_agent.invoke.called
+
+
+def test_s2_retrieve_id_tool():
+    """Test if the retrieve_semantic_scholar_paper_id tool is correctly utilized by the agent."""
+    thread_id = "test_thread"
+    mock_state = Talk2Scholars(
+        messages=[HumanMessage(content="Retrieve paper ID for AI research")]
+    )
+    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": [HumanMessage(content="Retrieve paper ID for AI research")],
+            "last_displayed_papers": {},
+            "papers": {
+                "paper_id": "MockPaper123"
+            },  # Ensure 'paper_id' is inside 'papers'
+            "multi_papers": {},
+        }
+        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 "paper_id" in result["papers"]
+        assert mock_agent.invoke.called
 
 
 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)