src/llm/prompts.py

from __future__ import annotations
"""Prompt management for the Ask-the-Web Agent."""

from pathlib import Path
from typing import Any

from src.utils.logging import get_logger

logger = get_logger(__name__)

# Prompts directory path
PROMPTS_DIR = Path(__file__).parent.parent.parent / "prompts"


def load_prompt(prompt_name: str) -> str:
    """Load a prompt template from the prompts directory.
    
    Args:
        prompt_name: Name of the prompt file (without .md extension)
        
    Returns:
        Prompt content as string
    """
    prompt_path = PROMPTS_DIR / f"{prompt_name}.md"

    if not prompt_path.exists():
        logger.error(f"Prompt file not found: {prompt_path}")
        raise FileNotFoundError(f"Prompt not found: {prompt_name}")

    with open(prompt_path, "r", encoding="utf-8") as f:
        content = f.read()

    # Extract the prompt template section (between ``` markers after "## Prompt Template")
    return _extract_template(content)


def _extract_template(content: str) -> str:
    """Extract the prompt template from markdown content.
    
    Args:
        content: Full markdown content
        
    Returns:
        Extracted template or full content if no template found
    """
    # Find the Prompt Template section
    template_marker = "## Prompt Template"
    if template_marker not in content:
        return content

    # Find content after the marker
    start_idx = content.find(template_marker) + len(template_marker)
    remaining = content[start_idx:]

    # Find the code block
    code_start = remaining.find("```")
    if code_start == -1:
        return content

    # Skip the opening ``` and any language marker
    code_start = remaining.find("\n", code_start) + 1
    code_end = remaining.find("```", code_start)

    if code_end == -1:
        return content

    return remaining[code_start:code_end].strip()


def format_prompt(prompt_name: str, **kwargs: Any) -> str:
    """Load and format a prompt template with variables.
    
    Args:
        prompt_name: Name of the prompt file
        **kwargs: Variables to substitute in the template
        
    Returns:
        Formatted prompt string
    """
    template = load_prompt(prompt_name)

    # Replace placeholders with provided values
    for key, value in kwargs.items():
        placeholder = "{" + key + "}"
        if placeholder in template:
            template = template.replace(placeholder, str(value))

    return template


# System prompt - loaded once and cached
_system_prompt_cache: str | None = None


def get_system_prompt() -> str:
    """Get the main system prompt.
    
    Returns:
        System prompt content
    """
    global _system_prompt_cache
    if _system_prompt_cache is None:
        _system_prompt_cache = load_prompt("00_system_prompt")
    return _system_prompt_cache


# Pre-defined prompt names for easy access
class PromptNames:
    """Available prompt template names."""

    SYSTEM = "00_system_prompt"
    INTENT_PARSER = "01_intent_parser_prompt"
    WORKFLOW_PLANNER = "02_workflow_planner_prompt"
    TOOL_SELECTION = "03_tool_selection_prompt"
    REACT_REASONING = "04_react_reasoning_prompt"
    CHAIN_OF_THOUGHT = "05_chain_of_thought_prompt"
    SEARCH_QUERY_GEN = "06_search_query_generation_prompt"
    SEARCH_RESULT_ANALYSIS = "07_search_result_analysis_prompt"
    RESPONSE_SYNTHESIS = "08_response_synthesis_prompt"
    REFLECTION = "09_reflection_prompt"
    ERROR_HANDLING = "10_error_handling_prompt"
    CLARIFICATION = "11_clarification_prompt"
    FOLLOW_UP_GEN = "12_follow_up_generation_prompt"