fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl

mcp_agent/mcp/prompts/prompt_server.py

@@ -5,187 +5,94 @@ A server that loads prompts from text files with simple delimiters and serves th
 Uses the prompt_template module for clean, testable handling of prompt templates.
 """
 
-import asyncio
 import argparse
+import asyncio
 import base64
 import logging
 import sys
 from pathlib import Path
-from typing import List, Dict, Optional, Callable, Awaitable, Literal, Any
-from mcp.server.fastmcp.resources import FileResource
-from pydantic import AnyUrl
-
-from mcp_agent.mcp import mime_utils, resource_utils
+from typing import Any, Awaitable, Callable, Dict, List, Optional
 
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.prompts.base import (
-    UserMessage,
     AssistantMessage,
     Message,
+    UserMessage,
 )
-from mcp.types import (
-    TextContent,
-)
+from mcp.server.fastmcp.resources import FileResource
+from mcp.types import PromptMessage
+from pydantic import AnyUrl
 
+from mcp_agent.mcp import mime_utils, resource_utils
+from mcp_agent.mcp.prompts.prompt_constants import (
+    ASSISTANT_DELIMITER as DEFAULT_ASSISTANT_DELIMITER,
+)
+from mcp_agent.mcp.prompts.prompt_constants import (
+    RESOURCE_DELIMITER as DEFAULT_RESOURCE_DELIMITER,
+)
+from mcp_agent.mcp.prompts.prompt_constants import (
+    USER_DELIMITER as DEFAULT_USER_DELIMITER,
+)
+from mcp_agent.mcp.prompts.prompt_load import create_messages_with_resources
 from mcp_agent.mcp.prompts.prompt_template import (
-    PromptTemplateLoader,
     PromptMetadata,
-    PromptContent,
-    PromptTemplate,
+    PromptTemplateLoader,
 )
 
 # Configure logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.ERROR)
 logger = logging.getLogger("prompt_server")
 
 # Create FastMCP server
 mcp = FastMCP("Prompt Server")
 
 
-class PromptConfig(PromptMetadata):
-    """Configuration for the prompt server"""
-
-    prompt_files: List[Path] = []
-    user_delimiter: str = "---USER"
-    assistant_delimiter: str = "---ASSISTANT"
-    resource_delimiter: str = "---RESOURCE"
-    http_timeout: float = 10.0
-    transport: str = "stdio"
-    port: int = 8000
-
-
-# Will be initialized with command line args
-config = None
-
-# We'll maintain registries of all exposed resources and prompts
-exposed_resources: Dict[str, Path] = {}
-prompt_registry: Dict[str, PromptMetadata] = {}
-
-# Define message role type
-MessageRole = Literal["user", "assistant"]
-
-
-def create_content_message(text: str, role: MessageRole) -> Message:
-    """Create a text content message with the specified role"""
-    message_class = UserMessage if role == "user" else AssistantMessage
-    return message_class(content=TextContent(type="text", text=text))
-
-
-def create_resource_message(
-    resource_path: str, content: str, mime_type: str, is_binary: bool, role: MessageRole
-) -> Message:
-    """Create a resource message with the specified content and role"""
-    message_class = UserMessage if role == "user" else AssistantMessage
-
-    if mime_utils.is_image_mime_type(mime_type):
-        # For images, create an ImageContent
-        image_content = resource_utils.create_image_content(
-            data=content, mime_type=mime_type
-        )
-        return message_class(content=image_content)
-    else:
-        # For other resources, create an EmbeddedResource
-        embedded_resource = resource_utils.create_embedded_resource(
-            resource_path, content, mime_type, is_binary
-        )
-        return message_class(content=embedded_resource)
-
-
-def create_messages_with_resources(
-    content_sections: List[PromptContent], prompt_files: List[Path]
-) -> List[Message]:
+def convert_to_fastmcp_messages(prompt_messages: List[PromptMessage]) -> List[Message]:
     """
-    Create a list of messages from content sections, with resources properly handled.
-
-    This implementation produces one message for each content section's text,
-    followed by separate messages for each resource (with the same role type
-    as the section they belong to).
+    Convert PromptMessage objects from prompt_load to FastMCP Message objects.
+    This adapter prevents double-wrapping of messages.
 
     Args:
-        content_sections: List of PromptContent objects
-        prompt_files: List of prompt files (to help locate resource files)
+        prompt_messages: List of PromptMessage objects from prompt_load
 
     Returns:
-        List of Message objects
+        List of FastMCP Message objects
     """
-    messages = []
+    result = []
 
-    for section in content_sections:
-        # Convert to our literal type for role
-        role = cast_message_role(section.role)
+    for msg in prompt_messages:
+        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))
 
-        # Add the text message
-        messages.append(create_content_message(section.text, role))
+    return result
 
-        # Add resource messages with the same role type as the section
-        for resource_path in section.resources:
-            try:
-                # Load resource with information about its type
-                resource_content, mime_type, is_binary = (
-                    resource_utils.load_resource_content(resource_path, prompt_files)
-                )
 
-                # Create and add the resource message
-                resource_message = create_resource_message(
-                    resource_path, resource_content, mime_type, is_binary, role
-                )
-                messages.append(resource_message)
-            except Exception as e:
-                logger.error(f"Error loading resource {resource_path}: {e}")
+class PromptConfig(PromptMetadata):
+    """Configuration for the prompt server"""
 
-    return messages
+    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
 
 
-def cast_message_role(role: str) -> MessageRole:
-    """Cast a string role to a MessageRole literal type"""
-    if role == "user" or role == "assistant":
-        return role  # type: ignore
-    # Default to user if the role is invalid
-    logger.warning(f"Invalid message role: {role}, defaulting to 'user'")
-    return "user"
+# 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]]]
 
 
-def create_prompt_handler(
-    template: "PromptTemplate", template_vars: List[str], prompt_files: List[Path]
-) -> PromptHandler:
-    """Create a prompt handler function for the given template"""
-    if template_vars:
-        # With template variables
-        docstring = f"Prompt with template variables: {', '.join(template_vars)}"
-
-        async def prompt_handler(**kwargs: Any) -> List[Message]:
-            # Build context from parameters
-            context = {
-                var: kwargs.get(var)
-                for var in template_vars
-                if var in kwargs and kwargs[var] is not None
-            }
-
-            # Apply substitutions to the template
-            content_sections = template.apply_substitutions(context)
-
-            # Convert to MCP Message objects, handling resources properly
-            return create_messages_with_resources(content_sections, prompt_files)
-    else:
-        # No template variables
-        docstring = "Get a prompt with no variable substitution"
-
-        async def prompt_handler(**kwargs: Any) -> List[Message]:
-            # Get the content sections
-            content_sections = template.content_sections
-
-            # Convert to MCP Message objects, handling resources properly
-            return create_messages_with_resources(content_sections, prompt_files)
-
-    # Set the docstring
-    prompt_handler.__doc__ = docstring
-    return prompt_handler
-
-
 # Type for resource handler
 ResourceHandler = Callable[[], Awaitable[str | bytes]]
 
@@ -208,13 +115,9 @@ def create_resource_handler(resource_path: Path, mime_type: str) -> ResourceHand
     return get_resource
 
 
-# Default delimiter values
-DEFAULT_USER_DELIMITER = "---USER"
-DEFAULT_ASSISTANT_DELIMITER = "---ASSISTANT"
-DEFAULT_RESOURCE_DELIMITER = "---RESOURCE"
-
-
-def get_delimiter_config(file_path: Optional[Path] = None) -> Dict[str, Any]:
+def get_delimiter_config(
+    config: Optional[PromptConfig] = None, file_path: Optional[Path] = None
+) -> Dict[str, Any]:
     """Get delimiter configuration, falling back to defaults if config is None"""
     # Set defaults
     config_values = {
@@ -234,11 +137,11 @@ def get_delimiter_config(file_path: Optional[Path] = None) -> Dict[str, Any]:
     return config_values
 
 
-def register_prompt(file_path: Path):
+def register_prompt(file_path: Path, config: Optional[PromptConfig] = None) -> None:
     """Register a prompt file"""
     try:
         # Get delimiter configuration
-        config_values = get_delimiter_config(file_path)
+        config_values = get_delimiter_config(config, file_path)
 
         # Use our prompt template loader to analyze the file
         loader = PromptTemplateLoader(
@@ -268,10 +171,61 @@ def register_prompt(file_path: Path):
 
         # Create and register prompt handler
         template_vars = list(metadata.template_variables)
-        handler = create_prompt_handler(
-            template, template_vars, config_values["prompt_files"]
-        )
-        mcp.prompt(name=metadata.name, description=metadata.description)(handler)
+
+        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:
@@ -295,9 +249,7 @@ def register_prompt(file_path: Path):
                             )
                         )
 
-                        logger.info(
-                            f"Registered resource: {resource_id} ({resource_file})"
-                        )
+                        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)
 
@@ -305,9 +257,7 @@ def register_prompt(file_path: Path):
 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("prompt_files", nargs="+", type=str, help="Prompt files to serve")
     parser.add_argument(
         "--user-delimiter",
         type=str,
@@ -352,7 +302,39 @@ def parse_args():
     return parser.parse_args()
 
 
-async def register_file_resource_handler():
+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,
+    )
+
+
+async def register_file_resource_handler(config: PromptConfig) -> None:
     """Register the general file resource handler"""
 
     @mcp.resource("file://{path}")
@@ -384,14 +366,14 @@ async def register_file_resource_handler():
             raise
 
 
-async def test_prompt(prompt_name: str) -> int:
+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_values = get_delimiter_config(config)
 
     metadata = prompt_registry[prompt_name]
     print(f"\nTesting prompt: {prompt_name}")
@@ -431,48 +413,24 @@ async def test_prompt(prompt_name: str) -> int:
     return 0
 
 
-async def async_main():
+async def async_main() -> int:
     """Run the FastMCP server (async version)"""
-    global config
-
     # Parse command line arguments
     args = parse_args()
 
-    # 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")
+    try:
+        # Initialize configuration
+        config = initialize_config(args)
+    except ValueError as e:
+        logger.error(str(e))
         return 1
 
-    # Initialize configuration
-    config = 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,
-    )
-
     # Register resource handlers
-    await register_file_resource_handler()
+    await register_file_resource_handler(config)
 
     # Register all prompts
     for file_path in config.prompt_files:
-        register_prompt(file_path)
+        register_prompt(file_path, config)
 
     # Print startup info
     logger.info("Starting prompt server")
@@ -484,13 +442,15 @@ async def async_main():
 
     # If a test prompt was specified, print it and exit
     if args.test:
-        return await test_prompt(args.test)
+        return await test_prompt(args.test, config)
 
     # Start the server with the specified transport
     if config.transport == "stdio":
         await mcp.run_stdio_async()
     else:  # sse
-        await mcp.run_sse_async(port=config.port)
+        # TODO update to 2025-03-26 specification and test config.
+        await mcp.run_sse_async()
+    return 0
 
 
 def main() -> int: