aiagents4pharma 1.31.0__py3-none-any.whl → 1.33.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_zotero_read.py

@@ -7,7 +7,12 @@ import unittest
 from unittest.mock import patch, MagicMock
 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,
+)
 
+# pylint: disable=protected-access
+# pylint: disable=protected-access, too-many-arguments, too-many-positional-arguments
 
 # Dummy Hydra configuration to be used in tests
 dummy_zotero_read_config = SimpleNamespace(
@@ -92,11 +97,11 @@ class TestZoteroSearchTool(unittest.TestCase):
         # Verify the Command update structure and contents
         self.assertIsInstance(result, Command)
         update = result.update
-        self.assertIn("zotero_read", update)
+        self.assertIn("article_data", update)
         self.assertIn("last_displayed_papers", update)
         self.assertIn("messages", update)
 
-        filtered_papers = update["zotero_read"]
+        filtered_papers = update["article_data"]
         self.assertIn("paper1", filtered_papers)
         self.assertIn("paper2", filtered_papers)
         message_content = update["messages"][0].content
@@ -149,7 +154,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         result = zotero_read.run(tool_input)
 
         update = result.update
-        filtered_papers = update["zotero_read"]
+        filtered_papers = update["article_data"]
         self.assertIn("paper1", filtered_papers)
         fake_zot.items.assert_called_with(
             limit=dummy_cfg.tools.zotero_read.zotero.max_limit
@@ -198,17 +203,19 @@ class TestZoteroSearchTool(unittest.TestCase):
     @patch(
         "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.hydra.initialize"
     )
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper."
+        "ZoteroSearchData._download_pdfs_in_parallel"
+    )
     def test_filtering_no_matching_papers(
         self,
+        mock_batch_download,
         mock_hydra_init,
         mock_hydra_compose,
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """
-        Test that when non-research items (e.g. attachments, notes) are returned,
-        they are still included since filtering is disabled.
-        """
+        """Testing filtering when no paper matching"""
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
 
@@ -222,6 +229,8 @@ class TestZoteroSearchTool(unittest.TestCase):
                     "date": "2021",
                     "url": "http://example.com",
                     "itemType": "attachment",
+                    "contentType": "application/pdf",  # orphaned
+                    "filename": "paper1.pdf",
                 }
             },
             {
@@ -242,20 +251,25 @@ class TestZoteroSearchTool(unittest.TestCase):
             "paper2": ["/Test Collection"],
         }
 
-        tool_call_id = "test_id_4"
+        mock_batch_download.return_value = {
+            "paper1": ("/tmp/fake_path.pdf", "paper1.pdf", "paper1")
+        }
+
         tool_input = {
             "query": "test",
-            "only_articles": True,
-            "tool_call_id": tool_call_id,
+            "only_articles": False,
+            "tool_call_id": "test_id_4",
             "limit": 2,
         }
-        # Instead of expecting a RuntimeError, we now expect both items to be returned.
+
         result = zotero_read.run(tool_input)
-        update = result.update
-        filtered_papers = update["zotero_read"]
+        filtered_papers = result.update["article_data"]
+
         self.assertIn("paper1", filtered_papers)
         self.assertIn("paper2", filtered_papers)
-        self.assertEqual(len(filtered_papers), 2)
+        self.assertEqual(filtered_papers["paper1"]["filename"], "paper1.pdf")
+        self.assertEqual(filtered_papers["paper1"]["pdf_url"], "/tmp/fake_path.pdf")
+        self.assertEqual(filtered_papers["paper1"]["source"], "zotero")
 
     @patch(
         "aiagents4pharma.talk2scholars.tools.zotero.utils.zotero_path.get_item_collections"
@@ -349,7 +363,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         result = zotero_read.run(tool_input)
 
         update = result.update
-        filtered_papers = update["zotero_read"]
+        filtered_papers = update["article_data"]
         self.assertIn("paper_valid", filtered_papers)
         self.assertEqual(len(filtered_papers), 1)
 
@@ -408,26 +422,390 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """
-        Test that if an item has a 'data' field that is not a dict, it is skipped.
-        """
+        """Test for no dict"""
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
 
         fake_zot = MagicMock()
-        # Supply one item whose "data" field is not a dict.
-        fake_items = [{"data": "this is not a dict"}]
+        # Make the item itself non-dict (not just `data`)
+        fake_items = ["this is not a dict"]
         fake_zot.items.return_value = fake_items
         mock_zotero_class.return_value = fake_zot
         mock_get_item_collections.return_value = {}
 
-        tool_call_id = "test_id_8"
         tool_input = {
             "query": "dummy",
             "only_articles": True,
-            "tool_call_id": tool_call_id,
+            "tool_call_id": "test_id_8",
             "limit": 2,
         }
+
         with self.assertRaises(RuntimeError) as context:
             zotero_read.run(tool_input)
         self.assertIn("No matching papers returned from Zotero", str(context.exception))
+
+    @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"
+    )
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper.requests.Session.get"
+    )
+    def test_pdf_attachment_success(
+        self,
+        mock_session_get,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Test for pdf attachment success"""
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        fake_zot = MagicMock()
+        fake_zot.items.return_value = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "Paper 1",
+                    "abstractNote": "Abstract 1",
+                    "date": "2021",
+                    "url": "http://example.com",
+                    "itemType": "journalArticle",
+                    "creators": [
+                        {
+                            "firstName": "John",
+                            "lastName": "Doe",
+                            "creatorType": "author",
+                        }
+                    ],
+                }
+            }
+        ]
+
+        fake_pdf_child = {
+            "data": {
+                "key": "attachment1",
+                "filename": "file1.pdf",
+                "contentType": "application/pdf",
+            }
+        }
+        fake_zot.children.return_value = [fake_pdf_child]
+        mock_zotero_class.return_value = fake_zot
+        mock_get_item_collections.return_value = {"paper1": ["/Test Collection"]}
+
+        # Mock successful PDF download via session
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.iter_content = lambda chunk_size: [b"fake pdf content"]
+        mock_response.headers = {
+            "Content-Disposition": 'attachment; filename="file1.pdf"'
+        }
+        mock_response.raise_for_status = lambda: None
+        mock_session_get.return_value = mock_response
+
+        tool_input = {
+            "query": "pdf test",
+            "only_articles": True,
+            "tool_call_id": "test_pdf_success",
+            "limit": 1,
+        }
+
+        result = zotero_read.run(tool_input)
+        paper = result.update["article_data"]["paper1"]
+
+        self.assertIn("pdf_url", paper)
+        self.assertTrue(paper["pdf_url"].endswith(".pdf"))
+        self.assertEqual(paper["filename"], "file1.pdf")
+        self.assertEqual(paper["attachment_key"], "attachment1")
+
+    @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_pdf_attachment_children_exception(
+        self,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Test that when children() raises an exception, PDF info is not added."""
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        fake_zot = MagicMock()
+        fake_items = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "Paper 1",
+                    "abstractNote": "Abstract 1",
+                    "date": "2021",
+                    "url": "http://example.com",
+                    "itemType": "journalArticle",
+                    "creators": [
+                        {
+                            "firstName": "John",
+                            "lastName": "Doe",
+                            "creatorType": "author",
+                        }
+                    ],
+                }
+            },
+        ]
+        fake_zot.items.return_value = fake_items
+
+        # Simulate children() raising an exception
+        fake_zot.children.side_effect = Exception("Child fetch error")
+        mock_zotero_class.return_value = fake_zot
+        mock_get_item_collections.return_value = {"paper1": ["/Test Collection"]}
+
+        tool_call_id = "test_pdf_children_exception"
+        tool_input = {
+            "query": "pdf test exception",
+            "only_articles": True,
+            "tool_call_id": tool_call_id,
+            "limit": 1,
+        }
+        result = zotero_read.run(tool_input)
+        filtered_papers = result.update["article_data"]
+
+        # Ensure no PDF-related keys are added
+        self.assertIn("paper1", filtered_papers)
+        self.assertNotIn("pdf_url", filtered_papers["paper1"])
+        self.assertNotIn("filename", filtered_papers["paper1"])
+        self.assertNotIn("attachment_key", filtered_papers["paper1"])
+
+    @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_pdf_attachment_missing_key(
+        self,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Test for pdf attachment missing"""
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        fake_zot = MagicMock()
+        fake_zot.items.return_value = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "Paper 1",
+                    "abstractNote": "Abstract 1",
+                    "date": "2021",
+                    "url": "http://example.com",
+                    "itemType": "journalArticle",
+                    "creators": [
+                        {
+                            "firstName": "Alice",
+                            "lastName": "Smith",
+                            "creatorType": "author",
+                        }
+                    ],
+                }
+            },
+        ]
+
+        fake_pdf_child = {
+            "data": {
+                "filename": "no_key.pdf",
+                "contentType": "application/pdf",
+            }
+        }
+        fake_zot.children.return_value = [fake_pdf_child]
+        mock_zotero_class.return_value = fake_zot
+        mock_get_item_collections.return_value = {"paper1": ["/Test Collection"]}
+
+        tool_input = {
+            "query": "missing key test",
+            "only_articles": True,
+            "tool_call_id": "test_pdf_missing_key",
+            "limit": 1,
+        }
+
+        result = zotero_read.run(tool_input)
+        paper = result.update["article_data"]["paper1"]
+
+        self.assertNotIn("pdf_url", paper)
+        self.assertNotIn("attachment_key", paper)
+
+    @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_pdf_attachment_outer_exception(
+        self,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Test that if children() returns a non-iterable (causing an exception),
+        PDF info is not added."""
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        fake_zot = MagicMock()
+        fake_items = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "Paper 1",
+                    "abstractNote": "Abstract 1",
+                    "date": "2021",
+                    "url": "http://example.com",
+                    "itemType": "journalArticle",
+                    "creators": [
+                        {
+                            "firstName": "Bob",
+                            "lastName": "Jones",
+                            "creatorType": "author",
+                        }
+                    ],
+                }
+            },
+        ]
+        fake_zot.items.return_value = fake_items
+
+        # Simulate children() returning None to trigger an exception in list comprehension.
+        fake_zot.children.return_value = None
+
+        mock_zotero_class.return_value = fake_zot
+        mock_get_item_collections.return_value = {"paper1": ["/Test Collection"]}
+
+        tool_call_id = "test_pdf_outer_exception"
+        tool_input = {
+            "query": "outer exception test",
+            "only_articles": True,
+            "tool_call_id": tool_call_id,
+            "limit": 1,
+        }
+        result = zotero_read.run(tool_input)
+        filtered_papers = result.update["article_data"]
+
+        # Ensure no PDF-related keys are added if an exception occurs
+        self.assertIn("paper1", filtered_papers)
+        self.assertNotIn("pdf_url", filtered_papers["paper1"])
+        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"
+    )
+    @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"
+        )
+
+        result = zotero_search._download_zotero_pdf("FAKE_ATTACHMENT_KEY")
+
+        self.assertIsNone(result)
+
+    @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_pdf_exception_logging(
+        self,
+        mock_hydra_init,
+        mock_hydra_compose,
+        mock_zotero_class,
+        mock_get_item_collections,
+    ):
+        """Test that a failed download logs the error and does not break the pipeline."""
+        mock_hydra_compose.return_value = dummy_cfg
+        mock_hydra_init.return_value.__enter__.return_value = None
+
+        fake_zot = MagicMock()
+        fake_zot.items.return_value = [
+            {
+                "data": {
+                    "key": "paper1",
+                    "title": "Fake Title",
+                    "itemType": "journalArticle",
+                }
+            }
+        ]
+        # Simulate an attachment
+        fake_zot.children.return_value = [
+            {
+                "data": {
+                    "key": "attachment1",
+                    "filename": "file1.pdf",
+                    "contentType": "application/pdf",
+                }
+            }
+        ]
+
+        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
+        with patch(
+            "aiagents4pharma.talk2scholars.tools.zotero.utils.read_helper."
+            "ZoteroSearchData._download_zotero_pdf"
+        ) as mock_download_pdf:
+            mock_download_pdf.side_effect = Exception("Simulated download error")
+
+            search = ZoteroSearchData(
+                query="failure test",
+                only_articles=True,
+                limit=1,
+                tool_call_id="fail_test",
+            )
+            search.process_search()
+
+            article_data = search.get_search_results()["article_data"]
+            assert "paper1" in article_data
+            assert "pdf_url" not in article_data["paper1"]  # download failed, no URL

aiagents4pharma/talk2scholars/tools/paper_download/__init__.py

@@ -4,14 +4,8 @@ This package provides modules for fetching and downloading academic papers from
 """
 
 # Import modules
-from . import abstract_downloader
-from . import arxiv_downloader
 from . import download_arxiv_input
-from .download_arxiv_input import download_arxiv_paper
 
 __all__ = [
-    "abstract_downloader",
-    "arxiv_downloader",
     "download_arxiv_input",
-    "download_arxiv_paper",
 ]

aiagents4pharma/talk2scholars/tools/paper_download/download_arxiv_input.py

@@ -1,64 +1,122 @@
-# File: aiagents4pharma/talk2scholars/tools/paper_download/download_arxiv_input.py
+#!/usr/bin/env python3
 """
-This module defines the `download_arxiv_paper` tool, which leverages the
-`ArxivPaperDownloader` class to fetch and download academic papers from arXiv
-based on their unique arXiv ID.
+Tool for downloading arXiv paper metadata and retrieving the PDF URL.
 """
+
+import logging
+import xml.etree.ElementTree as ET
 from typing import Annotated, Any
-from pydantic import BaseModel, Field
-from langchain_core.tools import tool
+
+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
 
-# Local import from the same package:
-from .arxiv_downloader import ArxivPaperDownloader
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 class DownloadArxivPaperInput(BaseModel):
-    """
-    Input schema for the arXiv paper download tool.
-    (Optional: if you decide to keep Pydantic validation in the future)
-    """
+    """Input schema for the arXiv paper download tool."""
 
     arxiv_id: str = Field(
-        description="The arXiv paper ID used to retrieve the paper details and PDF."
+        description="The arXiv paper ID used to retrieve the paper details and PDF URL."
     )
     tool_call_id: Annotated[str, InjectedToolCallId]
 
 
+def fetch_arxiv_metadata(
+    api_url: str, arxiv_id: str, request_timeout: int
+) -> ET.Element:
+    """Fetch and parse metadata from the arXiv API."""
+    query_url = f"{api_url}?search_query=id:{arxiv_id}&start=0&max_results=1"
+    response = requests.get(query_url, timeout=request_timeout)
+    response.raise_for_status()
+    return ET.fromstring(response.text)
+
+
+def extract_metadata(entry: ET.Element, ns: dict, arxiv_id: str) -> dict:
+    """Extract metadata from the XML entry."""
+    title_elem = entry.find("atom:title", ns)
+    title = title_elem.text.strip() if title_elem is not None else "N/A"
+
+    authors = [
+        author_elem.find("atom:name", ns).text.strip()
+        for author_elem in entry.findall("atom:author", ns)
+        if author_elem.find("atom:name", ns) is not None
+    ]
+
+    summary_elem = entry.find("atom:summary", ns)
+    abstract = summary_elem.text.strip() if summary_elem is not None else "N/A"
+
+    published_elem = entry.find("atom:published", ns)
+    pub_date = published_elem.text.strip() if published_elem is not None else "N/A"
+
+    pdf_url = next(
+        (
+            link.attrib.get("href")
+            for link in entry.findall("atom:link", ns)
+            if link.attrib.get("title") == "pdf"
+        ),
+        None,
+    )
+    if not pdf_url:
+        raise RuntimeError(f"Could not find PDF URL for arXiv ID {arxiv_id}")
+
+    return {
+        "Title": title,
+        "Authors": authors,
+        "Abstract": abstract,
+        "Publication Date": pub_date,
+        "URL": pdf_url,
+        "pdf_url": pdf_url,
+        "filename": f"{arxiv_id}.pdf",
+        "source": "arxiv",
+        "arxiv_id": arxiv_id,
+    }
+
+
 @tool(args_schema=DownloadArxivPaperInput, parse_docstring=True)
 def download_arxiv_paper(
     arxiv_id: str,
     tool_call_id: Annotated[str, InjectedToolCallId],
 ) -> Command[Any]:
     """
-    Download an arXiv paper's PDF using its unique arXiv ID.
+    Get metadata and PDF URL for an arXiv paper using its unique arXiv ID.
+    """
+    logger.info("Fetching metadata from arXiv for paper ID: %s", arxiv_id)
 
-    This function:
-      1. Creates an `ArxivPaperDownloader` instance.
-      2. Fetches metadata from arXiv using the provided `arxiv_id`.
-      3. Downloads the PDF from the returned link.
-      4. Returns a `Command` object containing the PDF data and a success message.
+    # Load configuration
+    with hydra.initialize(version_base=None, config_path="../../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["tools/download_arxiv_paper=default"]
+        )
+        api_url = cfg.tools.download_arxiv_paper.api_url
+        request_timeout = cfg.tools.download_arxiv_paper.request_timeout
 
-    Args:
-        arxiv_id (str): The unique arXiv paper ID.
-        tool_call_id (InjectedToolCallId): A unique identifier for tracking this tool call.
+    # Fetch and parse metadata
+    root = fetch_arxiv_metadata(api_url, arxiv_id, request_timeout)
+    ns = {"atom": "http://www.w3.org/2005/Atom"}
 
-    Returns:
-        Command[Any]: Contains metadata and messages about the success of the operation.
-    """
-    downloader = ArxivPaperDownloader()
+    entry = root.find("atom:entry", ns)
+    if entry is None:
+        raise ValueError(f"No entry found for arXiv ID {arxiv_id}")
+
+    # Extract metadata
+    metadata = extract_metadata(entry, ns, arxiv_id)
 
-    # If the downloader fails or the arxiv_id is invalid, this might raise an error
-    pdf_data = downloader.download_pdf(arxiv_id)
-    # print (pdf_data)
+    # Create article_data entry with the paper ID as the key
+    article_data = {arxiv_id: metadata}
 
-    content = f"Successfully downloaded PDF for arXiv ID {arxiv_id}"
+    content = f"Successfully retrieved metadata and PDF URL for arXiv ID {arxiv_id}"
 
     return Command(
         update={
-            "pdf_data": pdf_data,
+            "article_data": article_data,
             "messages": [ToolMessage(content=content, tool_call_id=tool_call_id)],
         }
     )