Update app.py

app.py

@@ -1,44 +1,31 @@
 import os
 import gradio as gr
 import requests
-import inspect
+import re  # Added missing import
 import pandas as pd
 from langgraph.graph import StateGraph, END
 from huggingface_hub import InferenceClient
+import time  # Added missing import
 
-
-# (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# --- Basic Agent Definition ---
-# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-# --- GAIA-Optimized Agent Implementation ---
-# Configure fallback models
+# --- Optimized Agent Implementation ---
 MODELS = [
     "Qwen/Qwen2-0.5B-Instruct",
     "google/flan-t5-xxl",
     "mistralai/Mistral-7B-Instruct-v0.2"
 ]
 
-# Initialize clients with automatic retry
 clients = [InferenceClient(model=model, token=os.environ["HF_TOKEN"]) for model in MODELS]
 
-# Define state structure using dictionary
-initial_state = {
-    "question": "",
-    "retries": 0,
-    "current_model": 0,
-    "answer": ""
-}
-
 def model_router(state: dict) -> dict:
     """Rotate through available models"""
     state["current_model"] = (state["current_model"] + 1) % len(MODELS)
     return state
 
 def query_model(state: dict) -> dict:
-    """Attempt to get answer from current model"""
+    """Generate answer with error handling"""
     try:
         response = clients[state["current_model"]].text_generation(
             prompt=f"""<|im_start|>system
@@ -51,48 +38,50 @@ Answer with ONLY the exact value requested.<|im_end|>
             max_new_tokens=50,
             stop_sequences=["<|im_end|>"]
         )
-        state["answer"] = response.split("<|im_start|>assistant")[-1].split("<|im_end|>")[0].strip()
-        state["answer"] = re.sub(r'[^a-zA-Z0-9]', '', state["answer"]).lower()
+        # Fixed answer extraction
+        answer_part = response.split("<|im_start|>assistant")[-1]
+        answer = answer_part.split("<|im_end|>")[0].strip()
+        state["answer"] = re.sub(r'[^a-zA-Z0-9]', '', answer).lower()
     except Exception as e:
         print(f"Model error: {str(e)}")
         state["answer"] = ""
     return state
 
-def validate_answer(state: dict) -> str:
-    """Check if we have a valid answer"""
-    return "final_answer" if state["answer"] else "retry"
+def should_continue(state: dict) -> str:
+    """Conditional edge function (not a node)"""
+    return END if state["answer"] else "route_model"
 
 # Build workflow
 workflow = StateGraph(dict)
 workflow.add_node("route_model", model_router)
 workflow.add_node("query", query_model)
-workflow.add_node("validate", validate_answer)
 
+workflow.set_entry_point("route_model")
 workflow.add_edge("route_model", "query")
-workflow.add_edge("query", "validate")
-
 workflow.add_conditional_edges(
-    "validate",
-    lambda x: "final_answer" if x["answer"] else "retry",
-    {"final_answer": END, "retry": "route_model"}
+    "query",
+    should_continue,
+    {END: END, "route_model": "route_model"}
 )
 
-workflow.set_entry_point("route_model")
 compiled_agent = workflow.compile()
 
-# GAIA Interface
 class BasicAgent:
     def __call__(self, question: str) -> str:
-        state = initial_state.copy()
-        state["question"] = question
+        state = {
+            "question": question,
+            "retries": 0,
+            "current_model": 0,
+            "answer": ""
+        }
         
         for _ in range(3):  # Max 3 attempts
             state = compiled_agent.invoke(state)
             if state["answer"]:
                 return state["answer"]
-            time.sleep(1)  # Backoff
+            time.sleep(1)
             
-        return ""  # Return empty to preserve scoring
+        return ""
         
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """