"""FastAPI HTTP-only vectorization service.""" import logging import os import secrets from pathlib import Path from contextlib import asynccontextmanager from typing import Literal from fastapi import Depends, FastAPI, HTTPException, Request, Security from fastapi.exception_handlers import request_validation_exception_handler from fastapi.exceptions import RequestValidationError from fastapi.responses import RedirectResponse from fastapi.security import APIKeyHeader from pydantic import BaseModel, Field, field_validator from app.embedding import MODEL_ID, TooManyTokens from app.models import ( ModelRegistry, QWEN_ID, COLBERT_ID, ModelUnavailable, ServiceBusy, memory_mb, ) log = logging.getLogger(__name__) api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) EmbeddingModel = Literal[ "e5", "qwen", "intfloat/multilingual-e5-small", "Qwen/Qwen3-Embedding-0.6B" ] def deployment_version(): path = Path(__file__).resolve().parent.parent / "VERSION" return path.read_text(encoding="utf-8").strip() if path.is_file() else "development" @asynccontextmanager async def lifespan(application: FastAPI): if not os.getenv("API_KEY"): raise RuntimeError("API_KEY must be configured") # Resident models (including BGE-M3 when enabled) load before readiness. application.state.registry = ModelRegistry.load() try: yield finally: application.state.registry = None from app.voyage_api import router as voyage_router from app.openai_api import router as openai_router from app.bge_m3_api import router as bge_m3_router app = FastAPI( title="Vectorization API", version="2.0.0", description="Resident E5, Qwen3, mLateOn and BGE-M3 CPU embedding API.", lifespan=lifespan, ) app.include_router(voyage_router) app.include_router(openai_router) app.include_router(bge_m3_router) @app.exception_handler(RequestValidationError) async def log_openai_validation_error(request: Request, exc: RequestValidationError): """Explain 422 for OpenAI embeddings without exposing credentials.""" if request.url.path == "/openai/v1/embeddings": failures = [ {"field": ".".join(map(str, error.get("loc", ()))), "code": error.get("type"), "message": error.get("msg")} for error in exc.errors()[:8] ] auth = request.headers.get("authorization", "") scheme, separator, provided = auth.partition(" ") expected = os.getenv("API_KEY") trusted = bool(expected and separator and scheme.lower() == "bearer" and secrets.compare_digest(provided, expected)) if trusted: body = exc.body if isinstance(exc.body, dict) else {} raw_input = body.get("input") inputs = raw_input if isinstance(raw_input, list) else [raw_input] previews = [ (index, value[:240] if isinstance(value, str) else f"<{type(value).__name__}>") for index, value in enumerate(inputs[:3]) ] model = body.get("model") log.warning( "OpenAI embeddings rejected: reason=request_validation_error " "fields=%r model=%r batch_size=%d input_previews=%r", failures, model if isinstance(model, str) else type(model).__name__, len(inputs), previews, ) else: log.warning( "OpenAI embeddings rejected: reason=request_validation_error " "authentication=unverified fields=%r", failures, ) return await request_validation_exception_handler(request, exc) class TextRequest(BaseModel): text: str = Field(min_length=1, description="Full input, no silent truncation") type: Literal["query", "passage"] = Field(examples=["query"]) @field_validator("text") @classmethod def nonblank(cls, text): if not text.strip(): raise ValueError("Text must not be blank") return text class EmbedRequest(TextRequest): model: EmbeddingModel | None = Field( default=None, examples=["qwen"], description="E5 by default" ) class EmbedResponse(BaseModel): text: str type: Literal["query", "passage"] tokens_count: int embedding_model: str dimensions: int embedding: list[float] class ColbertResponse(BaseModel): embedding_model: str tokens_count: int dimensions: int token_embeddings: list[list[float]] class ModelResponse(BaseModel): embedding_model: str dimensions: int max_input_tokens: int normalization: Literal["none"] = "none" version: str def require_api_key(provided: str | None = Security(api_key_header)): expected = os.getenv("API_KEY") if not expected or provided is None or not secrets.compare_digest(provided, expected): raise HTTPException(status_code=401, detail="Invalid or missing API key", headers={"WWW-Authenticate": "ApiKey"}) def get_registry(): registry = getattr(app.state, "registry", None) if registry is None: raise HTTPException(status_code=503, detail="Models are not ready") return registry def guard(callback): try: return callback() except TooManyTokens as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc except ModelUnavailable as exc: raise HTTPException(status_code=503, detail=str(exc)) from exc except ServiceBusy as exc: raise HTTPException(status_code=429, detail=str(exc), headers={"Retry-After": "2"}) from exc except Exception as exc: log.exception("Inference failed") raise HTTPException(status_code=500, detail="Inference failed") from exc @app.get("/", include_in_schema=False) def root(): return RedirectResponse(url="/docs") @app.get("/health") def health(): if getattr(app.state, "registry", None) is None: raise HTTPException(status_code=503, detail="Models are not ready") return {"status": "ok", "version": deployment_version(), "loaded_models": app.state.registry.loaded_models()} @app.get("/model", response_model=ModelResponse, dependencies=[Depends(require_api_key)]) def model_info(registry=Depends(get_registry)): if registry.e5 is None: raise HTTPException(status_code=503, detail="Model e5 is not enabled") return ModelResponse( embedding_model=MODEL_ID, dimensions=registry.e5.dimensions, max_input_tokens=registry.e5.max_tokens, version=deployment_version(), ) @app.get("/models", dependencies=[Depends(require_api_key)]) def models_info(registry=Depends(get_registry)): models = [] if registry.e5 is not None: models.append({"embedding_model": MODEL_ID, "purpose": "standard", "dimensions": registry.e5.dimensions, "max_input_tokens": registry.e5.max_tokens, "normalization": "none", "ready": True}) if registry.qwen is not None: models.append({"embedding_model": QWEN_ID, "purpose": "literature", "dimensions": registry.qwen.dimensions, "max_input_tokens": registry.qwen.max_tokens, "normalization": "l2", "ready": True}) if registry.colbert is not None: models.append({"embedding_model": COLBERT_ID, "purpose": "colbert", "dimensions": registry.colbert.dimensions, "max_input_tokens": registry.colbert.max_tokens, "normalization": "model-native", "ready": True}) if registry.bge_m3 is not None: models.append({"embedding_model": "BAAI/bge-m3", "purpose": "triple-vector", "dimensions": registry.bge_m3.dimensions, "colbert_dimensions": registry.bge_m3.colbert_dimensions, "max_input_tokens": registry.bge_m3.max_tokens, "normalization": "model-native", "ready": True}) return { "version": deployment_version(), "models": models, "loaded_models": registry.loaded_models(), "resources": memory_mb(), "startup_memory": registry.startup_memory, } @app.post("/embed", response_model=EmbedResponse, dependencies=[Depends(require_api_key)]) def embed(payload: EmbedRequest, registry=Depends(get_registry)): identifier, dimensions, result = guard( lambda: registry.embed(payload.text, payload.type, payload.model) ) return EmbedResponse( text=payload.text, type=payload.type, tokens_count=result.tokens_count, embedding_model=identifier, dimensions=dimensions, embedding=result.embedding, ) @app.post("/colbert/encode", response_model=ColbertResponse, dependencies=[Depends(require_api_key)]) def colbert_encode(payload: TextRequest, registry=Depends(get_registry)): result = guard(lambda: registry.colbert_encode(payload.text, payload.type)) return ColbertResponse( embedding_model=COLBERT_ID, dimensions=registry.colbert.dimensions, tokens_count=result.tokens_count, token_embeddings=result.token_embeddings, )