from openai import OpenAI

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

MODEL_ID = "gpt-4o"

def run_agent(level, question, file_name, ground_truth):
    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),
            planning_interval=3,
            tools = [DuckDuckGoSearchTool(), 
                     WikipediaSearchTool()],
            verbosity_level = 2
        )

    answer = agent.run(question)
    
    return answer, str(answer == ground_truth)

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

    prompt = """
        You are given a question and a preliminary answer:
        Question: {question}
        Preliminary answer: {answer}
        You must give the precise answer requested. Do not include explanations, steps, reasoning, or additional text. Be direct and specific.
        For example, if asked "What is the capital of France?", respond simply with "Paris".
    """

    completion = client.chat.completions.create(
        max_tokens = 1000,
        messages = [{"role": "user",
                     "content": [{"type": "text", "text": prompt}],
        model = MODEL_ID,
        temperature = 0.0
    )

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

    print("###")
    print(content)
    print("###")

    return content