Update app.py

app.py

@@ -3,8 +3,8 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from huggingface_hub import HfApi, InferenceClient, login
 from langgraph.graph import StateGraph, END
+from huggingface_hub import InferenceClient
 
 
 # (Keep Constants as is)
@@ -14,54 +14,56 @@ 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 ---
-from huggingface_hub import InferenceClient, login
-
-# Configure models
+# Configure fallback models
 MODELS = [
     "Qwen/Qwen2-0.5B-Instruct",
     "google/flan-t5-xxl",
     "mistralai/Mistral-7B-Instruct-v0.2"
 ]
 
-class AgentState:
-    def __init__(self):
-        self.question = ""
-        self.retries = 0
-        self.current_model = 0
-        self.answer = ""
-
-# Initialize clients
+# Initialize clients with automatic retry
 clients = [InferenceClient(model=model, token=os.environ["HF_TOKEN"]) for model in MODELS]
 
-def model_router(state):
+# 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)
+    state["current_model"] = (state["current_model"] + 1) % len(MODELS)
     return state
 
-def query_model(state):
+def query_model(state: dict) -> dict:
     """Attempt to get answer from current model"""
     try:
-        response = clients[state.current_model].text_generation(
-            prompt=f"GAIA Question: {state.question}\nAnswer:",
+        response = clients[state["current_model"]].text_generation(
+            prompt=f"""<|im_start|>system
+Answer with ONLY the exact value requested.<|im_end|>
+<|im_start|>user
+{state['question']}<|im_end|>
+<|im_start|>assistant
+""",
+            temperature=0.01,
             max_new_tokens=50,
-            temperature=0.01
+            stop_sequences=["<|im_end|>"]
         )
-        state.answer = response.split("Answer:")[-1].strip()
-        return state
+        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()
     except Exception as e:
-        print(f"Error with {MODELS[state.current_model]}: {str(e)}")
-        state.retries += 1
-        return state
+        print(f"Model error: {str(e)}")
+        state["answer"] = ""
+    return state
 
-def validate_answer(state):
-    """Basic GAIA answer validation"""
-    if len(state.answer) > 0 and 2 <= len(state.answer) <= 100:
-        return "final_answer"
-    return "retry"
+def validate_answer(state: dict) -> str:
+    """Check if we have a valid answer"""
+    return "final_answer" if state["answer"] else "retry"
 
 # Build workflow
-workflow = StateGraph(AgentState)
-
+workflow = StateGraph(dict)
 workflow.add_node("route_model", model_router)
 workflow.add_node("query", query_model)
 workflow.add_node("validate", validate_answer)
@@ -71,29 +73,26 @@ 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"
-    }
+    lambda x: "final_answer" if x["answer"] else "retry",
+    {"final_answer": END, "retry": "route_model"}
 )
 
 workflow.set_entry_point("route_model")
-agent = workflow.compile()
+compiled_agent = workflow.compile()
 
+# GAIA Interface
 class BasicAgent:
     def __call__(self, question: str) -> str:
-        state = AgentState()
-        state.question = question
+        state = initial_state.copy()
+        state["question"] = question
         
         for _ in range(3):  # Max 3 attempts
-            state = agent.invoke(state)
-            if state.answer:
-                return state.answer
+            state = compiled_agent.invoke(state)
+            if state["answer"]:
+                return state["answer"]
             time.sleep(1)  # Backoff
             
         return ""  # Return empty to preserve scoring
-
         
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """