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

aiagents4pharma/talk2knowledgegraphs/configs/config.yaml

@@ -5,4 +5,5 @@ defaults:
   - tools/subgraph_summarization: default
   - tools/graphrag_reasoning: default
   - utils/pubchem_utils: default
+  - utils/enrichments/uniprot_proteins: default
   - app/frontend: default

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_uniprot.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+"""
+Test cases for utils/enrichments/uniprot_proteins.py
+"""
+
+import pytest
+from ..utils.enrichments.uniprot_proteins import EnrichmentWithUniProt
+
+# In this test, we will consider 2 examples:
+# 1. Gene Name: TP53
+# 2. Gene Name: TP5 (Incomplete; must return empty results)
+# 2. Gene Name: XZ (Shorter than 3 characters; must return empty results)
+# The expected description of TP53 startswith:
+START_DESCP = "Multifunctional transcription factor"
+# The expected amino acid sequence of TP53 startswith:
+START_SEQ = "MEEPQSDPSV"
+
+@pytest.fixture(name="enrich_obj")
+def fixture_uniprot_config():
+    """Return a dictionary with the configuration for UniProt enrichment."""
+    return EnrichmentWithUniProt()
+
+def test_enrich_documents(enrich_obj):
+    """Test the enrich_documents method."""
+    gene_names = ["TP53", "TP5", "XZ"]
+    descriptions, sequences = enrich_obj.enrich_documents(gene_names)
+    assert descriptions[0].startswith(START_DESCP)
+    assert sequences[0].startswith(START_SEQ)
+    assert descriptions[1] is None
+    assert sequences[1] is None
+    assert descriptions[2] is None
+    assert sequences[2] is None
+
+def test_enrich_documents_with_rag(enrich_obj):
+    """Test the enrich_documents_with_rag method."""
+    gene_names = ["TP53", "TP5", "XZ"]
+    descriptions, sequences = enrich_obj.enrich_documents_with_rag(gene_names, None)
+    assert descriptions[0].startswith(START_DESCP)
+    assert sequences[0].startswith(START_SEQ)
+    assert descriptions[1] is None
+    assert sequences[1] is None
+    assert descriptions[2] is None
+    assert sequences[2] is None

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py

@@ -4,3 +4,4 @@ This package contains modules to use the enrichment model
 from . import enrichments
 from . import ollama
 from . import pubchem_strings
+from . import uniprot_proteins

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/uniprot_proteins.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+
+"""
+Enrichment class for enriching Gene names with their function and sequence using UniProt.
+"""
+
+from typing import List
+import logging
+import json
+import hydra
+import requests
+from .enrichments import Enrichments
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class EnrichmentWithUniProt(Enrichments):
+    """
+    Enrichment class using UniProt
+    """
+    def enrich_documents(self, texts: List[str]) -> List[str]:
+        """
+        Enrich a list of input UniProt gene names with their function and sequence.
+
+        Args:
+            texts: The list of gene names to be enriched.
+
+        Returns:
+            The list of enriched functions and sequences
+        """
+
+        enriched_gene_names = texts
+
+        logger.log(logging.INFO,
+                   "Load Hydra configuration for Gene enrichment with description and sequence.")
+        with hydra.initialize(version_base=None, config_path="../../configs"):
+            cfg = hydra.compose(config_name='config',
+                                overrides=['utils/enrichments/uniprot_proteins=default'])
+            cfg = cfg.utils.enrichments.uniprot_proteins
+
+
+        descriptions = []
+        sequences = []
+        for gene in enriched_gene_names:
+            params = {
+                "reviewed": cfg.reviewed,
+                "isoform": cfg.isoform,
+                "exact_gene": gene,
+                "organism": cfg.organism,
+                # You can get the list of all available organisms here:
+                # https://www.uniprot.org/help/taxonomy
+            }
+
+            r = requests.get(cfg.uniprot_url,
+                             headers={ "Accept" : "application/json"},
+                             params=params,
+                             timeout=cfg.timeout)
+            # if the response is not ok
+            if not r.ok:
+                descriptions.append(None)
+                sequences.append(None)
+                continue
+            response_body = json.loads(r.text)
+            # if the response body is empty
+            if not response_body:
+                descriptions.append(None)
+                sequences.append(None)
+                continue
+            description = ''
+            for comment in response_body[0]['comments']:
+                if comment['type'] == 'FUNCTION':
+                    for value in comment['text']:
+                        description += value['value']
+            sequence = response_body[0]['sequence']['sequence']
+            descriptions.append(description)
+            sequences.append(sequence)
+        return descriptions, sequences
+
+    def enrich_documents_with_rag(self, texts, docs):
+        """
+        Enrich a list of input UniProt gene names with their function and sequence.
+
+        Args:
+            texts: The list of gene names to be enriched.
+
+        Returns:
+            The list of enriched functions and sequences
+        """
+        return self.enrich_documents(texts)

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