""" 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 langchain_community.cache import SQLiteCache from typing import TypedDict import langchain_core #from langchain_ollama import ChatOllama CACHE_DB_PATH = ".langchain_llm_cache.db" #langchain_core.globals.set_debug(True) langchain_core.globals.set_llm_cache(SQLiteCache(database_path=CACHE_DB_PATH)) 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() print(f"LangChain LLM cache enabled: {CACHE_DB_PATH}") memory = MemorySaver() llm = None planner_llm = None def connect_models(): """Initialize and return the LLM instance.""" try: global llm, planner_llm # llm = ChatOllama( # model="gemma4:e2b", # base_url="http://localhost:11434/", # temperature=0, # ) # Reset derived clients whenever base model is reconnected. planner_llm = None if NVIDIA: llm = ChatNVIDIA( model=NVIDIA_MODEL, api_key= NVIDIA_API_KEY, temperature=0.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 get_planner_llm(): """Create structured planner client once and reuse it across questions.""" global planner_llm, llm if llm is None: llm = connect_models() if planner_llm is None: planner_llm = llm.with_structured_output(Plan, method="json_schema") return planner_llm def planner_node(state: AgentState): """Planner node: breaks down question into steps.""" prompt = PLANNER_PROMPT_TEMPLATE.format(question=state['question']) response = get_planner_llm().invoke(prompt) print(f"Planner generated {len(response.steps)} steps") return { **state, "plan": [step.model_dump() 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.""" # Compact context to reduce token usage sent to finalizer. results_text = "\n".join([ f"S{i+1}: {r['step'].get('description', '')} | O: {str(r['output'])[:80]}" 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__": # Example questions to process questions = [ """ Where were the Vietnamese specimens described by Kuznetzov in Nedoshivina's 2010 paper eventually deposited? Just give me the city name without abbreviations. """ #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']}")