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

aiagents4pharma/talk2scholars/configs/tools/zotero_read/default.yaml

@@ -2,6 +2,7 @@
 library_type: "user" # Type of library ('user' or 'group')
 default_limit: 2
 request_timeout: 10
+chunk_size: 16384  # Size (in bytes) for streaming PDF download chunks
 user_id: ${oc.env:ZOTERO_USER_ID} # Load from environment variable
 api_key: ${oc.env:ZOTERO_API_KEY} # Load from environment variable
 

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -7,6 +7,7 @@ across agent interactions.
 """
 
 import logging
+from collections.abc import Mapping
 from typing import Annotated, Any, Dict
 
 from langchain_core.embeddings import Embeddings
@@ -18,7 +19,24 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any]:
+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.
 
@@ -39,9 +57,13 @@ def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any
         >>> print(updated_state)
         {"papers": {"id2": "Paper 2"}}
     """
-    # No-op operation to use the 'existing' variable
-    _ = len(existing)
-    return new
+    # 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):
@@ -63,10 +85,14 @@ class Talk2Scholars(AgentState):
     """
 
     # 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]
-    papers: Annotated[Dict[str, Any], replace_dict]
-    multi_papers: Annotated[Dict[str, Any], replace_dict]
-    article_data: 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/test_question_and_answer_tool.py

@@ -3,11 +3,14 @@ Unit tests for question_and_answer tool functionality.
 """
 
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.messages import ToolMessage
 
+import aiagents4pharma.talk2scholars.tools.pdf.question_and_answer as qa_module
 from aiagents4pharma.talk2scholars.tools.pdf.question_and_answer import (
     Vectorstore,
     generate_answer,
@@ -145,8 +148,9 @@ class TestQuestionAndAnswerTool(unittest.TestCase):
 
         vector_store = Vectorstore(embedding_model=mock_embedding_model)
         vector_store.vector_store = True
+        # Add a document chunk with required metadata including chunk_id
         vector_store.documents["test_doc"] = Document(
-            page_content="Test content", metadata={"paper_id": "test_paper"}
+            page_content="Test content", metadata={"paper_id": "test_paper", "chunk_id": 0}
         )
 
         results = vector_store.retrieve_relevant_chunks(query="test query")
@@ -793,8 +797,9 @@ class TestMissingState(unittest.TestCase):
 
         vector_store = Vectorstore(embedding_model=mock_embedding_model)
         vector_store.vector_store = True
-        doc1 = Document(page_content="Doc 1", metadata={"paper_id": "paper1"})
-        doc2 = Document(page_content="Doc 2", metadata={"paper_id": "paper2"})
+        # Add document chunks with necessary metadata including chunk_ids
+        doc1 = Document(page_content="Doc 1", metadata={"paper_id": "paper1", "chunk_id": 0})
+        doc2 = Document(page_content="Doc 2", metadata={"paper_id": "paper2", "chunk_id": 1})
         vector_store.documents = {"doc1": doc1, "doc2": doc2}
 
         results = vector_store.retrieve_relevant_chunks(
@@ -820,3 +825,54 @@ class TestMissingState(unittest.TestCase):
             query="test", paper_ids=["nonexistent_id"]
         )
         assert results == []
+
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.pdf.question_and_answer.load_hydra_config"
+    )
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.pdf.question_and_answer.generate_answer"
+    )
+    def test_prebuilt_vector_store_branch(self, mock_generate, mock_load_config):
+        """Test question_and_answer tool with a shared pre-built vector store branch."""
+        # Mock configuration for tool-level thresholds
+        config = SimpleNamespace(top_k_papers=1, top_k_chunks=1)
+        mock_load_config.return_value = config
+        # Mock generate_answer to return a simple response
+        mock_generate.return_value = {"output_text": "Answer", "papers_used": ["p1"]}
+
+        # Prepare a dummy pre-built vector store
+        dummy_vs = SimpleNamespace(
+            loaded_papers=set(),
+            vector_store=True,
+            retrieve_relevant_chunks=lambda *_args, **_kwargs: [
+                Document(page_content="chunk", metadata={"paper_id": "p1"})
+            ],
+        )
+        # Override the module-level prebuilt_vector_store
+        qa_module.prebuilt_vector_store = dummy_vs
+
+        # Prepare state with required models and article_data
+        state = {
+            "text_embedding_model": MagicMock(),
+            "llm_model": MagicMock(),
+            "article_data": {"p1": {"source": "upload"}},
+        }
+
+        # Invoke the tool-level function via .run with appropriate input schema
+        input_data = {
+            "question": "What?",
+            "paper_ids": None,
+            "use_all_papers": False,
+            "tool_call_id": "testid",
+            "state": state,
+        }
+        result = qa_module.question_and_answer.run(input_data)
+
+        # Ensure the prebuilt branch was used and a Command is returned
+        self.assertTrue(hasattr(result, "update"))
+        messages = result.update.get("messages", [])
+        self.assertEqual(len(messages), 1)
+        self.assertIsInstance(messages[0], ToolMessage)
+
+        # Clean up global override
+        qa_module.prebuilt_vector_store = None

aiagents4pharma/talk2scholars/tests/test_read_helper_utils.py

@@ -0,0 +1,110 @@
+"""
+Unit tests for Zotero read helper download branches.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper import (
+    ZoteroSearchData,
+)
+
+# Dummy Hydra configuration for tests
+dummy_zotero_read_config = SimpleNamespace(
+    user_id="dummy_user",
+    library_type="user",
+    api_key="dummy_api_key",
+    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))
+
+
+class TestReadHelperDownloadsFalse(unittest.TestCase):
+    """Tests for read_helper download_pdfs=False branches."""
+
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_path.get_item_collections"
+    )
+    @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_pdfs_false_branches(
+        self,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Ensure attachment_key and filename are set when download_pdfs=False."""
+        # Setup Hydra mocks
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        # Fake Zotero items: one paper with child PDF, one orphaned PDF
+        fake_zot = MagicMock()
+        fake_items = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "P1",
+                    "abstractNote": "A1",
+                    "date": "2021",
+                    "url": "u1",
+                    "itemType": "journalArticle",
+                }
+            },
+            {
+                "data": {
+                    "key": "attach2",
+                    "itemType": "attachment",
+                    "contentType": "application/pdf",
+                    "filename": "file2.pdf",
+                }
+            },
+        ]
+        fake_zot.items.return_value = fake_items
+        # children for paper1
+        fake_child = {
+            "data": {
+                "key": "attach1",
+                "filename": "file1.pdf",
+                "contentType": "application/pdf",
+            }
+        }
+
+        def children_side_effect(key):
+            return [fake_child] if key == "paper1" else []
+
+        fake_zot.children.side_effect = children_side_effect
+        mock_zotero_class.return_value = fake_zot
+        mock_get_item_collections.return_value = {"paper1": ["/C1"], "attach2": ["/C2"]}
+
+        # Instantiate with download_pdfs=False
+        search = ZoteroSearchData(
+            query="test",
+            only_articles=False,
+            limit=2,
+            tool_call_id="id",
+            download_pdfs=False,
+        )
+        search.process_search()
+        data = search.get_search_results()["article_data"]
+
+        # Regular paper1 should have attachment_key and filename, no pdf_url
+        self.assertIn("paper1", data)
+        self.assertEqual(data["paper1"]["attachment_key"], "attach1")
+        self.assertEqual(data["paper1"]["filename"], "file1.pdf")
+        self.assertNotIn("pdf_url", data["paper1"])
+
+        # Orphan attach2 should have attachment_key and filename, no pdf_url
+        self.assertIn("attach2", data)
+        self.assertEqual(data["attach2"]["attachment_key"], "attach2")
+        self.assertEqual(data["attach2"]["filename"], "file2.pdf")
+        self.assertNotIn("pdf_url", data["attach2"])

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",