Feat/talk2scholar hydra migration

aiagents4pharma/talk2scholars/__init__.py

@@ -2,4 +2,4 @@
 This file is used to import all the modules in the package.
 """
 
-from . import agents, config, state, tests, tools
+from . import agents, configs, state, tests, tools

aiagents4pharma/talk2scholars/agents/main_agent.py

@@ -5,36 +5,35 @@
 """
 
 import logging
-from typing import Literal
-from dotenv import load_dotenv
+from typing import Literal, Any
+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.types import Command
 from ..agents import s2_agent
-from ..config.config import config
 from ..state.state_talk2scholars import Talk2Scholars
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-load_dotenv()
 
-def make_supervisor_node(llm: BaseChatModel) -> str:
+def make_supervisor_node(llm: BaseChatModel, cfg: Any) -> str:
     """
     Creates a supervisor node following LangGraph patterns.
 
     Args:
         llm (BaseChatModel): The language model to use for generating responses.
+        cfg (Any): The configuration object.
 
     Returns:
         str: The supervisor node function.
     """
-    # options = ["FINISH", "s2_agent"]
-
-    def supervisor_node(state: Talk2Scholars) -> Command[Literal["s2_agent", "__end__"]]:
+    def supervisor_node(
+        state: Talk2Scholars,
+    ) -> Command[Literal["s2_agent", "__end__"]]:
         """
         Supervisor node that routes to appropriate sub-agents.
 
@@ -44,9 +43,13 @@ def supervisor_node(state: Talk2Scholars) -> Command[Literal["s2_agent", "__end_
         Returns:
             Command[Literal["s2_agent", "__end__"]]: The command to execute next.
         """
-        logger.info("Supervisor node called")
+        logger.info(
+            "Supervisor node called - Messages count: %d, Current Agent: %s",
+            len(state["messages"]),
+            state.get("current_agent", "None"),
+        )
 
-        messages = [{"role": "system", "content": config.MAIN_AGENT_PROMPT}] + state[
+        messages = [{"role": "system", "content": cfg.state_modifier}] + state[
             "messages"
         ]
         response = llm.invoke(messages)
@@ -81,7 +84,8 @@ def supervisor_node(state: Talk2Scholars) -> Command[Literal["s2_agent", "__end_
 
     return supervisor_node
 
-def get_app(thread_id: str, llm_model ='gpt-4o-mini') -> StateGraph:
+
+def get_app(thread_id: str, llm_model="gpt-4o-mini") -> StateGraph:
     """
     Returns the langraph app with hierarchical structure.
 
@@ -91,6 +95,16 @@ def get_app(thread_id: str, llm_model ='gpt-4o-mini') -> StateGraph:
     Returns:
         The compiled langraph app.
     """
+
+    # 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)
+
     def call_s2_agent(state: Talk2Scholars) -> Command[Literal["__end__"]]:
         """
         Node for calling the S2 agent.
@@ -101,10 +115,10 @@ def call_s2_agent(state: Talk2Scholars) -> Command[Literal["__end__"]]:
         Returns:
             Command[Literal["__end__"]]: The command to execute next.
         """
-        logger.info("Calling S2 agent")
+        logger.info("Calling S2 agent with state: %s", state)
         app = s2_agent.get_app(thread_id, llm_model)
         response = app.invoke(state)
-        logger.info("S2 agent completed")
+        logger.info("S2 agent completed with response: %s", response)
         return Command(
             goto=END,
             update={
@@ -114,10 +128,17 @@ def call_s2_agent(state: Talk2Scholars) -> Command[Literal["__end__"]]:
                 "current_agent": "s2_agent",
             },
         )
-    llm = ChatOpenAI(model=llm_model, temperature=0)
+
+    logger.log(
+        logging.INFO,
+        "Using OpenAI model %s with temperature %s",
+        llm_model,
+        cfg.temperature
+    )
+    llm = ChatOpenAI(model=llm_model, temperature=cfg.temperature)
     workflow = StateGraph(Talk2Scholars)
 
-    supervisor = make_supervisor_node(llm)
+    supervisor = make_supervisor_node(llm, cfg)
     workflow.add_node("supervisor", supervisor)
     workflow.add_node("s2_agent", call_s2_agent)
 

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -1,56 +1,65 @@
-#/usr/bin/env python3
+# /usr/bin/env python3
 
-'''
+"""
 Agent for interacting with Semantic Scholar
-'''
+"""
 
 import logging
-from dotenv import load_dotenv
+import hydra
 from langchain_openai import ChatOpenAI
 from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import create_react_agent
+from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.checkpoint.memory import MemorySaver
-from ..config.config import config
 from ..state.state_talk2scholars import Talk2Scholars
-# from ..tools.s2 import s2_tools
-from ..tools.s2.search import search_tool
-from ..tools.s2.display_results import display_results
-from ..tools.s2.single_paper_rec import get_single_paper_recommendations
-from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations
+from ..tools.s2.search import search_tool as s2_search
+from ..tools.s2.display_results import display_results as s2_display
+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
 
-load_dotenv()
 
 # 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="gpt-4o-mini"):
+    """
     This function returns the langraph app.
-    '''
+    """
+
     def agent_s2_node(state: Talk2Scholars):
-        '''
+        """
         This function calls the model.
-        '''
+        """
         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
 
+    # Load hydra configuration
+    logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars S2 agent.")
+    with hydra.initialize(version_base=None, config_path="../../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["agents/talk2scholars/s2_agent=default"]
+        )
+        cfg = cfg.agents.talk2scholars.s2_agent
+
     # Define the tools
-    tools = [search_tool,
-            display_results,
-            get_single_paper_recommendations,
-            get_multi_paper_recommendations]
+    tools = ToolNode([s2_search, s2_display, 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 LLM
-    llm = ChatOpenAI(model=llm_model, temperature=0)
+    # Create the agent
     model = create_react_agent(
-                            llm,
-                            tools=tools,
-                            state_schema=Talk2Scholars,
-                            state_modifier=config.S2_AGENT_PROMPT,
-                            checkpointer=MemorySaver()
-                        )
+        llm,
+        tools=tools,
+        state_schema=Talk2Scholars,
+        state_modifier=cfg.s2_agent,
+        checkpointer=MemorySaver(),
+    )
 
     # Define a new graph
     workflow = StateGraph(Talk2Scholars)

aiagents4pharma/talk2scholars/config/__init__.py → aiagents4pharma/talk2scholars/configs/__init__.py

@@ -2,4 +2,6 @@
 This package contains configuration settings and prompts used by various AI agents
 """
 
-from . import config
+from . import agents
+from . import tools
+from . import app

aiagents4pharma/talk2scholars/configs/agents/__init__.py

@@ -0,0 +1,5 @@
+"""
+Import all the modules in the package
+"""
+
+from . import talk2scholars

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

@@ -0,0 +1,6 @@
+"""
+Import all the modules in the package
+"""
+
+from . import s2_agent
+from . import main_agent

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

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

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

@@ -0,0 +1,39 @@
+_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
+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"