gaia_agent.py

"""
GAIA Agent - Multi-step reasoning agent for complex tasks.
Uses LanggraphStateGraph for workflow orchestration and multiple specialized tools.
"""

import os
import json
from typing import List, Dict, Any, Optional, Literal
from pathlib import Path

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage
from langchain_openrouter import ChatOpenRouter
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from typing import TypedDict

from customtools import (
    load_and_analyze_excel_file,
    extract_text_from_image,
    web_search,
    wikisearch,
    youtube_transcript,
    addition_tool,
    subtraction_tool,
    multiplication_tool,
    transcribe_audio,
)
from config import (
    OPENROUTER_API_KEY,
    LLM_MODEL,
    LLM_TEMPERATURE,
    OUTPUT_FILE,
    FINAL_ANSWER_MAX_LENGTH,
    REASONING_TRACE_MAX_LENGTH,
)
from prompts import (
    PLANNER_PROMPT_TEMPLATE,
    FINALIZER_PROMPT_TEMPLATE,
)

load_dotenv()

memory = MemorySaver()


def connect_models():
    """Initialize and return the LLM instance."""
    try:
        print(f"Connecting to LLM: {LLM_MODEL}")
        llm = ChatOpenRouter(
            model=LLM_MODEL,
            temperature=LLM_TEMPERATURE,
            api_key=OPENROUTER_API_KEY,
        )
        return llm
    except Exception as e:
        print(f"Error initializing LLM: {e}")
        raise


# Tool registry
TOOLS = {
    "web_search": web_search,
    "addition_tool": addition_tool,
    "subtraction_tool": subtraction_tool,
    "multiplication_tool": multiplication_tool,
    "youtube_transcript": youtube_transcript,
    "load_and_analyze_excel_file": load_and_analyze_excel_file,
    "extract_text_from_image": extract_text_from_image,
    "wikisearch": wikisearch,
    "transcribe_audio": transcribe_audio,
}


class AgentState(TypedDict):
    """State structure for the agent workflow."""
    question: str
    plan: List[Dict[str, Any]]
    current_step: int
    selected_tool: Optional[str]
    tool_input: Optional[str]
    tool_output: Optional[str]
    intermediate_results: List[Dict[str, Any]]
    final_answer: Optional[str]
    done: bool


class Step(BaseModel):
    """Represents a single step in the plan."""
    step_number: int
    description: str
    tool: Literal[
        "web_search",
        "wikisearch",
        "youtube_transcript",
        "load_and_analyze_excel_file",
        "extract_text_from_image",
        "transcribe_audio",
        "addition_tool",
        "subtraction_tool",
        "multiplication_tool",
        "none",
    ]
    tool_input: str


class Plan(BaseModel):
    """Structured plan with multiple steps."""
    steps: List[Step]
def planner_node(state: AgentState):
    """Planner node: breaks down question into steps."""
    prompt = PLANNER_PROMPT_TEMPLATE.format(question=state['question'])
    
    planner_llm = llm.with_structured_output(Plan, method="json_schema")
    response = planner_llm.invoke(prompt)
    
    print(f"Planner generated {len(response.steps)} steps")
    
    return {
        **state,
        "plan": [step.dict() for step in response.steps],
        "current_step": 0,
        "intermediate_results": [],
        "done": False,
    }



def execute_step_node(state: AgentState):
    """Execute step node: prepares tool invocation."""
    step = state["plan"][state["current_step"]]
    tool_name = step.get("tool", "none")
    
    print(f"Executing step {state['current_step'] + 1}/{len(state['plan'])}: {tool_name}")
    
    return {
        **state,
        "tool_input": step.get("tool_input"),
        "selected_tool": tool_name,
    }


def tool_node(state: AgentState):
    """Tool execution node: invokes the selected tool."""
    tool_name = state.get("selected_tool")
    tool_input = state.get("tool_input")

    if tool_name == "none":
        return {**state, "tool_output": tool_input}
    
    print(f"Invoking tool: {tool_name}")
    tool = TOOLS.get(tool_name)
    
    # Special handling for load_and_analyze_excel_file: parse query|file_path format
    if tool_name == "load_and_analyze_excel_file" and isinstance(tool_input, str) and "|" in tool_input:
        parts = tool_input.split("|", 1)
        query = parts[0].strip()
        file_path = parts[1].strip()
        tool_input = {"query": query, "file_path": file_path}
        print(f"Parsed Excel input - Query: '{query[:50]}...', File: '{file_path}'")
    
    # Special handling for math tools: parse "a,b" format
    if tool in (addition_tool, subtraction_tool, multiplication_tool):
        try:
            a, b = tool_input.split(",")
            tool_input = {"a": a.strip(), "b": b.strip()}
        except Exception as e:
            print(f"Error parsing math tool input: {e}")
            return {**state, "tool_output": f"Error parsing input: {e}"}
    
    if not tool:
        return {**state, "tool_output": f"Unknown tool: {tool_name}"}

    try:
        result = tool.invoke(tool_input)
    except Exception as e:
        print(f"Error invoking tool {tool_name}: {e}")
        result = f"Tool error: {str(e)}"

    return {**state, "tool_output": result}


def update_state_node(state: AgentState):
    """Update state node: records tool output and progresses to next step."""
    step = state["plan"][state["current_step"]]

    state["intermediate_results"].append({
        "step": step,
        "output": state["tool_output"]
    })

    next_step = state["current_step"] + 1
    done = next_step >= len(state["plan"])

    return {
        **state,
        "current_step": next_step,
        "done": done,
    }




def should_continue(state: AgentState):
    """Conditional edge: determines if workflow should continue or finalize."""
    return "finalize" if state["done"] else "continue"


def finalizer_node(state: AgentState):
    """Finalizer node: summarizes results and generates final answer."""
    # Format intermediate results for the finalizer
    results_text = "\n".join([
        f"Step {i+1}: {r['step'].get('description', '')}\n  Output: {str(r['output'])[:100]}..."
        for i, r in enumerate(state["intermediate_results"])
    ])
    
    prompt = FINALIZER_PROMPT_TEMPLATE.format(
        question=state['question'],
        intermediate_results=results_text
    )
    
    response = llm.invoke(prompt)

    return {
        **state,
        "final_answer": response.content,
    }



def create_agent_workflow():

    graph = StateGraph(AgentState)

    # Nodes
    graph.add_node("planner", planner_node)
    graph.add_node("executor", execute_step_node)
    graph.add_node("tool", tool_node)
    graph.add_node("updater", update_state_node)
    graph.add_node("finalizer", finalizer_node) 
    # Entry
    graph.set_entry_point("planner")

    # Flow
    graph.add_edge("planner", "executor")
    graph.add_edge("executor", "tool")
    graph.add_edge("tool", "updater")

    # Loop
    graph.add_conditional_edges(
        "updater",
        should_continue,
        {
            "continue": "executor",
            "finalize": "finalizer"
        }
    )

    # End
    graph.add_edge("finalizer", END)

    return graph.compile()


def format_reasoning_trace(intermediate_results: List[Dict[str, Any]]) -> str:
    """Format intermediate results into a readable reasoning trace"""
    trace_lines = []
    for result in intermediate_results:
        step = result.get("step", {})
        output = result.get("output", "")
        description = step.get("description", "Unknown step")
        tool = step.get("tool", "none")
        
        trace_lines.append(f"Step: {description}")
        trace_lines.append(f" Tool: {tool}")
        trace_lines.append(f" Output: {output[:200]}{'...' if len(str(output)) > 200 else ''}")
    
    return "\n".join(trace_lines)


##################################################################################################################
# For Local env.
##################################################################################################################

def process_questions(questions_file: str = None, questions_list: List[str] = None) -> str:
    """
    Process multiple questions and save results to a file
    
    Args:
        questions_file: Path to a file containing questions (one per line)
        questions_list: List of questions to process
    
    Returns:
        Path to the output file with results
    """
    global llm
    llm = connect_models()
    print(f"LLM available: {llm}")
    agent = create_agent_workflow()
    
    # Get questions from either file or list
    if questions_file:
        with open(questions_file, 'r') as f:
            questions = [q.strip() for q in f.readlines() if q.strip()]
    elif questions_list:
        questions = questions_list
    else:
        raise ValueError("Either questions_file or questions_list must be provided")
    
    results = []
    
    for idx, question in enumerate(questions, 1):
        task_id = f"task_id_{idx}"
        print(f"\n{'='*80}")
        print(f"Processing {task_id}: {question[:80]}...")
        print(f"{'='*80}")
        
        try:
            # Run the agent
            result = agent.invoke({
                "question": question
            })
            
            # Extract the final answer and reasoning trace
            final_answer = result.get("final_answer", "No answer generated")
            intermediate_results = result.get("intermediate_results", [])
            
            # Format the reasoning trace
            reasoning_trace = format_reasoning_trace(intermediate_results)
            
            # Create the result object
            task_result = {
                "task_id": task_id,
                "model_answer": final_answer,
                "reasoning_trace": reasoning_trace
            }
            
            results.append(task_result)
            
            print(f"Completed {task_id}")
            print(f"Answer: {final_answer[:100]}...")
            
        except Exception as e:
            print(f"✗ Error processing {task_id}: {str(e)}")
            task_result = {
                "task_id": task_id,
                "model_answer": f"Error: {str(e)}",
                "reasoning_trace": "Failed to execute agent"
            }
            results.append(task_result)
    
    # Save results to file
    output_file = "/home/nitin/AI/hfagent/results.jsonl"
    with open(output_file, 'w') as f:
        for result in results:
            f.write(json.dumps(result) + '\n')
    
    print(f"\n{'='*80}")
    print(f"All tasks completed. Results saved to: {output_file}")
    print(f"{'='*80}")
    
    return output_file





#if __name__ == "__main__":

#    global llm