from smolagents import CodeAgent, HfApiModel, FinalAnswerTool

class LLMOnlyAgent:
    def __init__(self):

    # Basic inference model
        model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
        custom_role_conversions=None,
        )

        # Code Agent
        self.agent = CodeAgent(
            model=model,
            tools=[FinalAnswerTool()],
            max_steps=2
        )
        
        print("BasicAgent initialized.")
    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        answer = self.agent.run(question)
        print(f"Agent returning answer: {answer}")
        return answer

    def 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 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."""