python-services/Retrieve/main.py

"""生物医学RAG服务主程序入口。"""

import importlib
import pkgutil
import time
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

import uvicorn
from asgi_correlation_id import CorrelationIdMiddleware, correlation_id
from fastapi import FastAPI, Request
from fastapi_mcp import FastApiMCP
from fastapi.middleware.cors import CORSMiddleware

from routers import sensor, mcp_sensor
from utils.bio_logger import bio_logger as logger

# 调试：验证环境变量是否加载
logger.info(f"SERPER_API_KEY loaded: {'Yes' if os.getenv('SERPER_API_KEY') else 'No'}")


app = FastAPI(
    docs_url=None,  # 关闭 Swagger UI 文档
    redoc_url=None,  # 关闭 ReDoc 文档
    openapi_url=None,  # 关闭 OpenAPI 规范文件
    debug=False,  # 关闭调试模式
)

# 第一个添加的中间件
app.add_middleware(CorrelationIdMiddleware)
# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 路由
app.include_router(sensor.router)
app.include_router(mcp_sensor.router)  # 包含 MCP 路由


@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """HTTP中间件，记录请求处理时间和状态。"""
    start_time = time.time()

    logger.info(f"Request started  | URL: {request.url}")

    response = await call_next(request)
    process_time = time.time() - start_time

    logger.info(
        f"Request completed | "
        f"Status: {response.status_code} | "
        f"Time: {process_time:.2f}s"
    )

    return response


def dynamic_import_subclasses(parent_dir: str) -> None:
    """动态导入指定目录下的所有Python模块。

    Args:
        parent_dir: 要导入的目录路径
    """
    for _, module_name, _ in pkgutil.iter_modules([parent_dir]):
        module = importlib.import_module(f"{parent_dir}.{module_name}")
        logger.info(f"Imported: {module.__name__}")


# Add MCP server to the FastAPI app
mcp = FastApiMCP(
    app, 
    name="bio qa mcp", 
    include_operations=["bio_qa_stream_chat"]
)

# Mount the MCP server to the FastAPI app
# 挂载SSE端点到 /mcp/sse
mcp.mount_sse()


if __name__ == "__main__":
    logger.info("Starting Bio RAG Server...")
    dynamic_import_subclasses("search_service")
    uvicorn.run(app, host="0.0.0.0", port=9487)