prompthandler.py•2.88 kB
#
# Copyright (C) 2025 wisevision
#
# SPDX-License-Identifier: MPL-2.0
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
#
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Literal
from mcp.types import (
Prompt,
PromptArgument,
PromptMessage,
GetPromptResult,
TextContent
)
@dataclass
class ArgSpec:
name: str
description: str = ""
required: bool = True
type: Literal["string", "boolean", "number"] = "string"
enum: list[Any] | None = None
default: Any = None
@dataclass
class BasePromptHandler:
"""
Implement prompts by specifying:
- name, description
- args: list of ArgSpec
- messages_template: list of (role, text_template) tuples
where text_template can use `{var}` placeholders.
"""
namespace: str
name: str
description: str
args: list[ArgSpec] = field(default_factory=list)
messages_template: list[tuple[str, str]] = field(default_factory=list)
def get_prompt_description(self) -> Prompt:
return Prompt(
name=self.namespace + "." + self.name,
description=self.description,
arguments=[
PromptArgument(
name=a.name,
description=a.description,
required=a.required,
)
for a in self.args
],
)
def render(self, arguments: dict[str, str] | None) -> GetPromptResult:
arguments = arguments or {}
# Validate required args + fill defaults for optional
missing = [a.name for a in self.args if a.required and a.name not in arguments]
if missing:
raise ValueError(
f"Missing required argument(s): {', '.join(missing)} "
f"for prompt '{self.name}'"
)
for a in self.args:
if not a.required and a.name not in arguments and a.default is not None:
arguments[a.name] = a.default
# Render each message with simple str.format
messages: list[PromptMessage] = []
for role, template in self.messages_template:
try:
text = template.format(**arguments)
except KeyError as e:
raise ValueError(
f"Template for '{self.name}' references missing arg: {e}"
) from e
messages.append(
PromptMessage(
role=role,
content=TextContent(type="text", text=text),
)
)
return GetPromptResult(
description=self.description,
messages=messages,
)