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'}:
        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'}:
        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'}:
            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
)