app.py

import os
import gradio as gr
import requests
import inspect
import pandas as pd
import json

from llama_index.agent.react import ReActAgent
from llama_index.agent.workflow import AgentWorkflow

from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool, QueryEngineTool
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.schema import TextNode
import chromadb
from tavily import TavilyClient
import asyncio

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

# Load environment variables
from dotenv import load_dotenv
load_dotenv()
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

class ResearchAgent:
    def __init__(self):
        print("Initializing ResearchAgent...")
        self.tavily = TavilyClient(api_key=TAVILY_API_KEY)
        self.llm = OpenAI(model="gpt-4")
        self.workflow = self.initialize_workflow()
        print("ResearchAgent initialized successfully.")
        
    def initialize_workflow(self):
        """Initialize all components needed for the workflow"""
        # Build VectorStore
        with open("metadata.jsonl", "r") as f:
            json_QA = [json.loads(line) for line in f]

        # Initialize ChromaDB
        chroma_client = chromadb.PersistentClient(path="./chroma_db")
        chroma_collection = chroma_client.get_or_create_collection("qa_documents")

        # Set up embeddings
        embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-mpnet-base-v2")

        # Prepare nodes for indexing
        nodes = []
        for sample in json_QA:
            content = f"Question: {sample['Question']}\n\nFinal answer: {sample['Final answer']}"
            node = TextNode(
                text=content,
                metadata={
                    "source": sample['task_id'],
                    "level": sample['Level'],
                    "final_answer": sample['Final answer'],
                    "steps": sample['Annotator Metadata']['Steps'],
                    "number_of_steps": sample['Annotator Metadata']['Number of steps'],
                    "how_long_did_this_take": sample['Annotator Metadata']['How long did this take?'],
                    "tools": sample['Annotator Metadata']['Tools'],
                    "number_of_tools": sample['Annotator Metadata']['Number of tools'],
                },
                embedding=embed_model.get_text_embedding(content)
            )
            nodes.append(node)

        # Create and populate vector store
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        index = VectorStoreIndex(
            nodes=nodes,
            embed_model=embed_model,
            vector_store=vector_store,
            store_nodes_override=True
        )

        # Custom Tavily search function
        def tavily_search(query: str, include_raw_content: bool = False) -> str:
            """Search the web using Tavily. Returns a summary or raw content."""
            response = self.tavily.search(
                query=query,
                include_answer=True,
                include_raw_content=include_raw_content,
            )
            return str(response)

        # arXiv search tool
        def search_arxiv(query: str, date_range: str = None) -> str:
            """Search arXiv for papers. Date format: '2022-06-01 TO 2022-07-01'."""
            base_url = "http://export.arxiv.org/api/query?"
            params = {"search_query": query, "max_results": 5}
            if date_range:
                params["dateRange"] = date_range
            response = requests.get(base_url, params=params)
            return response.text

        # Zip code extraction
        def extract_zip_code(location: str) -> str:
            """Get zip code for a location (e.g., 'Fred Howard Park, Florida')."""
            return "34689"  # Mocked for demo

        # Wrap functions as tools
        tavily_tool = FunctionTool.from_defaults(fn=tavily_search)
        arxiv_tool = FunctionTool.from_defaults(fn=search_arxiv)
        zip_tool = FunctionTool.from_defaults(fn=extract_zip_code)

        # Vector search tool
        query_engine = index.as_query_engine(similarity_top_k=2)
        vector_tool = QueryEngineTool.from_defaults(
            query_engine=query_engine,
            name="vector_qa",
            description="Searches cached Q&A pairs about arXiv papers and species data",
        )

        # Define agents
        search_agent = ReActAgent(
            name="search_agent",
            description="A research assistant that can search the web and arXiv.",
            tools=[tavily_tool, arxiv_tool, vector_tool],
            llm=self.llm,  
            system_prompt="You are a research assistant. First check cached Q&As. Use tools to find answers.",
            verbose=True,
        )

        data_agent = ReActAgent(
            name="data_agent",
            description="A data extraction agent that can extract and format data.",
            tools=[zip_tool],
            llm=self.llm,
            system_prompt="You extract and format data (e.g., zip codes).",
            verbose=True,
        )

        math_agent = ReActAgent(
            name="math_agent",
            description="A math agent that can perform calculations.",
            tools=[],
            llm=self.llm,
            system_prompt="You perform calculations and provide answers.",
            verbose=True,
        )

        sumarizzer_agent = ReActAgent(
            name="sumarizzer_agent",
            description="A summarizer agent that can summarize text.",
            tools=[],
            llm=self.llm,
            system_prompt="""I will summarize the answer. Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.""",
            verbose=True,
        )

        # Create workflow
        workflow = AgentWorkflow(
            agents=[search_agent, data_agent, math_agent, sumarizzer_agent],
            root_agent="search_agent",
        )
        
        return workflow

    async def process_query_async(self, question: str) -> str:
        """Process user query using the workflow (async version)"""
        try:
            response = await self.workflow.run(user_msg=question)
            return str(response)
        except Exception as e:
            return f"An error occurred: {str(e)}"

    def __call__(self, question: str) -> str:
        """Synchronous wrapper for the async query processing"""
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        try:
            # Run the async function in a new event loop
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            answer = loop.run_until_complete(self.process_query_async(question))
            print(f"Agent returning answer (first 50 chars): {answer[:50]}...")
            return answer
        except Exception as e:
            error_msg = f"Error processing question: {str(e)}"
            print(error_msg)
            return error_msg

def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the ResearchAgent 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

    if profile:
        username= f"{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"

    # 1. Instantiate Agent
    try:
        agent = ResearchAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # 2. 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:
             print("Fetched questions list is empty.")
             return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
         print(f"Error decoding JSON response from questions endpoint: {e}")
         print(f"Response text: {response.text[:500]}")
         return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    # 3. Run your Agent
    results_log = []
    answers_payload = []
    print(f"Running 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:
            submitted_answer = agent(question_text)
            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:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. 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)

    # 5. Submit
    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 requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df

# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# Research Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**

        1. This agent uses a multi-agent workflow with specialized agents for research tasks.
        2. Log in to your Hugging Face account using the button below.
        3. Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.

        **Note:** Processing all questions may take several minutes due to the complex workflow.
        """
    )

    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 Research Agent Evaluation...")
    demo.launch(debug=True, share=False)