import os
from smolagents import CodeAgent, DuckDuckGoSearchTool
from smolagents import TransformersModel

class GaiaAgent:
    def __init__(self, model_id: str = "google/gemma-2b-it"): # <-- CHANGE MODEL HERE
        self.llm_model = TransformersModel(
            model_id=model_id,
            task="text-generation",
            # device_map="auto" # Can keep this, accelerate will manage
        )

        self.agent = CodeAgent(
            model=self.llm_model,
            tools=[DuckDuckGoSearchTool()],
            add_base_tools=False,
            verbose=True
        )

    def process_task(self, task_description: str) -> str:
        try:
            response = self.agent.run(task_description)
            return response
        except Exception as e:
            return f"An error occurred during agent processing: {e}"