# tools.py
import os
import openai
from smolagents.tools import Tool

# Carichiamo la tua chiave
openai.api_key = os.getenv("OPENAI_API_KEY")

SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL 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 commas, symbols or units (e.g. %, $, km) unless explicitly asked.
- If you are asked for a string, don't use articles ("a", "the"), abbreviations (e.g. "NYC"), or extra words; write digits in plain text unless specified otherwise.
- If you are asked for a comma separated list, apply the above rules to each element.

Examples:
Q: How many legs does a spider have?
Thought: It's basic biology.
FINAL ANSWER: 8

Q: What are the primary colors in light?
Thought: Physics concept.
FINAL ANSWER: red, green, blue

Q: Who won the Best Picture Oscar in 1994?
Thought: Film history.
FINAL ANSWER: Forrest Gump

Now it's your turn.
"""

class AnswerTool(Tool):
    name = "answer_tool"
    description = "Answer GAIA Level 1 questions in exact-match format."
    inputs = {"question": {"type": "string", "description": "The GAIA question text."}}
    output_type = "string"

    def forward(self, question: str) -> str:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": question},
        ]
        # → qui usiamo la nuova API v1
        resp = openai.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            temperature=0.0,
            max_tokens=64,
        )
        text = resp.choices[0].message.content.strip()
        if "FINAL ANSWER:" in text:
            return text.split("FINAL ANSWER:")[-1].strip()
        return text