import os
from smolagents import (
    CodeAgent, 
    LiteLLMModel,
    PythonInterpreterTool,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
)
from rate_limiter import RateLimiter

class BasicAgent:
    def __init__(self):
        model = LiteLLMModel(
            model_id=os.getenv("MODEL_ID"), 
            api_key=os.getenv("MODEL_API_KEY")
        )
        self.agent = CodeAgent(
            tools=[
                PythonInterpreterTool(),
                DuckDuckGoSearchTool(),
                VisitWebpageTool(),
            ], 
            model=model,
            step_callbacks=[RateLimiter().increment_and_sleep_if_needed]
        )
        
    def __call__(self, question: str) -> str:
        final_answer = self.agent.run(task=question)
        print(f"Agent returning final answer: {final_answer}")
        return final_answer