aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_agents_s2_agent.py

@@ -0,0 +1,198 @@
+"""
+Updated Unit Tests for the S2 agent (Semantic Scholar sub-agent).
+"""
+
+from unittest import mock
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.s2_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+@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_dataframe.display_dataframe"
+        ) 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.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.query_dataframe.query_dataframe"
+        ) as mock_s2_query_dataframe,
+        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 = {"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_dataframe.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_dataframe,
+            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, llm_model=LLM_MODEL)
+        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, llm_model=LLM_MODEL)
+        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"
+    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_tool_instance = mock.Mock()
+        mock_tool_instance.tools = mock_tools
+        mock_toolnode.return_value = mock_tool_instance
+        get_app(thread_id, llm_model=LLM_MODEL)
+        assert mock_toolnode.called
+        assert len(mock_tool_instance.tools) == 6
+
+
+def test_s2_query_dataframe_tool():
+    """Test if the query_dataframe 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_dataframe": "Mock Query Result"
+            },  # Ensure the expected key is inside 'papers'
+            "multi_papers": {},
+        }
+        app = get_app(thread_id, llm_model=LLM_MODEL)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+        assert "query_dataframe" 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, llm_model=LLM_MODEL)
+        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, llm_model=LLM_MODEL)
+        assert "Hydra error" in str(exc_info.value)

aiagents4pharma/talk2scholars/tests/test_agents_zotero_agent.py

@@ -0,0 +1,160 @@
+"""
+Updated Unit Tests for the Zotero agent (Zotero Library Managent sub-agent).
+"""
+
+from unittest import mock
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.zotero_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+@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.zotero_agent.temperature = 0
+        cfg_mock.agents.talk2scholars.zotero_agent.zotero_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.display_dataframe.display_dataframe"
+        ) as mock_s2_display,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2.query_dataframe.query_dataframe"
+        ) as mock_s2_query_dataframe,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.s2."
+            "retrieve_semantic_scholar_paper_id."
+            "retrieve_semantic_scholar_paper_id"
+        ) as mock_s2_retrieve_id,
+        mock.patch(
+            "aiagents4pharma.talk2scholars.tools.zotero.zotero_read.zotero_read"
+        ) as mock_zotero_query_dataframe,
+    ):
+        mock_s2_display.return_value = {"result": "Mock Display Result"}
+        mock_s2_query_dataframe.return_value = {"result": "Mock Query Result"}
+        mock_s2_retrieve_id.return_value = {"paper_id": "MockPaper123"}
+        mock_zotero_query_dataframe.return_value = {"result": "Mock Search Result"}
+
+        yield [
+            mock_s2_display,
+            mock_s2_query_dataframe,
+            mock_s2_retrieve_id,
+            mock_zotero_query_dataframe,
+        ]
+
+
+@pytest.mark.usefixtures("mock_hydra_fixture")
+def test_zotero_agent_initialization():
+    """Test that S2 agent initializes correctly with mock configuration."""
+    thread_id = "test_thread"
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.zotero_agent.create_react_agent"
+    ) as mock_create:
+        mock_create.return_value = mock.Mock()
+        app = get_app(thread_id, llm_model=LLM_MODEL)
+        assert app is not None
+        assert mock_create.called
+
+
+def test_zotero_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.zotero_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, llm_model=LLM_MODEL)
+        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_zotero_agent_tools_assignment(request):
+    """Ensure that the correct tools are assigned to the agent."""
+    thread_id = "test_thread"
+    mock_tools = request.getfixturevalue("mock_tools_fixture")
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.zotero_agent.create_react_agent"
+        ) as mock_create,
+        mock.patch("aiagents4pharma.talk2scholars.agents.zotero_agent.ToolNode") as mock_toolnode,
+    ):
+        mock_agent = mock.Mock()
+        mock_create.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_model=LLM_MODEL)
+        assert mock_toolnode.called
+        assert len(mock_tool_instance.tools) == 4
+
+
+def test_s2_query_dataframe_tool():
+    """Test if the query_dataframe 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.zotero_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_dataframe": "Mock Query Result"
+            },  # Ensure the expected key is inside 'papers'
+            "multi_papers": {},
+        }
+        app = get_app(thread_id, llm_model=LLM_MODEL)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+        assert "query_dataframe" in result["papers"]
+        assert mock_agent.invoke.called
+
+
+def test_zotero_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, llm_model=LLM_MODEL)
+        assert "Hydra error" in str(exc_info.value)

aiagents4pharma/talk2scholars/tests/test_s2_tools_display_dataframe.py

@@ -0,0 +1,91 @@
+"""
+Unit tests for S2 tools functionality.
+"""
+
+import pytest
+from langgraph.types import Command
+
+from ..tools.s2.display_dataframe import (
+    NoPapersFoundError as raised_error,
+)
+from ..tools.s2.display_dataframe import (
+    display_dataframe,
+)
+
+
+@pytest.fixture(name="initial_state")
+def initial_state_fixture():
+    """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_dataframe_empty_state(self, initial_state):
+        """Verifies display_dataframe 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_dataframe.invoke({"state": initial_state, "tool_call_id": "test123"})
+
+    def test_display_dataframe_shows_papers(self, initial_state):
+        """Verifies display_dataframe tool correctly returns papers from state"""
+        state = initial_state.copy()
+        state["last_displayed_papers"] = "papers"
+        state["papers"] = MOCK_STATE_PAPER
+
+        result = display_dataframe.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
+        )
+
+    def test_display_dataframe_direct_mapping(self, initial_state):
+        """Verifies display_dataframe handles direct dict mapping in last_displayed_papers."""
+        # Prepare state with direct mapping of papers
+        state = initial_state.copy()
+        state["last_displayed_papers"] = MOCK_STATE_PAPER
+        # Invoke display tool
+        result = display_dataframe.invoke({"state": state, "tool_call_id": "test123"})
+        assert isinstance(result, Command)
+        update = result.update
+        # Artifact should be the direct mapping
+        messages = update.get("messages", [])
+        assert len(messages) == 1
+        artifact = messages[0].artifact
+        assert artifact == MOCK_STATE_PAPER
+        # Content count should match mapping length
+        assert "1 papers found" in messages[0].content

aiagents4pharma/talk2scholars/tests/test_s2_tools_query_dataframe.py

@@ -0,0 +1,191 @@
+"""
+Unit tests for S2 tools functionality.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from langchain_core.messages import ToolMessage
+
+from ..tools.s2.query_dataframe import NoPapersFoundError, query_dataframe
+
+
+@pytest.fixture(name="initial_state")
+def initial_state_fixture():
+    """Provides an empty initial state for tests with a dummy llm_model."""
+    return {"papers": {}, "multi_papers": {}, "llm_model": MagicMock()}
+
+
+# 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_query_dataframe_empty_state(self, initial_state):
+        """Tests query_dataframe tool behavior when no papers are found."""
+        # Calling without any papers should raise NoPapersFoundError
+        tool_input = {
+            "question": "List all papers",
+            "state": initial_state,
+            "tool_call_id": "test_id",
+        }
+        with pytest.raises(
+            NoPapersFoundError,
+            match="No papers found. A search needs to be performed first.",
+        ):
+            query_dataframe.run(tool_input)
+
+    @patch("aiagents4pharma.talk2scholars.tools.s2.query_dataframe.create_pandas_dataframe_agent")
+    def test_query_dataframe_with_papers(self, mock_create_agent, initial_state):
+        """Tests querying papers when data is available."""
+        state = initial_state.copy()
+        state["last_displayed_papers"] = "papers"
+        state["papers"] = MOCK_STATE_PAPER
+
+        # Mock the dataframe agent instead of the LLM
+        mock_agent = MagicMock()
+        mock_agent.invoke.return_value = {"output": "Mocked response"}
+
+        mock_create_agent.return_value = mock_agent  # Mock the function returning the agent
+
+        # Ensure that the output of query_dataframe is correctly structured
+        # Invoke the tool with a test tool_call_id
+        tool_input = {
+            "question": "List all papers",
+            "state": state,
+            "tool_call_id": "test_id",
+        }
+        result = query_dataframe.run(tool_input)
+        # The tool returns a Command with messages
+        assert hasattr(result, "update")
+        update = result.update
+        assert "messages" in update
+        msgs = update["messages"]
+        assert len(msgs) == 1
+        msg = msgs[0]
+        assert isinstance(msg, ToolMessage)
+        assert msg.content == "Mocked response"
+
+    @patch("aiagents4pharma.talk2scholars.tools.s2.query_dataframe.create_pandas_dataframe_agent")
+    def test_query_dataframe_direct_mapping(self, mock_create_agent, initial_state):
+        """Tests query_dataframe when last_displayed_papers is a direct dict mapping."""
+        # Prepare state with direct mapping
+        state = initial_state.copy()
+        state["last_displayed_papers"] = MOCK_STATE_PAPER
+        # Mock the dataframe agent
+        mock_agent = MagicMock()
+        mock_agent.invoke.return_value = {"output": "Direct mapping response"}
+        mock_create_agent.return_value = mock_agent
+        # Invoke tool
+        # Invoke the tool with direct mapping and test tool_call_id
+        tool_input = {
+            "question": "Filter papers",
+            "state": state,
+            "tool_call_id": "test_id",
+        }
+        result = query_dataframe.run(tool_input)
+        update = result.update
+        assert "messages" in update
+        msgs = update["messages"]
+        assert len(msgs) == 1
+        msg = msgs[0]
+        assert isinstance(msg, ToolMessage)
+        assert msg.content == "Direct mapping response"
+
+    def test_query_dataframe_missing_llm(self, initial_state):
+        """Test that missing llm_model raises ValueError."""
+        # Remove llm_model
+        state = {k: v for k, v in initial_state.items() if k != "llm_model"}
+        state["last_displayed_papers"] = MOCK_STATE_PAPER
+        tool_input = {"question": "Test", "state": state, "tool_call_id": "test_id"}
+        with pytest.raises(ValueError) as exc:
+            query_dataframe.run(tool_input)
+        assert "Missing 'llm_model' in state." in str(exc.value)
+
+    def test_query_dataframe_invalid_mapping(self, initial_state):
+        """Test that invalid last_displayed_papers mapping raises ValueError."""
+        # Provide invalid mapping key
+        state = initial_state.copy()
+        state["last_displayed_papers"] = "nonexistent_key"
+        # llm_model present
+        tool_input = {"question": "Test", "state": state, "tool_call_id": "test_id"}
+        with pytest.raises(ValueError) as exc:
+            query_dataframe.run(tool_input)
+        assert "Could not resolve a valid metadata dictionary" in str(exc.value)
+
+    @patch("aiagents4pharma.talk2scholars.tools.s2.query_dataframe.create_pandas_dataframe_agent")
+    def test_query_dataframe_extract_ids(self, mock_create_agent):
+        """Test extract_ids returns the raw list or single element correctly."""
+        # Prepare state with fake paper_ids column
+        state = {"llm_model": MagicMock()}
+        state_key = "papers"
+        dic = {
+            "p1": {"paper_ids": ["id1", "id2"]},
+            "p2": {"paper_ids": ["id3"]},
+        }
+        state["last_displayed_papers"] = dic
+        state[state_key] = dic  # simulate indirect mapping
+        # Mock agent to echo the Python expression
+        mock_agent = MagicMock()
+        mock_agent.invoke.side_effect = lambda args, stream_mode=None: {"output": args["input"]}
+        mock_create_agent.return_value = mock_agent
+        # Test full list
+        tool_input = {
+            "question": "",
+            "state": state,
+            "tool_call_id": "tid",
+            "extract_ids": True,
+            "id_column": "paper_ids",
+        }
+        result = query_dataframe.run(tool_input)
+        output = result.update["messages"][0].content
+        # Should be the base list expression
+        expected = "df['paper_ids'].dropna().str[0].tolist()"
+        assert output == expected
+        # Test single element
+        tool_input["row_number"] = 2
+        result2 = query_dataframe.run(tool_input)
+        output2 = result2.update["messages"][0].content
+        expected2 = "df['paper_ids'].dropna().str[0].tolist()[1]"
+        assert output2 == expected2
+
+    def test_query_dataframe_extract_ids_missing_column(self, initial_state):
+        """Test that missing id_column raises ValueError when extract_ids=True."""
+        state = initial_state.copy()
+        state["last_displayed_papers"] = {"p1": {"paper_ids": ["id1"]}}
+        state["papers"] = state["last_displayed_papers"]
+        with pytest.raises(ValueError) as exc:
+            query_dataframe.run(
+                {
+                    "question": "",
+                    "state": state,
+                    "tool_call_id": "tid",
+                    "extract_ids": True,
+                    "id_column": "",
+                }
+            )
+        assert "Must specify 'id_column' when extract_ids=True." in str(exc.value)

aiagents4pharma/talk2scholars/tests/test_states_state.py

@@ -0,0 +1,38 @@
+"""
+Tests for state management functionality.
+"""
+
+from ..state.state_talk2scholars import merge_dict, replace_dict
+
+
+def test_state_replace_dict():
+    """Verifies state dictionary replacement works correctly"""
+    existing = {"key1": "value1", "key2": "value2"}
+    new = {"key3": "value3"}
+    result = replace_dict(existing, new)
+    assert result == new
+    assert isinstance(result, dict)
+
+
+def test_state_merge_dict():
+    """Verifies state dictionary merging works correctly"""
+    existing = {"a": 1, "b": 2}
+    new = {"b": 3, "c": 4}
+    result = merge_dict(existing, new)
+    # result should contain merged keys, with new values overriding existing ones
+    assert result == {"a": 1, "b": 3, "c": 4}
+    assert isinstance(result, dict)
+    # original existing dict should be unchanged
+    assert existing == {"a": 1, "b": 2}
+
+
+def test_replace_dict_non_mapping():
+    """Verifies replace_dict returns non-mapping values directly"""
+
+    existing = {"key": "value"}
+    # When new is not a dict, replace_dict should return new value unchanged
+    new_value = "not_a_dict"
+    result = replace_dict(existing, new_value)
+    assert result == new_value
+    # existing should remain unmodified when returning new directly
+    assert existing == {"key": "value"}