from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI 

# 1. Initialize the search tool (for internal use by our custom tool only)
duckduckgo_tool = DuckDuckGoSearchTool()

# 2. Define Custom Tool
# --- Herramientas personalizadas (Solo queda la de compositores) ---
@tool
def get_composer_lifespan_from_web(composer_lastname: str) -> str:
    """
    Busca información sobre un compositor y devuelve los datos crudos para que el agente los procese.
    
    Args:
        composer_lastname: El apellido del compositor (p. ej., 'Beethoven', 'Mozart').
    """
    global duckduckgo_tool
    
    # CRITICAL CHANGE: Use a shorter, less restrictive query.
    # The original was: "classical composer {composer_lastname} full name birth and death year"
    # The new one is simply the composer's name.
    query = f"{composer_lastname} composer lifespan"
    
    # Use .forward() instead of the direct call if available, or just the direct call
    # depending on the smolagents version. Assuming .forward() from the previous correction.
    search_result = duckduckgo_tool.forward(query) 
    
    return f"""
    RESULTADOS DE BÚSQUEDA:
    {search_result}
    
    INSTRUCCIÓN OBLIGATORIA PARA EL AGENTE:
    Usa la información de arriba para responder ÚNICAMENTE con este formato exacto:
    
    Nombre: [Nombre Completo]
    Fecha de nacimiento: [Año]
    Fecha de fallecimiento: [Año]
    """

# 3. Setup Agent
final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Load prompts
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    # CRITICAL FIX: Removed 'duckduckgo_tool' from this list. 
    # Now the agent MUST use 'get_composer_lifespan_from_web'.
    tools=[final_answer, get_composer_lifespan_from_web], 
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# 4. Launch UI
# Removed 'title' argument to prevent the previous crash
GradioUI(agent).launch()