MCP-RAG Server

mcp-rag
src
mcp_rag
from mcp import types
from mcp.server import FastMCP
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.routing import Route

import logging

logger = logging.getLogger()
logging.basicConfig(level=logging.DEBUG)
app = Server("example-server")
sse = SseServerTransport("/messages/")

async def handle_sse(request):
    async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
        await app.run(streams[0], streams[1], app.create_initialization_options())

@app.list_prompts()
async def handle_list_prompts() -> list[types.Prompt]:
    return [
        types.Prompt(
            name="example-prompt",
            description="An example prompt template",
            arguments=[
                types.PromptArgument(
                    name="arg1",
                    description="Example argument",
                    required=True
                )
            ]
        )
    ]

@app.get_prompt()
async def handle_get_prompt(
    name: str,
    arguments: dict[str, str] | None
) -> types.GetPromptResult:
    if name != "example-prompt":
        raise ValueError(f"Unknown prompt: {name}")

    return types.GetPromptResult(
        description="Example prompt",
        messages=[
            types.PromptMessage(
                role="user",
                content=types.TextContent(
                    type="text",
                    text="Example prompt text"
                )
            )
        ]
    )

@app.list_tools()
async def list_tools() -> list[types.Tool]:
    ...

# first argument is tool name, the second is a dict of arguments
@app.call_tool()
async def handle_call_tool(tool_name: str, arguments):
    print(tool_name)
    print(arguments)
    return [types.TextContent(text=f"Hello {tool_name}", type="text")]


starlette_app = Starlette(
    routes=[
        Route("/sse", endpoint=handle_sse),
        Mount("/messages/", app=sse.handle_post_message),
    ]
)

print(starlette_app.routes)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--host")
    parser.add_argument("--port", type=int)

    options = parser.parse_args()

    import uvicorn

    uvicorn.run(starlette_app, host=options.host, port=options.port)