app.py

import os
import gradio as gr
from PIL import Image
import requests
import pandas as pd
import io

# Import LangChain and LangGraph components
from langgraph.prebuilt import create_react_agent
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.tools import tool

# Constants (unchanged)
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# ------------------------------------------------------------------
# 1. Define your LangGraph ReAct Agent with Gemini 2.0 Flash
# ------------------------------------------------------------------
def get_agent():
    """Initialize the LangGraph ReAct agent with Gemini 2.0 Flash and tools."""
    # Gemini API key from Hugging Face secrets
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        raise ValueError("GEMINI_API_KEY environment variable not set. Please add it to Hugging Face Space secrets.")

    # Initialize the Gemini model
    model = ChatGoogleGenerativeAI(
        model="gemini-2.0-flash",
        api_key=gemini_api_key,
        temperature=0.7,
        timeout=60,
        max_retries=2
    )

    # --------------------------------------------------------------
    # 2. Define your custom tools using LangChain's @tool decorator
    # --------------------------------------------------------------

    # Tool: DuckDuckGo Search (free web search)
    ddg_search = DuckDuckGoSearchRun()

    # Tool: Image generation using FLUX.1 Schnell on Hugging Face
    @tool
    def image_generator(prompt: str) -> dict:
        """
        Generates a high-quality visual image based on a descriptive text prompt.
        Uses FLUX.1 Schnell from Hugging Face.
        Returns the image as a base64 string or URL.
        """
        HF_TOKEN = os.getenv("HF_TOKEN")  # Optional but helps with rate limits
        API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}

        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
        if response.status_code == 200:
            return {
                "success": True,
                "image": response.content,
                "message": "Image generated successfully."
            }
        else:
            return {
                "success": False,
                "message": f"Image generation failed: {response.status_code} - {response.text}"
            }

    # Combine tools into a list
    tools = [ddg_search, image_generator]

    # --------------------------------------------------------------
    # 3. Create the ReAct agent with LangGraph
    # --------------------------------------------------------------
    agent = create_react_agent(
        model=model,
        tools=tools,
        prompt=(
            "You are a helpful AI assistant with access to web search and image generation. "
            "Always think step by step before using tools. For image requests, use the image_generator tool. "
            "For web information, use the duckduckgo_search tool."
        )
    )
    
    return agent

def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the LangGraph ReAct agent, submits all answers,
    and displays the results.
    """
    space_id = os.getenv("SPACE_ID")
    if profile:
        username = profile.username
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # ------------------------------------------------------------------
    # Initialize LangGraph ReAct Agent
    # ------------------------------------------------------------------
    try:
        agent = get_agent()
    except Exception as e:
        print(f"Error initializing LangGraph agent: {e}")
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # Fetch questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except Exception as e:
        return f"Error fetching questions: {e}", None

    # Run agent on each question
    results_log = []
    answers_payload = []
    print(f"Running LangGraph ReAct agent on {len(questions_data)} questions...")
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            # LangGraph agent invocation returns a dict with 'output' key
            result = agent.invoke({
                "messages": [("user", question_text)]
            })
            submitted_answer = result["output"]
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append({
                "Task ID": task_id,
                "Question": question_text,
                "Submitted Answer": submitted_answer
            })
        except Exception as e:
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append({
                "Task ID": task_id,
                "Question": question_text,
                "Submitted Answer": f"AGENT ERROR: {e}"
            })

    if not answers_payload:
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # Prepare submission
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)

    # Submit answers
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except Exception as e:
        status_message = f"Submission Failed: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df

# ------------------------------------------------------------------
# Gradio Interface (unchanged)
# ------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**
        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.

        ---
        **Disclaimers:**
        Once clicking on the "submit" button, it can take quite some time (this is the time for the agent to go through all the questions).
        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution.
        """
    )

    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table]
    )

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID")
    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")
    if space_id_startup:
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")
    print("-"*(60 + len(" App Starting ")) + "\n")
    print("Launching Gradio Interface for LangGraph ReAct Agent Evaluation...")
    demo.launch(debug=True, share=False)