9fd680c366
Co-authored-by: luhualin
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
import contextlib
|
|
import sys
|
|
from collections.abc import AsyncIterator
|
|
from mcp.server.fastmcp import FastMCP
|
|
from starlette.applications import Starlette
|
|
from mcp.server.sse import SseServerTransport
|
|
from starlette.requests import Request
|
|
from starlette.routing import Mount, Route
|
|
from starlette.types import Receive, Scope, Send
|
|
from mcp.server import Server
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
from markitdown import MarkItDown
|
|
import uvicorn
|
|
|
|
# Initialize FastMCP server for MarkItDown (SSE)
|
|
mcp = FastMCP("markitdown")
|
|
|
|
|
|
@mcp.tool()
|
|
async def convert_to_markdown(uri: str) -> str:
|
|
"""Convert a resource described by an http:, https:, file: or data: URI to markdown"""
|
|
return MarkItDown().convert_uri(uri).markdown
|
|
|
|
|
|
def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
|
|
sse = SseServerTransport("/messages/")
|
|
session_manager = StreamableHTTPSessionManager(
|
|
app=mcp_server,
|
|
event_store=None,
|
|
json_response=True,
|
|
stateless=True,
|
|
)
|
|
|
|
async def handle_sse(request: Request) -> None:
|
|
async with sse.connect_sse(
|
|
request.scope,
|
|
request.receive,
|
|
request._send,
|
|
) as (read_stream, write_stream):
|
|
await mcp_server.run(
|
|
read_stream,
|
|
write_stream,
|
|
mcp_server.create_initialization_options(),
|
|
)
|
|
|
|
async def handle_streamable_http(
|
|
scope: Scope, receive: Receive, send: Send
|
|
) -> None:
|
|
await session_manager.handle_request(scope, receive, send)
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(app: Starlette) -> AsyncIterator[None]:
|
|
"""Context manager for session manager."""
|
|
async with session_manager.run():
|
|
print("Application started with StreamableHTTP session manager!")
|
|
try:
|
|
yield
|
|
finally:
|
|
print("Application shutting down...")
|
|
|
|
return Starlette(
|
|
debug=debug,
|
|
routes=[
|
|
Route("/sse", endpoint=handle_sse),
|
|
Mount("/mcp", app=handle_streamable_http),
|
|
Mount("/messages/", app=sse.handle_post_message),
|
|
],
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
|
|
# Main entry point
|
|
def main():
|
|
import argparse
|
|
|
|
mcp_server = mcp._mcp_server
|
|
|
|
parser = argparse.ArgumentParser(description="Run MCP SSE-based MarkItDown server")
|
|
|
|
parser.add_argument(
|
|
"--sse",
|
|
action="store_true",
|
|
help="Run the server with SSE transport rather than STDIO (default: False)",
|
|
)
|
|
parser.add_argument(
|
|
"--host", default=None, help="Host to bind to (default: 127.0.0.1)"
|
|
)
|
|
parser.add_argument(
|
|
"--port", type=int, default=None, help="Port to listen on (default: 3001)"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not args.sse and (args.host or args.port):
|
|
parser.error("Host and port arguments are only valid when using SSE transport.")
|
|
sys.exit(1)
|
|
|
|
if args.sse:
|
|
starlette_app = create_starlette_app(mcp_server, debug=True)
|
|
uvicorn.run(
|
|
starlette_app,
|
|
host=args.host if args.host else "127.0.0.1",
|
|
port=args.port if args.port else 3001,
|
|
)
|
|
else:
|
|
mcp.run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|