rag_server/mcp_app.py

import contextlib
import logging
from collections.abc import AsyncIterator

import click
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Receive, Scope, Send

from applications.config import (
    ES_HOSTS,
    ELASTIC_SEARCH_INDEX,
    ES_PASSWORD,
    MILVUS_CONFIG,
)
from applications.resource import init_resource_manager
from mcp_server.server import create_mcp_server

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


# 初始化资源管理器
resource_manager = init_resource_manager(
    es_hosts=ES_HOSTS,
    es_index=ELASTIC_SEARCH_INDEX,
    es_password=ES_PASSWORD,
    milvus_config=MILVUS_CONFIG,
)


@click.command()
@click.option("--port", default=8002, help="服务器监听端口")
@click.option("--host", default="0.0.0.0", help="服务器监听地址")
@click.option("--json-response", is_flag=True, help="使用JSON响应而不是SSE流")
def main(port: int, host: str, json_response: bool) -> int:
    """启动MCP"""
    # 创建MCP服务器
    app = create_mcp_server()

    # 创建会话管理器
    session_manager = StreamableHTTPSessionManager(
        app=app,
        event_store=None,
        json_response=json_response,
        stateless=True,
    )

    # 处理Streamable HTTP请求
    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]:
        await resource_manager.startup()
        async with session_manager.run():
            yield

    # 创建ASGI应用
    starlette_app = Starlette(
        debug=True,
        routes=[Mount("/mcp", app=handle_streamable_http)],
        lifespan=lifespan,
    )

    # 启动服务器
    import uvicorn

    uvicorn.run(starlette_app, host=host, port=port)
    return 0


if __name__ == "__main__":
    main()