aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -0,0 +1,98 @@
+"""
+State management for the Talk2Scholars agent.
+
+This module defines the state class `Talk2Scholars`, which maintains the conversation
+context, retrieved papers, and other relevant metadata. The state ensures consistency
+across agent interactions.
+"""
+
+import logging
+from collections.abc import Mapping
+from typing import Annotated, Any
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models import BaseChatModel
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def merge_dict(existing: dict[str, Any], new: dict[str, Any]) -> dict[str, Any]:
+    """
+    Merges the existing dictionary with a new dictionary.
+
+    This function logs the state merge and ensures that the new values
+    are appended to the existing state without overwriting other entries.
+    Args:
+        existing (Dict[str, Any]): The current dictionary state.
+        new (Dict[str, Any]): The new dictionary state to merge.
+    Returns:
+        Dict[str, Any]: The merged dictionary state.
+    """
+    merged = dict(existing) if existing else {}
+    merged.update(new or {})
+    return merged
+
+
+def replace_dict(existing: dict[str, Any], new: Any) -> Any:
+    """
+    Replaces the existing dictionary with a new dictionary.
+
+    This function logs the state update and ensures that the old state is replaced
+    with the new one.
+
+    Args:
+        existing (Dict[str, Any]): The current dictionary state.
+        new (Dict[str, Any]): The new dictionary state to replace the existing one.
+
+    Returns:
+        Dict[str, Any]: The updated dictionary state.
+
+    Example:
+        >>> old_state = {"papers": {"id1": "Paper 1"}}
+        >>> new_state = {"papers": {"id2": "Paper 2"}}
+        >>> updated_state = replace_dict(old_state, new_state)
+        >>> print(updated_state)
+        {"papers": {"id2": "Paper 2"}}
+    """
+    # If new is not a mapping, just replace existing value outright
+    if not isinstance(new, Mapping):
+        return new
+    # In-place replace: clear existing mapping and update with new entries
+    existing.clear()
+    existing.update(new)
+    return existing
+
+
+class Talk2Scholars(AgentState):
+    """
+    Represents the state of the Talk2Scholars agent.
+    This class extends `AgentState` to maintain conversation history, retrieved papers,
+    and interactions with the language model.
+    Attributes:
+        last_displayed_papers (Dict[str, Any]): Stores the most recently displayed papers.
+        papers (Dict[str, Any]): Stores the research papers retrieved from the agent's queries.
+        multi_papers (Dict[str, Any]): Stores multiple recommended papers from various sources.
+        article_data (Dict[str, Any]): Stores the papers retrieved from Zotero and the pdf
+        download agent with their metadata.
+        zotero_write_approval_status (Dict[str, Any]): Stores the approval status and collection
+        path for Zotero save operations.
+        llm_model (BaseChatModel): The language model instance used for generating responses.
+        text_embedding_model (Embeddings): The text embedding model used for
+        similarity calculations.
+    """
+
+    # Agent state fields
+    # Key controlling UI display: always replace to reference latest output
+    # Stores the most recently displayed papers metadata
+    last_displayed_papers: Annotated[dict[str, Any], replace_dict]
+    # Accumulative keys: merge new entries into existing state
+    papers: Annotated[dict[str, Any], merge_dict]
+    multi_papers: Annotated[dict[str, Any], merge_dict]
+    article_data: Annotated[dict[str, Any], merge_dict]
+    # Approval status: always replace to reflect latest operation
+    zotero_write_approval_status: Annotated[dict[str, Any], replace_dict]
+    llm_model: BaseChatModel
+    text_embedding_model: Embeddings

aiagents4pharma/talk2scholars/tests/__init__.py

@@ -0,0 +1,3 @@
+"""
+This module contains the test cases.
+"""

aiagents4pharma/talk2scholars/tests/test_agents_main_agent.py

@@ -0,0 +1,256 @@
+"""
+Unit tests for main agent functionality.
+Tests the supervisor agent's routing logic and state management.
+"""
+
+from types import SimpleNamespace
+
+import hydra
+import pytest
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.outputs import ChatGeneration, ChatResult
+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, messages, stop=None, run_manager=None, **kwargs):
+        """generate a dummy response based on the input messages."""
+        DummyLLM.called_prompt = messages[0].content
+        message = AIMessage(content="dummy output")
+        generation = ChatGeneration(message=message)
+        return ChatResult(generations=[generation])
+
+    @property
+    def _llm_type(self):
+        """Return the type of the language model."""
+        return "dummy"
+
+    # NEW: public shim to avoid protected access in tests
+    def public_llm_type(self) -> str:
+        """public method to access the LLM type."""
+        return self._llm_type
+
+
+# --- 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 get_supervisor_args(self):
+        """Return the supervisor arguments stored in this workflow."""
+        return self.supervisor_args
+
+
+def dummy_s2_agent(uniq_id, llm_model):
+    """Return a DummyWorkflow for the S2 agent."""
+    dummy_s2_agent.called_uniq_id = uniq_id
+    dummy_s2_agent.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "s2", "uniq_id": uniq_id})
+
+
+def dummy_zotero_agent(uniq_id, llm_model):
+    """Return a DummyWorkflow for the Zotero agent."""
+    dummy_zotero_agent.called_uniq_id = uniq_id
+    dummy_zotero_agent.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "zotero", "uniq_id": uniq_id})
+
+
+def dummy_question_and_answer_agent(uniq_id, llm_model):
+    """Return a DummyWorkflow for the PDF agent."""
+    dummy_question_and_answer_agent.called_uniq_id = uniq_id
+    dummy_question_and_answer_agent.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "pdf", "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, {}))
+
+    def get_config(self):
+        """Get the raw dummy configuration dictionary."""
+        return self.config
+
+
+# --- 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)
+    )
+
+
+def dummy_paper_download_agent(uniq_id, llm_model):
+    """Return a DummyWorkflow for the paper download agent."""
+    dummy_paper_download_agent.called_uniq_id = uniq_id
+    dummy_paper_download_agent.called_llm_model = llm_model
+    return DummyWorkflow(supervisor_args={"agent": "paper_download", "uniq_id": uniq_id})
+
+
+@pytest.fixture(autouse=True)
+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_s2_agent
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_zotero",
+        dummy_zotero_agent,
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_pdf",
+        dummy_question_and_answer_agent,
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_paper_download",
+        dummy_paper_download_agent,
+    )
+    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 through public interface."""
+    dummy = DummyLLM(model_name="test-model")
+    # Test that the dummy LLM can be used (testing the class works)
+    assert dummy.model_name == "test-model"
+    # Test through public interface that internally calls _generate (covers lines 26-27)
+    # Use invoke which internally calls _generate
+    messages = [HumanMessage(content="test prompt")]
+    result = dummy.invoke(messages)
+    # Verify the internal state was set
+    assert hasattr(DummyLLM, "called_prompt")
+    assert result is not None
+    assert DummyLLM.called_prompt == "test prompt"
+
+
+def test_dummy_llm_llm_type():
+    """Test the dummy LLM's type identification."""
+    dummy = DummyLLM(model_name="test-model")
+
+    # Use public shim instead of protected attribute access
+    llm_type = dummy.public_llm_type()
+    assert llm_type == "dummy"
+
+    # Also test the public string representation
+    assert "DummyLLM" in str(dummy.__class__.__name__)
+
+
+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"
+
+
+def test_dummy_workflow_get_supervisor_args():
+    """Test that DummyWorkflow.get_supervisor_args returns the stored args."""
+    dummy_args = {"agent": "test", "uniq_id": "id123"}
+    wf = DummyWorkflow(supervisor_args=dummy_args)
+    assert wf.get_supervisor_args() is dummy_args
+
+
+def test_dummy_hydra_compose_get_config():
+    """Test that DummyHydraCompose.get_config returns the raw config."""
+    config_dict = {"agents": {"test": {"key": "value"}}}
+    compose = DummyHydraCompose(config_dict)
+    assert compose.get_config() is config_dict

aiagents4pharma/talk2scholars/tests/test_agents_paper_agents_download_agent.py

@@ -0,0 +1,139 @@
+"""Unit tests for the paper download agent in Talk2Scholars."""
+
+from unittest import mock
+
+import pytest
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage
+
+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.paper_download_agent.paper_download_agent = "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.paper_downloader.download_papers"
+    ) as mock_download_papers:
+        mock_download_papers.return_value = {"article_data": {"dummy_key": "dummy_value"}}
+        yield [mock_download_papers]
+
+
+@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")],
+            "article_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 "article_data" in result
+
+
+def test_paper_download_agent_tools_assignment(
+    request,
+):
+    """Checks correct tool assignment (download_papers tool)."""
+    thread_id = "test_thread_paper_dl"
+    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_toolnode.return_value = mock_tool_instance
+
+        get_app(thread_id, llm_mock)
+        # Verify ToolNode was called with download_papers function
+        assert mock_toolnode.called
+        # Check that ToolNode was called with a list containing the download_papers tool
+        call_args = mock_toolnode.call_args[0][0]  # Get first positional argument (the tools list)
+        assert len(call_args) == 1
+        # The tool should be a StructuredTool with name 'download_papers'
+        tool = call_args[0]
+        assert hasattr(tool, "name")
+        assert tool.name == "download_papers"
+
+
+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_agents_pdf_agent.py

@@ -0,0 +1,114 @@
+"""
+Unit Tests for the PDF agent.
+"""
+
+from unittest import mock
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+
+from ..agents.pdf_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:
+        # Create a mock configuration with a pdf_agent section.
+        cfg_mock = mock.MagicMock()
+        # The pdf_agent config will be accessed as cfg.agents.talk2scholars.pdf_agent in get_app.
+        cfg_mock.agents.talk2scholars.pdf_agent.some_property = "Test prompt"
+        mock_compose.return_value = cfg_mock
+        yield mock_compose
+
+
+@pytest.fixture
+def mock_tools_fixture():
+    """Mock PDF agent tools to prevent execution of real API calls."""
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.pdf_agent.question_and_answer"
+        ) as mock_question_and_answer,
+    ):
+        mock_question_and_answer.return_value = {"result": "Mock Question and Answer Result"}
+        yield [mock_question_and_answer]
+
+
+@pytest.fixture(name="mock_llm")
+def llm_fixture():
+    """Provide a dummy language model to pass into get_app."""
+    return mock.Mock()
+
+
+@pytest.mark.usefixtures("mock_hydra_fixture")
+def test_pdf_agent_initialization(mock_llm):
+    """Test that PDF agent initializes correctly with mock configuration."""
+    thread_id = "test_thread"
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.pdf_agent.create_react_agent"
+    ) as mock_create:
+        mock_create.return_value = mock.Mock()
+        app = get_app(thread_id, mock_llm)
+        assert app is not None
+        assert mock_create.called
+
+
+def test_pdf_agent_invocation(mock_llm):
+    """Test that the PDF agent processes user input and returns a valid response."""
+    thread_id = "test_thread"
+    # Create a sample state with a human message.
+    mock_state = Talk2Scholars(messages=[HumanMessage(content="Extract key data from PDF")])
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.agents.pdf_agent.create_react_agent"
+    ) as mock_create:
+        mock_agent = mock.Mock()
+        mock_create.return_value = mock_agent
+        # Simulate a response from the PDF agent.
+        mock_agent.invoke.return_value = {
+            "messages": [AIMessage(content="PDF content extracted successfully")],
+            "article_data": {"page": 1, "text": "Sample PDF text"},
+        }
+        app = get_app(thread_id, mock_llm)
+        result = app.invoke(
+            mock_state,
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": "test_ns",
+                    "checkpoint_id": "test_checkpoint",
+                }
+            },
+        )
+        assert "messages" in result
+        assert "article_data" in result
+        assert result["article_data"]["page"] == 1
+
+
+def test_pdf_agent_tools_assignment(request, mock_llm):
+    """Ensure that the correct tools are assigned to the PDF agent."""
+    thread_id = "test_thread"
+    mock_tools = request.getfixturevalue("mock_tools_fixture")
+    with (
+        mock.patch(
+            "aiagents4pharma.talk2scholars.agents.pdf_agent.create_react_agent"
+        ) as mock_create,
+        mock.patch("aiagents4pharma.talk2scholars.agents.pdf_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, mock_llm)
+        assert mock_toolnode.called
+        assert len(mock_tool_instance.tools) == 1
+
+
+def test_pdf_agent_hydra_failure(mock_llm):
+    """Test exception handling when Hydra fails to load config for PDF agent."""
+    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, mock_llm)
+        assert "Hydra error" in str(exc_info.value)