aiagents4pharma 1.30.4__py3-none-any.whl → 1.32.0__py3-none-any.whl

aiagents4pharma/talk2scholars/agents/main_agent.py

@@ -68,15 +68,16 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         [
             get_app_s2(uniq_id, llm_model),  # semantic scholar
             get_app_zotero(uniq_id, llm_model),  # zotero
-            get_app_pdf(uniq_id, llm_model),  # pdf
-            get_app_paper_download(uniq_id, llm_model),  # paper download
+            get_app_paper_download(uniq_id, llm_model),  # pdf
+            get_app_pdf(uniq_id, llm_model),  # paper download
         ],
         model=llm_model,
         state_schema=Talk2Scholars,
         # Full history is needed to extract
         # the tool artifacts
         output_mode="full_history",
-        add_handoff_back_messages=False,
+        # Allow supervisor to resume control and chain multiple sub-agent calls
+        add_handoff_back_messages=True,
         prompt=cfg.system_prompt,
     )
 

aiagents4pharma/talk2scholars/agents/paper_download_agent.py

@@ -13,8 +13,7 @@ from langgraph.prebuilt.chat_agent_executor import create_react_agent
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.checkpoint.memory import MemorySaver
 from ..state.state_talk2scholars import Talk2Scholars
-from ..tools.paper_download import download_arxiv_paper
-from ..tools.s2.query_results import query_results
+from ..tools.paper_download.download_arxiv_input import download_arxiv_paper
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -45,7 +44,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         cfg = cfg.agents.talk2scholars.paper_download_agent
 
     # Define tools properly
-    tools = ToolNode([download_arxiv_paper, query_results])
+    tools = ToolNode([download_arxiv_paper])
 
     # Define the model
     logger.info("Using OpenAI model %s", llm_model)
@@ -78,7 +77,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     checkpointer = MemorySaver()
 
     # Compile the graph
-    app = workflow.compile(checkpointer=checkpointer, name="agent_paper_download")
+    app = workflow.compile(checkpointer=checkpointer, name="paper_download_agent")
 
     # Logging the information and returning the app
     logger.info("Compiled the graph")

aiagents4pharma/talk2scholars/agents/pdf_agent.py

@@ -18,8 +18,7 @@ from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.checkpoint.memory import MemorySaver
 from ..state.state_talk2scholars import Talk2Scholars
-from ..tools.pdf.question_and_answer import question_and_answer_tool
-from ..tools.s2.query_results import query_results
+from ..tools.pdf.question_and_answer import question_and_answer
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -56,7 +55,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         cfg = cfg.agents.talk2scholars.pdf_agent
         logger.info("Loaded pdf_agent configuration.")
 
-    def agent_pdf_node(state: Talk2Scholars):
+    def pdf_agent_node(state: Talk2Scholars):
         """
         Processes the current state by invoking the language model for PDF question and answer.
 
@@ -72,7 +71,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         return response
 
     # Define the tool node that includes the PDF QnA tool.
-    tools = ToolNode([question_and_answer_tool, query_results])
+    tools = ToolNode([question_and_answer])
 
     logger.info("Using OpenAI model %s", llm_model)
 
@@ -87,14 +86,14 @@ def get_app(uniq_id, llm_model: BaseChatModel):
 
     # Define a new workflow graph with the state schema.
     workflow = StateGraph(Talk2Scholars)
-    workflow.add_node("agent_pdf", agent_pdf_node)
-    workflow.add_edge(START, "agent_pdf")
+    workflow.add_node("pdf_agent", pdf_agent_node)
+    workflow.add_edge(START, "pdf_agent")
 
     # Initialize memory to persist state between runs.
     checkpointer = MemorySaver()
 
     # Compile the graph into a runnable app.
-    app = workflow.compile(checkpointer=checkpointer, name="agent_pdf")
+    app = workflow.compile(checkpointer=checkpointer, name="pdf_agent")
     logger.info("Compiled the PDF agent graph.")
 
     return app

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -6,22 +6,26 @@ Agent for interacting with Semantic Scholar
 
 import logging
 from typing import Any, Dict
+
 import hydra
 from langchain_core.language_models.chat_models import BaseChatModel
-from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import ToolNode, create_react_agent
+
 from ..state.state_talk2scholars import Talk2Scholars
-from ..tools.s2.search import search_tool as s2_search
-from ..tools.s2.display_results import display_results as s2_display
-from ..tools.s2.query_results import query_results as s2_query_results
+from ..tools.s2.display_dataframe import display_dataframe
+from ..tools.s2.multi_paper_rec import (
+    get_multi_paper_recommendations,
+)
+from ..tools.s2.query_dataframe import query_dataframe
 from ..tools.s2.retrieve_semantic_scholar_paper_id import (
-    retrieve_semantic_scholar_paper_id as s2_retrieve_id,
+    retrieve_semantic_scholar_paper_id,
 )
+from ..tools.s2.search import search_tool
 from ..tools.s2.single_paper_rec import (
-    get_single_paper_recommendations as s2_single_rec,
+    get_single_paper_recommendations,
 )
-from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations as s2_multi_rec
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -50,8 +54,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         >>> result = app.invoke(initial_state)
     """
 
-    # def agent_s2_node(state: Talk2Scholars) -> Command[Literal["supervisor"]]:
-    def agent_s2_node(state: Talk2Scholars) -> Dict[str, Any]:
+    def s2_agent_node(state: Talk2Scholars) -> Dict[str, Any]:
         """
         Processes the user query and retrieves relevant research papers.
 
@@ -67,7 +70,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
             Dict[str, Any]: A dictionary containing the updated conversation state.
 
         Example:
-            >>> result = agent_s2_node(current_state)
+            >>> result = s2_agent_node(current_state)
             >>> papers = result.get("papers", [])
         """
         logger.log(logging.INFO, "Creating Agent_S2 node with thread_id %s", uniq_id)
@@ -89,12 +92,12 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     # Define the tools
     tools = ToolNode(
         [
-            s2_search,
-            s2_display,
-            s2_query_results,
-            s2_retrieve_id,
-            s2_single_rec,
-            s2_multi_rec,
+            search_tool,
+            display_dataframe,
+            query_dataframe,
+            retrieve_semantic_scholar_paper_id,
+            get_single_paper_recommendations,
+            get_multi_paper_recommendations,
         ]
     )
 
@@ -111,8 +114,8 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     )
 
     workflow = StateGraph(Talk2Scholars)
-    workflow.add_node("agent_s2", agent_s2_node)
-    workflow.add_edge(START, "agent_s2")
+    workflow.add_node("s2_agent", s2_agent_node)
+    workflow.add_edge(START, "s2_agent")
 
     # Initialize memory to persist state between graph runs
     checkpointer = MemorySaver()
@@ -121,7 +124,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     # This compiles it into a LangChain Runnable,
     # meaning you can use it as you would any other runnable.
     # Note that we're (optionally) passing the memory when compiling the graph
-    app = workflow.compile(checkpointer=checkpointer, name="agent_s2")
+    app = workflow.compile(checkpointer=checkpointer, name="s2_agent")
     logger.log(
         logging.INFO,
         "Compiled the graph with thread_id %s and llm_model %s",

aiagents4pharma/talk2scholars/agents/zotero_agent.py

@@ -16,8 +16,8 @@ from ..state.state_talk2scholars import Talk2Scholars
 from ..tools.zotero.zotero_read import zotero_read
 from ..tools.zotero.zotero_review import zotero_review
 from ..tools.zotero.zotero_write import zotero_write
-from ..tools.s2.display_results import display_results as s2_display
-from ..tools.s2.query_results import query_results as s2_query_results
+from ..tools.s2.display_dataframe import display_dataframe
+from ..tools.s2.query_dataframe import query_dataframe
 from ..tools.s2.retrieve_semantic_scholar_paper_id import (
     retrieve_semantic_scholar_paper_id,
 )
@@ -49,7 +49,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         >>> result = app.invoke(initial_state)
     """
 
-    def agent_zotero_node(state: Talk2Scholars) -> Dict[str, Any]:
+    def zotero_agent_node(state: Talk2Scholars) -> Dict[str, Any]:
         """
         Processes the user query and retrieves relevant research papers from Zotero.
 
@@ -65,7 +65,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
             Dict[str, Any]: A dictionary containing the updated conversation state.
 
         Example:
-            >>> result = agent_zotero_node(current_state)
+            >>> result = zotero_agent_node(current_state)
             >>> papers = result.get("papers", [])
         """
         logger.log(
@@ -89,11 +89,11 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     tools = ToolNode(
         [
             zotero_read,
-            s2_display,
-            s2_query_results,
+            display_dataframe,
+            query_dataframe,
             retrieve_semantic_scholar_paper_id,
-            zotero_review,  # First review
-            zotero_write,  # Then save with user confirmation
+            zotero_review,
+            zotero_write,
         ]
     )
 
@@ -110,14 +110,14 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     )
 
     workflow = StateGraph(Talk2Scholars)
-    workflow.add_node("agent_zotero", agent_zotero_node)
-    workflow.add_edge(START, "agent_zotero")
+    workflow.add_node("zotero_agent", zotero_agent_node)
+    workflow.add_edge(START, "zotero_agent")
 
     # Initialize memory to persist state between graph runs
     checkpointer = MemorySaver()
 
     # Compile the graph
-    app = workflow.compile(checkpointer=checkpointer, name="agent_zotero")
+    app = workflow.compile(checkpointer=checkpointer, name="zotero_agent")
     logger.log(
         logging.INFO,
         "Compiled the graph with thread_id %s and llm_model %s",

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/default.yaml

@@ -1,22 +1,22 @@
 _target_: agents.main_agent.get_app
 temperature: 0
-system_prompt: >
-  You are Talk2Scholars agent coordinating academic paper discovery
-  and analysis with help of the following agents:
-   1. Agent S2: This agent can be used to search and recommend papers
-      from Semantic Scholar. Use this agent when the user asks for
-      general paper/article searches and recommendations, or to retrieve information
-      from the last displayed results table or query abstract of last
-      displayed results.
-   2. Agent Zotero: This agent can be used to retrieve, display, and query
-      papers/articles from the Zotero library. Use this agent only when the user
-      explicitly asks for papers from Zotero. This tool can also be used to
-      save papers in the zotero library.
-   3. Agent PaperFetch: This agent can be used to download papers/articles
-      from ArXiv.
-   4. Agent PDFQuery: This agent can be used to query contents of an
-      uploaded or downloaded PDF/paper/article.
+system_prompt: |
+  You are the Main Supervisor Agent.
 
-   Your final response should be a one sentence summary of the information
-   retrieved from the agents above. Do not repeat the information already
-   displayed to the user in the response of the agents.
+  You have access to four tools, each represented by a sub-agent:
+
+  - s2_agent(Use this to search for or recommend academic papers. This agent
+    should be used when the user requests general paper or article
+    searches, recommendations, or wants to retrieve information—such as
+    abstracts, from the most last displayed or searched results table.),
+  - zotero_agent(Use to Read or Write academic papers to zotero account,
+    This agent can also be used to save papers in the zotero library only
+    with explicit approval from the user),
+  - pdf_agent(Use this to perform question-and-answer tasks on downloaded, uploaded, or Zotero papers/PDFs.), and 
+  - paper_download_agent(Use to download PDFs).
+
+  Each sub-agent is specialized for a different task.
+
+  You can call multiple sub-agents at the same time, or sequentially. After receiving output from one agent, you can call another based on the user’s query.
+  Your goal is to analyze the user’s request carefully, decide which sub-agent(s) should be used, and coordinate their execution efficiently.
+  Always prioritize delegating tasks correctly. Think step-by-step before acting. Avoid answering by yourself unless explicitly necessary.

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/default.yaml

@@ -1,19 +1,24 @@
 _target_: agents.s2_agent.get_app
-s2_agent: >
-  You are an academic research assistant with access to the
-  Semantic Scholar API for paper discovery and analysis.
+s2_agent: |
+  You are the S2 Agent.
+
+  You are responsible for searching academic papers using the Semantic Scholar API.
+
+  Your capabilities include:
+
+  - Retrieving papers based on user queries.
+  - Recommending papers based on a single paper or multiple papers provided by the user.
+  - Retrieving the Semantic Scholar ID of a paper based on its title.
+  - This ID can later be used by other tools (search or recommend) based on the user’s needs.
+  - Always respond accurately based on Semantic Scholar search and recommendation features.
+  - Use `query_dataframe` tool query over the last displayed papers or the search table.
+  - Always call `display_dataframe` tool at the end.
 
-  AVAILABLE TOOLS:
-  1. search - Search for academic papers by query string
-  2. display_results - Display the papers retrieved by other tools
-  3. single_paper_rec - Get recommendations based on a SINGLE paper
-  4. multi_paper_rec - Get recommendations based on MULTIPLE papers
-  5. query_results - Ask questions about the current set of papers
-  6. retrieve_semantic_scholar_paper_id - Get Semantic Scholar ID for a paper title
 
   WORKFLOW STEPS:
-  1. When user requests papers, use search/recommendation tools to find papers
-  2. Use `display_results` tool to display the response from the search/recommendation tools 
-  3. Use `query_results` tool to query over the selected paper only when the user asks to
-  4. When the user wants recommendations, you can get the "semantic_scholar_paper_id" using `query_results` tool in the "last_displayed_results" key, then
-     pass the "semantic_scholar_paper_id" to `search`, `single_paper_rec` or `multi_paper_rec` tools depending on the user's query. Do not use "arxiv_id"
+  1. When user requests papers, use search/recommendation tools to find papers.
+  2. Use `display_dataframe` tool to display the response from the search/recommendation tools.
+  3. Use `query_dataframe` tool to query over the selected paper only when the user asks to.
+  4. When the user only wants recommendations, you can get the "semantic_scholar_paper_id"
+     using `query_dataframe` tool, then pass the "semantic_scholar_paper_id" to `search`,
+     `single_paper_rec` or `multi_paper_rec` tools depending on the user's query. Do not use "arxiv_id"(It is used to download pdfs)

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/zotero_agent/default.yaml

@@ -1,12 +1,33 @@
-target: agents.zotero_agent.get_app
-zotero_agent: >
-  You are a knowledgeable research assistant with expertise in academic literature.
-  Your goal is to help users find, manage, and save relevant research papers and display them.
+_target_: agents.zotero_agent.get_app
+zotero_agent: |
+  You are the Zotero Agent.
+
+  You are responsible for read and writing of papers to user's Zotero library.
+  Behavior:
+
+  - Once you have successfully read the papers, you must immediately stop, return a clear 'Search complete' message along with a summary of the articles, call the
+    `display_dataframe` tool, and return to the main supervisor for further processing based on the user's query.
+  - Do not continue any further processing or re-enter into reading steps.
+  - You can write papers to user's library but only after explicit user confirmation.
+  - Do not attempt to answer any scientific or content-related questions yourself.
+  - You can retrieve all articles or search based on the user's query, inferring whether to return the full collection or filter by title, keywords, or other details.
+  - Never call `query_dataframe` tool regarding any question or any information retrival only if the user explicitly asks for metadata.
+
+  In multi-step workflows:
+
+  - Your job is only to read the requested paper or all the papers in user's library and return the successful search output.
+  - After that, the Main Supervisor Agent will decide the next step (such as passing the paper to the pdf_agent).
+  - Always call `display_dataframe` tool at the end before transfering to Main Supervisor Agent.
+  - Never attempt to call other agents yourself.
+
+  Stopping Condition:
+
+  - After successful search, indicate completion clearly and terminate your action.
 
   When saving papers to Zotero:
-  1. First use zotero_review with the collection path.
+  1. First use `zotero_review` tool with the collection path.
   2. Wait for user confirmation (they must say "Yes" or "Approve").
-  3. Use zotero_write with both the collection_path and user_confirmation and call `s2_display` tool after the papers as saved.
+  3. Use `zotero_write` tool with both the collection_path and user_confirmation and call `display_dataframe` tool after the papers as saved.
 
   IMPORTANT: Human approval is required for saving papers to Zotero. Never save papers
   without explicit approval from the user. Always respect the user's decision if they

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -8,8 +8,9 @@ across agent interactions.
 
 import logging
 from typing import Annotated, Any, Dict
-from langchain_core.language_models import BaseChatModel
+
 from langchain_core.embeddings import Embeddings
+from langchain_core.language_models import BaseChatModel
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
 # Configure logging
@@ -38,22 +39,22 @@ def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any
         >>> print(updated_state)
         {"papers": {"id2": "Paper 2"}}
     """
-    logger.info("Updating existing state %s with the state dict: %s", existing, new)
+    # No-op operation to use the 'existing' variable
+    _ = len(existing)
     return new
 
 
 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.
-        zotero_read (Dict[str, Any]): Stores the papers retrieved from Zotero.
+        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.
@@ -65,8 +66,7 @@ class Talk2Scholars(AgentState):
     last_displayed_papers: Annotated[Dict[str, Any], replace_dict]
     papers: Annotated[Dict[str, Any], replace_dict]
     multi_papers: Annotated[Dict[str, Any], replace_dict]
-    pdf_data: Annotated[Dict[str, Any], replace_dict]
-    zotero_read: Annotated[Dict[str, Any], replace_dict]
+    article_data: Annotated[Dict[str, Any], replace_dict]
     zotero_write_approval_status: Annotated[Dict[str, Any], replace_dict]
     llm_model: BaseChatModel
     text_embedding_model: Embeddings

aiagents4pharma/talk2scholars/tests/test_main_agent.py

@@ -51,24 +51,24 @@ class DummyWorkflow:
         return self
 
 
-def dummy_get_app_s2(uniq_id, llm_model):
+def dummy_s2_agent(uniq_id, llm_model):
     """Return a DummyWorkflow for the S2 agent."""
-    dummy_get_app_s2.called_uniq_id = uniq_id
-    dummy_get_app_s2.called_llm_model = llm_model
+    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_get_app_zotero(uniq_id, llm_model):
+def dummy_zotero_agent(uniq_id, llm_model):
     """Return a DummyWorkflow for the Zotero agent."""
-    dummy_get_app_zotero.called_uniq_id = uniq_id
-    dummy_get_app_zotero.called_llm_model = llm_model
+    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_get_app_pdf(uniq_id, llm_model):
+def dummy_question_and_answer_agent(uniq_id, llm_model):
     """Return a DummyWorkflow for the PDF agent."""
-    dummy_get_app_pdf.called_uniq_id = uniq_id
-    dummy_get_app_pdf.called_llm_model = llm_model
+    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})
 
 
@@ -143,10 +143,10 @@ def patch_hydra(monkeypatch):
     )
 
 
-def dummy_get_app_paper_download(uniq_id, llm_model):
+def dummy_paper_download_agent(uniq_id, llm_model):
     """Return a DummyWorkflow for the paper download agent."""
-    dummy_get_app_paper_download.called_uniq_id = uniq_id
-    dummy_get_app_paper_download.called_llm_model = llm_model
+    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}
     )
@@ -156,19 +156,19 @@ def dummy_get_app_paper_download(uniq_id, llm_model):
 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_get_app_s2
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_s2", dummy_s2_agent
     )
     monkeypatch.setattr(
         "aiagents4pharma.talk2scholars.agents.main_agent.get_app_zotero",
-        dummy_get_app_zotero,
+        dummy_zotero_agent,
     )
     monkeypatch.setattr(
         "aiagents4pharma.talk2scholars.agents.main_agent.get_app_pdf",
-        dummy_get_app_pdf,
+        dummy_question_and_answer_agent,
     )
     monkeypatch.setattr(
         "aiagents4pharma.talk2scholars.agents.main_agent.get_app_paper_download",
-        dummy_get_app_paper_download,
+        dummy_paper_download_agent,
     )
     monkeypatch.setattr(
         "aiagents4pharma.talk2scholars.agents.main_agent.create_supervisor",

aiagents4pharma/talk2scholars/tests/test_paper_download_agent.py

@@ -22,22 +22,15 @@ def mock_hydra_fixture():
 @pytest.fixture
 def mock_tools_fixture():
     """Mocks paper download tools to prevent real HTTP calls."""
-    with (
-        mock.patch(
-            "aiagents4pharma.talk2scholars.tools.paper_download."
-            "download_arxiv_input.download_arxiv_paper"
-        ) as mock_download_arxiv_paper,
-        mock.patch(
-            "aiagents4pharma.talk2scholars.tools.s2.query_results.query_results"
-        ) as mock_query_results,
-    ):
+    with mock.patch(
+        "aiagents4pharma.talk2scholars.tools.paper_download."
+        "download_arxiv_input.download_arxiv_paper"
+    ) as mock_download_arxiv_paper:
         mock_download_arxiv_paper.return_value = {
-            "pdf_data": {"dummy_key": "dummy_value"}
+            "article_data": {"dummy_key": "dummy_value"}
         }
-        mock_query_results.return_value = {
-        "result": "Mocked Query Result"
-        }
-        yield [mock_download_arxiv_paper, mock_query_results]
+        yield [mock_download_arxiv_paper]
+
 
 @pytest.mark.usefixtures("mock_hydra_fixture")
 def test_paper_download_agent_initialization():
@@ -54,6 +47,7 @@ def test_paper_download_agent_initialization():
         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
@@ -70,10 +64,9 @@ def test_paper_download_agent_invocation():
         mock_create_agent.return_value = mock_agent
         mock_agent.invoke.return_value = {
             "messages": [AIMessage(content="Here is the paper")],
-            "pdf_data": {"file_bytes": b"FAKE_PDF_CONTENTS"},
+            "article_data": {"file_bytes": b"FAKE_PDF_CONTENTS"},
         }
 
-
         app = get_app(thread_id, llm_mock)
         result = app.invoke(
             mock_state,
@@ -87,11 +80,11 @@ def test_paper_download_agent_invocation():
         )
 
         assert "messages" in result
-        assert "pdf_data" in result
+        assert "article_data" in result
 
 
 def test_paper_download_agent_tools_assignment(request):  # Keep fixture name
-    """Checks correct tool assignment (download_arxiv_paper, query_results)."""
+    """Checks correct tool assignment (download_arxiv_paper, query_dataframe)."""
     thread_id = "test_thread_paper_dl"
     mock_tools = request.getfixturevalue("mock_tools_fixture")
     llm_mock = mock.Mock(spec=BaseChatModel)
@@ -107,12 +100,12 @@ def test_paper_download_agent_tools_assignment(request):  # Keep fixture name
         mock_agent = mock.Mock()
         mock_create_agent.return_value = mock_agent
         mock_tool_instance = mock.Mock()
-        mock_tool_instance.tools = mock_tools
-        mock_toolnode.return_value= mock_tool_instance
+        mock_tool_instance.tools = mock_tools if mock_tools else []
+        mock_toolnode.return_value = mock_tool_instance
 
         get_app(thread_id, llm_mock)
         assert mock_toolnode.called
-        assert len(mock_tool_instance.tools) == 2
+        assert len(mock_tool_instance.tools) == 1
 
 
 def test_paper_download_agent_hydra_failure():
@@ -137,6 +130,6 @@ def test_paper_download_agent_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."
-        )
+        assert "Mock model failure" in str(
+            exc_info.value
+        ), "Model initialization failure should raise an exception."