File size: 1,898 Bytes
56749a4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""
FastAPI wrapper around the SigLIP text encoder for HuggingFace Spaces.

  POST /encode {query} -> {vector:[768]}   (bearer-token auth)
  GET  /health         -> {status:"ok"}    (open, for the keep-warm ping)

The model + tokenizer load once at startup (lifespan) and stay resident, so every
request after wake is warm. The 283 MB model is baked into the Docker image, so a
post-sleep restart doesn't re-download.
"""
from __future__ import annotations

import os
import secrets
from contextlib import asynccontextmanager

from fastapi import Depends, FastAPI, Header, HTTPException
from pydantic import BaseModel

from text_encoder import EMBED_DIM, encode, warmup

ENCODER_TOKEN = os.environ.get("SIGLIP_ENCODER_TOKEN", "")


@asynccontextmanager
async def lifespan(_app: FastAPI):
    warmup()
    yield


app = FastAPI(lifespan=lifespan)


class EncodeRequest(BaseModel):
    query: str


class EncodeResponse(BaseModel):
    vector: list[float]


def require_token(authorization: str = Header(default="")) -> None:
    if not ENCODER_TOKEN:
        raise HTTPException(status_code=500, detail="encoder token not configured")
    # Constant-time compare so the endpoint isn't a timing oracle for the token.
    if not secrets.compare_digest(authorization, f"Bearer {ENCODER_TOKEN}"):
        raise HTTPException(status_code=401, detail="unauthorized")


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


@app.post("/encode", response_model=EncodeResponse, dependencies=[Depends(require_token)])
def encode_query(body: EncodeRequest) -> EncodeResponse:
    query = body.query.strip()
    if not query:
        raise HTTPException(status_code=422, detail="empty query")
    vec = encode(query)
    if len(vec) != EMBED_DIM:
        raise HTTPException(status_code=500, detail=f"expected {EMBED_DIM} dims, got {len(vec)}")
    return EncodeResponse(vector=vec)