aiagents4pharma 1.19.1__py3-none-any.whl → 1.20.1__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

@@ -1,151 +1,207 @@
 #!/usr/bin/env python3
 
 """
-Main agent for the talk2scholars app.
+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, Any
+from typing import Literal, Callable
 import hydra
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import 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
 
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def make_supervisor_node(llm: BaseChatModel, cfg: Any) -> str:
+def get_hydra_config():
     """
-    Creates a supervisor node following LangGraph patterns.
+    Loads and returns 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.
+
+    Returns:
+        Any: The configuration object 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: BaseChatModel, thread_id: str) -> Callable:
+    """
+    Creates and returns a supervisor node for intelligent routing using the ReAct pattern.
+
+    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.
 
     Args:
-        llm (BaseChatModel): The language model to use for generating responses.
-        cfg (Any): The configuration object.
+        llm (BaseChatModel): The language model used by the supervisor agent.
+        thread_id (str): Unique identifier for the conversation session.
 
     Returns:
-        str: The supervisor node function.
+        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)
     """
+    logger.info("Loading Hydra configuration for Talk2Scholars main agent.")
+    cfg = get_hydra_config()
+    logger.info("Hydra configuration loaded with values: %s", cfg)
+
+    # 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(),
+    )
+
     def supervisor_node(
         state: Talk2Scholars,
     ) -> Command[Literal["s2_agent", "__end__"]]:
         """
-        Supervisor node that routes to appropriate sub-agents.
+        Processes user queries and determines the next step in the conversation flow.
+
+        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.
 
         Args:
-            state (Talk2Scholars): The current state of the conversation.
+            state (Talk2Scholars): The current state of the conversation, containing
+                messages, papers, and metadata.
 
         Returns:
-            Command[Literal["s2_agent", "__end__"]]: The command to execute next.
+            Command: The next action to be executed, along with updated state data.
+
+        Example:
+            result = supervisor_node(current_state)
+            next_step = result.goto
         """
         logger.info(
-            "Supervisor node called - Messages count: %d, Current Agent: %s",
+            "Supervisor node called - Messages count: %d",
             len(state["messages"]),
-            state.get("current_agent", "None"),
         )
 
-        messages = [{"role": "system", "content": cfg.state_modifier}] + state[
-            "messages"
-        ]
-        response = llm.invoke(messages)
-        goto = (
-            "FINISH"
-            if not any(
-                kw in state["messages"][-1].content.lower()
-                for kw in ["search", "paper", "find"]
-            )
-            else "s2_agent"
+        # 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)
 
-        if goto == "FINISH":
-            return Command(
-                goto=END,
-                update={
-                    "messages": state["messages"]
-                    + [AIMessage(content=response.content)],
-                    "is_last_step": True,
-                    "current_agent": None,
-                },
-            )
-
-        return Command(
-            goto="s2_agent",
-            update={
-                "messages": state["messages"],
-                "is_last_step": False,
-                "current_agent": "s2_agent",
-            },
-        )
+        return Command(goto=goto)
 
     return supervisor_node
 
 
-def get_app(thread_id: str, llm_model="gpt-4o-mini") -> StateGraph:
+def get_app(thread_id: str, llm_model: str = "gpt-4o-mini") -> StateGraph:
     """
-    Returns the langraph app with hierarchical structure.
+    Initializes and returns the LangGraph application with a 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.
 
     Args:
-        thread_id (str): The thread ID for the conversation.
+        thread_id (str): Unique identifier for the conversation session.
+        llm_model (str, optional): The language model to be used. Defaults to "gpt-4o-mini".
 
     Returns:
-        The compiled langraph app.
-    """
+        StateGraph: A compiled LangGraph application ready for query invocation.
 
-    # Load hydra configuration
-    logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars main agent.")
-    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.info("Hydra configuration loaded with values: %s", cfg)
+    Example:
+        app = get_app("thread_123")
+        result = app.invoke(initial_state)
+    """
+    cfg = get_hydra_config()
 
-    def call_s2_agent(state: Talk2Scholars) -> Command[Literal["__end__"]]:
+    def call_s2_agent(
+        state: Talk2Scholars,
+    ) -> Command[Literal["supervisor", "__end__"]]:
         """
-        Node for calling the S2 agent.
+        Calls the Semantic Scholar (S2) agent to process academic paper queries.
+
+        This function invokes the S2 agent, retrieves relevant research papers,
+        and updates the conversation state accordingly.
 
         Args:
-            state (Talk2Scholars): The current state of the conversation.
+            state (Talk2Scholars): The current conversation state, including user queries
+                and any previously retrieved papers.
 
         Returns:
-            Command[Literal["__end__"]]: The command to execute next.
+            Command: The next action to execute, along with updated messages and papers.
+
+        Example:
+            result = call_s2_agent(current_state)
+            next_step = result.goto
         """
         logger.info("Calling S2 agent with state: %s", state)
         app = s2_agent.get_app(thread_id, llm_model)
-        response = app.invoke(state)
+
+        # 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,
+                }
+            },
+        )
         logger.info("S2 agent completed with response: %s", response)
+
         return Command(
             goto=END,
             update={
                 "messages": response["messages"],
-                "papers": response.get("papers", []),
-                "is_last_step": True,
-                "current_agent": "s2_agent",
+                "papers": response.get("papers", {}),
+                "multi_papers": response.get("multi_papers", {}),
             },
         )
 
-    logger.log(
-        logging.INFO,
-        "Using OpenAI model %s with temperature %s",
-        llm_model,
-        cfg.temperature
-    )
+    # Initialize LLM
+    logger.info("Using OpenAI model %s with temperature %s", llm_model, cfg.temperature)
     llm = ChatOpenAI(model=llm_model, temperature=cfg.temperature)
+
+    # Build the graph
     workflow = StateGraph(Talk2Scholars)
+    supervisor = make_supervisor_node(llm, thread_id)
 
-    supervisor = make_supervisor_node(llm, cfg)
     workflow.add_node("supervisor", supervisor)
     workflow.add_node("s2_agent", call_s2_agent)
-
-    # Define edges
     workflow.add_edge(START, "supervisor")
     workflow.add_edge("s2_agent", END)
 
+    # Compile the graph without initial state
     app = workflow.compile(checkpointer=MemorySaver())
     logger.info("Main agent workflow compiled")
     return app

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -39,7 +39,7 @@ def get_app(uniq_id, llm_model="gpt-4o-mini"):
 
     # Load hydra configuration
     logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars S2 agent.")
-    with hydra.initialize(version_base=None, config_path="../../configs"):
+    with hydra.initialize(version_base=None, config_path="../configs"):
         cfg = hydra.compose(
             config_name="config", overrides=["agents/talk2scholars/s2_agent=default"]
         )
@@ -57,6 +57,7 @@ def get_app(uniq_id, llm_model="gpt-4o-mini"):
         llm,
         tools=tools,
         state_schema=Talk2Scholars,
+        # prompt=cfg.s2_agent,
         state_modifier=cfg.s2_agent,
         checkpointer=MemorySaver(),
     )

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

@@ -6,34 +6,13 @@ openai_llms:
   - "gpt-3.5-turbo"
 temperature: 0
 main_agent: >
-  "You are a supervisory AI agent that routes user queries to specialized tools.\n"
-  "Your task is to select the most appropriate tool based on the user's request.\n\n"
-  "Available tools and their capabilities:\n\n"
-  "1. semantic_scholar_agent:\n"
-  "   - Search for academic papers and research\n"
-  "   - Get paper recommendations\n"
-  "   - Find similar papers\n"
-  "   USE FOR: Any queries about finding papers, academic research, "
-  "or getting paper recommendations\n\n"
-  "ROUTING GUIDELINES:\n\n"
-  "ALWAYS route to semantic_scholar_agent for:\n"
-  "- Finding academic papers\n"
-  "- Searching research topics\n"
-  "- Getting paper recommendations\n"
-  "- Finding similar papers\n"
-  "- Any query about academic literature\n\n"
-  "Approach:\n"
-  "1. Identify the core need in the user's query\n"
-  "2. Select the most appropriate tool based on the guidelines above\n"
-  "3. If unclear, ask for clarification\n"
-  "4. For multi-step tasks, focus on the immediate next step\n\n"
-  "Remember:\n"
-  "- Be decisive in your tool selection\n"
-  "- Focus on the immediate task\n"
-  "- Default to semantic_scholar_agent for any paper-finding tasks\n"
-  "- Ask for clarification if the request is ambiguous\n\n"
-  "When presenting paper search results, always use this exact format:\n\n"
-  "Remember to:\n"
-  "- To always add the url\n"
-  "- Put URLs on the title line itself as markdown\n"
-  "- Maintain consistent spacing and formatting"
+  You are an intelligent research assistant 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

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

@@ -6,63 +6,19 @@ openai_llms:
   - "gpt-3.5-turbo"
 temperature: 0
 s2_agent: >
-  "You are a specialized academic research assistant with access to the following tools:\n\n"
-  "1. search_papers:\n"
-  "   USE FOR: General paper searches\n"
-  "   - Enhances search terms automatically\n"
-  "   - Adds relevant academic keywords\n"
-  "   - Focuses on recent research when appropriate\n\n"
-  "2. get_single_paper_recommendations:\n"
-  "   USE FOR: Finding papers similar to a specific paper\n"
-  "   - Takes a single paper ID\n"
-  "   - Returns related papers\n\n"
-  "3. get_multi_paper_recommendations:\n"
-  "   USE FOR: Finding papers similar to multiple papers\n"
-  "   - Takes multiple paper IDs\n"
-  "   - Finds papers related to all inputs\n\n"
-  "GUIDELINES:\n\n"
-  "For paper searches:\n"
-  "- Enhance search terms with academic language\n"
-  "- Include field-specific terminology\n"
-  '- Add "recent" or "latest" when appropriate\n'
-  "- Keep queries focused and relevant\n\n"
-  "For paper recommendations:\n"
-  "- Identify paper IDs (40-character hexadecimal strings)\n"
-  "- Use single_paper_recommendations for one ID\n"
-  "- Use multi_paper_recommendations for multiple IDs\n\n"
-  "Best practices:\n"
-  "1. Start with a broad search if no paper IDs are provided\n"
-  "2. Look for paper IDs in user input\n"
-  "3. Enhance search terms for better results\n"
-  "4. Consider the academic context\n"
-  "5. Be prepared to refine searches based on feedback\n\n"
-  "Remember:\n"
-  "- Always select the most appropriate tool\n"
-  "- Enhance search queries naturally\n"
-  "- Consider academic context\n"
-  "- Focus on delivering relevant results\n\n"
-  "IMPORTANT GUIDELINES FOR PAPER RECOMMENDATIONS:\n\n"
-  "For Multiple Papers:\n"
-  "- When getting recommendations for multiple papers, always use "
-  "get_multi_paper_recommendations tool\n"
-  "- DO NOT call get_single_paper_recommendations multiple times\n"
-  "- Always pass all paper IDs in a single call to get_multi_paper_recommendations\n"
-  '- Use for queries like "find papers related to both/all papers" or '
-  '"find similar papers to these papers"\n\n'
-  "For Single Paper:\n"
-  "- Use get_single_paper_recommendations when focusing on one specific paper\n"
-  "- Pass only one paper ID at a time\n"
-  '- Use for queries like "find papers similar to this paper" or '
-  '"get recommendations for paper X"\n'
-  "- Do not use for multiple papers\n\n"
-  "Examples:\n"
-  '- For "find related papers for both papers":\n'
-  "  ✓ Use get_multi_paper_recommendations with both paper IDs\n"
-  "  × Don't make multiple calls to get_single_paper_recommendations\n\n"
-  '- For "find papers related to the first paper":\n'
-  "  ✓ Use get_single_paper_recommendations with just that paper's ID\n"
-  "  × Don't use get_multi_paper_recommendations\n\n"
-  "Remember:\n"
-  "- Be precise in identifying which paper ID to use for single recommendations\n"
-  "- Don't reuse previous paper IDs unless specifically requested\n"
-  "- For fresh paper recommendations, always use the original paper ID"
+  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.

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -3,10 +3,8 @@ This is the state file for the talk2scholars agent.
 """
 
 import logging
-from typing import Annotated, Any, Dict, Optional
-
+from typing import Annotated, Any, Dict
 from langgraph.prebuilt.chat_agent_executor import AgentState
-from typing_extensions import NotRequired, Required
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -22,11 +20,14 @@ def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any
 class Talk2Scholars(AgentState):
     """
     The state for the talk2scholars agent, inheriting from AgentState.
+
+    Attributes:
+        papers: Dictionary of papers from search results
+        multi_papers: Dictionary of papers from multi-paper recommendations
+        llm_model: Model being used
     """
 
-    papers: Annotated[Dict[str, Any], replace_dict]  # Changed from List to Dict
-    search_table: NotRequired[str]
-    next: str  # Required for routing in LangGraph
-    current_agent: NotRequired[Optional[str]]
-    is_last_step: Required[bool]  # Required field for LangGraph
+    # Agent state fields
+    papers: Annotated[Dict[str, Any], replace_dict]
+    multi_papers: Annotated[Dict[str, Any], replace_dict]
     llm_model: str