knowledge-adapter-service/src/knowledge_adapter/app.py

from typing import Annotated, Any
from uuid import UUID

from fastapi import Body, Depends, FastAPI, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, RedirectResponse, Response
from pydantic import ValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.types import Lifespan

from knowledge_adapter.asset_content_processor import ContentProcessingError
from knowledge_adapter.asset_redirect import (
    AssetNotFoundError,
    AssetRedirector,
    AssetRedirectUnavailableError,
)
from knowledge_adapter.auth import DifyKeyAuthenticator
from knowledge_adapter.models import RetrievalRequest
from knowledge_adapter.observability import (
    bind_request_id,
    log_http_request,
    normalize_request_id,
    request_started_at,
    reset_request_id,
)
from knowledge_adapter.service import InvalidRequestError, RetrievalService
from knowledge_adapter.weknora_client import WeKnoraBadResponseError, WeKnoraUnavailableError

_SAFE_HTTP_DETAILS = frozenset(
    {
        "content processing failed",
        "invalid API key",
        "invalid request",
        "invalid upstream response",
        "upstream unavailable",
    }
)
_DEFAULT_ERROR_DETAILS = {
    400: "invalid request",
    401: "invalid API key",
    403: "forbidden",
    405: "method not allowed",
    413: "invalid request",
    422: "invalid request",
    500: "internal server error",
    502: "bad gateway",
    503: "service unavailable",
}


def create_app(
    *,
    authenticator: DifyKeyAuthenticator | None = None,
    retrieval_service: RetrievalService | None = None,
    asset_redirect_service: AssetRedirector | None = None,
    max_query_chars: int = 2000,
    lifespan: Lifespan[FastAPI] | None = None,
) -> FastAPI:
    app = FastAPI(
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
        lifespan=lifespan,
    )

    @app.middleware("http")
    async def request_id_middleware(request: Request, call_next: Any) -> Response:
        request_id = normalize_request_id(request.headers.get("X-Request-ID"))
        request.state.request_id = request_id
        request_id_token = bind_request_id(request_id)
        started_at_ns = request_started_at()
        try:
            try:
                response = await call_next(request)
            except Exception:
                response = JSONResponse(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    content={"detail": "internal server error"},
                )
            response.headers["X-Request-ID"] = request_id
            route = getattr(request.scope.get("route"), "path", "<unmatched>")
            log_http_request(
                request_id=request_id,
                method=request.method,
                route=route,
                status_code=response.status_code,
                started_at_ns=started_at_ns,
            )
            return response
        finally:
            reset_request_id(request_id_token)

    @app.exception_handler(RequestValidationError)
    async def invalid_request_handler(
        _request: Request,
        _exc: RequestValidationError,
    ) -> JSONResponse:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"detail": "invalid request"},
        )

    @app.exception_handler(StarletteHTTPException)
    async def safe_http_error_handler(
        request: Request,
        exc: StarletteHTTPException,
    ) -> Response:
        if exc.status_code == status.HTTP_404_NOT_FOUND:
            if request.url.path.startswith("/i/"):
                return Response(status_code=status.HTTP_404_NOT_FOUND)
            return JSONResponse(
                status_code=status.HTTP_404_NOT_FOUND,
                content={"detail": "not found"},
            )
        detail = (
            exc.detail
            if isinstance(exc.detail, str) and exc.detail in _SAFE_HTTP_DETAILS
            else _DEFAULT_ERROR_DETAILS.get(exc.status_code, "request failed")
        )
        headers = (
            {"WWW-Authenticate": "Bearer"}
            if exc.status_code == status.HTTP_401_UNAUTHORIZED
            else None
        )
        return JSONResponse(
            status_code=exc.status_code,
            content={"detail": detail},
            headers=headers,
        )

    @app.get("/healthz")
    async def healthz() -> dict[str, str]:
        return {"status": "ok"}

    if asset_redirect_service is not None:

        @app.get("/i/{asset_id:uuid}", response_class=RedirectResponse)
        async def asset_redirect(asset_id: UUID) -> Response:
            try:
                location = await asset_redirect_service.get_redirect_url(asset_id)
            except AssetNotFoundError:
                return Response(status_code=status.HTTP_404_NOT_FOUND)
            except AssetRedirectUnavailableError:
                return Response(status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
            return RedirectResponse(
                url=location,
                status_code=status.HTTP_302_FOUND,
                headers={
                    "Cache-Control": "no-store",
                    "Referrer-Policy": "no-referrer",
                },
            )

        @app.get("/i/{_invalid_asset_id}")
        async def invalid_asset_id(_invalid_asset_id: str) -> Response:
            return Response(status_code=status.HTTP_404_NOT_FOUND)

    if authenticator is not None:

        @app.post("/retrieval", dependencies=[Depends(authenticator)])
        async def retrieval(
            payload: Annotated[Any, Body()],
        ) -> dict[str, list[dict[str, object]]]:
            try:
                request_model = RetrievalRequest.model_validate(payload)
            except ValidationError as exc:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="invalid request",
                ) from exc
            if len(request_model.query) > max_query_chars:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="invalid request",
                )
            if retrieval_service is not None:
                try:
                    return await retrieval_service.retrieve(request_model)
                except InvalidRequestError as exc:
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail="invalid request",
                    ) from exc
                except WeKnoraUnavailableError as exc:
                    raise HTTPException(
                        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                        detail="upstream unavailable",
                    ) from exc
                except WeKnoraBadResponseError as exc:
                    raise HTTPException(
                        status_code=status.HTTP_502_BAD_GATEWAY,
                        detail="invalid upstream response",
                    ) from exc
                except ContentProcessingError as exc:
                    raise HTTPException(
                        status_code=status.HTTP_502_BAD_GATEWAY,
                        detail="content processing failed",
                    ) from exc
            return {"records": []}

    return app