Building Useful Prompts

"""
prompts.py - Reusable prompt registry for MCP servers.

This module builds on jsonrpc.py from Lesson 3.
"""

from typing import Callable

# Import error handling from Lesson 3
from jsonrpc import JsonRpcError, JsonRpcErrorCode


# =============================================================================
# Section 1: Prompt Registry
# =============================================================================

class PromptRegistry:
    """
    Registry for MCP prompts with decorator-based registration.

    Usage:
        prompts = PromptRegistry()

        @prompts.register(
            name="greet",
            description="Generate a greeting",
            arguments=[{"name": "name", "description": "Name to greet", "required": True}]
        )
        def greet_prompt(args: dict) -> dict:
            name = args.get("name", "World")
            return {
                "description": f"Greeting for {name}",
                "messages": [{"role": "user", "content": {"type": "text", "text": f"Hello, {name}!"}}]
            }
    """

    def __init__(self):
        self._prompts: dict[str, dict] = {}
        self._generators: dict[str, Callable] = {}

    def register(
        self,
        name: str,
        description: str,
        arguments: list[dict] | None = None
    ) -> Callable:
        """Decorator to register a prompt generator."""
        def decorator(func: Callable[[dict], dict]) -> Callable:
            self._prompts[name] = {
                "name": name,
                "description": description,
                "arguments": arguments or []
            }
            self._generators[name] = func
            return func
        return decorator

    def list_prompts(self) -> list[dict]:
        """Return all prompt definitions for prompts/list."""
        return list(self._prompts.values())

    def get_prompt(self, name: str, arguments: dict) -> dict:
        """Generate a prompt by name with arguments."""
        if name not in self._generators:
            raise JsonRpcError(
                JsonRpcErrorCode.INVALID_PARAMS,
                f"Unknown prompt: {name}"
            )

        # Validate required arguments
        prompt_def = self._prompts[name]
        for arg in prompt_def.get("arguments", []):
            if arg.get("required") and arg["name"] not in arguments:
                raise JsonRpcError(
                    JsonRpcErrorCode.INVALID_PARAMS,
                    f"Missing required argument: {arg['name']}"
                )

        return self._generators[name](arguments)

# =============================================================================
# Section 2: MCP Router Integration
# =============================================================================

def register_prompt_handlers(router, prompts: PromptRegistry) -> None:
    """
    Register prompts/list and prompts/get handlers with an MCP router.

    Args:
        router: A JsonRpcRouter from mcp_core.py
        prompts: A PromptRegistry with registered prompts
    """

    @router.method("prompts/list")
    def handle_prompts_list(params: dict) -> dict:
        return {"prompts": prompts.list_prompts()}

    @router.method("prompts/get")
    def handle_prompts_get(params: dict) -> dict:
        name = params.get("name", "")
        arguments = params.get("arguments", {})
        return prompts.get_prompt(name, arguments)

# =============================================================================
# Test the module
# =============================================================================

if __name__ == "__main__":
    prompts = PromptRegistry()

    @prompts.register(
        name="greet",
        description="Generate a greeting",
        arguments=[{"name": "name", "description": "Name to greet", "required": True}]
    )
    def greet_prompt(args: dict) -> dict:
        name = args.get("name", "World")
        return {
            "description": f"Greeting for {name}",
            "messages": [{"role": "user", "content": {"type": "text", "text": f"Hello, {name}!"}}]
        }

    print("Registered prompts:")
    for p in prompts.list_prompts():
        print(f"  - {p['name']}: {p['description']}")

    print("\nGenerating prompt:")
    result = prompts.get_prompt("greet", {"name": "MCP"})
    print(f"  {result}")

"""
prompts_example.py - Server with useful prompts

Run: python prompts_example.py
"""

# Import from our modules
from mcp_core import router
from prompts import PromptRegistry, register_prompt_handlers
from http_sse_transport import run_http_server

# =============================================================================
# Register prompts
# =============================================================================

prompt_registry = PromptRegistry()


@prompt_registry.register(
    name="code-review",
    description="Review code for issues",
    arguments=[
        {"name": "code", "description": "Code to review", "required": True},
        {"name": "language", "description": "Programming language", "required": True},
        {"name": "focus", "description": "Focus: security, performance, style", "required": False}
    ]
)
def code_review(args: dict) -> dict:
    code = args.get("code", "")
    language = args.get("language", "")
    focus = args.get("focus", "all")

    text = f"""Please review this {language} code, focusing on {focus}:

```{language}
{code}
```

Identify issues and suggest improvements."""

    return {
        "description": f"Code review ({focus})",
        "messages": [{"role": "user", "content": {"type": "text", "text": text}}]
    }


@prompt_registry.register(
    name="summarize",
    description="Summarize text",
    arguments=[
        {"name": "text", "description": "Text to summarize", "required": True},
        {"name": "style", "description": "Style: tldr, bullets, paragraph", "required": False}
    ]
)
def summarize(args: dict) -> dict:
    text = args.get("text", "")
    style = args.get("style", "paragraph")

    styles = {
        "tldr": "Provide a one-sentence TL;DR.",
        "bullets": "Summarize as bullet points.",
        "paragraph": "Write a paragraph summary."
    }

    prompt_text = f"""Summarize this text:

{text}

{styles.get(style, styles['paragraph'])}"""

    return {
        "description": f"Summary ({style})",
        "messages": [{"role": "user", "content": {"type": "text", "text": prompt_text}}]
    }


@prompt_registry.register(
    name="explain-code",
    description="Explain what code does",
    arguments=[
        {"name": "code", "description": "Code to explain", "required": True},
        {"name": "audience", "description": "Audience: beginner, developer, expert", "required": False}
    ]
)
def explain_code(args: dict) -> dict:
    code = args.get("code", "")
    audience = args.get("audience", "developer")

    text = f"""Explain this code for a {audience}:

```
{code}
```

Cover what it does and how it works."""

    return {
        "description": f"Code explanation for {audience}",
        "messages": [{"role": "user", "content": {"type": "text", "text": text}}]
    }


# Register with the MCP router
register_prompt_handlers(router, prompt_registry)


# =============================================================================
# Run the server
# =============================================================================

if __name__ == "__main__":
    run_http_server()