from smolagents import HfApiModel, ToolCallingAgent, CodeAgent
from tools import DuckDuckGoSearchTool, get_current_time_in_timezone
import os

class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]} ...")

        # Using Qwen
        model = HfApiModel(
            max_tokens=2096,
            temperature=0.5,
            model_id = os.getenv("QWEN_MODEL"),# it is possible that this model may be overloaded
            custom_role_conversions=None,
        )
        agent = CodeAgent(
            model=model,
            tools=[DuckDuckGoSearchTool(),get_current_time_in_timezone]
        )
        prompt = f"""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 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.
{question}
"""
        try:
            answer = agent.run(prompt)
        except Exception as e:
            print(f"Agent failed: {e}")
            answer = "Error: Agent could not complete the request."
        
        print(f"Agent returning answer: {answer}")
        return answer