from openai import OpenAI

from smolagents import (
    CodeAgent, 
    OpenAIServerModel, 
    DuckDuckGoSearchTool, 
    WikipediaSearchTool
)

MODEL_ID_1 = "gpt-4o"
MODEL_ID_2 = "o4-mini"

def run_agent(level, question, file_name):
    if level == 1:
        max_steps = 5
    elif level == 2:
        max_steps = 10
    elif level == 3:
        max_steps = 20
   
    agent = CodeAgent(
            add_base_tools = True,
            additional_authorized_imports = ["pandas", "numpy"],
            max_steps = max_steps,
            model = OpenAIServerModel(model_id = MODEL_ID_1),
            #planning_interval=3,
            tools = [DuckDuckGoSearchTool(), 
                     WikipediaSearchTool()],
            #verbosity_level = 2
        )

    answer = agent.run(question)

    return get_final_answer(question, answer)

def get_final_answer(question, answer):
    client = OpenAI()

    prompt_template = """
You are a GAIA benchmark assistant, who is given a question and an answer. Read the question carefully. You must only give the final answer requested. Do not include explanations, steps, reasoning, or additional text. Be direct and specific.
**Question:** """ + str(question) + """
**Answer:** """ + str(answer) + """
**Example:** If asked "What is the capital of France?", respond with "Paris".
**Final answer:**:

"""

    completion = client.chat.completions.create(
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt_template}]}],
        model = MODEL_ID_2
    )

    final_answer = completion.choices[0].message.content

    return final_answer