Update app.py

app.py

@@ -247,26 +247,20 @@
 # =================================================================================================
 
 # 
-
 import os
 import io
-import json
 import requests
 import pandas as pd
 import gradio as gr
 from contextlib import redirect_stdout
-from typing import TypedDict, Annotated, List
-import operator
+from typing import List
 
 # --- LangChain & LangGraph Imports ---
-from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from langchain_core.messages import BaseMessage
 from langchain_core.tools import tool
-# <<<--- CHANGE: Import ChatCohere and Cohere's Tool annd ToolCall classes--->>>
-from langchain_cohere import ChatCohere
-from langchain_cohere.cohere_agent import CohereAgent
-from langchain.agents import AgentExecutor
-from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import ToolNode
+from langchain_cohere.chat_models import ChatCohere
+from langchain.agents import AgentExecutor, create_cohere_tools_agent
+from langchain_core.prompts import ChatPromptTemplate
 from tavily import TavilyClient
 import pypdf
 
@@ -276,7 +270,6 @@ FILES_DIR = "./files"
 os.makedirs(FILES_DIR, exist_ok=True)
 
 # --- System Prompt (Unchanged) ---
-# This prompt is excellent and requires no changes.
 AGENT_SYSTEM_PROMPT = """You are a world-class AI agent, specialized in solving complex problems from the GAIA benchmark.
 Your task is to analyze the user's question, think step-by-step, and use the provided tools to find the correct answer.
 CRITICAL INSTRUCTIONS:
@@ -341,25 +334,27 @@ def python_interpreter(code: str) -> str:
 
 #
 # ================================================================================================
-#  ✅ 2. CONFIGURE AND BUILD THE AGENT (Using ChatCohere)
+#  ✅ 2. CONFIGURE AND BUILD THE AGENT (Standard, Robust Method)
 # ================================================================================================
 #
-class AgentState(TypedDict):
-    input: str
-    chat_history: List[BaseMessage]
-    agent_outcome: dict | None
-
 def build_agent_graph():
-    """Builds the agent using the direct ChatCohere integration."""
+    """Builds the agent using the standard create_cohere_tools_agent function."""
     tools = [tavily_search, read_file, python_interpreter]
 
-    # <<<--- CHANGE: Instantiate ChatCohere directly --->>>
-    # It will use the COHERE_API_KEY from your secrets.
-    # We use command-r-plus, Cohere's most powerful model.
-    llm = ChatCohere(model="command-r-plus", temperature=0)
+    # Instantiate the ChatCohere model
+    llm = ChatCohere(model="command-r-plus", temperature=0, cohere_api_key=os.getenv("COHERE_API_KEY"))
+
+    # Create the prompt template. The 'agent_scratchpad' is crucial.
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", AGENT_SYSTEM_PROMPT),
+        ("human", "{input}"),
+        ("placeholder", "{agent_scratchpad}"),
+    ])
 
-    # <<<--- This is much simpler now, as ChatCohere has built-in agent capabilities --->>>
-    agent = CohereAgent(llm=llm, tools=tools)
+    # Use the standard agent creation function
+    agent = create_cohere_tools_agent(llm, tools, prompt)
+    
+    # Create the executor to run the agent loop
     agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
     return agent_executor
@@ -371,16 +366,16 @@ def build_agent_graph():
 #
 class GaiaAgent:
     def __init__(self):
-        print("GaiaAgent initialized. Building agent with direct ChatCohere integration...")
+        print("GaiaAgent initialized. Building agent with create_cohere_tools_agent...")
         self.agent_app = build_agent_graph()
 
     def __call__(self, question: str) -> str:
         print(f"\n{'='*60}\nAgent received question: {question[:100]}...\n{'='*60}")
         try:
-            # The Cohere agent executor expects 'input' and a 'preamble' for the system message.
+            # The standard agent executor expects 'input' and 'chat_history'.
             response = self.agent_app.invoke({
                 "input": question,
-                "preamble": AGENT_SYSTEM_PROMPT
+                "chat_history": [] # Each question is independent, so history is empty.
             })
             final_answer = str(response.get("output", "")).strip()
             print(f"\n--- Agent finished. Final Answer: {final_answer} ---\n")
@@ -389,7 +384,7 @@ class GaiaAgent:
             print(f"An error occurred during agent execution: {e}")
             return f"AGENT_EXECUTION_ERROR: {e}"
 
-# --- The rest of the file is mostly the same ---
+# --- The rest of the file is unchanged ---
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")
     if not profile: return "Please Login to Hugging Face with the button.", None
@@ -441,9 +436,9 @@ with gr.Blocks() as demo:
     gr.Markdown("# GAIA Agent Final Assessment (Direct Cohere Integration)")
     gr.Markdown(
         """
-        **Instructor's Note:** We are now using the direct `langchain-cohere` integration. This is the most reliable way to use the Command R+ model.
+        **Instructor's Note:** This version uses the standard `create_cohere_tools_agent` function, which is the most robust way to create a tool-calling agent with Cohere.
         1.  Ensure you have a **`COHERE_API_KEY`** and a **`TAVILY_API_KEY`** set in your Space secrets.
-        2.  Ensure your `requirements.txt` includes `langchain-cohere`.
+        2.  Ensure your `requirements.txt` includes `langchain-cohere` and `langchain`.
         """
     )
     gr.LoginButton()
@@ -454,5 +449,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     print("\n" + "-"*30 + " App Starting " + "-"*30)
-    # Disable experimental SSR to prevent startup crashes
     demo.launch(debug=True, share=False, ssr_mode=False)