""" 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 langchain_nvidia_ai_endpoints import ChatNVIDIA 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, modulus_tool, power_tool, ) from config import ( OPENROUTER_API_KEY, LLM_MODEL, NVIDIA, NVIDIA_API_KEY, NVIDIA_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: global llm if NVIDIA: llm = ChatNVIDIA( model=NVIDIA_MODEL, api_key= NVIDIA_API_KEY, temperature=1, top_p=1, ) else: 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, "modulus_tool": modulus_tool, "power_tool": power_tool, } 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", "modulus_tool", "power_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"]] print(f"Current Step:{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) 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 # Example questions to process questions = [ """Task ID: 52e8ce1c-09bd-4537-8e2d-67d1648779b9 ; Question: The attached .csv file shows precipitation amounts, in inches, for the five boroughs of New York City in a certain year. How many inches of precipitation did the city receive in total for that year? Don’t use commas if the number has four or more digits. ; file_name: /home/nitin/.cache/huggingface/hub/datasets--gaia-benchmark--GAIA/snapshots/682dd723ee1e1697e00360edccf2366dc8418dd9/2023/test/52e8ce1c-09bd-4537-8e2d-67d1648779b9.csv """ #"What is the square of the population of France in millions?", #"What is 50 plus 75?" ] # Process all questions output_file = process_questions(questions_list=questions) # Print the results print("\nResults from file:") with open(output_file, 'r') as f: for line in f: result = json.loads(line) print(f"\nTask ID: {result['task_id']}") print(f"Answer: {result['model_answer']}") print(f"Reasoning:\n{result['reasoning_trace']}")