Update app.py

app.py

@@ -5,83 +5,111 @@ 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
+from langchain_community.vectorstores import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from typing import TypedDict, Annotated
+import os
+import re
+
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# --- Optimized Agent Implementation ---
-MODELS = [
-    "Qwen/Qwen2-0.5B-Instruct",
-    "google/flan-t5-xxl",
-    "mistralai/Mistral-7B-Instruct-v0.2"
-]
+# Configuration
+MODEL_ID = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"  # 4.2GB quantized
+EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
+FALLBACK_MODELS = ["google/flan-t5-base", "mistralai/Mistral-7B-Instruct-v0.2"]
 
-clients = [InferenceClient(model=model, token=os.environ["HF_TOKEN"]) for model in MODELS]
+class AgentState(TypedDict):
+    question: str
+    context: str
+    answer: str
+    attempts: Annotated[int, lambda x, y: x + 1]
 
-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:
-    """Generate answer with error handling"""
-    try:
-        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,
-            stop_sequences=["<|im_end|>"]
+class BasicAgent:
+    def __init__(self):
+        # Initialize components
+        self.client = InferenceClient(
+            model=MODEL_ID,
+            token=os.environ["HF_TOKEN"],
+            timeout=120
         )
-        # 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
+        
+        # Initialize vector store (add your documents here)
+        self.vectorstore = Chroma.from_texts(
+            texts=["GAIA knowledge content..."],  # Replace with your documents
+            embedding=EMBEDDING_MODEL,
+            persist_directory="./chroma_db"
+        )
+        
+        # Build LangGraph workflow
+        self.workflow = self._build_graph()
 
-def should_continue(state: dict) -> str:
-    """Conditional edge function (not a node)"""
-    return END if state["answer"] else "route_model"
+    def _build_graph(self):
+        # Define nodes
+        def retrieve(state: AgentState):
+            docs = self.vectorstore.similarity_search(state["question"], k=3)
+            state["context"] = "\n".join([d.page_content for d in docs])
+            return state
 
-# Build workflow
-workflow = StateGraph(dict)
-workflow.add_node("route_model", model_router)
-workflow.add_node("query", query_model)
+        def generate(state: AgentState):
+            try:
+                response = self.client.text_generation(
+                    f"""<s>[INST]Answer using ONLY this context:
+                    {state['context']}
+                    Question: {state['question']}
+                    Answer: [/INST]""",
+                    temperature=0.1,
+                    max_new_tokens=100,
+                    stop_sequences=["</s>"]
+                )
+                state["answer"] = response.split("[/INST]")[-1].strip()
+            except Exception:
+                state["answer"] = ""
+            return state
 
-workflow.set_entry_point("route_model")
-workflow.add_edge("route_model", "query")
-workflow.add_conditional_edges(
-    "query",
-    should_continue,
-    {END: END, "route_model": "route_model"}
-)
+        def validate(state: AgentState):
+            if len(state["answer"]) > 5 and state["attempts"] < 3:
+                return "final"
+            return "retry"
 
-compiled_agent = workflow.compile()
+        # Build workflow
+        workflow = StateGraph(AgentState)
+        workflow.add_node("retrieve", retrieve)
+        workflow.add_node("generate", generate)
+        workflow.add_node("validate", validate)
+
+        workflow.set_entry_point("retrieve")
+        workflow.add_edge("retrieve", "generate")
+        workflow.add_edge("generate", "validate")
+        
+        workflow.add_conditional_edges(
+            "validate",
+            lambda x: "retry" if x["answer"] == "" else "final",
+            {
+                "retry": "retrieve",
+                "final": END
+            }
+        )
+
+        return workflow.compile()
 
-class BasicAgent:
     def __call__(self, question: str) -> str:
+        # GAIA-compliant formatting
         state = {
             "question": question,
-            "retries": 0,
-            "current_model": 0,
-            "answer": ""
+            "context": "",
+            "answer": "",
+            "attempts": 0
         }
         
         for _ in range(3):  # Max 3 attempts
-            state = compiled_agent.invoke(state)
+            state = self.workflow.invoke(state)
             if state["answer"]:
-                return state["answer"]
-            time.sleep(1)
-            
-        return ""
+                answer = re.sub(r'[^a-zA-Z0-9]', '', state["answer"]).lower()
+                return answer[:100]  # GAIA length constraint
+        
+        return ""  # Preserve scoring eligibility
         
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """