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

aiagents4pharma/talk2scholars/agents/main_agent.py

@@ -9,117 +9,21 @@ for multi-agent systems and implements proper state management.
 """
 
 import logging
-from typing import Literal, Callable
-from pydantic import BaseModel
 import hydra
+from langgraph_supervisor import create_supervisor
+from langchain_openai import ChatOpenAI
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.graph import END, START, StateGraph
-from langgraph.types import Command
-from ..agents import s2_agent
-from ..agents import zotero_agent
+from ..agents.s2_agent import get_app as get_app_s2
+from ..agents.zotero_agent import get_app as get_app_zotero
 from ..state.state_talk2scholars import Talk2Scholars
 
-# Configure logging
+# Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def get_hydra_config():
-    """
-    Loads the Hydra configuration for the main agent.
-
-    This function initializes the Hydra configuration system and retrieves the settings
-    for the `Talk2Scholars` agent, ensuring that all required parameters are loaded.
-
-    Returns:
-        DictConfig: The configuration object containing parameters for the main agent.
-    """
-    with hydra.initialize(version_base=None, config_path="../configs"):
-        cfg = hydra.compose(
-            config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
-        )
-    return cfg.agents.talk2scholars.main_agent
-
-
-def make_supervisor_node(llm_model: BaseChatModel, thread_id: str) -> Callable:
-    """
-    Creates the supervisor node responsible for routing user queries to the appropriate sub-agents.
-
-    This function initializes the routing logic by leveraging the system and router prompts defined
-    in the Hydra configuration. The supervisor determines whether to
-    call a sub-agent (like `s2_agent`, `zotero_agent`)
-    or directly generate a response using the language model.
-
-    Args:
-        llm_model (BaseChatModel): The language model used for decision-making.
-        thread_id (str): Unique identifier for the current conversation session.
-
-    Returns:
-        Callable: The supervisor node function that processes user queries and
-        decides the next step.
-    """
-    cfg = get_hydra_config()
-    logger.info("Hydra configuration for Talk2Scholars main agent loaded: %s", cfg)
-    members = ["s2_agent", "zotero_agent"]
-    options = ["FINISH"] + members
-    # Define system prompt for general interactions
-    system_prompt = cfg.system_prompt
-    # Define router prompt for routing to sub-agents
-    router_prompt = cfg.router_prompt + " " + " ".join(members)
-
-    class Router(BaseModel):
-        """Worker to route to next. If no workers needed, route to FINISH."""
-
-        next: Literal[*options]
-
-    def supervisor_node(
-        state: Talk2Scholars,
-    ) -> Command:
-        """
-        Handles the routing logic for the supervisor agent.
-
-        This function determines the next agent to invoke based on the router prompt response.
-        If no further processing is required, it generates an AI response using the system prompt.
-
-        Args:
-            state (Talk2Scholars): The current conversation state, including messages
-            exchanged so far.
-
-        Returns:
-            Command: A command dictating whether to invoke a sub-agent or generate a final response.
-        """
-        messages = [SystemMessage(content=router_prompt)] + list(state["messages"])
-        structured_llm = llm_model.with_structured_output(Router)
-        response = structured_llm.invoke(messages)
-        goto = response.next
-        logger.info("Routing to: %s, Thread ID: %s", goto, thread_id)
-        if goto == "FINISH":
-            goto = END  # Using END from langgraph.graph
-            # If no agents were called, and the last message was
-            # from the user, call the LLM to respond to the user
-            # with a slightly different system prompt.
-            if isinstance(messages[-1], HumanMessage):
-                response = llm_model.invoke(
-                    [
-                        SystemMessage(content=system_prompt),
-                    ]
-                    + messages[1:]
-                )
-                return Command(
-                    goto=goto, update={"messages": AIMessage(content=response.content)}
-                )
-        # Go to the requested agent
-        return Command(goto=goto)
-
-    return supervisor_node
-
-
-def get_app(
-    thread_id: str,
-    llm_model: BaseChatModel
-):
+def get_app(uniq_id, llm_model: BaseChatModel):
     """
     Initializes and returns the LangGraph-based hierarchical agent system.
 
@@ -139,114 +43,36 @@ def get_app(
         >>> app = get_app("thread_123")
         >>> result = app.invoke(initial_state)
     """
-    cfg = get_hydra_config()
-
-    def call_s2_agent(
-        state: Talk2Scholars,
-    ) -> Command[Literal["supervisor"]]:
-        """
-        Invokes the Semantic Scholar (S2) agent to retrieve relevant research papers.
-
-        This function calls the `s2_agent` and updates the conversation state with retrieved
-        academic papers. The agent uses Semantic Scholar's API to find papers based on
-        user queries.
-
-        Args:
-            state (Talk2Scholars): The current state of the conversation, containing messages
-                and any previous search results.
-
-        Returns:
-            Command: A command to update the conversation state with the retrieved papers
-                and return control to the supervisor node.
-
-        Example:
-            >>> result = call_s2_agent(current_state)
-            >>> next_step = result.goto
-        """
-        logger.info("Calling S2 agent")
-        app = s2_agent.get_app(thread_id, llm_model)
-
-        # Invoke the S2 agent, passing state,
-        # Pass both config_id and thread_id
-        response = app.invoke(
-            state,
-            {
-                "configurable": {
-                    "config_id": thread_id,
-                    "thread_id": thread_id,
-                }
-            },
+    if llm_model.model_name == "gpt-4o-mini":
+        llm_model = ChatOpenAI(
+            model="gpt-4o-mini",
+            temperature=0,
+            model_kwargs={"parallel_tool_calls": False},
         )
-        logger.info("S2 agent completed with response")
-        return Command(
-            update={
-                "messages": response["messages"],
-                "papers": response.get("papers", {}),
-                "multi_papers": response.get("multi_papers", {}),
-                "last_displayed_papers": response.get("last_displayed_papers", {}),
-            },
-            # Always return to supervisor
-            goto="supervisor",
-        )
-
-    def call_zotero_agent(
-        state: Talk2Scholars,
-    ) -> Command[Literal["supervisor"]]:
-        """
-        Invokes the Zotero agent to retrieve and process papers from the user's Zotero library.
-
-        This function calls the Zotero agent, which interacts with the user's Zotero database
-        to retrieve relevant papers based on the conversation context. It updates the
-        conversation state with the retrieved papers and relevant metadata.
-
-        Args:
-            state (Talk2Scholars): The current conversation state, containing user messages
-                and any previously retrieved Zotero data.
-
-        Returns:
-            Command: A command that updates the conversation state with retrieved Zotero
-                papers and metadata before returning control to the supervisor node.
-
-        Example:
-            >>> result = call_zotero_agent(current_state)
-            >>> next_step = result.goto
-        """
-        logger.info("Calling Zotero agent")
-        app = zotero_agent.get_app(thread_id, llm_model)
-        # Invoke the Zotero agent, passing state
-        response = app.invoke(
-            state,
-            {
-                "configurable": {
-                    "config_id": thread_id,
-                    "thread_id": thread_id,
-                }
-            },
-        )
-        logger.info("Zotero agent completed with response")
-        return Command(
-            update={
-                "messages": response["messages"],
-                "zotero_read": response.get("zotero_read", {}),
-                "last_displayed_papers": response.get("last_displayed_papers", {}),
-            },
-            # Always return to supervisor
-            goto="supervisor",
+    # Load hydra configuration
+    logger.log(logging.INFO, "Launching Talk2Scholars with thread_id %s", uniq_id)
+    with hydra.initialize(version_base=None, config_path="../configs/"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
         )
+        cfg = cfg.agents.talk2scholars.main_agent
+    logger.log(logging.INFO, "System_prompt of Talk2Scholars: %s", cfg.system_prompt)
+    # Create supervisor workflow
+    workflow = create_supervisor(
+        [
+            get_app_s2(uniq_id, llm_model),  # semantic scholar
+            get_app_zotero(uniq_id, llm_model),  # zotero
+        ],
+        model=llm_model,
+        state_schema=Talk2Scholars,
+        # Full history is needed to extract
+        # the tool artifacts
+        output_mode="full_history",
+        add_handoff_back_messages=False,
+        prompt=cfg.system_prompt,
+    )
+
+    # Compile and run
+    app = workflow.compile(checkpointer=MemorySaver(), name="Talk2Scholars_MainAgent")
 
-    # Initialize LLM
-    logger.info("Using model %s with temperature %s", llm_model, cfg.temperature)
-
-    # Build the graph
-    workflow = StateGraph(Talk2Scholars)
-    supervisor = make_supervisor_node(llm_model, thread_id)
-    # Add nodes
-    workflow.add_node("supervisor", supervisor)
-    workflow.add_node("s2_agent", call_s2_agent)
-    workflow.add_node("zotero_agent", call_zotero_agent)
-    # Add edges
-    workflow.add_edge(START, "supervisor")
-    # Compile the workflow
-    app = workflow.compile(checkpointer=MemorySaver())
-    logger.info("Main agent workflow compiled")
     return app

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -28,9 +28,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 Semantic Scholar (S2) agent.
 
@@ -86,6 +84,7 @@ def get_app(
             config_name="config", overrides=["agents/talk2scholars/s2_agent=default"]
         )
         cfg = cfg.agents.talk2scholars.s2_agent
+        logger.log(logging.INFO, "Loaded configuration for S2 agent")
 
     # Define the tools
     tools = ToolNode(
@@ -107,7 +106,7 @@ def get_app(
         llm_model,
         tools=tools,
         state_schema=Talk2Scholars,
-        state_modifier=cfg.s2_agent,
+        prompt=cfg.s2_agent,
         checkpointer=MemorySaver(),
     )
 
@@ -122,7 +121,12 @@ def get_app(
     # 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)
-    logger.log(logging.INFO, "Compiled the graph")
+    app = workflow.compile(checkpointer=checkpointer, name="agent_s2")
+    logger.log(
+        logging.INFO,
+        "Compiled the graph with thread_id %s and llm_model %s",
+        uniq_id,
+        llm_model,
+    )
 
     return app

aiagents4pharma/talk2scholars/agents/zotero_agent.py

@@ -14,6 +14,7 @@ from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.checkpoint.memory import MemorySaver
 from ..state.state_talk2scholars import Talk2Scholars
 from ..tools.zotero.zotero_read import zotero_search_tool
+from ..tools.zotero.zotero_write import zotero_save_tool
 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.retrieve_semantic_scholar_paper_id import (
@@ -25,9 +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 Zotero agent.
 
@@ -83,6 +82,7 @@ def get_app(
             overrides=["agents/talk2scholars/zotero_agent=default"],
         )
         cfg = cfg.agents.talk2scholars.zotero_agent
+        logger.log(logging.INFO, "Loaded configuration for Zotero agent")
 
     # Define the tools
     tools = ToolNode(
@@ -91,6 +91,7 @@ def get_app(
             s2_display,
             s2_query_results,
             retrieve_semantic_scholar_paper_id,
+            zotero_save_tool,
         ]
     )
 
@@ -102,7 +103,7 @@ def get_app(
         llm_model,
         tools=tools,
         state_schema=Talk2Scholars,
-        state_modifier=cfg.zotero_agent,
+        prompt=cfg.zotero_agent,
         checkpointer=MemorySaver(),
     )
 
@@ -114,7 +115,12 @@ def get_app(
     checkpointer = MemorySaver()
 
     # Compile the graph
-    app = workflow.compile(checkpointer=checkpointer)
-    logger.log(logging.INFO, "Compiled the graph")
+    app = workflow.compile(checkpointer=checkpointer, name="agent_zotero")
+    logger.log(
+        logging.INFO,
+        "Compiled the graph with thread_id %s and llm_model %s",
+        uniq_id,
+        llm_model,
+    )
 
     return app

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

@@ -1,9 +1,4 @@
 _target_: agents.main_agent.get_app
-openai_api_key: ${oc.env:OPENAI_API_KEY}
-openai_llms:
-  - "gpt-4o-mini"
-  - "gpt-4-turbo"
-  - "gpt-3.5-turbo"
 temperature: 0
 system_prompt: >
   You are the Talk2Scholars agent coordinating academic paper discovery and analysis.
@@ -14,46 +9,5 @@ system_prompt: >
      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.
-
-router_prompt: >
-  You are a supervisor tasked with managing a conversation between the
-  following workers/members: ["s2_agent", "zotero_agent"] Given the user request, respond with the
-  worker to act next. Each worker will perform a task and respond with
-  their results and status. When finished, respond with FINISH.
-
-  Here is a description of the workers:
-  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. This agent can also
-     retrieve the Semantic Scholar ID of a paper. It can also be used to
-     provide more information about a paper.
-  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 agent can also
-     retrieve the Semantic Scholar ID of a paper.
-
-  CRITICAL RULES:
-  1. Do not generate any content or modify worker outputs
-  2. Route to FINISH ONLY when a worker has COMPLETELY finished their task
-  3. For the S2_agent and zotero_agent, ensure it has both SEARCHED and DISPLAYED results before FINISH
-
-  Available workers: members
-
-  Worker descriptions:
-  1. S2_agent: Routes to this agent ONLY for:
-     - Initial paper searches
-     - Getting paper recommendations
-     - Retrieving paper IDs based on the title of a paper
-     - Displaying search/recommendation results
-     - Query over papers
-  2. Zotero_agent: Routes to this agent ONLY for:
-     - Paper/journals searches from Zotero library
-     - Displaying search results
-     - Retrieving paper IDs based on the title of a paper
-     - Query over papers
-
-  Respond with FINISH when and ONLY when:
-  1. A worker has COMPLETELY finished their task (including display)
-  2. The results have been displayed to the user using display_results
-  3. No further action is needed
+     explicitly asks for papers from Zotero. This tool can also be used to
+     save papers in under collections in the zotero library

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

@@ -1,42 +1,19 @@
 _target_: agents.s2_agent.get_app
-openai_api_key: ${oc.env:OPENAI_API_KEY}
-openai_llms:
-  - "gpt-4o-mini"
-  - "gpt-4-turbo"
-  - "gpt-3.5-turbo"
-temperature: 0
 s2_agent: >
   You are an academic research assistant with access to the
   Semantic Scholar API for paper discovery and analysis.
 
   AVAILABLE TOOLS:
-  1. search_tool - Search for academic papers by query string
+  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
 
-  You also have tools to gain more insights on the papers and
-  display them.
-  You must strictly rely on retrieved information and avoid
-  generating unsupported content. Do not generate hallucinations
-  or fabricate details of any article. Stay focused on accurate,
-  sourced academic insights.
-
-  CRITICAL INSTRUCTIONS:
-  1. You must ONLY use information retrieved directly from the API
-  2. NEVER generate or fabricate paper details
-  3. NEVER modify or enhance the API responses
-  4. If information is missing from the API response, state that it's not available
-  5. ALWAYS CALL THE DISPLAY_RESULTS TOOL after completing a search
-
-  WORKFLOW STEPS (ALWAYS FOLLOW THIS EXACT SEQUENCE):
+  WORKFLOW STEPS:
   1. When user requests papers, use search/recommendation tools to find papers
-  2. IMMEDIATELY AFTER `search_tool`, `single_paper_rec`, `multi_paper_rec` completes, ALWAYS call display_results tool
+  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. NEVER skip the display_results step - it is REQUIRED for showing results to the user
-  5. Do not try to analyze or summarize papers yourself
-
-  Remember: The display_results tool is MANDATORY after every search - 
-  without it, users cannot see the search results.
+  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"

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

@@ -1,10 +1,4 @@
 target: agents.zotero_agent.get_app
-openai_api_key: ${oc.env:OPENAI_API_KEY}
-openai_llms:
-  - "gpt-4o-mini"
-  - "gpt-4-turbo"
-  - "gpt-3.5-turbo"
-temperature: 0
 zotero_agent: >
   You are a specialized Zotero library agent with access to tools for paper retrieval and management.
 
@@ -13,23 +7,13 @@ zotero_agent: >
   2. display_results - Display the papers retrieved by other tools
   3. query_results - Ask questions about the current set of papers
   4. retrieve_semantic_scholar_paper_id - Get Semantic Scholar ID for a paper title for the papers from zotero library
+  5. zotero_write - Save paper to users zotero library under specfied collections
 
-  You also have tools to gain more insights on the papers and display them.
-  You must strictly rely on retrieved information and avoid generating unsupported content. Do not generate hallucinations or fabricate details of any article. Stay focused on accurate, sourced academic insights.
 
-  CRITICAL INSTRUCTIONS:
-  1. You must ONLY use information retrieved directly from the Zotero Library
-  2. NEVER generate or fabricate paper details
-  3. NEVER modify or enhance the responses
-  4. If information is missing from the response, state that it's not available
-  5. ALWAYS CALL THE DISPLAY_RESULTS TOOL after completing a search
-
-  WORKFLOW STEPS (ALWAYS FOLLOW THIS EXACT SEQUENCE):
+  WORKFLOW STEPS
   1. When user requests papers, use `zotero_search_tool` to find papers
-  2. IMMEDIATELY AFTER `zotero_search_tool` completes, ALWAYS call display_results tool
+  2. Use `display_results` tool to display the response
   3. Use `query_results` tool to query over the selected paper only when the user asks to
   4. Use `retrieve_semantic_scholar_paper_id` to get the semantic scholar id of a paper title for the papers from zotero library
-  5. NEVER skip the display_results step - it is REQUIRED for showing results to the user
-  6. Do not try to analyze or summarize papers yourself
-
-  Remember: The display_results tool is MANDATORY after every search - without it, users cannot see the search results.
+  5. Use `zotero_write` to save the papers to users zotero library under collections and call `display_results` only after you recive
+     the save was successfull

aiagents4pharma/talk2scholars/configs/config.yaml

@@ -11,3 +11,4 @@ defaults:
   - tools/retrieve_semantic_scholar_paper_id: default
   - tools/question_and_answer: default
   - tools/zotero_read: default
+  - tools/zotero_write: default

aiagents4pharma/talk2scholars/configs/tools/__init__.py

@@ -7,3 +7,4 @@ from . import single_paper_recommendation
 from . import multi_paper_recommendation
 from . import question_and_answer
 from . import zotero_read
+from . import zotero_write

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

@@ -9,8 +9,8 @@ api_fields:
   - "authors"
   - "citationCount"
   - "url"
+  - "externalIds"
 # Commented fields that could be added later if needed
-# - "externalIds"
 
 # Default headers and params
 headers:

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

@@ -9,8 +9,8 @@ api_fields:
   - "authors"
   - "citationCount"
   - "url"
+  - "externalIds"
 # Commented fields that could be added later if needed
-# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

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

@@ -9,8 +9,8 @@ api_fields:
   - "authors"
   - "citationCount"
   - "url"
+  - "externalIds"
 # Commented fields that could be added later if needed
-# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

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

@@ -12,4 +12,45 @@ search_params:
 # Item Types and Limit
 zotero:
   max_limit: 100
-  filter_item_types: ["journalArticle", "conferencePaper", "preprint"]
+  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/configs/tools/zotero_write/__inti__.py

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