ai-robot-core/ai-service/app/main.py

"""
Main FastAPI application for AI Service.
[AC-AISVC-01] Entry point with middleware and exception handlers.
"""

import logging
import os
from logging.handlers import RotatingFileHandler
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, status
from fastapi.exceptions import HTTPException, RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response

from app.api import chat_router, health_router
from app.api.mid import router as mid_router
from app.api.openapi import router as openapi_router
from app.api.admin import (
    api_key_router,
    dashboard_router,
    decomposition_template_router,
    embedding_router,
    flow_test_router,
    guardrails_router,
    intent_rules_router,
    kb_router,
    llm_router,
    metadata_field_definition_router,
    metadata_schema_router,
    monitoring_router,
    prompt_templates_router,
    rag_router,
    retrieval_strategy_router,
    scene_slot_bundle_router,
    script_flows_router,
    sessions_router,
    slot_definition_router,
    tenants_router,
)
from app.api.admin.kb_optimized import router as kb_optimized_router
from app.core.config import get_settings
from app.core.database import close_db, init_db
from app.core.exceptions import (
    AIServiceException,
    ErrorCode,
    ErrorResponse,
    ai_service_exception_handler,
    generic_exception_handler,
    http_exception_handler,
)
from app.core.middleware import ApiKeyMiddleware, TenantContextMiddleware
from app.core.qdrant_client import close_qdrant_client

settings = get_settings()


def setup_logging():
    """
    配置滚动日志文件。
    - 日志文件存储在 logs/ 目录
    - 单文件最大 2MB，超过则切分
    - 保留最近 7 天的日志（约 70 个备份文件）
    - 同时输出到控制台
    """
    log_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "logs")
    os.makedirs(log_dir, exist_ok=True)
    
    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    formatter = logging.Formatter(log_format)
    
    root_logger = logging.getLogger()
    root_logger.setLevel(getattr(logging, settings.log_level.upper()))
    
    root_logger.handlers.clear()
    
    console_handler = logging.StreamHandler()
    console_handler.setLevel(getattr(logging, settings.log_level.upper()))
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)
    
    log_file = os.path.join(log_dir, "ai-service.log")
    file_handler = RotatingFileHandler(
        filename=log_file,
        maxBytes=2 * 1024 * 1024,
        backupCount=70,
        encoding="utf-8",
    )
    file_handler.setLevel(getattr(logging, settings.log_level.upper()))
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)
    
    logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
    logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING)
    logging.getLogger("sqlalchemy.dialects").setLevel(logging.WARNING)
    logging.getLogger("sqlalchemy.orm").setLevel(logging.WARNING)
    
    return log_dir


setup_logging()

logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    [AC-AISVC-01, AC-AISVC-11, AC-AISVC-50] Application lifespan manager.
    Handles startup and shutdown of database and external connections.
    """
    logger.info(f"[AC-AISVC-01] Starting {settings.app_name} v{settings.app_version}")

    try:
        await init_db()
        logger.info("[AC-AISVC-11] Database initialized successfully")
    except Exception as e:
        logger.warning(f"[AC-AISVC-11] Database initialization skipped: {e}")

    try:
        from app.core.database import async_session_maker
        from app.services.api_key import get_api_key_service

        logger.info("[AC-AISVC-50] Starting API key initialization...")
        async with async_session_maker() as session:
            api_key_service = get_api_key_service()
            logger.info(f"[AC-AISVC-50] Got API key service instance, initializing...")
            await api_key_service.initialize(session)
            logger.info(f"[AC-AISVC-50] API key service initialized, cache size: {len(api_key_service._keys_cache)}")
            default_key = await api_key_service.create_default_key(session)
            if default_key:
                logger.info(f"[AC-AISVC-50] Default API key created: {default_key.key}")
    except Exception as e:
        logger.error(f"[AC-AISVC-50] API key initialization FAILED: {e}", exc_info=True)

    try:
        from app.services.mid.tool_guide_registry import init_tool_guide_registry

        logger.info("[ToolGuideRegistry] Starting tool guides initialization...")
        tool_guide_registry = init_tool_guide_registry()
        logger.info(f"[ToolGuideRegistry] Tool guides loaded: {tool_guide_registry.list_tools()}")
    except Exception as e:
        logger.error(f"[ToolRegistry] Tools initialization FAILED: {e}", exc_info=True)

    # [AC-AISVC-29] 预初始化 Embedding 服务，避免首次查询时的延迟
    try:
        from app.services.embedding import get_embedding_provider

        logger.info("[AC-AISVC-29] Pre-initializing embedding service...")
        embedding_provider = await get_embedding_provider()
        logger.info(
            f"[AC-AISVC-29] Embedding service pre-initialized: "
            f"provider={embedding_provider.PROVIDER_NAME}"
        )
    except Exception as e:
        logger.error(f"[AC-AISVC-29] Embedding service pre-initialization FAILED: {e}", exc_info=True)

    yield

    await close_db()
    await close_qdrant_client()
    logger.info(f"Shutting down {settings.app_name}")


app = FastAPI(
    title=settings.app_name,
    version=settings.app_version,
    description="""
    Python AI Service for intelligent chat with RAG support.

    ## Features
    - Multi-tenant isolation via X-Tenant-Id header
    - SSE streaming support via Accept: text/event-stream
    - RAG-powered responses with confidence scoring

    ## Response Modes
    - **JSON**: Default response mode (Accept: application/json or no Accept header)
    - **SSE Streaming**: Set Accept: text/event-stream for streaming responses
    """,
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(TenantContextMiddleware)
app.add_middleware(ApiKeyMiddleware)

app.add_exception_handler(AIServiceException, ai_service_exception_handler)
app.add_exception_handler(HTTPException, http_exception_handler)
app.add_exception_handler(Exception, generic_exception_handler)


@app.get("/favicon.ico", include_in_schema=False)
async def favicon() -> Response:
    return Response(status_code=204)


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """
    [AC-AISVC-03] Handle request validation errors with structured response.
    """
    logger.warning(f"[AC-AISVC-03] Request validation error: {exc.errors()}")
    error_response = ErrorResponse(
        code=ErrorCode.INVALID_REQUEST.value,
        message="Request validation failed",
        details=[{"loc": list(err["loc"]), "msg": err["msg"], "type": err["type"]} for err in exc.errors()],
    )
    return JSONResponse(
        status_code=status.HTTP_400_BAD_REQUEST,
        content=error_response.model_dump(exclude_none=True),
    )


app.include_router(health_router)
app.include_router(chat_router)

app.include_router(api_key_router)
app.include_router(dashboard_router)
app.include_router(decomposition_template_router)
app.include_router(embedding_router)
app.include_router(flow_test_router)
app.include_router(guardrails_router)
app.include_router(intent_rules_router)
app.include_router(kb_router)
app.include_router(kb_optimized_router)
app.include_router(llm_router)
app.include_router(metadata_field_definition_router)
app.include_router(metadata_schema_router)
app.include_router(monitoring_router)
app.include_router(prompt_templates_router)
app.include_router(rag_router)
app.include_router(retrieval_strategy_router)
app.include_router(scene_slot_bundle_router)
app.include_router(script_flows_router)
app.include_router(sessions_router)
app.include_router(slot_definition_router)
app.include_router(tenants_router)

app.include_router(mid_router)
app.include_router(openapi_router)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
    )