fast-agent-mcp 0.4.7__py3-none-any.whl

fast_agent/mcp/prompts/prompt_server.py

@@ -0,0 +1,552 @@
+"""
+FastMCP Prompt Server V2
+
+A server that loads prompts from text files with simple delimiters and serves them via MCP.
+Uses the prompt_template module for clean, testable handling of prompt templates.
+"""
+
+import argparse
+import asyncio
+import base64
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Union
+
+from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp.prompts.base import (
+    AssistantMessage,
+    Message,
+    UserMessage,
+)
+from mcp.server.fastmcp.resources import FileResource
+from mcp.types import PromptMessage
+from pydantic import AnyUrl
+
+from fast_agent.mcp import mime_utils, resource_utils
+from fast_agent.mcp.prompts.prompt_constants import (
+    ASSISTANT_DELIMITER as DEFAULT_ASSISTANT_DELIMITER,
+)
+from fast_agent.mcp.prompts.prompt_constants import (
+    RESOURCE_DELIMITER as DEFAULT_RESOURCE_DELIMITER,
+)
+from fast_agent.mcp.prompts.prompt_constants import (
+    USER_DELIMITER as DEFAULT_USER_DELIMITER,
+)
+from fast_agent.mcp.prompts.prompt_load import create_messages_with_resources
+from fast_agent.mcp.prompts.prompt_template import (
+    PromptMetadata,
+    PromptTemplateLoader,
+)
+from fast_agent.types import PromptMessageExtended
+
+# Configure logging
+logging.basicConfig(level=logging.ERROR)
+logger = logging.getLogger("prompt_server")
+
+# Create FastMCP server
+mcp = FastMCP("Prompt Server")
+
+
+def convert_to_fastmcp_messages(prompt_messages: list[Union[PromptMessage, PromptMessageExtended]]) -> list[Message]:
+    """
+    Convert PromptMessage or PromptMessageExtended objects to FastMCP Message objects.
+    This adapter prevents double-wrapping of messages and handles both types.
+
+    Args:
+        prompt_messages: List of PromptMessage or PromptMessageExtended objects
+
+    Returns:
+        List of FastMCP Message objects
+    """
+    result = []
+
+    for msg in prompt_messages:
+        # Handle both PromptMessage and PromptMessageExtended
+        if hasattr(msg, 'from_multipart'):
+            # PromptMessageExtended - convert to regular PromptMessage format
+            flat_messages = msg.from_multipart()
+            for flat_msg in flat_messages:
+                if flat_msg.role == "user":
+                    result.append(UserMessage(content=flat_msg.content))
+                elif flat_msg.role == "assistant":
+                    result.append(AssistantMessage(content=flat_msg.content))
+                else:
+                    logger.warning(f"Unknown message role: {flat_msg.role}, defaulting to user")
+                    result.append(UserMessage(content=flat_msg.content))
+        else:
+            # Regular PromptMessage - use directly
+            if msg.role == "user":
+                result.append(UserMessage(content=msg.content))
+            elif msg.role == "assistant":
+                result.append(AssistantMessage(content=msg.content))
+            else:
+                logger.warning(f"Unknown message role: {msg.role}, defaulting to user")
+                result.append(UserMessage(content=msg.content))
+
+    return result
+
+
+class PromptConfig(PromptMetadata):
+    """Configuration for the prompt server"""
+
+    prompt_files: list[Path] = []
+    user_delimiter: str = DEFAULT_USER_DELIMITER
+    assistant_delimiter: str = DEFAULT_ASSISTANT_DELIMITER
+    resource_delimiter: str = DEFAULT_RESOURCE_DELIMITER
+    http_timeout: float = 10.0
+    transport: str = "stdio"
+    port: int = 8000
+    host: str = "0.0.0.0"
+
+
+# We'll maintain registries of all exposed resources and prompts
+exposed_resources: dict[str, Path] = {}
+prompt_registry: dict[str, PromptMetadata] = {}
+
+
+# Define a single type for prompt handlers to avoid mypy issues
+PromptHandler = Callable[..., Awaitable[list[Message]]]
+
+
+# Type for resource handler
+ResourceHandler = Callable[[], Awaitable[str | bytes]]
+
+
+def create_resource_handler(resource_path: Path, mime_type: str) -> ResourceHandler:
+    """Create a resource handler function for the given resource"""
+
+    async def get_resource() -> str | bytes:
+        is_binary = mime_utils.is_binary_content(mime_type)
+
+        if is_binary:
+            # For binary files, read in binary mode and base64 encode
+            with open(resource_path, "rb") as f:
+                return f.read()
+        else:
+            # For text files, read as utf-8 text
+            with open(resource_path, "r", encoding="utf-8") as f:
+                return f.read()
+
+    return get_resource
+
+
+def get_delimiter_config(
+    config: PromptConfig | None = None, file_path: Path | None = None
+) -> dict[str, Any]:
+    """Get delimiter configuration, falling back to defaults if config is None"""
+    # Set defaults
+    config_values = {
+        "user_delimiter": DEFAULT_USER_DELIMITER,
+        "assistant_delimiter": DEFAULT_ASSISTANT_DELIMITER,
+        "resource_delimiter": DEFAULT_RESOURCE_DELIMITER,
+        "prompt_files": [file_path] if file_path else [],
+    }
+
+    # Override with config values if available
+    if config is not None:
+        config_values["user_delimiter"] = config.user_delimiter
+        config_values["assistant_delimiter"] = config.assistant_delimiter
+        config_values["resource_delimiter"] = config.resource_delimiter
+        config_values["prompt_files"] = config.prompt_files
+
+    return config_values
+
+
+def register_prompt(file_path: Path, config: PromptConfig | None = None) -> None:
+    """Register a prompt file"""
+    try:
+        # Check if it's a JSON file for ultra-minimal path
+        file_str = str(file_path).lower()
+        if file_str.endswith(".json"):
+            # Simple JSON handling - just load and register directly
+            from mcp.server.fastmcp.prompts.base import Prompt, PromptArgument
+
+            from fast_agent.mcp.prompts.prompt_load import load_prompt
+
+            # Create metadata with minimal information
+            metadata = PromptMetadata(
+                name=file_path.stem,
+                description=f"JSON prompt: {file_path.stem}",
+                template_variables=set(),
+                resource_paths=[],  # Skip resource handling
+                file_path=file_path,
+            )
+
+            # Ensure unique name
+            prompt_name = metadata.name
+            if prompt_name in prompt_registry:
+                base_name = prompt_name
+                suffix = 1
+                while prompt_name in prompt_registry:
+                    prompt_name = f"{base_name}_{suffix}"
+                    suffix += 1
+                metadata.name = prompt_name
+
+            prompt_registry[metadata.name] = metadata
+
+            # Create a simple handler that directly loads the JSON file each time
+            async def json_prompt_handler():
+                # Load the messages from the JSON file
+                messages = load_prompt(file_path)
+                # Convert to FastMCP format
+                return convert_to_fastmcp_messages(messages)
+
+            # Register directly with MCP
+            prompt = Prompt(
+                name=metadata.name,
+                description=metadata.description,
+                arguments=[],  # No arguments for JSON prompts
+                fn=json_prompt_handler,
+            )
+            mcp._prompt_manager.add_prompt(prompt)
+
+            logger.info(f"Registered JSON prompt: {metadata.name} ({file_path})")
+            return  # Early return - we're done with JSON files
+
+        # For non-JSON files, continue with the standard approach
+        # Get delimiter configuration
+        config_values = get_delimiter_config(config, file_path)
+
+        # Use standard template loader for text files
+        loader = PromptTemplateLoader(
+            {
+                config_values["user_delimiter"]: "user",
+                config_values["assistant_delimiter"]: "assistant",
+                config_values["resource_delimiter"]: "resource",
+            }
+        )
+
+        # Get metadata and load the template
+        metadata = loader.get_metadata(file_path)
+        template = loader.load_from_file(file_path)
+
+        # Ensure unique name
+        prompt_name = metadata.name
+        if prompt_name in prompt_registry:
+            base_name = prompt_name
+            suffix = 1
+            while prompt_name in prompt_registry:
+                prompt_name = f"{base_name}_{suffix}"
+                suffix += 1
+            metadata.name = prompt_name
+
+        prompt_registry[metadata.name] = metadata
+        logger.info(f"Registered prompt: {metadata.name} ({file_path})")
+
+        # Create and register prompt handler
+        template_vars = list(metadata.template_variables)
+
+        from mcp.server.fastmcp.prompts.base import Prompt, PromptArgument
+
+        # For prompts with variables, create arguments list for FastMCP
+        if template_vars:
+            # Create a function with properly typed parameters
+            async def template_handler_with_vars(**kwargs):
+                # Extract template variables from kwargs
+                context = {var: kwargs.get(var) for var in template_vars if var in kwargs}
+
+                # Check for missing variables
+                missing_vars = [var for var in template_vars if var not in context]
+                if missing_vars:
+                    raise ValueError(
+                        f"Missing required template variables: {', '.join(missing_vars)}"
+                    )
+
+                # Apply template and create messages
+                content_sections = template.apply_substitutions(context)
+                prompt_messages = create_messages_with_resources(
+                    content_sections, config_values["prompt_files"]
+                )
+                return convert_to_fastmcp_messages(prompt_messages)
+
+            # Create a Prompt directly
+            arguments = [
+                PromptArgument(name=var, description=f"Template variable: {var}", required=True)
+                for var in template_vars
+            ]
+
+            # Create and add the prompt directly to the prompt manager
+            prompt = Prompt(
+                name=metadata.name,
+                description=metadata.description,
+                arguments=arguments,
+                fn=template_handler_with_vars,
+            )
+            mcp._prompt_manager.add_prompt(prompt)
+        else:
+            # Create a simple prompt without variables
+            async def template_handler_without_vars() -> list[Message]:
+                content_sections = template.content_sections
+                prompt_messages = create_messages_with_resources(
+                    content_sections, config_values["prompt_files"]
+                )
+                return convert_to_fastmcp_messages(prompt_messages)
+
+            # Create a Prompt object directly instead of using the decorator
+            prompt = Prompt(
+                name=metadata.name,
+                description=metadata.description,
+                arguments=[],
+                fn=template_handler_without_vars,
+            )
+            mcp._prompt_manager.add_prompt(prompt)
+
+        # Register any referenced resources in the prompt
+        for resource_path in metadata.resource_paths:
+            if not resource_path.startswith(("http://", "https://")):
+                # It's a local resource
+                resource_file = file_path.parent / resource_path
+                if resource_file.exists():
+                    resource_id = f"resource://fast-agent/{resource_file.name}"
+
+                    # Register the resource if not already registered
+                    if resource_id not in exposed_resources:
+                        exposed_resources[resource_id] = resource_file
+                        mime_type = mime_utils.guess_mime_type(str(resource_file))
+
+                        mcp.add_resource(
+                            FileResource(
+                                uri=AnyUrl(resource_id),
+                                path=resource_file,
+                                mime_type=mime_type,
+                                is_binary=mime_utils.is_binary_content(mime_type),
+                            )
+                        )
+
+                        logger.info(f"Registered resource: {resource_id} ({resource_file})")
+    except Exception as e:
+        logger.error(f"Error registering prompt {file_path}: {e}", exc_info=True)
+
+
+def parse_args():
+    """Parse command line arguments"""
+    parser = argparse.ArgumentParser(description="FastMCP Prompt Server")
+    parser.add_argument("prompt_files", nargs="+", type=str, help="Prompt files to serve")
+    parser.add_argument(
+        "--user-delimiter",
+        type=str,
+        default="---USER",
+        help="Delimiter for user messages (default: ---USER)",
+    )
+    parser.add_argument(
+        "--assistant-delimiter",
+        type=str,
+        default="---ASSISTANT",
+        help="Delimiter for assistant messages (default: ---ASSISTANT)",
+    )
+    parser.add_argument(
+        "--resource-delimiter",
+        type=str,
+        default="---RESOURCE",
+        help="Delimiter for resource references (default: ---RESOURCE)",
+    )
+    parser.add_argument(
+        "--http-timeout",
+        type=float,
+        default=10.0,
+        help="Timeout for HTTP requests in seconds (default: 10.0)",
+    )
+    parser.add_argument(
+        "--transport",
+        type=str,
+        choices=["stdio", "sse", "http"],
+        default="stdio",
+        help="Transport to use (default: stdio)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to use for SSE transport (default: 8000)",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to bind to for SSE transport (default: 0.0.0.0)",
+    )
+    parser.add_argument(
+        "--test", type=str, help="Test a specific prompt without starting the server"
+    )
+
+    return parser.parse_args()
+
+
+def initialize_config(args) -> PromptConfig:
+    """Initialize configuration from command line arguments"""
+    # Resolve file paths
+    prompt_files = []
+    for file_path in args.prompt_files:
+        path = Path(file_path)
+        if not path.exists():
+            logger.warning(f"File not found: {path}")
+            continue
+        prompt_files.append(path.resolve())
+
+    if not prompt_files:
+        logger.error("No valid prompt files specified")
+        raise ValueError("No valid prompt files specified")
+
+    # Initialize configuration
+    return PromptConfig(
+        name="prompt_server",
+        description="FastMCP Prompt Server",
+        template_variables=set(),
+        resource_paths=[],
+        file_path=Path(__file__),
+        prompt_files=prompt_files,
+        user_delimiter=args.user_delimiter,
+        assistant_delimiter=args.assistant_delimiter,
+        resource_delimiter=args.resource_delimiter,
+        http_timeout=args.http_timeout,
+        transport=args.transport,
+        port=args.port,
+        host=args.host,
+    )
+
+
+async def register_file_resource_handler(config: PromptConfig) -> None:
+    """Register the general file resource handler"""
+
+    @mcp.resource("file://{path}")
+    async def get_file_resource(path: str):
+        """Read a file from the given path."""
+        try:
+            # Find the file, checking relative paths first
+            file_path = resource_utils.find_resource_file(path, config.prompt_files)
+            if file_path is None:
+                # If not found as relative path, try absolute path
+                file_path = Path(path)
+                if not file_path.exists():
+                    raise FileNotFoundError(f"Resource file not found: {path}")
+
+            mime_type = mime_utils.guess_mime_type(str(file_path))
+            is_binary = mime_utils.is_binary_content(mime_type)
+
+            if is_binary:
+                # For binary files, read as binary and base64 encode
+                with open(file_path, "rb") as f:
+                    return base64.b64encode(f.read()).decode("utf-8")
+            else:
+                # For text files, read as text with UTF-8 encoding
+                with open(file_path, "r", encoding="utf-8") as f:
+                    return f.read()
+        except Exception as e:
+            # Log the error and re-raise
+            logger.error(f"Error accessing resource at '{path}': {e}")
+            raise
+
+
+async def test_prompt(prompt_name: str, config: PromptConfig) -> int:
+    """Test a prompt and print its details"""
+    if prompt_name not in prompt_registry:
+        logger.error(f"Test prompt not found: {prompt_name}")
+        return 1
+
+    # Get delimiter configuration with reasonable defaults
+    config_values = get_delimiter_config(config)
+
+    metadata = prompt_registry[prompt_name]
+    print(f"\nTesting prompt: {prompt_name}")
+    print(f"Description: {metadata.description}")
+    print(f"Template variables: {', '.join(metadata.template_variables)}")
+
+    # Load and print the template
+    loader = PromptTemplateLoader(
+        {
+            config_values["user_delimiter"]: "user",
+            config_values["assistant_delimiter"]: "assistant",
+            config_values["resource_delimiter"]: "resource",
+        }
+    )
+    template = loader.load_from_file(metadata.file_path)
+
+    # Print each content section
+    print("\nContent sections:")
+    for i, section in enumerate(template.content_sections):
+        print(f"\n[{i + 1}] Role: {section.role}")
+        print(f"Content: {section.text}")
+        if section.resources:
+            print(f"Resources: {', '.join(section.resources)}")
+
+    # If there are template variables, test with dummy values
+    if metadata.template_variables:
+        print("\nTemplate substitution test:")
+        test_context = {var: f"[TEST-{var}]" for var in metadata.template_variables}
+        applied = template.apply_substitutions(test_context)
+
+        for i, section in enumerate(applied):
+            print(f"\n[{i + 1}] Role: {section.role}")
+            print(f"Content with substitutions: {section.text}")
+            if section.resources:
+                print(f"Resources with substitutions: {', '.join(section.resources)}")
+
+    return 0
+
+
+async def async_main() -> int:
+    """Run the FastMCP server (async version)"""
+    # Parse command line arguments
+    args = parse_args()
+
+    try:
+        # Initialize configuration
+        config = initialize_config(args)
+    except ValueError as e:
+        logger.error(str(e))
+        return 1
+
+    # Register resource handlers
+    await register_file_resource_handler(config)
+
+    # Register all prompts
+    for file_path in config.prompt_files:
+        register_prompt(file_path, config)
+
+    # Print startup info
+    logger.info("Starting prompt server")
+    logger.info(f"Registered {len(prompt_registry)} prompts")
+    logger.info(f"Registered {len(exposed_resources)} resources")
+    logger.info(
+        f"Using delimiters: {config.user_delimiter}, {config.assistant_delimiter}, {config.resource_delimiter}"
+    )
+
+    # If a test prompt was specified, print it and exit
+    if args.test:
+        return await test_prompt(args.test, config)
+
+    # Start the server with the specified transport
+    if config.transport == "sse":  # sse
+        # Set the host and port in settings before running the server
+        mcp.settings.host = config.host
+        mcp.settings.port = config.port
+        logger.info(f"Starting SSE server on {config.host}:{config.port}")
+        await mcp.run_sse_async()
+    elif config.transport == "http":
+        mcp.settings.host = config.host
+        mcp.settings.port = config.port
+        logger.info(f"Starting SSE server on {config.host}:{config.port}")
+        await mcp.run_streamable_http_async()
+    elif config.transport == "stdio":
+        await mcp.run_stdio_async()
+    else:
+        logger.error(f"Unknown transport: {config.transport}")
+        return 1
+    return 0
+
+
+def main() -> int:
+    """Run the FastMCP server"""
+    try:
+        return asyncio.run(async_main())
+    except KeyboardInterrupt:
+        logger.info("\nServer stopped by user")
+    except Exception as e:
+        logger.error(f"\nError: {e}", exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())