Update gaia_agent.py

gaia_agent.py

@@ -15,8 +15,14 @@ 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,
@@ -29,7 +35,6 @@ from customtools import (
     transcribe_audio,
     modulus_tool,
     power_tool,
-    
 )
 from config import (
     OPENROUTER_API_KEY,
@@ -48,31 +53,39 @@ from prompts import (
 )
 
 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
+        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=1,
-                                top_p=1,
+             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,
-            )
+             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}")
@@ -91,7 +104,7 @@ TOOLS = {
     "wikisearch": wikisearch,
     "transcribe_audio": transcribe_audio,
     "modulus_tool": modulus_tool,
-    "power_tool": power_tool,
+    "power_tool":power_tool,
 }
 
 
@@ -132,18 +145,29 @@ class Step(BaseModel):
 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'])
-    
-    planner_llm = llm.with_structured_output(Plan, method="json_schema")
-    response = planner_llm.invoke(prompt)
+
+    response = get_planner_llm().invoke(prompt)
     
     print(f"Planner generated {len(response.steps)} steps")
     
     return {
         **state,
-        "plan": [step.dict() for step in response.steps],
+        "plan": [step.model_dump() for step in response.steps],
         "current_step": 0,
         "intermediate_results": [],
         "done": False,
@@ -234,9 +258,9 @@ def should_continue(state: AgentState):
 
 def finalizer_node(state: AgentState):
     """Finalizer node: summarizes results and generates final answer."""
-    # Format intermediate results for the finalizer
+    # Compact context to reduce token usage sent to finalizer.
     results_text = "\n".join([
-        f"Step {i+1}: {r['step'].get('description', '')}\n  Output: {str(r['output'])[:100]}..."
+        f"S{i+1}: {r['step'].get('description', '')} | O: {str(r['output'])[:80]}"
         for i, r in enumerate(state["intermediate_results"])
     ])
     
@@ -388,12 +412,12 @@ def process_questions(questions_file: str = None, questions_list: List[str] = No
 
 
 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
-"""
+        """
+          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?"