Update app.py

app.py

@@ -4,6 +4,7 @@ import requests
 import inspect
 import pandas as pd
 from huggingface_hub import HfApi, InferenceClient, login
+from langgraph.graph import StateGraph, END
 
 
 # (Keep Constants as is)
@@ -15,52 +16,85 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- GAIA-Optimized Agent Implementation ---
 from huggingface_hub import InferenceClient, login
 
+# Configure models
+MODELS = [
+    "Qwen/Qwen2-0.5B-Instruct",
+    "google/flan-t5-xxl",
+    "mistralai/Mistral-7B-Instruct-v0.2"
+]
 
-class BasicAgent:
+class AgentState:
     def __init__(self):
-        login(token=os.environ["HF_TOKEN"])  # Required authentication
-        
-        # Primary model (7B quantized)
-        self.client = InferenceClient(
-            model="Qwen/Qwen2-0.5B-Instruct",  # 1.2GB, free-tier compatible
-            token=os.environ["HF_TOKEN"],
-            timeout=60
-        )
-        
-        # Verify model access
-        test_response = self._call_model("2+2=")
-        if "4" not in test_response:
-            raise RuntimeError("Model initialization failed")
-
-    def _call_model(self, question: str) -> str:
-        """Optimized prompt engineering for GAIA"""
-        prompt = f"""<|im_start|>system
-Answer with ONLY the exact value requested. No explanations.<|im_end|>
-<|im_start|>user
-{question}<|im_end|>
-<|im_start|>assistant
-"""
-        return self.client.text_generation(
-            prompt=prompt,
-            temperature=0.01,
+        self.question = ""
+        self.retries = 0
+        self.current_model = 0
+        self.answer = ""
+
+# Initialize clients
+clients = [InferenceClient(model=model, token=os.environ["HF_TOKEN"]) for model in MODELS]
+
+def model_router(state):
+    """Rotate through available models"""
+    state.current_model = (state.current_model + 1) % len(MODELS)
+    return state
+
+def query_model(state):
+    """Attempt to get answer from current model"""
+    try:
+        response = clients[state.current_model].text_generation(
+            prompt=f"GAIA Question: {state.question}\nAnswer:",
             max_new_tokens=50,
-            stop_sequences=["<|im_end|>"]
+            temperature=0.01
         )
+        state.answer = response.split("Answer:")[-1].strip()
+        return state
+    except Exception as e:
+        print(f"Error with {MODELS[state.current_model]}: {str(e)}")
+        state.retries += 1
+        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"
+
+# Build workflow
+workflow = StateGraph(AgentState)
 
+workflow.add_node("route_model", model_router)
+workflow.add_node("query", query_model)
+workflow.add_node("validate", validate_answer)
+
+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"
+    }
+)
+
+workflow.set_entry_point("route_model")
+agent = workflow.compile()
+
+class BasicAgent:
     def __call__(self, question: str) -> str:
-        try:
-            raw_response = self._call_model(question)
-            
-            # Robust answer extraction
-            answer = raw_response.split("<|im_start|>assistant")[-1]
-            answer = answer.split("<|im_end|>")[0].strip()
+        state = AgentState()
+        state.question = question
+        
+        for _ in range(3):  # Max 3 attempts
+            state = agent.invoke(state)
+            if state.answer:
+                return state.answer
+            time.sleep(1)  # Backoff
             
-            # GAIA-compliant normalization
-            return re.sub(r'[^a-zA-Z0-9]', '', answer).lower()
-        except Exception as e:
-            print(f"Error: {str(e)}")
-            return ""
+        return ""  # Return empty to preserve scoring
 
+        
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,