aiagents4pharma 1.37.0__py3-none-any.whl → 1.39.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_s2_display.py

@@ -52,7 +52,9 @@ class TestS2Tools:
             raised_error,
             match="No papers found. A search/rec needs to be performed first.",
         ):
-            display_dataframe.invoke({"state": initial_state, "tool_call_id": "test123"})
+            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"""
@@ -72,3 +74,20 @@ class TestS2Tools:
             "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_query.py

@@ -76,3 +76,20 @@ class TestS2Tools:
 
         assert isinstance(result, str)  # Ensure output is a string
         assert result == "Mocked response"  # Validate the expected 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
+        result = query_dataframe.invoke({"question": "Filter papers", "state": state})
+        assert isinstance(result, str)
+        assert result == "Direct mapping response"

aiagents4pharma/talk2scholars/tests/test_state.py

@@ -2,7 +2,7 @@
 Tests for state management functionality.
 """
 
-from ..state.state_talk2scholars import replace_dict
+from ..state.state_talk2scholars import merge_dict, replace_dict
 
 
 def test_state_replace_dict():
@@ -12,3 +12,27 @@ def test_state_replace_dict():
     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"}

aiagents4pharma/talk2scholars/tests/test_zotero_pdf_downloader_utils.py

@@ -0,0 +1,46 @@
+"""
+Unit tests for Zotero PDF downloader utilities.
+"""
+
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+import requests
+
+from aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_pdf_downloader import (
+    download_pdfs_in_parallel,
+    download_zotero_pdf,
+)
+
+
+class TestZoteroPDFDownloaderUtils(unittest.TestCase):
+    """Tests for zotero_pdf_downloader module."""
+
+    @patch("requests.Session.get")
+    def test_download_zotero_pdf_default_filename(self, mock_get):
+        """Test download_zotero_pdf returns default filename when header has no filename."""
+        # Mock response without Content-Disposition filename
+        mock_response = MagicMock()
+        mock_response.raise_for_status = lambda: None
+        mock_response.iter_content = lambda chunk_size: [b"fakepdf"]
+        mock_response.headers = {}
+        mock_get.return_value = mock_response
+
+        session = requests.Session()
+        result = download_zotero_pdf(session, "user123", "apikey", "attach123")
+        # Should return a tuple (file_path, filename)
+        self.assertIsNotNone(result)
+        file_path, filename = result
+        # File should exist
+        self.assertTrue(os.path.isfile(file_path))
+        # Filename should default to 'downloaded.pdf'
+        self.assertEqual(filename, "downloaded.pdf")
+        # Clean up temp file
+        os.remove(file_path)
+
+    def test_download_pdfs_in_parallel_empty(self):
+        """Test that download_pdfs_in_parallel returns empty dict on empty input."""
+        session = requests.Session()
+        result = download_pdfs_in_parallel(session, "user123", "apikey", {})
+        self.assertEqual(result, {})

aiagents4pharma/talk2scholars/tests/test_zotero_read.py

@@ -2,14 +2,20 @@
 Unit tests for Zotero search tool in zotero_read.py.
 """
 
-from types import SimpleNamespace
 import unittest
-from unittest.mock import patch, MagicMock
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import requests
 from langgraph.types import Command
-from aiagents4pharma.talk2scholars.tools.zotero.zotero_read import zotero_read
+
 from aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper import (
     ZoteroSearchData,
 )
+from aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_pdf_downloader import (
+    download_zotero_pdf,
+)
+from aiagents4pharma.talk2scholars.tools.zotero.zotero_read import zotero_read
 
 # pylint: disable=protected-access
 # pylint: disable=protected-access, too-many-arguments, too-many-positional-arguments
@@ -22,7 +28,6 @@ dummy_zotero_read_config = SimpleNamespace(
     zotero=SimpleNamespace(
         max_limit=5,
         filter_item_types=["journalArticle", "conferencePaper"],
-        filter_excluded_types=["attachment", "note"],
     ),
 )
 dummy_cfg = SimpleNamespace(tools=SimpleNamespace(zotero_read=dummy_zotero_read_config))
@@ -204,8 +209,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.hydra.initialize"
     )
     @patch(
-        "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper."
-        "ZoteroSearchData._download_pdfs_in_parallel"
+        "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.download_pdfs_in_parallel"
     )
     def test_filtering_no_matching_papers(
         self,
@@ -260,6 +264,7 @@ class TestZoteroSearchTool(unittest.TestCase):
             "only_articles": False,
             "tool_call_id": "test_id_4",
             "limit": 2,
+            "download_pdfs": True,
         }
 
         result = zotero_read.run(tool_input)
@@ -514,6 +519,7 @@ class TestZoteroSearchTool(unittest.TestCase):
             "only_articles": True,
             "tool_call_id": "test_pdf_success",
             "limit": 1,
+            "download_pdfs": True,
         }
 
         result = zotero_read.run(tool_input)
@@ -713,39 +719,26 @@ class TestZoteroSearchTool(unittest.TestCase):
         self.assertNotIn("filename", filtered_papers["paper1"])
         self.assertNotIn("attachment_key", filtered_papers["paper1"])
 
-    @patch("aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.requests.get")
     @patch(
-        "aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_path.get_item_collections"
+        "aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_pdf_downloader."
+        "requests.Session.get"
     )
-    @patch("aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.zotero.Zotero")
-    @patch("aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.hydra.compose")
-    @patch(
-        "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.hydra.initialize"
-    )
-    def test_download_zotero_pdf_exception(
-        self,
-        mock_hydra_init,
-        mock_hydra_compose,
-        mock_zotero_class,
-        mock_get_item_collections,
-        mock_requests_get,
-    ):
-        """Test that _download_zotero_pdf returns None and logs error on request exception."""
-        # Setup mocks for config and Zotero client
-        mock_hydra_compose.return_value = dummy_cfg
-        mock_hydra_init.return_value.__enter__.return_value = None
-        mock_zotero_class.return_value = MagicMock()
-        mock_get_item_collections.return_value = {}
-
-        # Simulate a request exception during PDF download
-        mock_requests_get.side_effect = Exception("Simulated download failure")
-
-        zotero_search = ZoteroSearchData(
-            query="test", only_articles=False, limit=1, tool_call_id="test123"
+    def test_download_zotero_pdf_exception(self, mock_session_get):
+        """Test that download_zotero_pdf returns None and logs error on request exception."""
+        # Simulate a session.get exception during PDF download
+        mock_session_get.side_effect = requests.exceptions.RequestException(
+            "Simulated download failure"
         )
-
-        result = zotero_search._download_zotero_pdf("FAKE_ATTACHMENT_KEY")
-
+        # Create a session for testing
+        session = requests.Session()
+        # Call the module-level download function
+        result = download_zotero_pdf(
+            session,
+            dummy_cfg.tools.zotero_read.user_id,
+            dummy_cfg.tools.zotero_read.api_key,
+            "FAKE_ATTACHMENT_KEY",
+        )
+        # Should return None on failure
         self.assertIsNone(result)
 
     @patch(
@@ -791,12 +784,14 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class.return_value = fake_zot
         mock_get_item_collections.return_value = {"paper1": ["/Fake Collection"]}
 
-        # Patch just the internal _download_zotero_pdf to raise an exception
+        # Patch the module-level download_zotero_pdf to raise an exception
         with patch(
-            "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper."
-            "ZoteroSearchData._download_zotero_pdf"
+            "aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_pdf_downloader."
+            "download_zotero_pdf"
         ) as mock_download_pdf:
-            mock_download_pdf.side_effect = Exception("Simulated download error")
+            mock_download_pdf.side_effect = requests.exceptions.RequestException(
+                "Simulated download error"
+            )
 
             search = ZoteroSearchData(
                 query="failure test",

aiagents4pharma/talk2scholars/tools/paper_download/__init__.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """
-This package provides modules for fetching and downloading academic papers from arXiv.
+This package provides modules for fetching and downloading academic papers from arXiv,
+biorxiv and medrxiv.
 """
 
 # Import modules
@@ -8,4 +9,6 @@ from . import download_arxiv_input
 
 __all__ = [
     "download_arxiv_input",
+    "download_biorxiv_input",
+    "download_medrxiv_input",
 ]

aiagents4pharma/talk2scholars/tools/paper_download/download_biorxiv_input.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+Tool for downloading bioRxiv paper metadata and retrieving the PDF URL.
+"""
+
+import logging
+from typing import Annotated, Any
+
+import hydra
+import requests
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DownloadBiorxivPaperInput(BaseModel):
+    """Input schema for the bioRxiv paper download tool."""
+
+    doi: str = Field(description=
+    """The bioRxiv DOI, from search_helper or multi_helper or single_helper, 
+    used to retrieve the paper details and PDF URL."""
+    )
+    logger.info("DOI Received: %s", doi)
+    tool_call_id: Annotated[str, InjectedToolCallId]
+
+def fetch_biorxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict:
+    """
+    Fetch metadata for a bioRxiv paper using its DOI and extract relevant fields.
+    
+    Parameters:
+        doi (str): The DOI of the bioRxiv paper.
+    
+    Returns:
+        dict: A dictionary containing the title, authors, abstract, publication date, and URLs.
+    """
+    # Strip any version suffix (e.g., v1) since bioRxiv's API is version-sensitive
+    clean_doi = doi.split("v")[0]
+
+    api_url = f"{api_url}{clean_doi}"
+    logger.info("Fetching metadata from api url: %s", api_url)
+    response = requests.get(api_url, timeout=request_timeout)
+    response.raise_for_status()
+
+    data = response.json()
+    if not data.get("collection"):
+        raise ValueError(f"No metadata found for DOI: {doi}")
+
+    data = response.json()
+
+    return data["collection"][0]
+
+def extract_metadata(paper: dict, doi: str) -> dict:
+    """
+    Extract relevant metadata fields from a bioRxiv paper entry.
+    """
+    title = paper.get("title", "")
+    authors = paper.get("authors", "")
+    abstract = paper.get("abstract", "")
+    pub_date = paper.get("date", "")
+    doi_suffix = paper.get("doi", "").split("10.1101/")[-1]
+    pdf_url = f"https://www.biorxiv.org/content/10.1101/{doi_suffix}.full.pdf"
+    logger.info("PDF URL: %s", pdf_url)
+    return {
+        "Title": title,
+        "Authors": authors,
+        "Abstract": abstract,
+        "Publication Date": pub_date,
+        "URL": pdf_url,
+        "pdf_url": pdf_url,
+        "filename": f"{doi_suffix}.pdf",
+        "source": "biorxiv",
+        "biorxiv_id": doi
+    }
+
+@tool(args_schema=DownloadBiorxivPaperInput, parse_docstring=True)
+def download_biorxiv_paper(
+    doi: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+) -> Command[Any]:
+    """
+    Get metadata and PDF URL for a bioRxiv paper using its DOI.
+    """
+    logger.info("Fetching metadata from bioRxiv for DOI: %s", doi)
+
+    # Load configuration
+    with hydra.initialize(version_base=None, config_path="../../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["tools/download_biorxiv_paper=default"]
+        )
+        api_url = cfg.tools.download_biorxiv_paper.api_url
+        request_timeout = cfg.tools.download_biorxiv_paper.request_timeout
+        logger.info("API URL: %s", api_url)
+        logger.info("Request Timeout: %s", request_timeout)
+
+    # Fetch metadata
+    raw_data = fetch_biorxiv_metadata(doi, api_url, request_timeout)
+    metadata = extract_metadata(raw_data, doi)
+    article_data = {doi: metadata}
+    content = f"Successfully retrieved metadata and PDF URL for bioRxiv DOI {doi}"
+
+    return Command(
+        update={
+            "article_data": article_data,
+            "messages": [ToolMessage(content=content, tool_call_id=tool_call_id)],
+        }
+    )

aiagents4pharma/talk2scholars/tools/paper_download/download_medrxiv_input.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+Tool for downloading medRxiv paper metadata and retrieving the PDF URL.
+"""
+
+import logging
+from typing import Annotated, Any
+
+import hydra
+import requests
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DownloadMedrxivPaperInput(BaseModel):
+    """Input schema for the medRxiv paper download tool."""
+
+    doi: str = Field(description=
+    """The medRxiv DOI, from search_helper or multi_helper or single_helper, 
+    used to retrieve the paper details and PDF URL."""
+    )
+    logger.info("DOI Received: %s", doi)
+    tool_call_id: Annotated[str, InjectedToolCallId]
+
+# Fetching raw metadata from medRxiv API for a given DOI
+def fetch_medrxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict:
+    """
+    Fetch metadata for a medRxiv paper using its DOI and extract relevant fields.
+
+    Parameters:
+        doi (str): The DOI of the medRxiv paper.
+
+    Returns:
+        dict: A dictionary containing the title, authors, abstract, publication date, and URLs.
+    """
+    # Strip any version suffix (e.g., v1) since bioRxiv's API is version-sensitive
+    clean_doi = doi.split("v")[0]
+
+    api_url = f"{api_url}{clean_doi}"
+    logger.info("Fetching metadata from api url: %s", api_url)
+    response = requests.get(api_url, timeout=request_timeout)
+    response.raise_for_status()
+
+    data = response.json()
+    if not data.get("collection"):
+        raise ValueError(f"No entry found for medRxiv ID {doi}")
+
+    return data["collection"][0]
+
+# Extracting relevant metadata fields from the raw data
+def extract_metadata(paper: dict, doi: str) -> dict:
+    """
+    Extract relevant metadata fields from a medRxiv paper entry.
+    """
+    title = paper.get("title", "")
+    authors = paper.get("authors", "")
+    abstract = paper.get("abstract", "")
+    pub_date = paper.get("date", "")
+    doi_suffix = paper.get("doi", "").split("10.1101/")[-1]
+    pdf_url = f"https://www.medrxiv.org/content/10.1101/{doi_suffix}.full.pdf"
+    logger.info("PDF URL: %s", pdf_url)
+    return {
+        "Title": title,
+        "Authors": authors,
+        "Abstract": abstract,
+        "Publication Date": pub_date,
+        "URL": pdf_url,
+        "pdf_url": pdf_url,
+        "filename": f"{doi_suffix}.pdf",
+        "source": "medrxiv",
+        "medrxiv_id": doi
+    }
+
+# Tool to download medRxiv paper metadata and PDF URL
+@tool(args_schema=DownloadMedrxivPaperInput, parse_docstring=True)
+def download_medrxiv_paper(
+    doi: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+) -> Command[Any]:
+    """
+    Get metadata and PDF URL for a medRxiv paper using its doi or medrxiv id.
+    """
+    logger.info("Fetching metadata from medRxiv for DOI: %s", doi)
+
+    # Load configuration
+    with hydra.initialize(version_base=None, config_path="../../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["tools/download_medrxiv_paper=default"]
+        )
+        api_url = cfg.tools.download_medrxiv_paper.api_url
+        request_timeout = cfg.tools.download_medrxiv_paper.request_timeout
+        logger.info("API URL: %s", api_url)
+
+    raw_data = fetch_medrxiv_metadata(doi, api_url, request_timeout)
+    metadata = extract_metadata(raw_data, doi)
+    article_data = {doi: metadata}
+
+    content = f"Successfully retrieved metadata and PDF URL for medRxiv DOI {doi}"
+
+    return Command(
+        update={
+            "article_data": article_data,
+            "messages": [ToolMessage(content=content, tool_call_id=tool_call_id)],
+        }
+    )