Update app.py

app.py

@@ -8,6 +8,7 @@ import inspect
 import pandas as pd
 import tools
 from smolagents import CodeAgent
+# Resolve the correct LLM model class across smolagents versions
 try:
     from smolagents import InferenceClientModel as _HFModel  # smolagents >= 1.0
 except ImportError:
@@ -15,10 +16,13 @@ except ImportError:
         from smolagents.models import HfApiModel as _HFModel
     except ImportError:
         from smolagents import HfApiModel as _HFModel
-
 from typing import TypedDict, List, Dict, Any, Optional
 from langgraph.graph import StateGraph, START, END
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage  # kept for LangGraph compatibility
+
+# Helper to build a smolagents-compatible message list
+def _msg(content: str) -> list:
+    return [{"role": "user", "content": content}]
 
 
 # (Keep Constants as is)
@@ -29,19 +33,21 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 def _build_hf_model(model_name: str):
     """Build a text model across smolagents versions."""
     for kwargs in (
+        {"model_id": model_name, "max_tokens": 2048, "temperature": 0.3},
         {"model_id": model_name, "max_new_tokens": 2048, "temperature": 0.3},
+        {"repo_id": model_name, "max_tokens": 2048, "temperature": 0.3},
         {"repo_id": model_name, "max_new_tokens": 2048, "temperature": 0.3},
     ):
         try:
             return _HFModel(**kwargs)
-        except (TypeError, Exception):
+        except TypeError:
             continue
     raise RuntimeError(f"Cannot instantiate model {model_name} with available smolagents version")
 
 
 # Text/math models via smolagents
-model = _build_hf_model("Qwen3.5-35B-A3B")
-math_model = _build_hf_model("Qwen/Qwen2.5-Math-1.5B")
+model = _build_hf_model("meta-llama/Llama-3.2-3B-Instruct")  # General model for classification and final answer synthesis
+math_model = _build_hf_model("deepseek-ai/deepseek-math-7b-instruct")
 
 # FireRed OCR (Transformers) loaded lazily to avoid startup crashes
 _fire_red_model = None
@@ -82,8 +88,6 @@ def _extract_text_from_response(response: Any) -> str:
         return str(content)
     return str(response)
 
-
-    
 #define the state
 class AgentState(TypedDict):
     question: str
@@ -94,6 +98,7 @@ class AgentState(TypedDict):
     is_math: Optional[bool]
     have_image: Optional[bool]
     final_answer: Optional[str]  # The final answer produced by the agent
+    retry_count: Optional[int]   # Number of retries so far
     messages: List[Dict[str, Any]]  # Track conversation with LLM for analysis
 #define nodes
 
@@ -119,8 +124,9 @@ def classify(state: AgentState) -> str:
                 "have_image": false
             }}
             """ 
-    messages = [HumanMessage(content=prompt)]
-    response = model.invoke(messages)
+    messages = _msg(prompt)
+    response = model(messages)
+    raw = _extract_text_from_response(response)
     # Parse JSON from the model's response
     import json, re
     match = re.search(r'\{.*?\}', raw, re.DOTALL)
@@ -134,7 +140,6 @@ def classify(state: AgentState) -> str:
     have_file = bool(data.get("have_file", False))
     is_math = bool(data.get("is_math", False))
     have_image = bool(data.get("have_image", False))
-    
     print(f"Classification result: is_searching={is_searching}, have_file={have_file}, is_math={is_math}, have_image={have_image}")
     mew_messages = state.get("messages", []) + [
         {"role": "system", "content": "Classify the question to determine which tools to use."},
@@ -178,7 +183,7 @@ def handle_image(state: AgentState) -> str:
 
     # Use ImageReaderTool to download the image as base64
     image_reader = tools.ImageReaderTool()
-    image_data_uri = image_reader(task_id) if task_id and file_name else ""
+    image_data_uri = image_reader(task_id, file_name) if task_id and file_name else ""
 
     if not image_data_uri or image_data_uri.startswith("Failed"):
         print(f"Could not download image for task {task_id}")
@@ -203,8 +208,7 @@ Return a JSON object with the following fields:
     "transcribed_text": "All text visible in the image transcribed here."
 }}"""
 
-    
-    
+    # Run OCR through FireRed-OCR using Transformers
     try:
         # Decode base64 data URI into bytes/PIL image
         _, b64_data = image_data_uri.split(",", 1)
@@ -275,7 +279,7 @@ def handle_file(state: AgentState) -> str:
 
     # Use the file_reader tool to fetch the file content
     file_reader = tools.FileReaderTool()
-    file_content = file_reader(task_id) if task_id and file_name else ""
+    file_content = file_reader(task_id, file_name) if task_id and file_name else ""
 
     # Build prompt with the retrieved file content
     file_context = ""
@@ -293,8 +297,8 @@ Return a JSON object with the following field:
 {{
     "extracted_info": "The relevant extracted information from the file."
 }}"""
-    messages = [HumanMessage(content=prompt)]
-    response = model.invoke(messages)
+    messages = _msg(prompt)
+    response = model(messages)
     extracted_info = _extract_text_from_response(response)
     print(f"Extracted file info: {extracted_info[:100]}...")
     new_messages = state.get("messages", []) + [
@@ -311,8 +315,8 @@ def handle_math(state: AgentState) -> str:
     """Agent handles a math problem if classified as a math problem."""
     question = state["question"]
     print(f"Agent is handling a math problem: {question[:50]}...")
-    messages = [HumanMessage(content=f"Solve the following math problem step by step:\n\n{question}")]
-    response = math_model.invoke(messages)
+    messages = _msg(f"Solve the following math problem step by step:\n\n{question}")
+    response = math_model(messages)
     solution = _extract_text_from_response(response)
     print(f"Math solution: {solution[:100]}...")
     new_messages = state.get("messages", []) + [
@@ -345,11 +349,11 @@ Question: {question}
 Context gathered:
 {context}
 """
-    messages = [HumanMessage(content=prompt)]
+    messages = _msg(prompt)
     # Use the general model for final answer synthesis
-    response = model.invoke(messages)
+    response = model(messages)
     raw_response = _extract_text_from_response(response)
-    
+
     # Extract the final answer after "FINAL ANSWER:" if present
     if "FINAL ANSWER:" in raw_response:
         final_answer = raw_response.split("FINAL ANSWER:")[-1].strip()
@@ -360,6 +364,52 @@ Context gathered:
     return {"final_answer": final_answer}
 
 
+def evaluate(state: AgentState) -> dict:
+    """LLM evaluates whether the current final_answer is adequate.
+    If not, increments retry_count so the graph can loop back."""
+    import json, re
+    question = state["question"]
+    current_answer = state.get("final_answer", "")
+    retry_count = state.get("retry_count", 0) or 0
+
+    prompt = f"""You are a strict evaluator. Given the question and a candidate answer, decide if the answer is complete, relevant, and not an error message.
+
+Question: {question}
+Candidate answer: {current_answer}
+
+Return ONLY a JSON object:
+{{"is_adequate": true}} if the answer looks correct and complete,
+{{"is_adequate": false}} if the answer is wrong, incomplete, an error, or just says it could not find information."""
+
+    response = model(_msg(prompt))
+    raw = _extract_text_from_response(response)
+    match = re.search(r'\{.*?\}', raw, re.DOTALL)
+    data = {}
+    if match:
+        try:
+            data = json.loads(match.group())
+        except json.JSONDecodeError:
+            pass
+    is_adequate = bool(data.get("is_adequate", True))  # default: accept
+    print(f"Evaluation: is_adequate={is_adequate}, retry_count={retry_count}")
+    return {
+        "retry_count": retry_count + (0 if is_adequate else 1),
+        "is_searching": False if not is_adequate else state.get("is_searching"),
+        "have_file": False if not is_adequate else state.get("have_file"),
+        "is_math": False if not is_adequate else state.get("is_math"),
+        "have_image": False if not is_adequate else state.get("have_image"),
+    }
+
+
+def route_after_evaluate(state: AgentState) -> str:
+    """If answer was inadequate and retries remain, search web for more context."""
+    retry_count = state.get("retry_count", 0) or 0
+    if retry_count > 0 and retry_count <= 2:
+        print(f"Answer inadequate — retry {retry_count}/2, routing to web search")
+        return "handle_search"
+    return END
+
+
 def route_after_classify(state: AgentState) -> str:
     """Routing function: decide which handler to invoke based on classification."""
     if state.get("have_image"):
@@ -383,6 +433,7 @@ agent_graph.add_node("handle_image", handle_image)
 agent_graph.add_node("handle_file", handle_file)
 agent_graph.add_node("handle_math", handle_math)
 agent_graph.add_node("answer", answer)
+agent_graph.add_node("evaluate", evaluate)
 
 agent_graph.add_edge(START, "read")
 agent_graph.add_edge("read", "classify")
@@ -395,7 +446,11 @@ agent_graph.add_edge("handle_search", "answer")
 agent_graph.add_edge("handle_image", "answer")
 agent_graph.add_edge("handle_file", "answer")
 agent_graph.add_edge("handle_math", "answer")
-agent_graph.add_edge("answer", END)
+agent_graph.add_edge("answer", "evaluate")
+agent_graph.add_conditional_edges(
+    "evaluate",
+    route_after_evaluate,
+)
 
 compiled_agent = agent_graph.compile()
 
@@ -424,7 +479,8 @@ class BasicAgent:
             "have_file": False,
             "is_math": False,
             "have_image": False,
-            "final_answer": ""
+            "final_answer": "",
+            "retry_count": 0
         })
 
         # Extract the final answer from the state