import os
from smolagents import CodeAgent, DuckDuckGoSearchTool, GradioUI, InferenceClientModel, ManagedAgent

# 1. THE BRAIN
model = InferenceClientModel(
    model_id="deepseek-ai/DeepSeek-R1",
    token=os.getenv("HF_TOKEN")
)

# 2. THE RESEARCHER (Specialist)
# We wrap the researcher in ManagedAgent so the Boss can "call" it.
search_agent = CodeAgent(
    tools=[DuckDuckGoSearchTool()],
    model=model,
    name="research_specialist",
    description="Finds technical data and web information. Use for all research tasks."
)

managed_researcher = ManagedAgent(
    agent=search_agent,
    name="research_specialist",
    description="A web-search specialist. Call this agent when you need current data from the internet."
)

# 3. THE BOSS LOGIC (The Template Fix)
# This dictionary overrides the read-only system_prompt property.
custom_templates = {
    "system_prompt": """You are the BOSS OPERATOR (Elite PhD Level). 
    1. Delegate sub-tasks to 'research_specialist' for web-based data.
    2. Use your internal Python tools for math, simulation, and data analysis.
    3. Always verify data before presenting a final answer.
    4. Format all math in LaTeX ($E=mc^2$) and all data in Markdown tables."""
}

# 4. THE BOSS (Manager)
manager = CodeAgent(
    tools=[], 
    model=model,
    managed_agents=[managed_researcher],
    add_base_tools=True,
    prompt_templates=custom_templates # <--- THIS IS THE CRITICAL FIX
)

if __name__ == "__main__":
    GradioUI(manager).launch()