alfred/agent/prompts.py

"""Prompt builder for the agent system."""

import json

from infrastructure.persistence import get_memory

from .parameters import format_parameters_for_prompt, get_missing_required_parameters
from .registry import Tool


class PromptBuilder:
    """Builds system prompts for the agent with memory context.

    Attributes:
        tools: Dictionary of available tools.
    """

    def __init__(self, tools: dict[str, Tool]):
        """
        Initialize the prompt builder.

        Args:
            tools: Dictionary mapping tool names to Tool instances.
        """
        self.tools = tools

    def _format_tools_description(self) -> str:
        """Format tools with their descriptions and parameters."""
        return "\n".join(
            f"- {tool.name}: {tool.description}\n"
            f"  Parameters: {json.dumps(tool.parameters, ensure_ascii=False)}"
            for tool in self.tools.values()
        )

    def _format_episodic_context(self) -> str:
        """Format episodic memory context for the prompt."""
        memory = get_memory()
        lines = []

        # Last search results
        if memory.episodic.last_search_results:
            search = memory.episodic.last_search_results
            lines.append(f"LAST SEARCH: '{search.get('query')}'")
            results = search.get("results", [])
            if results:
                lines.append(f"   {len(results)} results available:")
                for r in results[:5]:
                    name = r.get("name", r.get("title", "Unknown"))
                    lines.append(f"   {r.get('index')}. {name}")
                if len(results) > 5:
                    lines.append(f"   ... and {len(results) - 5} more")

        # Pending question
        if memory.episodic.pending_question:
            q = memory.episodic.pending_question
            lines.append(f"\nPENDING QUESTION: {q.get('question')}")
            for opt in q.get("options", []):
                lines.append(f"   {opt.get('index')}. {opt.get('label')}")

        # Active downloads
        if memory.episodic.active_downloads:
            lines.append(f"\nACTIVE DOWNLOADS: {len(memory.episodic.active_downloads)}")
            for dl in memory.episodic.active_downloads[:3]:
                lines.append(f"   - {dl.get('name')}: {dl.get('progress', 0)}%")

        # Recent errors
        if memory.episodic.recent_errors:
            last_error = memory.episodic.recent_errors[-1]
            lines.append(
                f"\nLAST ERROR: {last_error.get('error')} "
                f"(action: {last_error.get('action')})"
            )

        # Unread events
        unread = [e for e in memory.episodic.background_events if not e.get("read")]
        if unread:
            lines.append(f"\nUNREAD EVENTS: {len(unread)}")
            for e in unread[:3]:
                lines.append(f"   - {e.get('type')}: {e.get('data', {})}")

        return "\n".join(lines) if lines else ""

    def _format_stm_context(self) -> str:
        """Format short-term memory context for the prompt."""
        memory = get_memory()
        lines = []

        # Current workflow
        if memory.stm.current_workflow:
            wf = memory.stm.current_workflow
            lines.append(f"CURRENT WORKFLOW: {wf.get('type')}")
            lines.append(f"   Target: {wf.get('target', {}).get('title', 'Unknown')}")
            lines.append(f"   Stage: {wf.get('stage')}")

        # Current topic
        if memory.stm.current_topic:
            lines.append(f"CURRENT TOPIC: {memory.stm.current_topic}")

        # Extracted entities
        if memory.stm.extracted_entities:
            entities_json = json.dumps(
                memory.stm.extracted_entities, ensure_ascii=False
            )
            lines.append(f"EXTRACTED ENTITIES: {entities_json}")

        return "\n".join(lines) if lines else ""

    def build_system_prompt(self) -> str:
        """
        Build the system prompt with context from memory.

        Returns:
            The complete system prompt string.
        """
        memory = get_memory()
        tools_desc = self._format_tools_description()
        params_desc = format_parameters_for_prompt()

        # Check for missing required parameters
        missing_params = get_missing_required_parameters({"config": memory.ltm.config})
        missing_info = ""
        if missing_params:
            missing_info = "\n\nMISSING REQUIRED PARAMETERS:\n"
            for param in missing_params:
                missing_info += f"- {param.key}: {param.description}\n"
                missing_info += f"  Why needed: {param.why_needed}\n"

        # Build context sections
        episodic_context = self._format_episodic_context()
        stm_context = self._format_stm_context()

        config_json = json.dumps(memory.ltm.config, indent=2, ensure_ascii=False)

        return f"""You are an AI agent helping a user manage their local media library.

{params_desc}

CURRENT CONFIGURATION:
{config_json}
{missing_info}

{f"SESSION CONTEXT:{chr(10)}{stm_context}" if stm_context else ""}

{f"CURRENT STATE:{chr(10)}{episodic_context}" if episodic_context else ""}

IMPORTANT RULES:
1. When the user refers to a number (e.g., "the 3rd one", "download number 2"), \
use `add_torrent_by_index` or `get_torrent_by_index` with that number.
2. If a torrent search was performed, results are numbered. \
The user can reference them by number.
3. To use a tool, respond STRICTLY with this JSON format:
   {{ "thought": "explanation", "action": {{ "name": "tool_name", "args": {{ }} }} }}
   - No text before or after the JSON
4. You can use MULTIPLE TOOLS IN SEQUENCE.
5. When you have all the information needed, respond in NATURAL TEXT (not JSON).
6. If a required parameter is missing, ask the user for it.
7. Respond in the same language as the user.

EXAMPLES:
- After a torrent search, if the user says "download the 3rd one":
  {{ "thought": "User wants torrent #3", "action": {{ "name": "add_torrent_by_index", \
"args": {{ "index": 3 }} }} }}

- To search for torrents:
  {{ "thought": "Searching torrents", "action": {{ "name": "find_torrents", \
"args": {{ "media_title": "Inception 1080p" }} }} }}

AVAILABLE TOOLS:
{tools_desc}
"""