chainlit / prompt_server.py
Gustav2811's picture
Better app logging
36619f7
# prompt_server.py
import asyncio
from mcp.server import Server
import mcp.types as types
from mcp.server.stdio import stdio_server
from typing import Any
PROMPTS = {
"simple_prompt": types.Prompt(
name="simple_prompt",
description="A basic prompt that asks the LLM to act like a pirate.",
),
"explain_code": types.Prompt(
name="explain-code",
description="Explain a snippet of code.",
arguments=[
types.PromptArgument(
name="code", description="The code snippet to explain", required=True
),
types.PromptArgument(
name="language", description="The programming language", required=False
),
],
),
}
app = Server("prompt-test-server")
@app.list_prompts()
async def list_prompts() -> list[types.Prompt]:
return list(PROMPTS.values())
@app.get_prompt()
async def get_prompt(
name: str, arguments: dict[str, Any] | None = None
) -> types.GetPromptResult:
arguments = arguments or {}
if name == "simple_prompt":
return types.GetPromptResult(
messages=[
types.PromptMessage(
role="system",
content=types.TextContent(
text="You are a helpful pirate assistant. You always answer in pirate slang."
),
),
types.PromptMessage(
role="user",
content=types.TextContent(
text="What are the benefits of using Python?"
),
),
]
)
if name == "explain-code":
code = arguments.get("code", "")
language = arguments.get("language", "unknown")
return types.GetPromptResult(
messages=[
types.PromptMessage(
role="user",
content=types.TextContent(
text=f"Explain this {language} code snippet to a five-year-old:\n\n```\n{code}\n```"
),
)
]
)
raise ValueError("Prompt not found")
async def main():
# This server only has prompts, no tools or resources.
async with stdio_server() as (reader, writer):
await app.run(reader, writer, app.create_initialization_options())
if __name__ == "__main__":
asyncio.run(main())