Spaces:
Running
Running
| """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" | |
| 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) | |
| 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"]) | |
| 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 | |
| def root(): | |
| return RedirectResponse(url="/docs") | |
| 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()} | |
| 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(), | |
| ) | |
| 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, | |
| } | |
| 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, | |
| ) | |
| 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, | |
| ) | |