dp-zp-agent/app/main.py

from __future__ import annotations

import asyncio
import logging
import sys
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import (
    FastAPI,
    Request,
)
from fastapi.middleware.cors import (
    CORSMiddleware,
)
from fastapi.exceptions import RequestValidationError
from fastapi.encoders import jsonable_encoder
from fastapi.exception_handlers import http_exception_handler, request_validation_exception_handler
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.responses import JSONResponse

PROJECT_ROOT = Path(__file__).resolve().parents[1]

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(
        0,
        str(PROJECT_ROOT),
    )


from app.routes import router
from app.security import (
    validate_security_configuration,
)
from tools.rag.embedding_utils import embed_query
from tools.retrieval.retrieval_contract import error_result
from tools.rag.common import DATA_DIR
from tools.sync.sync_jobs import SyncJobs
from tools.rag.rebuild_index import rebuild_index

logger = logging.getLogger(__name__)


OPENWEBUI_ORIGIN = "https://ui.tukekemt.xyz"


@asynccontextmanager
async def lifespan(
    application: FastAPI,
):
    validate_security_configuration()
    application.state.embedding_ready = False

    try:
        await asyncio.to_thread(
            embed_query,
            "warmup",
        )
        application.state.embedding_ready = True

    except Exception:
        logger.exception(
            "Predhriatie embedding modelu " "zlyhalo. API pokračuje bez " "úspešného warm-upu."
        )

    def run_sync(**kwargs):
        result = rebuild_index(**kwargs)
        application.state.embedding_ready = True
        return result

    application.state.sync_jobs = SyncJobs(DATA_DIR, runner=run_sync)
    application.state.sync_jobs.start()
    try:
        yield
    finally:
        await asyncio.to_thread(application.state.sync_jobs.stop)


app = FastAPI(
    title="ZP Agent API",
    description=(
        "RAG API pre vyhľadávanie "
        "a získavanie informácií "
        "z repozitára záverečných prác "
        "ZP Wiki."
    ),
    version="0.9.0",
    lifespan=lifespan,
    docs_url=None,
    redoc_url=None,
)


@app.exception_handler(StarletteHTTPException)
async def retrieval_http_error(request: Request, error: StarletteHTTPException):
    if request.url.path not in {"/rag", "/search", "/graphrag"}:
        return await http_exception_handler(request, error)
    reason = {
        400: "invalid_query",
        401: "unauthorized",
        403: "forbidden",
        404: "not_found",
        405: "method_not_allowed",
        422: "invalid_request",
        503: "retrieval_unavailable",
    }.get(error.status_code, "internal_error")
    return JSONResponse(
        status_code=error.status_code,
        headers=error.headers,
        content={**error_result(reason=reason), "detail": error.detail},
    )


@app.exception_handler(RequestValidationError)
async def retrieval_validation_error(request: Request, error: RequestValidationError):
    if request.url.path not in {"/rag", "/search", "/graphrag"}:
        return await request_validation_exception_handler(request, error)
    # Retain useful validation locations/types without echoing submitted secrets.
    details = [
        {"loc": row["loc"], "type": row["type"], "msg": row["msg"]} for row in error.errors()
    ]
    return JSONResponse(
        status_code=422,
        content=jsonable_encoder({**error_result(reason="invalid_request"), "detail": details}),
    )


app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        OPENWEBUI_ORIGIN,
    ],
    allow_credentials=True,
    allow_methods=[
        "GET",
        "POST",
        "OPTIONS",
    ],
    allow_headers=[
        "Authorization",
        "Content-Type",
        "X-API-Key",
        "X-Session-Id",
        "X-OpenWebUI-Chat-Id",
        "X-OpenWebUI-Message-Id",
        "X-OpenWebUI-User-Id",
        "X-OpenWebUI-User-Name",
        "X-OpenWebUI-User-Email",
        "X-OpenWebUI-User-Role",
        "X-OpenWebUI-User-Jwt",
        "X-OpenWebUI-User-Message-Id",
        "X-OpenWebUI-User-Message-Parent-Id",
        "X-OpenWebUI-Task",
    ],
    allow_private_network=True,
)


@app.middleware("http")
async def add_security_headers(
    request: Request,
    call_next,
):
    try:
        response = await call_next(request)
    except Exception:
        if request.url.path not in {"/rag", "/search", "/graphrag"}:
            raise
        # Serialization failures must obey the same contract as backend errors.
        logger.exception("Neočakávaná chyba výstupu retrieval API.")
        response = JSONResponse(
            status_code=500,
            content={
                **error_result(reason="internal_error"),
                "detail": "Nastala interná chyba servera.",
            },
        )

    response.headers["Cache-Control"] = "no-store"

    response.headers["X-Content-Type-Options"] = "nosniff"

    response.headers["Referrer-Policy"] = "no-referrer"

    response.headers["X-Frame-Options"] = "DENY"

    return response


app.include_router(router)