Update gaia_agent.py

gaia_agent.py

@@ -14,6 +14,7 @@ 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 (
@@ -26,10 +27,14 @@ from customtools import (
     subtraction_tool,
     multiplication_tool,
     transcribe_audio,
+    modulus_tool,
 )
 from config import (
     OPENROUTER_API_KEY,
     LLM_MODEL,
+    NVIDIA,
+    NVIDIA_API_KEY,
+    NVIDIA_MODEL,
     LLM_TEMPERATURE,
     OUTPUT_FILE,
     FINAL_ANSWER_MAX_LENGTH,
@@ -48,12 +53,24 @@ 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,
-        )
+        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}")
@@ -71,6 +88,7 @@ TOOLS = {
     "extract_text_from_image": extract_text_from_image,
     "wikisearch": wikisearch,
     "transcribe_audio": transcribe_audio,
+    "modulus_tool": modulus_tool,
 }
 
 
@@ -101,6 +119,7 @@ class Step(BaseModel):
         "addition_tool",
         "subtraction_tool",
         "multiplication_tool",
+        "modulus_tool",
         "none",
     ]
     tool_input: str
@@ -131,6 +150,7 @@ def planner_node(state: AgentState):
 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}")
@@ -280,10 +300,6 @@ def format_reasoning_trace(intermediate_results: List[Dict[str, Any]]) -> str:
     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
@@ -367,7 +383,26 @@ def process_questions(questions_file: str = None, questions_list: List[str] = No
 
 
 
-#if __name__ == "__main__":
+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
+"""
 
-#    global llm
+        #"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']}")