aiagents4pharma 1.29.0__py3-none-any.whl → 1.30.1__py3-none-any.whl

aiagents4pharma/talk2scholars/agents/__init__.py

@@ -4,5 +4,6 @@ This file is used to import all the modules in the package.
 
 from . import main_agent
 from . import s2_agent
+from . import paper_download_agent
 from . import zotero_agent
 from . import pdf_agent

aiagents4pharma/talk2scholars/agents/main_agent.py

@@ -1,11 +1,14 @@
 #!/usr/bin/env python3
 
 """
-Main agent for the talk2scholars app using ReAct pattern.
+Main agent module for initializing and running the Talk2Scholars application.
 
-This module implements a hierarchical agent system where a supervisor agent
-routes queries to specialized sub-agents. It follows the LangGraph patterns
-for multi-agent systems and implements proper state management.
+This module sets up the hierarchical agent system using LangGraph and integrates
+various sub-agents for handling different tasks such as semantic scholar, zotero,
+PDF processing, and paper downloading.
+
+Functions:
+- get_app: Initializes and returns the LangGraph-based hierarchical agent system.
 """
 
 import logging
@@ -16,6 +19,8 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from ..agents.s2_agent import get_app as get_app_s2
 from ..agents.zotero_agent import get_app as get_app_zotero
+from ..agents.pdf_agent import get_app as get_app_pdf
+from ..agents.paper_download_agent import get_app as get_app_paper_download
 from ..state.state_talk2scholars import Talk2Scholars
 
 # Initialize logger
@@ -43,12 +48,13 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         >>> app = get_app("thread_123")
         >>> result = app.invoke(initial_state)
     """
-    if llm_model.model_name == "gpt-4o-mini":
-        llm_model = ChatOpenAI(
-            model="gpt-4o-mini",
-            temperature=0,
-            model_kwargs={"parallel_tool_calls": False},
-        )
+    if hasattr(llm_model, "model_name"):
+        if llm_model.model_name == "gpt-4o-mini":
+            llm_model = ChatOpenAI(
+                model="gpt-4o-mini",
+                temperature=0,
+                model_kwargs={"parallel_tool_calls": False},
+            )
     # Load hydra configuration
     logger.log(logging.INFO, "Launching Talk2Scholars with thread_id %s", uniq_id)
     with hydra.initialize(version_base=None, config_path="../configs/"):
@@ -62,6 +68,8 @@ 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
         ],
         model=llm_model,
         state_schema=Talk2Scholars,

aiagents4pharma/talk2scholars/agents/paper_download_agent.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""
+This module defines the paper download agent that connects to the arXiv API to fetch
+paper details and PDFs. It is part of the Talk2Scholars project.
+"""
+
+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.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
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def get_app(uniq_id, llm_model: BaseChatModel):
+    """
+    Initializes and returns the LangGraph application for the Talk2Scholars paper download agent.
+
+    Args:
+        uniq_id (str): A unique identifier for tracking the current session.
+        llm_model (BaseChatModel, optional): The language model to be used by the agent.
+            Defaults to ChatOpenAI(model="gpt-4o-mini", temperature=0.5).
+
+    Returns:
+        StateGraph: A compiled LangGraph application that enables the paper download agent to
+            process user queries and retrieve arXiv papers.
+    """
+
+    # Load Hydra configuration
+    logger.info("Loading Hydra configuration for Talk2Scholars paper download agent")
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(
+            config_name="config",
+            overrides=["agents/talk2scholars/paper_download_agent=default"],
+        )
+        cfg = cfg.agents.talk2scholars.paper_download_agent
+
+    # Define tools properly
+    tools = ToolNode([download_arxiv_paper, query_results])
+
+    # Define the model
+    logger.info("Using OpenAI model %s", llm_model)
+    model = create_react_agent(
+        llm_model,
+        tools=tools,
+        state_schema=Talk2Scholars,
+        prompt=cfg.paper_download_agent,
+        checkpointer=MemorySaver(),
+    )
+
+    def paper_download_agent_node(state: Talk2Scholars) -> Dict[str, Any]:
+        """
+        Processes the current state to fetch the arXiv paper.
+        """
+        logger.info("Creating paper download agent node with thread_id: %s", uniq_id)
+        result = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+        return result
+
+    # Define new graph
+    workflow = StateGraph(Talk2Scholars)
+
+    # Adding node for paper download agent
+    workflow.add_node("paper_download_agent", paper_download_agent_node)
+
+    # Entering into the agent
+    workflow.add_edge(START, "paper_download_agent")
+
+    # Memory management for states between graph runs
+    checkpointer = MemorySaver()
+
+    # Compile the graph
+    app = workflow.compile(checkpointer=checkpointer, name="agent_paper_download")
+
+    # Logging the information and returning the app
+    logger.info("Compiled the graph")
+    return app

aiagents4pharma/talk2scholars/agents/pdf_agent.py

@@ -26,10 +26,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def get_app(
-    uniq_id,
-    llm_model: BaseChatModel
-):
+def get_app(uniq_id, llm_model: BaseChatModel):
     """
     Initializes and returns the LangGraph application for the PDF agent.
 
@@ -40,7 +37,7 @@ def get_app(
 
     Args:
         uniq_id (str): A unique identifier for the current conversation session or thread.
-        llm_model (BaseChatModel, optional): The language model instance to be used. 
+        llm_model (BaseChatModel, optional): The language model instance to be used.
             Defaults to ChatOpenAI(model="gpt-4o-mini", temperature=0).
 
     Returns:
@@ -71,10 +68,7 @@ def get_app(
             Any: The response generated by the language model after processing the state.
         """
         logger.info("Creating Agent_PDF node with thread_id %s", uniq_id)
-        response = model.invoke(
-            state,
-            {"configurable": {"thread_id": uniq_id}}
-        )
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
         return response
 
     # Define the tool node that includes the PDF QnA tool.
@@ -100,7 +94,7 @@ def get_app(
     checkpointer = MemorySaver()
 
     # Compile the graph into a runnable app.
-    app = workflow.compile(checkpointer=checkpointer)
+    app = workflow.compile(checkpointer=checkpointer, name="agent_pdf")
     logger.info("Compiled the PDF agent graph.")
 
     return app

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

@@ -1,13 +1,22 @@
 _target_: agents.main_agent.get_app
 temperature: 0
 system_prompt: >
-  You are the Talk2Scholars agent coordinating academic paper discovery and analysis.
+  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.
 
-  You have access to the following agents:
-  1. S2_agent: This agent can be used to search and recommend papers
-     from Semantic Scholar. Use this agent when the user asks for
-     general paper searches and recommendations.
-  2. Zotero_agent: This agent can be used to retrieve, display, and query
-     papers 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 under collections in the zotero library
+   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.

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/paper_download_agent/__init__.py

@@ -0,0 +1,3 @@
+"""
+Import all the modules in the package
+"""

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

@@ -15,5 +15,5 @@ s2_agent: >
   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 "paper_id" using `query_results` tool in the "last_displayed_results" key, then
-     pass the "paper_id" to `search`, `single_paper_rec` or `multi_paper_rec` tools depending on the user's query. Do not use "arxiv_id"
+  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"

aiagents4pharma/talk2scholars/configs/app/frontend/default.yaml

@@ -8,6 +8,7 @@ page:
 llms:
   available_models:
     - "OpenAI/gpt-4o-mini"
+    - "NVIDIA/llama-3.3-70b-instruct"
 # # Chat UI configuration
 # chat:
 #   assistant_avatar: "🤖"

aiagents4pharma/talk2scholars/configs/config.yaml

@@ -2,10 +2,12 @@ defaults:
   - _self_
   - agents/talk2scholars/main_agent: default
   - agents/talk2scholars/s2_agent: default
+  - agents/talk2scholars/paper_download_agent: default
   - agents/talk2scholars/zotero_agent: default
   - app/frontend: default
   - agents/talk2scholars/pdf_agent: default
   - tools/search: default
+  - tools/download_arxiv_paper: default
   - tools/single_paper_recommendation: default
   - tools/multi_paper_recommendation: default
   - tools/retrieve_semantic_scholar_paper_id: default

aiagents4pharma/talk2scholars/configs/tools/download_arxiv_paper/__init__.py

@@ -0,0 +1,3 @@
+"""
+Import all the modules in the package
+"""

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

@@ -6,10 +6,15 @@ api_fields:
   - "title"
   - "abstract"
   - "year"
-  - "authors"
+  - "authors.name"
+  - "authors.authorId"
   - "citationCount"
   - "url"
   - "externalIds"
+  - "venue"
+  - "publicationVenue" # Full object, instead of specific subfields
+  - "journal" # Full object, instead of specific subfields
+  - "publicationDate"
 # Commented fields that could be added later if needed
 
 # Default headers and params

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

@@ -6,10 +6,16 @@ api_fields:
   - "title"
   - "abstract"
   - "year"
-  - "authors"
+  - "authors.name"
+  - "authors.authorId"
   - "citationCount"
   - "url"
   - "externalIds"
+  - "venue"
+  - "publicationVenue" # Full object, instead of specific subfields
+  - "journal" # Full object, instead of specific subfields
+  - "publicationDate"
+
 # Commented fields that could be added later if needed
 #  - "publicationTypes"
 #  - "openAccessPdf"

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

@@ -6,10 +6,15 @@ api_fields:
   - "title"
   - "abstract"
   - "year"
-  - "authors"
+  - "authors.name"
+  - "authors.authorId"
   - "citationCount"
   - "url"
   - "externalIds"
+  - "venue"
+  - "publicationVenue" # Full object, instead of specific subfields
+  - "journal" # Full object, instead of specific subfields
+  - "publicationDate"
 # Commented fields that could be added later if needed
 #  - "publicationTypes"
 #  - "openAccessPdf"

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

@@ -53,4 +53,4 @@ zotero:
       "Web Page",
     ]
 
-  filter_excluded_types: ["attachment", "note", "annotation"]
+  # filter_excluded_types: ["attachment", "note", "annotation"]

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -9,6 +9,7 @@ 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 langgraph.prebuilt.chat_agent_executor import AgentState
 
 # Configure logging
@@ -54,6 +55,8 @@ class Talk2Scholars(AgentState):
         multi_papers (Dict[str, Any]): Stores multiple recommended papers from various sources.
         zotero_read (Dict[str, Any]): Stores the papers retrieved from Zotero.
         llm_model (BaseChatModel): The language model instance used for generating responses.
+        text_embedding_model (Embeddings): The text embedding model used for
+        similarity calculations.
     """
 
     # Agent state fields
@@ -63,3 +66,4 @@ class Talk2Scholars(AgentState):
     pdf_data: Annotated[Dict[str, Any], replace_dict]
     zotero_read: Annotated[Dict[str, Any], replace_dict]
     llm_model: BaseChatModel
+    text_embedding_model: Embeddings

aiagents4pharma/talk2scholars/tests/test_llm_main_integration.py

@@ -1,58 +1,89 @@
 """
 Integration tests for talk2scholars system with OpenAI.
+This test triggers all sub-agents by sending a conversation that covers:
+- Searching Semantic Scholar (S2 agent)
+- Retrieving Zotero results (Zotero agent)
+- Querying PDF content (PDF agent)
+- Downloading paper details from arXiv (Paper Download agent)
 """
 
-import os
-import pytest
-import hydra
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage
-from ..agents.main_agent import get_app
-from ..state.state_talk2scholars import Talk2Scholars
+# This will be covered in the next pr.
 
-# pylint: disable=redefined-outer-name
-
-
-@pytest.mark.skipif(
-    not os.getenv("OPENAI_API_KEY"), reason="Requires OpenAI API key to run"
-)
-def test_main_agent_real_llm():
-    """
-    Test that the main agent invokes S2 agent correctly
-    and updates the state with real LLM execution.
-    """
-
-    # Load Hydra Configuration EXACTLY like in main_agent.py
-    with hydra.initialize(version_base=None, config_path="../configs"):
-        cfg = hydra.compose(
-            config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
-        )
-    hydra_cfg = cfg.agents.talk2scholars.main_agent
-
-    assert hydra_cfg is not None, "Hydra config failed to load"
-
-    # Use the real OpenAI API (ensure env variable is set)
-    llm = ChatOpenAI(model="gpt-4o-mini", temperature=hydra_cfg.temperature)
-
-    # Initialize main agent workflow (WITH real Hydra config)
-    thread_id = "test_thread"
-    app = get_app(thread_id, llm)
-
-    # Provide an actual user query
-    initial_state = Talk2Scholars(
-        messages=[HumanMessage(content="Find AI papers on transformers")]
-    )
-
-    # Invoke the agent (triggers supervisor → s2_agent)
-    result = app.invoke(
-        initial_state,
-        {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
-    )
-
-    # Assert that the supervisor routed correctly
-    assert "messages" in result, "Expected messages in response"
-
-    # Fix: Accept AIMessage as a valid response type
-    assert isinstance(
-        result["messages"][-1], (HumanMessage, AIMessage, str)
-    ), "Last message should be a valid response"
+#
+# import os
+# import pytest
+# import hydra
+# from langchain_openai import ChatOpenAI
+# from langchain_core.messages import HumanMessage, AIMessage
+# from ..agents.main_agent import get_app
+# from ..state.state_talk2scholars import Talk2Scholars
+#
+# # pylint: disable=redefined-outer-name,too-few-public-methods
+#
+#
+# @pytest.mark.skipif(
+#     not os.getenv("OPENAI_API_KEY"), reason="Requires OpenAI API key to run"
+# )
+# def test_main_agent_real_llm():
+#     """
+#     Integration test for the Talk2Scholars system using a real OpenAI LLM.
+#     This test verifies that the supervisor correctly routes to all sub-agents by
+#     providing a conversation with queries intended to trigger each agent.
+#     """
+#     # Load Hydra configuration EXACTLY like in main_agent.py
+#     with hydra.initialize(version_base=None, config_path="../configs"):
+#         cfg = hydra.compose(
+#             config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
+#         )
+#     hydra_cfg = cfg.agents.talk2scholars.main_agent
+#     assert hydra_cfg is not None, "Hydra config failed to load"
+#
+#     # Use the real OpenAI API (ensure OPENAI_API_KEY is set in environment)
+#     llm = ChatOpenAI(model="gpt-4o-mini", temperature=hydra_cfg.temperature)
+#
+#     # Initialize the main agent workflow (with real Hydra config)
+#     thread_id = "test_thread"
+#     app = get_app(thread_id, llm)
+#
+#     # Provide a multi-turn conversation intended to trigger all sub-agents:
+#     # - S2 agent: "Search Semantic Scholar for AI papers on transformers."
+#     # - Zotero agent: "Retrieve Zotero results for these papers."
+#     # - PDF agent: "Analyze the attached PDF and summarize its key findings."
+#     # - Paper Download agent: "Download the paper details from arXiv."
+#     initial_state = Talk2Scholars(
+#         messages=[
+#             HumanMessage(
+#                 content="Search Semantic Scholar for AI papers on transformers."
+#             ),
+#             HumanMessage(content="Also, retrieve Zotero results for these papers."),
+#             HumanMessage(
+#                 content="I have attached a PDF; analyze it and tell me the key findings."
+#             ),
+#             HumanMessage(content="Finally, download the paper from arXiv."),
+#         ]
+#     )
+#
+#     # Invoke the agent (which routes to the appropriate sub-agents)
+#     result = app.invoke(
+#         initial_state,
+#         {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
+#     )
+#
+#     # Assert that the result contains messages and that the final message is valid.
+#     assert "messages" in result, "Expected 'messages' in the response"
+#     last_message = result["messages"][-1]
+#     assert isinstance(
+#         last_message, (HumanMessage, AIMessage, str)
+#     ), "Last message should be a valid response type"
+#
+#     # Concatenate message texts (if available) to perform keyword checks.
+#     output_text = " ".join(
+#         msg.content if hasattr(msg, "content") else str(msg)
+#         for msg in result["messages"]
+#     ).lower()
+#
+#     # Check for keywords that suggest each sub-agent was invoked.
+#     for keyword in ["semantic scholar", "zotero", "pdf", "arxiv"]:
+#         assert (
+#             keyword in output_text
+#         ), f"Expected keyword '{keyword}' in the output response"

aiagents4pharma/talk2scholars/tests/test_main_agent.py

@@ -65,6 +65,13 @@ def dummy_get_app_zotero(uniq_id, llm_model):
     return DummyWorkflow(supervisor_args={"agent": "zotero", "uniq_id": uniq_id})
 
 
+def dummy_get_app_pdf(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
+    return DummyWorkflow(supervisor_args={"agent": "pdf", "uniq_id": uniq_id})
+
+
 def dummy_create_supervisor(apps, model, state_schema, **kwargs):
     """Return a DummyWorkflow for the supervisor."""
     dummy_create_supervisor.called_kwargs = kwargs
@@ -136,6 +143,15 @@ def patch_hydra(monkeypatch):
     )
 
 
+def dummy_get_app_paper_download(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
+    return DummyWorkflow(
+        supervisor_args={"agent": "paper_download", "uniq_id": uniq_id}
+    )
+
+
 @pytest.fixture(autouse=True)
 def patch_sub_agents_and_supervisor(monkeypatch):
     """Patch the sub-agents and supervisor creation functions."""
@@ -146,6 +162,14 @@ def patch_sub_agents_and_supervisor(monkeypatch):
         "aiagents4pharma.talk2scholars.agents.main_agent.get_app_zotero",
         dummy_get_app_zotero,
     )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_pdf",
+        dummy_get_app_pdf,
+    )
+    monkeypatch.setattr(
+        "aiagents4pharma.talk2scholars.agents.main_agent.get_app_paper_download",
+        dummy_get_app_paper_download,
+    )
     monkeypatch.setattr(
         "aiagents4pharma.talk2scholars.agents.main_agent.create_supervisor",
         dummy_create_supervisor,