Optimized Memory MCP Server V2

"""Base classes for FastMCP prompts."""

import inspect
import json
from collections.abc import Callable
from typing import Any, Awaitable, Literal, Sequence

import pydantic_core
from pydantic import BaseModel, Field, TypeAdapter, validate_call

from mcp.types import EmbeddedResource, ImageContent, TextContent

CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource


class Message(BaseModel):
    """Base class for all prompt messages."""

    role: Literal["user", "assistant"]
    content: CONTENT_TYPES

    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
        if isinstance(content, str):
            content = TextContent(type="text", text=content)
        super().__init__(content=content, **kwargs)


class UserMessage(Message):
    """A message from the user."""

    role: Literal["user", "assistant"] = "user"

    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
        super().__init__(content=content, **kwargs)


class AssistantMessage(Message):
    """A message from the assistant."""

    role: Literal["user", "assistant"] = "assistant"

    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
        super().__init__(content=content, **kwargs)


message_validator = TypeAdapter(UserMessage | AssistantMessage)

SyncPromptResult = (
    str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]]
)
PromptResult = SyncPromptResult | Awaitable[SyncPromptResult]


class PromptArgument(BaseModel):
    """An argument that can be passed to a prompt."""

    name: str = Field(description="Name of the argument")
    description: str | None = Field(
        None, description="Description of what the argument does"
    )
    required: bool = Field(
        default=False, description="Whether the argument is required"
    )


class Prompt(BaseModel):
    """A prompt template that can be rendered with parameters."""

    name: str = Field(description="Name of the prompt")
    description: str | None = Field(
        None, description="Description of what the prompt does"
    )
    arguments: list[PromptArgument] | None = Field(
        None, description="Arguments that can be passed to the prompt"
    )
    fn: Callable = Field(exclude=True)

    @classmethod
    def from_function(
        cls,
        fn: Callable[..., PromptResult],
        name: str | None = None,
        description: str | None = None,
    ) -> "Prompt":
        """Create a Prompt from a function.

        The function can return:
        - A string (converted to a message)
        - A Message object
        - A dict (converted to a message)
        - A sequence of any of the above
        """
        func_name = name or fn.__name__

        if func_name == "<lambda>":
            raise ValueError("You must provide a name for lambda functions")

        # Get schema from TypeAdapter - will fail if function isn't properly typed
        parameters = TypeAdapter(fn).json_schema()

        # Convert parameters to PromptArguments
        arguments = []
        if "properties" in parameters:
            for param_name, param in parameters["properties"].items():
                required = param_name in parameters.get("required", [])
                arguments.append(
                    PromptArgument(
                        name=param_name,
                        description=param.get("description"),
                        required=required,
                    )
                )

        # ensure the arguments are properly cast
        fn = validate_call(fn)

        return cls(
            name=func_name,
            description=description or fn.__doc__ or "",
            arguments=arguments,
            fn=fn,
        )

    async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]:
        """Render the prompt with arguments."""
        # Validate required arguments
        if self.arguments:
            required = {arg.name for arg in self.arguments if arg.required}
            provided = set(arguments or {})
            missing = required - provided
            if missing:
                raise ValueError(f"Missing required arguments: {missing}")

        try:
            # Call function and check if result is a coroutine
            result = self.fn(**(arguments or {}))
            if inspect.iscoroutine(result):
                result = await result

            # Validate messages
            if not isinstance(result, (list, tuple)):
                result = [result]

            # Convert result to messages
            messages = []
            for msg in result:
                try:
                    if isinstance(msg, Message):
                        messages.append(msg)
                    elif isinstance(msg, dict):
                        msg = message_validator.validate_python(msg)
                        messages.append(msg)
                    elif isinstance(msg, str):
                        messages.append(
                            UserMessage(content=TextContent(type="text", text=msg))
                        )
                    else:
                        msg = json.dumps(pydantic_core.to_jsonable_python(msg))
                        messages.append(Message(role="user", content=msg))
                except Exception:
                    raise ValueError(
                        f"Could not convert prompt result to message: {msg}"
                    )

            return messages
        except Exception as e:
            raise ValueError(f"Error rendering prompt {self.name}: {e}")