Update app.py

app.py

@@ -259,10 +259,14 @@ from typing import TypedDict, Annotated, List
 import operator
 
 # --- LangChain & LangGraph Imports ---
-from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage, AIMessage, SystemMessage
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
 from langchain_core.tools import tool
-from langchain_huggingface import HuggingFaceEndpoint
+# <<<--- 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 tavily import TavilyClient
 import pypdf
 
@@ -271,32 +275,22 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 FILES_DIR = "./files"
 os.makedirs(FILES_DIR, exist_ok=True)
 
-# --- System Prompt (Updated for Manual Tool Calling) ---
+# --- 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.
-
-**TOOL USAGE INSTRUCTIONS:**
-When you need to use a tool, you MUST respond with a JSON object containing the tool name and its arguments. The JSON object should have two keys: "tool_name" and "parameters".
-
-Here is an example of how to call the `tavily_search` tool:
-```json
-{
-  "tool_name": "tavily_search",
-  "parameters": {
-    "query": "What was the score of the 2023 FIFA Women's World Cup final?"
-  }
-}```
-
-**CRITICAL FINAL ANSWER INSTRUCTIONS:**
-Once you have gathered all the necessary information and are absolutely certain of the answer, you MUST provide it directly and concisely.
-- Your final response must ONLY be the answer itself.
-- DO NOT wrap the final answer in a JSON object or include any conversational text like 'The answer is...'.
-
+CRITICAL INSTRUCTIONS:
+1.  **Analyze the Goal:** First, understand what the user is asking for.
+2.  **Plan & Execute:** Formulate a plan and use the available tools (`tavily_search`, `read_file`, `python_interpreter`) to gather information.
+3.  **Final Answer Format:** Once you are absolutely certain of the answer, you MUST provide it directly and concisely.
+    - DO NOT include your reasoning, thoughts, or any conversational text like 'The answer is...', 'Here is the result:', or 'Based on my search...'.
+    - Your final response must ONLY be the answer itself.
 EXAMPLES OF CORRECT FINAL ANSWERS:
-- `2023`
-- `John Doe`
-- `42`
-- `broccoli, celery, lettuce, sweet potatoes`
+- If the question asks for a year: `2023`
+- If it asks for a name: `John Doe`
+- If it asks for a number: `42`
+- If it asks for a comma-separated list: `item1, item2, item3`
+Think, use your tools, and then provide ONLY the final, precise answer.
 """
 
 #
@@ -347,73 +341,28 @@ def python_interpreter(code: str) -> str:
 
 #
 # ================================================================================================
-#  ✅ 2. CONFIGURE AND BUILD THE AGENT GRAPH
+#  ✅ 2. CONFIGURE AND BUILD THE AGENT (Using ChatCohere)
 # ================================================================================================
 #
 class AgentState(TypedDict):
-    messages: Annotated[List[BaseMessage], operator.add]
+    input: str
+    chat_history: List[BaseMessage]
+    agent_outcome: dict | None
 
 def build_agent_graph():
-    """Builds the LangGraph agent."""
+    """Builds the agent using the direct ChatCohere integration."""
     tools = [tavily_search, read_file, python_interpreter]
-    tool_map = {tool.name: tool for tool in tools}
-
-    repo_id = "CohereForAI/c4ai-command-r-plus"
-    
-    # <<<--- CHANGE 1: Explicitly set `task="conversational"` --->>>
-    # This is the crucial fix. We are telling the endpoint to use the correct API pipeline.
-    llm = HuggingFaceEndpoint(
-        repo_id=repo_id,
-        task="conversational", # This is the key fix!
-        max_new_tokens=1024,
-        temperature=0.1,
-        huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
-    )
 
-    def call_model(state: AgentState):
-        """Invokes the LLM using the conversational task."""
-        # <<<--- CHANGE 2: The conversational task takes a list of messages directly --->>>
-        # This is cleaner and the correct way to use this pipeline.
-        response = llm.invoke(state['messages'])
-        return {"messages": [response]}
-
-    def should_continue(state: AgentState) -> str:
-        """Determines whether to call a tool or end the loop."""
-        last_message_content = state['messages'][-1].content.strip()
-        if last_message_content.startswith('{') and last_message_content.endswith('}'):
-            try:
-                json.loads(last_message_content)
-                return "action"
-            except json.JSONDecodeError:
-                return "end"
-        else:
-            return "end"
-
-    def call_tool_node(state: AgentState):
-        """Parses the tool call from the LLM output and executes it."""
-        last_message_content = state['messages'][-1].content.strip()
-        try:
-            tool_call_data = json.loads(last_message_content)
-            tool_name = tool_call_data.get("tool_name")
-            parameters = tool_call_data.get("parameters", {})
+    # <<<--- 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)
 
-            if tool_name not in tool_map:
-                return {"messages": [ToolMessage(content=f"Error: Tool '{tool_name}' not found.", tool_call_id="error")]}
+    # <<<--- This is much simpler now, as ChatCohere has built-in agent capabilities --->>>
+    agent = CohereAgent(llm=llm, tools=tools)
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-            selected_tool = tool_map[tool_name]
-            tool_output = selected_tool.invoke(parameters)
-            return {"messages": [ToolMessage(content=str(tool_output), tool_call_id=tool_name)]}
-
-        except Exception as e:
-            return {"messages": [ToolMessage(content=f"Error processing tool call: {e}. Content: '{last_message_content}'", tool_call_id="error")]}
-
-    workflow = StateGraph(AgentState)
-    workflow.add_node("agent", call_model)
-    workflow.add_node("action", call_tool_node)
-    workflow.set_entry_point("agent")
-    workflow.add_conditional_edges("agent", should_continue, {"action": "action", "end": END})
-    workflow.add_edge('action', 'agent')
-    return workflow.compile()
+    return agent_executor
 
 #
 # ================================================================================================
@@ -422,23 +371,25 @@ def build_agent_graph():
 #
 class GaiaAgent:
     def __init__(self):
-        print("GaiaAgent initialized. Building Command R+ agent with 'conversational' task...")
+        print("GaiaAgent initialized. Building agent with direct ChatCohere integration...")
         self.agent_app = build_agent_graph()
 
     def __call__(self, question: str) -> str:
         print(f"\n{'='*60}\nAgent received question: {question[:100]}...\n{'='*60}")
-        initial_input = {"messages": [SystemMessage(content=AGENT_SYSTEM_PROMPT), HumanMessage(content=question)]}
-        final_state = None
-        for i, step in enumerate(self.agent_app.stream(initial_input, {"recursion_limit": 15})):
-            if i == 0: print("--- Starting Agentic Loop ---")
-            final_state = list(step.values())[0]
-
-        final_answer_message = final_state['messages'][-1]
-        final_answer = str(final_answer_message.content).strip()
-        print(f"\n--- Agent finished. Final Answer: {final_answer} ---\n")
-        return final_answer
-
-# --- The rest of the file (run_and_submit_all, Gradio UI) remains the same ---
+        try:
+            # The Cohere agent executor expects 'input' and a 'preamble' for the system message.
+            response = self.agent_app.invoke({
+                "input": question,
+                "preamble": AGENT_SYSTEM_PROMPT
+            })
+            final_answer = str(response.get("output", "")).strip()
+            print(f"\n--- Agent finished. Final Answer: {final_answer} ---\n")
+            return final_answer
+        except Exception as e:
+            print(f"An error occurred during agent execution: {e}")
+            return f"AGENT_EXECUTION_ERROR: {e}"
+
+# --- The rest of the file is mostly the same ---
 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
@@ -453,16 +404,13 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         response = requests.get(questions_url, timeout=15)
         response.raise_for_status()
         questions_data = response.json()
-    except Exception as e:
-        return f"An unexpected error occurred fetching questions: {e}", None
+    except Exception as e: return f"An unexpected error occurred fetching questions: {e}", None
 
-    results_log = []
-    answers_payload = []
+    results_log, answers_payload = [], []
     agent_instance = GaiaAgent()
 
     for item in questions_data:
-        task_id = item.get("task_id")
-        question_text = item.get("question")
+        task_id, question_text = item.get("task_id"), item.get("question")
         if not task_id or question_text is None: continue
         try:
             submitted_answer = agent_instance(question_text)
@@ -472,10 +420,9 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
              print(f"Error running agent on task {task_id}: {e}")
              results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
 
-    if not answers_payload:
-        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
-
+    if not answers_payload: return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
     submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
+
     try:
         response = requests.post(submit_url, json=submission_data, timeout=90)
         response.raise_for_status()
@@ -488,15 +435,15 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             f"Message: {result_data.get('message', 'No message received.')}"
         )
         return final_status, pd.DataFrame(results_log)
-    except Exception as e:
-        return f"An unexpected error occurred during submission: {e}", pd.DataFrame(results_log)
+    except Exception as e: return f"An unexpected error occurred during submission: {e}", pd.DataFrame(results_log)
 
 with gr.Blocks() as demo:
-    gr.Markdown("# GAIA Agent Final Assessment (Open Source: Command R+ - Corrected Task)")
+    gr.Markdown("# GAIA Agent Final Assessment (Direct Cohere Integration)")
     gr.Markdown(
         """
-        **Instructor's Note:** This version corrects the `HuggingFaceEndpoint` invocation by specifying `task="conversational"`.
-        This is the final key required to make the Command R+ model work correctly with the Hugging Face Inference API for our agent.
+        **Instructor's Note:** We are now using the direct `langchain-cohere` integration. This is the most reliable way to use the Command R+ model.
+        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`.
         """
     )
     gr.LoginButton()
@@ -507,5 +454,5 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     print("\n" + "-"*30 + " App Starting " + "-"*30)
-    demo.launch(debug=True, share=False, ssr_mode=False)
-    #demo.launch(debug=True, share=False)
+    # Disable experimental SSR to prevent startup crashes
+    demo.launch(debug=True, share=False, ssr_mode=False)