app.py

import os
import json
import time
import requests
import gradio as gr
import spaces
from groq import Groq

from prompt import SYSTEM_PROMPT
from tools import execute_tool


SCORING_API = "https://agents-course-unit4-scoring.hf.space"

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
TEXT_MODEL = os.getenv("GROQ_TEXT_MODEL", "openai/gpt-oss-120b")

MAX_STEPS = 6
MAX_TOOL_RESULT = 6000

client = Groq(api_key=GROQ_API_KEY)


# ---------------------------------------------------------
# ZeroGPU compatibility
# ---------------------------------------------------------

@spaces.GPU(duration=1)
def zerogpu_probe():
    return "ZeroGPU available"


# ---------------------------------------------------------
# Native Groq tools
# ---------------------------------------------------------

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": (
                "Search the public web for factual information. "
                "Use this first for questions requiring current, historical, "
                "Wikipedia, YouTube, NASA, IOC, news, or other web research."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The exact web search query."
                    }
                },
                "required": ["query"],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "read_webpage",
            "description": (
                "Read a webpage returned by web_search. "
                "Normally provide the full URL from a search result. "
                "If you do not know the URL, you may provide cursor as the "
                "zero-based search-result index and optionally pattern. "
                "The application will resolve the URL from the most recent search."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "Full webpage URL. Optional if cursor is provided."
                    },
                    "cursor": {
                        "type": "integer",
                        "description": "Zero-based index of the most recent web-search result."
                    },
                    "pattern": {
                        "type": "string",
                        "description": "Optional text/topic to focus on while reading."
                    }
                },
                "required": [],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "download_file",
            "description": (
                "Download the attachment associated with the current GAIA task. "
                "Use this before inspecting an Excel, PDF, image, audio, or Python attachment."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {
                        "type": "string",
                        "description": "The current GAIA task ID."
                    }
                },
                "required": ["task_id"],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "inspect_file",
            "description": (
                "Inspect the current GAIA attachment. "
                "Useful for Excel, CSV, PDF, TXT, Python, JSON and other files. "
                "The application automatically downloads the task attachment."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {
                        "type": "string",
                        "description": "The current GAIA task ID."
                    }
                },
                "required": ["task_id"],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": (
                "Run Python code for calculations, data processing, string operations, "
                "statistics, or analyzing a downloaded attachment. "
                "Do not use shell commands."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "Python code to execute."
                    }
                },
                "required": ["code"],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "transcribe_audio",
            "description": (
                "Transcribe the audio attachment for the current GAIA task. "
                "Use when the question depends on spoken audio."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {
                        "type": "string",
                        "description": "The current GAIA task ID."
                    }
                },
                "required": ["task_id"],
                "additionalProperties": False
            }
        }
    },

    {
        "type": "function",
        "function": {
            "name": "analyze_image",
            "description": (
                "Analyze the image attachment for the current GAIA task. "
                "Use this for visual questions, diagrams, screenshots, chess boards, etc."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "task_id": {
                        "type": "string",
                        "description": "The current GAIA task ID."
                    },
                    "question": {
                        "type": "string",
                        "description": "What specifically must be determined from the image."
                    }
                },
                "required": ["task_id", "question"],
                "additionalProperties": False
            }
        }
    }
]


# ---------------------------------------------------------
# Groq call
# ---------------------------------------------------------

def call_groq(messages, use_tools=True):
    last_error = None

    for attempt in range(3):
        try:
            kwargs = {
                "model": TEXT_MODEL,
                "messages": messages,
                "temperature": 0,
                "max_completion_tokens": 2500,
            }

            if use_tools:
                kwargs["tools"] = TOOLS
                kwargs["tool_choice"] = "auto"
            else:
                # No tools on final answer call.
                kwargs["tool_choice"] = "none"

            response = client.chat.completions.create(**kwargs)
            return response.choices[0].message

        except Exception as e:
            last_error = e
            print(f"Groq error (attempt {attempt + 1}/3): {type(e).__name__}: {e}")

            if attempt < 2:
                time.sleep(2 ** attempt)

    raise RuntimeError(
        f"Groq request failed after 3 attempts: {last_error}"
    )


# ---------------------------------------------------------
# Answer cleanup
# ---------------------------------------------------------

def clean_answer(answer):
    if answer is None:
        return ""

    answer = str(answer).strip()

    if answer.startswith("```"):
        lines = answer.splitlines()

        if lines and lines[0].startswith("```"):
            lines = lines[1:]

        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]

        answer = "\n".join(lines).strip()

    if answer.upper().startswith("FINAL ANSWER:"):
        answer = answer[len("FINAL ANSWER:"):].strip()

    # Handle accidental JSON final response
    try:
        obj = json.loads(answer)

        if isinstance(obj, dict) and "answer" in obj:
            answer = str(obj["answer"]).strip()
    except Exception:
        pass

    return answer


# ---------------------------------------------------------
# Solve one GAIA question
# ---------------------------------------------------------

def solve_question(task_id, question):

    print()
    print("=" * 60)
    print(f"Agent solving task {task_id}")
    print("=" * 60)

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": (
                f"GAIA TASK ID: {task_id}\n\n"
                f"QUESTION:\n{question}\n\n"
                "Solve this task and return only the exact answer when finished."
            )
        }
    ]

    for step in range(1, MAX_STEPS + 1):

        print(f"Agent step {step}/{MAX_STEPS}")

        try:
            message = call_groq(messages, use_tools=True)
        except Exception as e:
            return f"ERROR: {e}"

        tool_calls = message.tool_calls or []

        if not tool_calls:
            answer = clean_answer(message.content)
            print(f"FINAL ANSWER: {answer}")
            return answer

        print(
            "Tool calls:",
            [call.function.name for call in tool_calls]
        )

        # Important: append the assistant tool-call message exactly as returned.
        messages.append(message)

        for tool_call in tool_calls:

            name = tool_call.function.name

            try:
                args = json.loads(tool_call.function.arguments or "{}")
            except Exception:
                args = {}

            print(f"Executing: {name}")
            print(f"Arguments: {args}")

            # Force current task ID into task-dependent tools.
            if name in {
                "download_file",
                "inspect_file",
                "transcribe_audio",
                "analyze_image"
            }:
                args["task_id"] = task_id

            try:
                result = execute_tool(
                    name=name,
                    args=args,
                    task_id=task_id
                )

            except Exception as e:
                result = f"Tool error: {type(e).__name__}: {e}"

            result = str(result)

            if len(result) > MAX_TOOL_RESULT:
                result = (
                    result[:MAX_TOOL_RESULT]
                    + "\n[TOOL OUTPUT TRUNCATED]"
                )

            print(f"Tool result: {len(result)} characters")

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": name,
                    "content": result
                }
            )

    return "ERROR: Maximum reasoning steps reached."


# ---------------------------------------------------------
# Fetch questions
# ---------------------------------------------------------

def get_questions():

    response = requests.get(
        f"{SCORING_API}/questions",
        timeout=30
    )

    response.raise_for_status()

    data = response.json()

    if isinstance(data, list):
        return data

    if isinstance(data, dict):
        if "questions" in data:
            return data["questions"]

    raise RuntimeError("Unexpected /questions response format.")


# ---------------------------------------------------------
# Submit
# ---------------------------------------------------------

def submit_answers(username, answers):

    payload = {
        "username": username,
        "agent_code": (
            "https://huggingface.co/spaces/"
            f"{username}/Final_Assignment_Template/tree/main"
        ),
        "answers": answers
    }

    response = requests.post(
        f"{SCORING_API}/submit",
        json=payload,
        timeout=60
    )

    response.raise_for_status()

    return response.json()


# ---------------------------------------------------------
# Run all questions
# ---------------------------------------------------------

def run_agent(username):

    if not GROQ_API_KEY:
        return "ERROR: GROQ_API_KEY is not configured in Space Secrets."

    username = username.strip()

    if not username:
        return "ERROR: Enter your Hugging Face username."

    print("\nFetching GAIA questions...")

    questions = get_questions()

    print(f"Found {len(questions)} questions.")

    answers = []

    for i, item in enumerate(questions, start=1):

        task_id = item.get("task_id") or item.get("id")
        question = item.get("question")

        print()
        print("=" * 60)
        print(f"QUESTION {i}/{len(questions)}")
        print(f"TASK ID: {task_id}")
        print("=" * 60)

        answer = solve_question(
            task_id=task_id,
            question=question
        )

        answers.append(
            {
                "task_id": task_id,
                "submitted_answer": answer
            }
        )

        # Small delay between questions.
        time.sleep(0.5)

    print("\nSubmitting answers...")

    result = submit_answers(
        username=username,
        answers=answers
    )

    print("\nSUBMISSION RESULT:")
    print(result)

    return json.dumps(
        result,
        indent=2
    )


# ---------------------------------------------------------
# Simple UI
# ---------------------------------------------------------

with gr.Blocks(title="GAIA Agent") as demo:

    gr.Markdown("# GAIA Unit 4 Agent")

    username = gr.Textbox(
        label="Hugging Face Username",
        value="Viske"
    )

    run_button = gr.Button(
        "Run Agent"
    )

    output = gr.Textbox(
        label="Result",
        lines=15
    )

    run_button.click(
        fn=run_agent,
        inputs=username,
        outputs=output
    )


if __name__ == "__main__":
    demo.launch(
        debug=True,
        share=False,
        ssr_mode=False
    )