Update app.py

app.py

@@ -2,18 +2,133 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-from agent import BasicAgent
+from smolagents import LiteLLMModel, CodeAgent, Tool
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
+# --- Agent Tools ---
+class MathSolver(Tool):
+    name = "math_solver"
+    description = "Safely evaluate basic math expressions."
+    inputs = {"input": {"type": "string", "description": "Math expression to evaluate."}}
+    output_type = "string"
+
+    def forward(self, input: str) -> str:
+        try:
+            # Safe evaluation of math expressions
+            allowed_names = {
+                k: v for k, v in __builtins__.items() if k in [
+                    'abs', 'round', 'min', 'max', 'sum', 'pow'
+                ]
+            }
+            allowed_names.update({
+                'int': int, 'float': float, 'str': str,
+                '__builtins__': {}
+            })
+            return str(eval(input, allowed_names))
+        except Exception as e:
+            return f"Math error: {e}"
+
+class FileAttachmentQueryTool(Tool):
+    name = "run_query_with_file"
+    description = "Downloads a file mentioned in a user prompt, adds it to the context, and runs a query on it."
+    inputs = {
+        "task_id": {
+            "type": "string",
+            "description": "A unique identifier for the task related to this file, used to download it.",
+            "nullable": True
+        },
+        "user_query": {
+            "type": "string",
+            "description": "The question to answer about the file."
+        }
+    }
+    output_type = "string"
+
+    def forward(self, task_id: str | None, user_query: str) -> str:
+        if not task_id:
+            return "No task_id provided for file download."
+        
+        file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+        try:
+            file_response = requests.get(file_url)
+            if file_response.status_code != 200:
+                return f"Failed to download file: {file_response.status_code}"
+            
+            # For text-based files, return content directly
+            file_content = file_response.text[:2000]  # Limit content size
+            return f"Relevant information from file: {file_content}"
+        except Exception as e:
+            return f"File download error: {e}"
+
+# --- Agent Implementation ---
+def select_model(provider="groq"):
+    """Select and return a model based on the provider."""
+    GROQ_MODEL_NAME = "groq/llama3-70b-8192"
+    HF_MODEL_NAME = "huggingfaceh4/zephyr-7b-beta"
+    
+    if provider == "groq":
+        api_key = os.getenv("GROQ_API_KEY")
+        if not api_key:
+            raise ValueError("GROQ_API_KEY environment variable is not set")
+        return LiteLLMModel(model_id=GROQ_MODEL_NAME, api_key=api_key)
+    elif provider == "hf":
+        api_key = os.getenv("HF_TOKEN")
+        if not api_key:
+            raise ValueError("HF_TOKEN environment variable is not set")
+        return LiteLLMModel(model_id=HF_MODEL_NAME, api_key=api_key)
+    else:
+        # Default to Groq if no valid provider specified
+        api_key = os.getenv("GROQ_API_KEY")
+        if not api_key:
+            raise ValueError("GROQ_API_KEY environment variable is not set")
+        return LiteLLMModel(model_id=GROQ_MODEL_NAME, api_key=api_key)
+
+class BasicAgent:
+    def __init__(self, provider="groq"):
+        model = select_model(provider)
+        tools = [
+            MathSolver(),
+            FileAttachmentQueryTool(),
+        ]
+        self.agent = CodeAgent(
+            model=model,
+            tools=tools,
+            add_base_tools=False,
+            max_steps=15,
+        )
+        # System prompt to enforce exact answer format
+        self.agent.prompt_templates["system_prompt"] = (
+            "You are a GAIA benchmark AI assistant. Your sole purpose is to output the minimal, final answer. "
+            "You must NEVER output explanations, intermediate steps, reasoning, or comments — only the answer. "
+            "For numerical answers, use digits only, e.g., `4` not `four`. "
+            "For string answers, omit articles ('a', 'the') and use full words. "
+            "For lists, output in comma-separated format with no conjunctions. "
+            "If the answer is not found, say `- unknown`."
+        )
+
+    def __call__(self, question: str) -> str:
+        result = self.agent.run(question)
+        # Extract only the final answer without any wrappers
+        final_str = str(result).strip()
+        # Remove any potential prefixes
+        if final_str.startswith('[ANSWER]'):
+            final_str = final_str[8:].strip()
+        if final_str.startswith('Final answer:'):
+            final_str = final_str[13:].strip()
+        if final_str.startswith('Answer:'):
+            final_str = final_str[7:].strip()
+        return final_str
+
+# --- Main Application Functions ---
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
     and displays the results.
     """
     # --- Determine HF Space Runtime URL and Repo URL ---
-    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code
+    space_id = os.getenv("SPACE_ID")
 
     if profile:
         username = f"{profile.username}"