add agents

src/agents/langgraph_agent_v1.py

@@ -0,0 +1,289 @@
+from collections.abc import Sequence
+from typing import Annotated, Literal, TypedDict
+
+from langchain.chat_models import init_chat_model
+from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
+from langchain_community.tools.arxiv import ArxivQueryRun
+from langchain_community.tools.pubmed.tool import PubmedQueryRun
+from langchain_community.tools.semanticscholar.tool import SemanticScholarQueryRun
+from langchain_community.tools.wikidata.tool import WikidataAPIWrapper, WikidataQueryRun
+from langchain_community.utilities import WikipediaAPIWrapper
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import Tool
+from langchain_experimental.utilities import PythonREPL
+from langgraph.graph import END, StateGraph
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
+
+# Set up tools
+python_repl = PythonREPL()
+repl_tool = Tool(
+    name="python_repl",
+    description="A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.",
+    func=python_repl.run,
+)
+
+# Initialize all tools
+tools = [
+    DuckDuckGoSearchRun(),
+    PubmedQueryRun(),
+    SemanticScholarQueryRun(),
+    ArxivQueryRun(),
+    WikidataQueryRun(api_wrapper=WikidataAPIWrapper()),
+    WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
+    repl_tool,
+]
+
+# Initialize Gemini model
+model = init_chat_model("gemini-2.0-flash", model_provider="google_genai")
+model_with_tools = model.bind_tools(tools)
+
+# Create tools lookup
+tools_by_name = {tool.name: tool for tool in tools}
+
+
+# Pydantic models for structured output
+class ToolSufficiencyResponse(BaseModel):
+    """Response for tool sufficiency check"""
+
+    sufficient: bool = Field(description="Whether the available tools are sufficient to answer the question")
+    reasoning: str = Field(description="Brief reasoning for the decision")
+
+
+class FinalAnswer(BaseModel):
+    """Final answer structure"""
+
+    answer: str = Field(description="The comprehensive answer to the user's question")
+
+
+# Graph state
+class AgentState(TypedDict):
+    """The state of the agent."""
+
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+    llm_call_count: int
+    max_llm_calls: int
+    question: str | None = None
+    answer: FinalAnswer | None = None
+    tool_sufficiency: ToolSufficiencyResponse | None = None
+
+
+# Node functions
+def check_tool_sufficiency(state: AgentState, config: RunnableConfig):
+    """Check if available tools are sufficient to answer the question"""
+
+    question = state["question"]
+    question_message = HumanMessage(content=f"Question to analyze: {question}")
+
+    # Create system prompt for sufficiency check
+    available_tools_desc = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
+
+    system_prompt = f"""You are an AI assistant that needs to determine if the available tools are sufficient to answer a user's question.
+
+Available tools:
+{available_tools_desc}
+
+Your task is to analyze the user's question and determine if these tools provide sufficient capability to answer it comprehensively.
+
+Consider:
+- Can the question be answered with web search, academic papers, or computational tools?
+- Does the question require real-time data, personal information, or capabilities not available through these tools?
+- Can you break down the question into parts that these tools can handle?
+
+Be generous in your assessment - if there's a reasonable path to answer the question using these tools, respond with sufficient=True."""
+
+    messages = [SystemMessage(content=system_prompt), question_message]
+
+    structured_model = model.with_structured_output(ToolSufficiencyResponse)
+    response = structured_model.invoke(messages, config)
+
+    # Add response to messages for context
+    response_message = AIMessage(
+        content=f"Tool sufficiency check: {'Sufficient' if response.sufficient else 'Insufficient'}. Reasoning: {response.reasoning}"
+    )
+
+    return {"messages": [question_message, response_message], "tool_sufficiency": response}
+
+
+def call_model(state: AgentState, config: RunnableConfig):
+    """Call the model (ReAct agent LLM node)"""
+
+    system_prompt = SystemMessage(
+        content="""You are a helpful AI assistant with access to various tools. Use the tools available to you to answer the user's question comprehensively.
+
+    Think step by step:
+    1. Analyze what information you need.
+    2. Use appropriate tools to gather that information.
+    3. Synthesize the information to provide a complete answer.
+
+    IMPORTANT INSTRUCTIONS:
+    - Avoid repeating the same tool call with identical parameters.
+    - When calling tools, vary your queries and tool arguments to explore different aspects of the question.
+    - Use each tool intelligently and purposefully—avoid redundant or uninformative tool calls.
+    - Track which tools you've already used and how, so you don't repeat yourself.
+
+    Be thorough but efficient with your tool usage. Use tools only when needed, and prefer combining information from multiple sources."""
+    )
+
+    response = model_with_tools.invoke([system_prompt] + state["messages"], config)
+
+    # Increment LLM call count
+    new_count = state.get("llm_call_count", 0) + 1
+
+    return {"messages": [response], "llm_call_count": new_count}
+
+
+def tool_node(state: AgentState):
+    """Execute tools based on the last message's tool calls"""
+
+    outputs = []
+    last_message = state["messages"][-1]
+
+    for tool_call in last_message.tool_calls:
+        try:
+            tool_result = tools_by_name[tool_call["name"]].invoke(tool_call["args"])
+            outputs.append(
+                ToolMessage(
+                    content=str(tool_result),
+                    name=tool_call["name"],
+                    tool_call_id=tool_call["id"],
+                )
+            )
+        except Exception as e:
+            outputs.append(
+                ToolMessage(
+                    content=f"Error executing tool {tool_call['name']}: {str(e)}",
+                    name=tool_call["name"],
+                    tool_call_id=tool_call["id"],
+                )
+            )
+
+    return {"messages": outputs}
+
+
+def final_answer_node(state: AgentState, config: RunnableConfig):
+    """Generate final structured answer based on conversation history"""
+
+    system_prompt = SystemMessage(
+        content="""You are tasked with providing the most concise possible final answer to the user question based on the conversation history and tool usage.
+
+CRITICAL INSTRUCTIONS FOR CONCISENESS:
+- If the question asks for a number, provide ONLY the number (e.g., "5", "23", "147")
+- If the question asks for a name, provide ONLY the name (e.g., "John Smith", "Paris")
+- If the question asks for a yes/no, provide ONLY "Yes" or "No"
+- If the question asks for a date, provide ONLY the date (e.g., "2023-05-15", "March 2020")
+- Remove ALL unnecessary words, articles, explanations, or context
+- Do NOT include phrases like "The answer is", "Based on the research", "According to", etc.
+- Provide the absolute minimum text that directly answers the question
+- This is for a benchmark submission where brevity is crucial
+
+If the answer cannot be determined from the conversation history, respond with "Unable to determine" (nothing more)."""
+    )
+
+    filtered_messages = []
+    for msg in state["messages"]:
+        if hasattr(msg, "content") and msg.content and msg.content.strip():
+            filtered_messages.append(msg)
+
+    messages = [system_prompt] + filtered_messages
+
+    structured_model = model.with_structured_output(FinalAnswer)
+    response = structured_model.invoke(messages, config)
+
+    return {"messages": [AIMessage(content=f"Final Answer: {response.answer}")], "answer": response}
+
+
+# Edge functions
+def should_continue_sufficiency(state: AgentState) -> Literal["sufficient", "insufficient"]:
+    """Decide whether tools are sufficient"""
+    # Check if we have a tool sufficiency result
+    sufficiency = state["tool_sufficiency"]
+
+    if sufficiency and hasattr(sufficiency, "sufficient") and sufficiency.sufficient:
+        return "sufficient"
+    else:
+        return "insufficient"
+
+
+def should_continue_react(state: AgentState) -> Literal["tools", "final_answer"]:
+    """Decide whether to continue with ReAct loop or move to final answer"""
+
+    messages = state["messages"]
+    last_message = messages[-1]
+
+    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
+        return "tools"
+
+    return "final_answer"
+
+
+def should_continue_after_tools(state: AgentState) -> Literal["agent", "final_answer"]:
+    """Decide after tool execution whether to continue or finalize"""
+
+    llm_call_count = state.get("llm_call_count", 0)
+    max_calls = state.get("max_llm_calls", 4)
+
+    # If we've reached the maximum number of LLM calls, go to final answer
+    if llm_call_count >= max_calls:
+        return "final_answer"
+
+    # Otherwise, continue the ReAct loop (go back to agent)
+    return "agent"
+
+
+def create_react_agent_graph():
+    """Create and return the compiled ReAct agent graph"""
+
+    workflow = StateGraph(AgentState)
+
+    # Add nodes
+    workflow.add_node("check_sufficiency", check_tool_sufficiency)
+    workflow.add_node("agent", call_model)
+    workflow.add_node("tools", tool_node)
+    workflow.add_node("final_answer", final_answer_node)
+
+    # Set entry point
+    workflow.set_entry_point("check_sufficiency")
+
+    # Add conditional edges
+    workflow.add_conditional_edges(
+        "check_sufficiency", should_continue_sufficiency, {"sufficient": "agent", "insufficient": END}
+    )
+    workflow.add_conditional_edges("agent", should_continue_react, {"tools": "tools", "final_answer": "final_answer"})
+    workflow.add_conditional_edges(
+        "tools", should_continue_after_tools, {"agent": "agent", "final_answer": "final_answer"}
+    )
+
+    # Add edges
+    workflow.add_edge("final_answer", END)
+
+    return workflow.compile()
+
+
+def run_agent(question: str, max_llm_calls: int = 4):
+    """Run the ReAct agent with a question"""
+
+    graph = create_react_agent_graph()
+
+    initial_state = {"messages": [HumanMessage(content=question)], "llm_call_count": 0, "max_llm_calls": max_llm_calls}
+
+    # Stream the execution
+    print(f"Question: {question}")
+    print("=" * 50)
+
+    for step in graph.stream(initial_state):
+        for node, output in step.items():
+            print(f"\n--- {node.upper()} ---")
+            if "messages" in output and output["messages"]:
+                for msg in output["messages"]:
+                    if hasattr(msg, "content"):
+                        print(f"{msg.__class__.__name__}: {msg.content}")
+                    elif hasattr(msg, "tool_calls") and msg.tool_calls:
+                        print(f"Tool calls: {[tc['name'] for tc in msg.tool_calls]}")
+
+            if "final_answer" in output:
+                print("\nFINAL STRUCTURED ANSWER:")
+                print(f"Answer: {output['final_answer'].answer}")
+                print(f"Confidence: {output['final_answer'].confidence}")
+                print(f"Sources: {output['final_answer'].sources_used}")

src/agents/langgraph_agent_v2.py

@@ -0,0 +1,364 @@
+from collections.abc import Sequence
+from typing import Annotated, Literal
+
+from langchain.chat_models import init_chat_model
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph import END, StateGraph
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
+
+from src.agents.models import FeasibilityCheck, FinalAnswer, FinalConclusion, NextStep
+from src.agents.prompts import GAIAPrompts
+from src.agents.tools import tools
+
+# Initialize
+model = init_chat_model("gemini-2.0-flash", model_provider="google_genai")
+model_with_tools = model.bind_tools(tools)
+
+tools_by_name = {tool.name: tool for tool in tools}
+
+prompts = GAIAPrompts()
+
+
+# Graph state
+class GraphState(BaseModel):
+    """The state of the graph"""
+
+    # History
+    history: Annotated[Sequence[BaseMessage], add_messages] = Field(
+        default_factory=list
+    )  # Complete history with node info
+    coordinator_messages: Annotated[Sequence[BaseMessage], add_messages] = Field(
+        default_factory=list
+    )  # Coordinator-specific messages
+    executor_messages: Sequence[BaseMessage] = Field(default_factory=list)  # Executor-specific messages
+
+    # Input
+    question: str
+
+    # Feasibility check
+    feasibility: FeasibilityCheck | None = None
+
+    # Coordinator state
+    next_step: NextStep | None = None
+    coordinator_conclusion: FinalConclusion | None = None
+    coordinator_iterations: int
+    coordinator_max_iterations: int
+
+    # Executor state
+    executor_conclusion: FinalConclusion | None = None
+    executor_iterations: int
+    executor_max_iterations: int
+
+    # Final answer state
+    final_answer: FinalAnswer | None = None
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+
+# Nodes
+def check_feasibility(state: GraphState, config: RunnableConfig):
+    """Check if the question is feasible to answer with the available tools"""
+
+    question = state["question"]
+
+    system_message = SystemMessage(content=prompts.get_feasibility_check_prompt(tools), node="feasibility")
+    question_message = HumanMessage(content=question, node="feasibility")
+    messages = [system_message, question_message]
+
+    structured_model = model.with_structured_output(FeasibilityCheck)
+    response = structured_model.invoke(messages, config)
+
+    response_message = AIMessage(content=str(response), node="feasibility")
+    messages += [response_message]
+    return {
+        "history": messages,
+        "feasibility": response,
+    }
+
+
+def coordinator_node(state: GraphState, config: RunnableConfig):
+    """Determine the next step in the plan and select appropriate tools"""
+
+    coordinator_messages = state["coordinator_messages"]
+    new_messages = []
+
+    if not coordinator_messages:
+        system_message = SystemMessage(content=prompts.get_coordinator_system_prompt(tools), node="coordinator")
+        human_message = HumanMessage(
+            content=prompts.get_coordinator_context_prompt(state["question"]), node="coordinator"
+        )
+        coordinator_messages = [system_message, human_message]
+        new_messages = coordinator_messages
+
+    if state["executor_conclusion"]:
+        executor_message = AIMessage(
+            content=f"Executor conclusion: {state['executor_conclusion'].conclusion}. Complete text: {str(state['executor_conclusion'])}",
+            node="executor",
+        )
+        coordinator_messages += [executor_message]
+        new_messages += [executor_message]
+
+    # Check if we've reached max iterations
+    if (state["next_step"] and state["next_step"].is_final) or (
+        state["coordinator_iterations"] >= state["coordinator_max_iterations"]
+    ):
+        # Generate final conclusion instead of next step
+        human_message = HumanMessage(
+            content=prompts.get_coordinator_max_iterations_prompt(state["question"]), node="coordinator"
+        )
+
+        structured_model = model.with_structured_output(FinalConclusion)
+        response = structured_model.invoke(coordinator_messages + [human_message], config)
+        response_message = AIMessage(content=str(response), node="coordinator")
+
+        new_messages += [human_message, response_message]
+        return {
+            "history": new_messages,
+            "coordinator_messages": new_messages,
+            "coordinator_conclusion": response,
+            "coordinator_iterations": state["coordinator_iterations"] + 1,
+        }
+
+    structured_model = model.with_structured_output(NextStep)
+    response = structured_model.invoke(coordinator_messages, config)
+
+    response_message = AIMessage(content=str(response), node="coordinator")
+    new_messages += [response_message]
+
+    return {
+        "history": new_messages,
+        "coordinator_messages": new_messages,
+        "coordinator_iterations": state["coordinator_iterations"] + 1,
+        "next_step": response,
+        "executor_messages": [],
+        "executor_conclusion": None,
+        "executor_iterations": 0,
+    }
+
+
+def executor_node(state: GraphState, config: RunnableConfig):
+    """Plan the execution of the current step using ReAct pattern"""
+    if not state["next_step"]:
+        return {
+            "executor_conclusion": FinalConclusion(conclusion="No next step", partial_results=""),
+            "executor_iterations": state["executor_iterations"] + 1,
+        }
+
+    messages = state["executor_messages"]
+
+    if not messages:
+        system_message = SystemMessage(
+            content=prompts.get_executor_system_prompt(state["next_step"].tools),
+            node="executor",
+        )
+        human_message = HumanMessage(content=prompts.get_executor_task_prompt(state["next_step"].step), node="executor")
+        messages = [system_message, human_message]
+
+    if state["executor_iterations"] >= state["executor_max_iterations"]:
+        # Generate final conclusion and return to coordinator
+        human_message = HumanMessage(
+            content=prompts.get_executor_max_iterations_prompt(state["next_step"].step),
+            node="executor",
+        )
+
+        messages += [human_message]
+
+        structured_model = model.with_structured_output(FinalConclusion)
+        response = structured_model.invoke(messages, config)
+
+        response_message = AIMessage(
+            content=f"Executor conclusion: {str(response)}",
+            node="executor",
+        )
+
+        return {
+            "history": [human_message, response_message],
+            "executor_conclusion": response
+            or FinalConclusion(conclusion="Failed to generate conclusion", partial_results=""),
+            "executor_iterations": state["executor_iterations"] + 1,
+        }
+
+    selected_tools = [tool for tool in tools if tool.name in state["next_step"].tools]
+    model_with_selected_tools = model.bind_tools(selected_tools)
+
+    response_message = model_with_selected_tools.invoke(messages, config)
+    response_message.node = "executor"
+
+    return {
+        "history": response_message,
+        "executor_messages": messages + [response_message],
+        "executor_iterations": state["executor_iterations"] + 1,
+    }
+
+
+def tool_node(state: GraphState):
+    """Execute tools based on the last message's tool calls"""
+    outputs = []
+    messages = state["executor_messages"]
+    last_message = state["executor_messages"][-1]
+
+    for tool_call in last_message.tool_calls:
+        try:
+            tool_result = tools_by_name[tool_call["name"]].invoke(tool_call["args"])
+            tool_message = ToolMessage(
+                content=str(tool_result),
+                name=tool_call["name"],
+                tool_call_id=tool_call["id"],
+                node="tools",
+            )
+            outputs.append(tool_message)
+        except Exception as e:
+            tool_message = ToolMessage(
+                content=f"Error executing tool {tool_call['name']}: {str(e)}",
+                name=tool_call["name"],
+                tool_call_id=tool_call["id"],
+                node="tools",
+            )
+            outputs.append(tool_message)
+
+    return {
+        "history": outputs,
+        "executor_messages": messages + outputs,
+    }
+
+
+def finalise(state: GraphState, config: RunnableConfig):
+    """Generate the final answer based on coordinator history"""
+    system_message = SystemMessage(content=prompts.get_finalizer_prompt(), node="finalise")
+    messages = [system_message] + state["coordinator_messages"]
+
+    structured_model = model.with_structured_output(FinalAnswer)
+    response = structured_model.invoke(messages, config)
+    response_message = AIMessage(content=str(response), node="finalise")
+
+    return {"history": response_message, "final_answer": response}
+
+
+# Edges
+def should_continue_after_feasibility(state: GraphState) -> Literal["coordinator", END]:
+    """Decide whether to continue with coordination or end"""
+    if state["feasibility"] and state["feasibility"].feasible:
+        return "coordinator"
+    return END
+
+
+def should_continue_after_coordinator(state: GraphState) -> Literal["executor", "finalise"]:
+    """Decide whether to continue with execution or go to final answer"""
+    if state["coordinator_conclusion"] or (state["coordinator_iterations"] >= state["coordinator_max_iterations"]):
+        return "finalise"
+    return "executor"
+
+
+def should_continue_after_executor(state: GraphState) -> Literal["tools", "coordinator", "executor"]:
+    """Decide whether to continue with tools or go back to coordinator"""
+    last_message = state["executor_messages"][-1]
+    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
+        return "tools"
+
+    if state["executor_conclusion"]:
+        return "coordinator"
+
+    return "executor"
+
+
+def should_continue_after_tools(state: GraphState) -> Literal["executor"]:
+    """Tools always go back to executor"""
+    return "executor"
+
+
+# Graph
+def build_graph():
+    """Build the graph"""
+    graph = StateGraph(GraphState)
+
+    # Add nodes
+    graph.add_node("check_feasibility", check_feasibility)
+    graph.add_node("coordinator", coordinator_node)
+    graph.add_node("executor", executor_node)
+    graph.add_node("tools", tool_node)
+    graph.add_node("finalise", finalise)
+
+    # Set entry point
+    graph.set_entry_point("check_feasibility")
+
+    # Add edges
+    graph.add_conditional_edges(
+        "check_feasibility", should_continue_after_feasibility, {"coordinator": "coordinator", END: END}
+    )
+    graph.add_conditional_edges(
+        "coordinator", should_continue_after_coordinator, {"executor": "executor", "finalise": "finalise"}
+    )
+    graph.add_conditional_edges(
+        "executor",
+        should_continue_after_executor,
+        {"executor": "executor", "tools": "tools", "coordinator": "coordinator"},
+    )
+    graph.add_conditional_edges(
+        "tools",
+        should_continue_after_tools,
+        {"executor": "executor"},
+    )
+
+    # Finalise node goes to END
+    graph.add_edge("finalise", END)
+
+    return graph.compile()
+
+
+def run_agent(question: str, coordinator_max_iterations: int = 5, executor_max_iterations: int = 3):
+    """Run the agent with a question"""
+    graph = build_graph()
+
+    initial_state = {
+        "question": question,
+        "history": [],
+        "coordinator_messages": [],
+        "executor_messages": [],
+        "coordinator_iterations": 0,
+        "executor_iterations": 0,
+        "coordinator_max_iterations": coordinator_max_iterations,
+        "executor_max_iterations": executor_max_iterations,
+    }
+
+    # Stream the execution
+    print(f"Question: {question}")
+    print("=" * 50)
+
+    for step in graph.stream(initial_state):
+        for node, output in step.items():
+            print(f"\n--- {node.upper()} ---")
+
+            # Print history with node information
+            if "history" in output and output["history"]:
+                print("\nComplete History (with node info):")
+                for msg in output["history"]:
+                    node_info = getattr(msg, "node", "unknown") if hasattr(msg, "node") else "unknown"
+                    content = getattr(msg, "content", str(msg)) if hasattr(msg, "content") else str(msg)
+                    print(f"[{node_info}] {msg.__class__.__name__}: {content}")
+
+            if "coordinator_messages" in output and output["coordinator_messages"]:
+                print("\nCoordinator Messages:")
+                for msg in output["coordinator_messages"]:
+                    if hasattr(msg, "content"):
+                        print(f"{msg.__class__.__name__}: {msg.content}")
+
+            if "executor_messages" in output and output["executor_messages"]:
+                print("\nExecutor Messages:")
+                for msg in output["executor_messages"]:
+                    if hasattr(msg, "content"):
+                        print(f"{msg.__class__.__name__}: {msg.content}")
+
+            if "executor_conclusion" in output and output["executor_conclusion"]:
+                print("\n=== EXECUTOR CONCLUSION ===")
+                print(f"Conclusion: {output['executor_conclusion'].conclusion}")
+                print(f"Partial Results: {output['executor_conclusion'].partial_results}")
+                print(f"Confidence: {output['executor_conclusion'].confidence}")
+
+            if "final_answer" in output and output["final_answer"]:
+                print("\n=== FINAL ANSWER ===")
+                print(f"Answer: {output['final_answer'].answer}")
+                print(f"Confidence: {output['final_answer'].confidence}")
+                print(f"Reasoning: {output['final_answer'].reasoning}")