src/app.py - master - blanchefort/embedder_service

src/app.py
91 строка4 KB
Igor Shatalin
Merge remote-tracking branch 'origin/master'
29 июл 2026, 20:55
29 июл 2026, 20:5548f48f2
import asyncio
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator
from uuid import UUID

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from src.api.rerank_router import router as rerank_router
from src.api.router import router as embedding_router
from src.embedder.text import TextEmbedder
from src.reranker.gte import Reranker
from src.settings import Settings, get_settings


def create_app(settings: Settings | None = None, *, load_models: bool = True) -> FastAPI:
    configured = settings or get_settings()

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        if load_models:
            await asyncio.to_thread(app.state.embedder.load)
            await app.state.reranker.load()
            if configured.reranker_warmup:
                await app.state.reranker.warmup()
            else:
                app.state.reranker.ready = True
        yield
        await app.state.reranker.close()

    app = FastAPI(title="Embedder Service", version="0.2.0", lifespan=lifespan)
    app.state.settings = configured
    app.state.embedder = TextEmbedder(configured)
    app.state.reranker = Reranker(configured)
    app.dependency_overrides[get_settings] = lambda: configured
    app.include_router(embedding_router)
    app.include_router(rerank_router)

    @app.middleware("http")
    async def body_size_limit(request: Request, call_next):
        content_length = request.headers.get("content-length")
        if content_length:
            try:
                parsed_content_length = int(content_length)
            except ValueError:
                return JSONResponse(status_code=400, content={"detail": "Invalid Content-Length header"})
            if parsed_content_length < 0:
                return JSONResponse(status_code=400, content={"detail": "Invalid Content-Length header"})
            if parsed_content_length > configured.reranker_max_body_bytes:
                return JSONResponse(status_code=413, content={"detail": "Request body too large"})
        if request.method in {"POST", "PUT", "PATCH"}:
            body = await request.body()
            if len(body) > configured.reranker_max_body_bytes:
                return JSONResponse(status_code=413, content={"detail": "Request body too large"})
        significant = request.url.path.startswith("/api/v1/embeddings") or request.url.path == "/api/v1/rerank"
        if significant:
            supplied = request.headers.get("X-Request-ID")
            try:
                request_id = str(UUID(supplied)) if supplied else None
            except ValueError:
                request_id = None
            if request_id is None:
                return JSONResponse(
                    status_code=400,
                    content={"detail": "X-Request-ID header with a UUID is required"},
                )
            request.state.request_id = request_id
        response = await call_next(request)
        if significant:
            response.headers["X-Request-ID"] = request.state.request_id
        return response

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

    @app.get("/health/ready")
    async def ready():
        is_ready = app.state.embedder._embedder is not None and app.state.reranker.ready
        if not is_ready:
            return JSONResponse(status_code=503, content={"status": "not_ready"})
        return {"status": "ready"}

    @app.get("/v1/models")
    async def models() -> dict[str, object]:
        return {"reranker": app.state.reranker.status}

    return app


app = create_app()