import yaml

from smolagents import (
    PromptTemplates,
    PlanningPromptTemplate,
    FinalAnswerPromptTemplate,
    ManagedAgentPromptTemplate,
)

def load_prompts():

    """Builds a CodeAgent using the custom system and planning prompts from code_agent.yaml."""

    # 1) Load YAML
    with open("prompts/code_agent.yaml", "r", encoding="utf-8") as f:
        y = yaml.safe_load(f)

    # 2) Build PromptTemplates
    templates = PromptTemplates(
        system_prompt=y.get("system_prompt", ""),
        planning=PlanningPromptTemplate(**y.get("planning", {})),
        managed_agent=ManagedAgentPromptTemplate(**y.get("managed_agent", {})),
        final_answer=FinalAnswerPromptTemplate(**y.get("final_answer", {})),
    )

    return templates


def load_prompts2():
    system_prompt_override = "prompts/code_agent.yaml"
    
    try:
        with open(system_prompt_override, 'r', encoding='utf-8') as file:
            return file.read()
    except FileNotFoundError:
        raise FileNotFoundError(f"Prompt file '{system_prompt_override}' not found.")
    except Exception as e:
        raise Exception(f"Failed to read prompt file: {e}")