fix: T2S integration (#143)

* fix: refactor code

* fix: retry for the try/except block

* chores: added param

* chores: added param in the fetch

* fix: update write logic

* fix: update zotero read logic

* chores: updated

* chores: update system prompt

* chores: update tool config

* feat: added embedding model key

* chores: fix formating

* chores: fix formating

* fix: update packages

* fix: update utils

* chores: update app

* fix: formating

* fix: formating

* chores: Updated module level doc string

* chores: updated s2 prompt

* fix: updated tests

* fix: updated integartion and routing logic

* fix: no response test

* chores: added docstings for text embedding

* chores: added docstings for text embedding

* fix: pylinting

* chores: updated filter param

* chores: updated filter param

* chores: updated prompt

* fix: updated the trigger

* chores: update

---------

Co-authored-by: gurdeep330 <[email protected]>

Author: ansh-info

.github/workflows/ci.yml

@@ -11,9 +11,6 @@ on:
   push:
     branches:
       - main
-  pull_request:
-    branches:
-      - main
   workflow_dispatch:
 
 jobs:

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

@@ -20,6 +20,7 @@
 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.
@@ -39,22 +40,20 @@ def get_app(uniq_id, llm_model: BaseChatModel):
     with hydra.initialize(version_base=None, config_path="../configs"):
         cfg = hydra.compose(
             config_name="config",
-            overrides=["agents/talk2scholars/paper_download_agent=default"]
+            overrides=["agents/talk2scholars/paper_download_agent=default"],
         )
         cfg = cfg.agents.talk2scholars.paper_download_agent
 
     # Define tools properly
-    tools = ToolNode(
-    [download_arxiv_paper, query_results]
-    )
+    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.prompt,
+        prompt=cfg.paper_download_agent,
         checkpointer=MemorySaver(),
     )
 
@@ -79,7 +78,7 @@ def paper_download_agent_node(state: Talk2Scholars) -> Dict[str, Any]:
     checkpointer = MemorySaver()
 
     # Compile the graph
-    app = workflow.compile(checkpointer=checkpointer)
+    app = workflow.compile(checkpointer=checkpointer, name="agent_paper_download")
 
     # Logging the information and returning the app
     logger.info("Compiled the graph")

aiagents4pharma/talk2scholars/agents/pdf_agent.py

@@ -26,10 +26,7 @@
 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 agent_pdf_node(state: Talk2Scholars):
             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 agent_pdf_node(state: Talk2Scholars):
     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/default.yaml

@@ -36,4 +36,3 @@ paper_download_agent: >
   Do not attempt to analyze or summarize papers beyond what is retrieved from the API.
   Ensure responses are structured clearly and concisely, making
   it easy for the user to understand the retrieved information.
-

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/tools/download_arxiv_paper/default.yaml

@@ -1,4 +1,4 @@
 api_url: "http://export.arxiv.org/api/query"
 request_timeout: 10
 chunk_size: 1024
-pdf_base_url: "https://arxiv.org/pdf" 
+pdf_base_url: "https://arxiv.org/pdf"

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/question_and_answer/default.yaml

@@ -1,6 +1,6 @@
 # Default configuration for the PDF question_and_answer Tool
 openai_api_key: ${oc.env:OPENAI_API_KEY}
-chunk_size: 1000         # Number of characters per text chunk
-chunk_overlap: 200       # Overlap between adjacent chunks
-num_retrievals: 3        # Number of document chunks to retrieve for the QA chain
-qa_chain_type: "stuff"   # The type of QA chain to use
+chunk_size: 1000 # Number of characters per text chunk
+chunk_overlap: 200 # Overlap between adjacent chunks
+num_retrievals: 3 # Number of document chunks to retrieve for the QA chain
+qa_chain_type: "stuff" # The type of QA chain to use

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/configs/tools/zotero_write/default.yaml

@@ -12,5 +12,44 @@ search_params:
 # Item Types and Limit
 zotero:
   max_limit: 100
-  filter_item_types: ["journalArticle", "conferencePaper", "preprint"]
-  filter_excluded_types: ["attachment", "note", "annotation"]
+  filter_item_types:
+    [
+      "Artwork",
+      "Audio Recording",
+      "Bill",
+      "Blog Post",
+      "Book",
+      "Book Section",
+      "Case",
+      "Conference Paper",
+      "Dataset",
+      "Dictionary Entry",
+      "Document",
+      "E-mail",
+      "Encyclopedia Article",
+      "Film",
+      "Forum Post",
+      "Hearing",
+      "Instant Message",
+      "Interview",
+      "Journal Article",
+      "Letter",
+      "Magazine Article",
+      "Manuscript",
+      "Map",
+      "Newspaper Article",
+      "Patent",
+      "Podcast",
+      "Preprint",
+      "Presentation",
+      "Radio Broadcast",
+      "Report",
+      "Software",
+      "Standard",
+      "Statute",
+      "Thesis",
+      "TV Broadcast",
+      "Video Recording",
+      "Web Page",
+    ]
+  # filter_excluded_types: ["attachment", "note", "annotation"]

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -9,6 +9,7 @@
 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,4 +66,4 @@ class Talk2Scholars(AgentState):
     pdf_data: Annotated[Dict[str, Any], replace_dict]
     zotero_read: Annotated[Dict[str, Any], replace_dict]
     llm_model: BaseChatModel
-    pdf_data: Annotated[Dict[str, Any], replace_dict]
+    text_embedding_model: Embeddings