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

class GaiaAgent:
    def __init__(self, model_id: str = "HuggingFaceH4/zephyr-7b-beta"): # <-- CHANGE MODEL HERE
        self.llm_model = TransformersModel(
            model_id=model_id,
            # For Zephyr (a causal LM), the default `AutoModelForCausalLM` works,
            # and `task="text-generation"` is appropriate for the pipeline.
            task="text-generation", 
            # You might need device_map="auto" if you hit memory issues or have GPU:
            # device_map="auto" 
        )

        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}"