aiagents4pharma 1.20.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/configs/config.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - _self_
+  - agents/t2b_agent: default
+  - tools/ask_question: default
+  - tools/get_annotation: default

aiagents4pharma/talk2scholars/agents/main_agent.py

@@ -6,28 +6,17 @@ Main agent for the talk2scholars app using ReAct pattern.
 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.
-
-The main components are:
-1. Supervisor node with ReAct pattern for intelligent routing.
-2. S2 agent node for handling academic paper queries.
-3. Shared state management via Talk2Scholars.
-4. Hydra-based configuration system.
-
-Example:
-    app = get_app("thread_123", "gpt-4o-mini")
-    result = app.invoke({
-        "messages": [("human", "Find papers about AI agents")]
-    })
 """
 
 import logging
 from typing import Literal, Callable
+from pydantic import BaseModel
 import hydra
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, START, StateGraph
-from langgraph.prebuilt import create_react_agent
 from langgraph.types import Command
 from ..agents import s2_agent
 from ..state.state_talk2scholars import Talk2Scholars
@@ -39,13 +28,13 @@ logger = logging.getLogger(__name__)
 
 def get_hydra_config():
     """
-    Loads and returns the Hydra configuration for the main agent.
+    Loads the Hydra configuration for the main agent.
 
-    This function fetches the configuration settings for the Talk2Scholars
-    agent, ensuring that all required parameters are properly initialized.
+    This function initializes the Hydra configuration system and retrieves the settings
+    for the `Talk2Scholars` agent, ensuring that all required parameters are loaded.
 
     Returns:
-        Any: The configuration object for the main agent.
+        DictConfig: The configuration object containing parameters for the main agent.
     """
     with hydra.initialize(version_base=None, config_path="../configs"):
         cfg = hydra.compose(
@@ -54,116 +43,127 @@ def get_hydra_config():
     return cfg.agents.talk2scholars.main_agent
 
 
-def make_supervisor_node(llm: BaseChatModel, thread_id: str) -> Callable:
+def make_supervisor_node(llm_model: BaseChatModel, thread_id: str) -> Callable:
     """
-    Creates and returns a supervisor node for intelligent routing using the ReAct pattern.
+    Creates the supervisor node responsible for routing user queries to the appropriate sub-agents.
 
-    This function initializes a supervisor agent that processes user queries and
-    determines the appropriate sub-agent for further processing. It applies structured
-    reasoning to manage conversations and direct queries based on context.
+    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`)
+    or directly generate a response using the language model.
 
     Args:
-        llm (BaseChatModel): The language model used by the supervisor agent.
-        thread_id (str): Unique identifier for the conversation session.
+        llm_model (BaseChatModel): The language model used for decision-making.
+        thread_id (str): Unique identifier for the current conversation session.
 
     Returns:
-        Callable: A function that acts as the supervisor node in the LangGraph workflow.
-
-    Example:
-        supervisor = make_supervisor_node(llm, "thread_123")
-        workflow.add_node("supervisor", supervisor)
+        Callable: The supervisor node function that processes user queries and
+        decides the next step.
     """
-    logger.info("Loading Hydra configuration for Talk2Scholars main agent.")
     cfg = get_hydra_config()
-    logger.info("Hydra configuration loaded with values: %s", cfg)
+    logger.info("Hydra configuration for Talk2Scholars main agent loaded: %s", cfg)
+    members = ["s2_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
+
+    class Router(BaseModel):
+        """Worker to route to next. If no workers needed, route to FINISH."""
 
-    # Create the supervisor agent using the main agent's configuration
-    supervisor_agent = create_react_agent(
-        llm,
-        tools=[],  # Will add sub-agents later
-        state_modifier=cfg.main_agent,
-        state_schema=Talk2Scholars,
-        checkpointer=MemorySaver(),
-    )
+        next: Literal[*options]
 
     def supervisor_node(
         state: Talk2Scholars,
-    ) -> Command[Literal["s2_agent", "__end__"]]:
+    ) -> Command:
         """
-        Processes user queries and determines the next step in the conversation flow.
+        Handles the routing logic for the supervisor agent.
 
-        This function examines the conversation state and decides whether to forward
-        the query to a specialized sub-agent (e.g., S2 agent) or conclude the interaction.
+        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 state of the conversation, containing
-                messages, papers, and metadata.
+            state (Talk2Scholars): The current conversation state, including messages
+            exchanged so far.
 
         Returns:
-            Command: The next action to be executed, along with updated state data.
-
-        Example:
-            result = supervisor_node(current_state)
-            next_step = result.goto
+            Command: A command dictating whether to invoke a sub-agent or generate a final response.
         """
-        logger.info(
-            "Supervisor node called - Messages count: %d",
-            len(state["messages"]),
-        )
-
-        # Invoke the supervisor agent with configurable thread_id
-        result = supervisor_agent.invoke(
-            state, {"configurable": {"thread_id": thread_id}}
-        )
-        goto = "s2_agent"
-        logger.info("Supervisor agent completed with result: %s", result)
-
+        messages = [SystemMessage(content=router_prompt)] + 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: str = "gpt-4o-mini") -> StateGraph:
+def get_app(
+    thread_id: str,
+    llm_model: BaseChatModel = ChatOpenAI(model="gpt-4o-mini", temperature=0),
+):
     """
-    Initializes and returns the LangGraph application with a hierarchical agent system.
+    Initializes and returns the LangGraph-based hierarchical agent system.
 
-    This function sets up the full agent architecture, including the supervisor
-    and sub-agents, and compiles the LangGraph workflow for handling user queries.
+    This function constructs the agent workflow by defining nodes for the supervisor
+    and sub-agents. It compiles the graph using `StateGraph` to enable structured
+    conversational workflows.
 
     Args:
-        thread_id (str): Unique identifier for the conversation session.
-        llm_model (str, optional): The language model to be used. Defaults to "gpt-4o-mini".
+        thread_id (str): A unique session identifier for tracking conversation state.
+        llm_model (BaseChatModel, optional): The language model used for query processing.
+            Defaults to `ChatOpenAI(model="gpt-4o-mini", temperature=0)`.
 
     Returns:
-        StateGraph: A compiled LangGraph application ready for query invocation.
+        StateGraph: A compiled LangGraph application that can process user queries.
 
     Example:
-        app = get_app("thread_123")
-        result = app.invoke(initial_state)
+        >>> app = get_app("thread_123")
+        >>> result = app.invoke(initial_state)
     """
     cfg = get_hydra_config()
 
     def call_s2_agent(
         state: Talk2Scholars,
-    ) -> Command[Literal["supervisor", "__end__"]]:
+    ) -> Command[Literal["supervisor"]]:
         """
-        Calls the Semantic Scholar (S2) agent to process academic paper queries.
+        Invokes the Semantic Scholar (S2) agent to retrieve relevant research papers.
 
-        This function invokes the S2 agent, retrieves relevant research papers,
-        and updates the conversation state accordingly.
+        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 conversation state, including user queries
-                and any previously retrieved papers.
+            state (Talk2Scholars): The current state of the conversation, containing messages
+                and any previous search results.
 
         Returns:
-            Command: The next action to execute, along with updated messages and papers.
+            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
+            >>> result = call_s2_agent(current_state)
+            >>> next_step = result.goto
         """
-        logger.info("Calling S2 agent with state: %s", state)
+        logger.info("Calling S2 agent")
         app = s2_agent.get_app(thread_id, llm_model)
 
         # Invoke the S2 agent, passing state,
@@ -177,31 +177,30 @@ def get_app(thread_id: str, llm_model: str = "gpt-4o-mini") -> StateGraph:
                 }
             },
         )
-        logger.info("S2 agent completed with response: %s", response)
-
+        logger.info("S2 agent completed with response")
         return Command(
-            goto=END,
             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",
         )
 
     # Initialize LLM
-    logger.info("Using OpenAI model %s with temperature %s", llm_model, cfg.temperature)
-    llm = ChatOpenAI(model=llm_model, temperature=cfg.temperature)
+    logger.info("Using model %s with temperature %s", llm_model, cfg.temperature)
 
     # Build the graph
     workflow = StateGraph(Talk2Scholars)
-    supervisor = make_supervisor_node(llm, thread_id)
-
+    supervisor = make_supervisor_node(llm_model, thread_id)
+    # Add nodes
     workflow.add_node("supervisor", supervisor)
     workflow.add_node("s2_agent", call_s2_agent)
+    # Add edges
     workflow.add_edge(START, "supervisor")
-    workflow.add_edge("s2_agent", END)
-
-    # Compile the graph without initial state
+    # Compile the workflow
     app = workflow.compile(checkpointer=MemorySaver())
     logger.info("Main agent workflow compiled")
     return app

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -5,37 +5,80 @@ Agent for interacting with Semantic Scholar
 """
 
 import logging
+from typing import Any, Dict
 import hydra
 from langchain_openai import ChatOpenAI
+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 ..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.retrieve_semantic_scholar_paper_id import (
+    retrieve_semantic_scholar_paper_id as s2_retrieve_id,
+)
 from ..tools.s2.single_paper_rec import (
     get_single_paper_recommendations as s2_single_rec,
 )
 from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations as s2_multi_rec
 
-
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def get_app(uniq_id, llm_model="gpt-4o-mini"):
+def get_app(
+    uniq_id, llm_model: BaseChatModel = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+):
     """
-    This function returns the langraph app.
+    Initializes and returns the LangGraph application for the Semantic Scholar (S2) agent.
+
+    This function sets up the S2 agent, which integrates various tools to search, retrieve,
+    and display research papers from Semantic Scholar. The agent follows the ReAct pattern
+    for structured interaction.
+
+    Args:
+        uniq_id (str): Unique identifier for the current conversation session.
+        llm_model (BaseChatModel, optional): The language model to be used by the agent.
+            Defaults to `ChatOpenAI(model="gpt-4o-mini", temperature=0)`.
+
+    Returns:
+        StateGraph: A compiled LangGraph application that enables the S2 agent to process
+            user queries and retrieve research papers.
+
+    Example:
+        >>> app = get_app("thread_123")
+        >>> result = app.invoke(initial_state)
     """
 
-    def agent_s2_node(state: Talk2Scholars):
+    # def agent_s2_node(state: Talk2Scholars) -> Command[Literal["supervisor"]]:
+    def agent_s2_node(state: Talk2Scholars) -> Dict[str, Any]:
         """
-        This function calls the model.
+        Processes the user query and retrieves relevant research papers.
+
+        This function calls the language model using the configured `ReAct` agent to analyze
+        the state and generate an appropriate response. The function then returns control
+        to the main supervisor.
+
+        Args:
+            state (Talk2Scholars): The current conversation state, including messages exchanged
+                and any previously retrieved research papers.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the updated conversation state.
+
+        Example:
+            >>> result = agent_s2_node(current_state)
+            >>> papers = result.get("papers", [])
         """
         logger.log(logging.INFO, "Creating Agent_S2 node with thread_id %s", uniq_id)
-        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
-        return response
+        result = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+
+        return result
+
+    logger.log(logging.INFO, "thread_id, llm_model: %s, %s", uniq_id, llm_model)
 
     # Load hydra configuration
     logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars S2 agent.")
@@ -46,30 +89,31 @@ def get_app(uniq_id, llm_model="gpt-4o-mini"):
         cfg = cfg.agents.talk2scholars.s2_agent
 
     # Define the tools
-    tools = ToolNode([s2_search, s2_display, s2_single_rec, s2_multi_rec])
+    tools = ToolNode(
+        [
+            s2_search,
+            s2_display,
+            s2_query_results,
+            s2_retrieve_id,
+            s2_single_rec,
+            s2_multi_rec,
+        ]
+    )
 
     # Define the model
     logger.log(logging.INFO, "Using OpenAI model %s", llm_model)
-    llm = ChatOpenAI(model=llm_model, temperature=cfg.temperature)
 
     # Create the agent
     model = create_react_agent(
-        llm,
+        llm_model,
         tools=tools,
         state_schema=Talk2Scholars,
-        # prompt=cfg.s2_agent,
         state_modifier=cfg.s2_agent,
         checkpointer=MemorySaver(),
     )
 
-    # Define a new graph
     workflow = StateGraph(Talk2Scholars)
-
-    # Define the two nodes we will cycle between
     workflow.add_node("agent_s2", agent_s2_node)
-
-    # Set the entrypoint as `agent`
-    # This means that this node is the first one called
     workflow.add_edge(START, "agent_s2")
 
     # Initialize memory to persist state between graph runs

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

@@ -5,14 +5,35 @@ openai_llms:
   - "gpt-4-turbo"
   - "gpt-3.5-turbo"
 temperature: 0
-main_agent: >
-  You are an intelligent research assistant coordinating academic paper discovery and analysis.
+system_prompt: >
+  You are the Talk2Scholars agent coordinating academic paper discovery and analysis.
 
-  AVAILABLE TOOLS AND ROUTING:
-  1. semantic_scholar_agent:
-     Access to tools:
-     - search_tool: For paper discovery
-     - display_results: For showing paper results
-     - get_single_paper_recommendations: For single paper recommendations
-     - get_multi_paper_recommendations: For multi-paper recommendations
-     → ROUTE TO THIS AGENT FOR: Any query about academic papers, research, or articles
+  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. This agent can also
+    retrieve the Semantic Scholar ID of a paper.
+router_prompt: >
+  You are a supervisor tasked with managing a conversation between the
+  following workers: {members}. 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.
+
+  Here are some instructions for the workers:
+  1. Call the S2 agent for general paper searches and recommendations.
+  2. The S2 agent has access to tools for querying and displaying papers.
+  3. If the user wants suggestions for papers and you don’t have 
+    a Semantic Scholar ID for it but do have the title from
+    the last displayed results, use the S2 agent to retrieve the
+    Semantic Scholar ID of the paper. Then, use the S2 agent again to display
+    recommendations for the paper.
+  4. You can call the S2 agent to get more information about a paper based
+    on the context of the conversation.
+  5. Respond with FINISH when all tasks are completed.

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

@@ -6,19 +6,11 @@ openai_llms:
   - "gpt-3.5-turbo"
 temperature: 0
 s2_agent: >
-  You are a specialized academic research agent with access to tools for paper discovery and analysis.
-
-  YOUR TOOLS:
-  1. search_tool:
-     - Finds research papers based on user queries.
-     - If no papers are found, it performs a new search.
-
-  2. display_results:
-     - Shows the current research papers.
-     - If no papers are found, it will instruct you to perform a search.
-
-  3. get_single_paper_recommendations:
-     - Provides recommendations based on a single selected paper.
-
-  4. get_multi_paper_recommendations:
-     - Provides recommendations based on multiple selected papers.
+  You are an academic research assistant with access to the
+  Semantic Scholar API for paper discovery and analysis.
+  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.

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

@@ -1,14 +1,13 @@
-# # Page configuration
-# page:
-#   title: "Talk2Scholars"
-#   icon: "🤖"
-#   layout: "wide"
+# Page configuration
+page:
+  title: "Talk2Scholars"
+  icon: "🤖"
+  layout: "wide"
 
 # Available LLM models
-llm_models:
-  - "gpt-4o-mini"
-  - "gpt-4-turbo"
-  - "gpt-3.5-turbo"
+llms:
+  available_models:
+    - "OpenAI/gpt-4o-mini"
 # # Chat UI configuration
 # chat:
 #   assistant_avatar: "🤖"
@@ -16,6 +15,9 @@ llm_models:
 #   input_placeholder: "Say something ..."
 #   spinner_text: "Fetching response ..."
 
+api_keys:
+  openai_key: "OPENAI_API_KEY"
+  nvidia_key: "NVIDIA_API_KEY"
 # # Feedback configuration
 # feedback:
 #   type: "thumbs"

aiagents4pharma/talk2scholars/configs/config.yaml

@@ -5,4 +5,5 @@ defaults:
   - tools/search: default
   - tools/single_paper_recommendation: default
   - tools/multi_paper_recommendation: default
+  - tools/retrieve_semantic_scholar_paper_id: default
   - app/frontend: default

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

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

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

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

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

@@ -10,6 +10,7 @@ api_fields:
   - "citationCount"
   - "url"
 # Commented fields that could be added later if needed
+# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

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

@@ -10,6 +10,7 @@ api_fields:
   - "citationCount"
   - "url"
 # Commented fields that could be added later if needed
+# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -1,9 +1,14 @@
 """
-This is the state file for the talk2scholars agent.
+State management for the Talk2Scholars agent.
+
+This module defines the state class `Talk2Scholars`, which maintains the conversation
+context, retrieved papers, and other relevant metadata. The state ensures consistency
+across agent interactions.
 """
 
 import logging
 from typing import Annotated, Any, Dict
+from langchain_core.language_models import BaseChatModel
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
 # Configure logging
@@ -12,22 +17,46 @@ logger = logging.getLogger(__name__)
 
 
 def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any]:
-    """Replace the existing dict with the new one."""
+    """
+    Replaces the existing dictionary with a new dictionary.
+
+    This function logs the state update and ensures that the old state is replaced
+    with the new one.
+
+    Args:
+        existing (Dict[str, Any]): The current dictionary state.
+        new (Dict[str, Any]): The new dictionary state to replace the existing one.
+
+    Returns:
+        Dict[str, Any]: The updated dictionary state.
+
+    Example:
+        >>> old_state = {"papers": {"id1": "Paper 1"}}
+        >>> new_state = {"papers": {"id2": "Paper 2"}}
+        >>> updated_state = replace_dict(old_state, new_state)
+        >>> print(updated_state)
+        {"papers": {"id2": "Paper 2"}}
+    """
     logger.info("Updating existing state %s with the state dict: %s", existing, new)
     return new
 
 
 class Talk2Scholars(AgentState):
     """
-    The state for the talk2scholars agent, inheriting from 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:
-        papers: Dictionary of papers from search results
-        multi_papers: Dictionary of papers from multi-paper recommendations
-        llm_model: Model being used
+        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.
+        llm_model (BaseChatModel): The language model instance used for generating responses.
     """
 
     # Agent state fields
+    last_displayed_papers: Annotated[Dict[str, Any], replace_dict]
     papers: Annotated[Dict[str, Any], replace_dict]
     multi_papers: Annotated[Dict[str, Any], replace_dict]
-    llm_model: str
+    llm_model: BaseChatModel