import os
from smolagents import OpenAIServerModel


system_text = """
Revise the answer to match exactly what the question is asking. Keep it concise and focused, removing irrelevant details, redundant phrases, or formatting. 
Do not change the meaning or add new information. Ensure the output directly answers the question, with the correct level of specificity and format requested.
Only return the revised answer. Do not include explanations or headers. If the question asks a number of something return only the number, no extra words.
"""


def check_final_answer(question: str, answer: str) -> str:
    """
    Pass the question and answer to a LLM, to make it proper for GAIA comparison.

    Args:
        question: Question.
        answer: Original final answer.
    """
    OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

    answer_model = OpenAIServerModel(
        api_key=OPENAI_API_KEY,
        model_id='gpt-4o-mini',
        temperature=0,
    )

    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": system_text
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": f"Question:\n{question}\n\nOriginal Answer:\n{answer}",
                }
            ]
        }
    ]
    response = answer_model(messages).content

    return response