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
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

from app.api import chat_router, health_router
from app.api.admin import (
    api_key_router,
    dashboard_router,
    embedding_router,
    flow_test_router,
    guardrails_router,
    intent_rules_router,
    kb_router,
    llm_router,
    monitoring_router,
    prompt_templates_router,
    rag_router,
    script_flows_router,
    sessions_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()

logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)

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)

    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.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(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(monitoring_router)
app.include_router(prompt_templates_router)
app.include_router(rag_router)
app.include_router(script_flows_router)
app.include_router(sessions_router)
app.include_router(tenants_router)


if __name__ == "__main__":
    import uvicorn

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