File size: 7,249 Bytes
fd4e724 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """Private dispatcher ASGI app. Does not bind a listener or change production."""
from contextlib import asynccontextmanager
from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from .config import Settings
from .http_security import Guard, owner
from .jobs import Jobs, status
from .routing import Dispatcher
from .schemas import OCROptions, SafeError
from .security import sensitivity
CODES = {
"unauthorized": 401,
"job_not_found": 404,
"no_authorized_destination": 403,
"idempotency_conflict": 409,
"result_not_ready": 409,
"job_capacity": 429,
"queue_full": 429,
"upload_limit": 413,
"model_not_ready": 503,
}
def create_app(settings=None, cpu=None, gpu=None):
settings = settings or Settings.from_env()
@asynccontextmanager
async def lifespan(app):
nonlocal cpu, gpu
probe_task = None
if cpu is None:
if settings.cpu_url:
from .adapters.cpu_http import RemoteCPU
cpu = RemoteCPU(settings.cpu_url, settings.worker_token, settings.timeout)
else:
from .adapters.rapidocr_cpu import RapidCPU
cpu = RapidCPU(settings.models_dir)
try:
await cpu.initialize()
except SafeError:
pass # /ready reports failure; never pretend warmed readiness.
if gpu is None and settings.gpu_enabled and settings.freellm_url and settings.freellm_key:
from .adapters.freellm_alias import FreeLLMAlias
gpu = FreeLLMAlias(settings.freellm_url, settings.freellm_key, settings.timeout)
async def probe():
try:
await gpu.probe()
except Exception:
gpu.qualified = False
import asyncio
probe_task = asyncio.create_task(probe())
dispatcher = Dispatcher(cpu, gpu, settings)
dispatcher.start()
jobs = Jobs(dispatcher, settings)
jobs.start()
app.state.jobs = jobs
app.state.dispatcher = dispatcher
try:
yield
finally:
if probe_task:
probe_task.cancel()
await asyncio.gather(probe_task, return_exceptions=True)
await jobs.close()
await dispatcher.close()
app = FastAPI(title="Private Hybrid OCR", version="1.0.0", lifespan=lifespan)
app.add_middleware(
Guard,
token=settings.token,
max_bytes=settings.max_upload,
rate=settings.rate_per_minute,
)
@app.exception_handler(SafeError)
async def safe_error(request, exc):
return JSONResponse({"code": exc.code}, status_code=CODES.get(exc.code, 422))
@app.exception_handler(RequestValidationError)
async def invalid(request, exc):
return JSONResponse({"code": "invalid_request"}, status_code=422)
@app.exception_handler(Exception)
async def unexpected(request, exc):
return JSONResponse({"code": "internal_error"}, status_code=500)
def parse_options(raw):
try:
if len(raw) > 8192:
raise ValueError()
opts = OCROptions.model_validate_json(raw)
opts.sensitivity = sensitivity(opts.sensitivity, settings.min_sensitivity)
opts.max_pages = min(opts.max_pages, settings.max_pages)
return opts
except Exception:
raise SafeError("invalid_options") from None
async def upload(file):
try:
data = await file.read(settings.max_upload + 1)
finally:
await file.close()
if len(data) > settings.max_upload:
raise SafeError("upload_limit")
return data
@app.get("/health")
async def health():
return {"status": "ok"}
@app.get("/ready")
async def ready():
cpu_health = await app.state.dispatcher.cpu.backend.health()
return JSONResponse(
{"ready": cpu_health.ready}, status_code=200 if cpu_health.ready else 503
)
@app.post("/v1/ocr/jobs", status_code=202)
async def submit(request: Request, file: UploadFile = File(...), options: str = Form("{}")):
opts = parse_options(options)
data = await upload(file)
key = request.headers.get("idempotency-key")
if key and len(key) > 128:
raise SafeError("invalid_idempotency_key")
j = app.state.jobs.submit(data, file.filename or "", opts, owner(request), key)
return dict(
job_id=j.id,
status=j.status,
request_id=j.request_id,
status_url=f"/v1/ocr/jobs/{j.id}",
result_url=f"/v1/ocr/jobs/{j.id}/result",
)
@app.get("/v1/ocr/jobs/{job_id}")
async def get(job_id: str, request: Request):
return status(app.state.jobs.get(job_id, owner(request)))
@app.get("/v1/ocr/jobs/{job_id}/result")
async def result(job_id: str, request: Request):
j = app.state.jobs.get(job_id, owner(request))
if j.result is None:
raise SafeError("result_not_ready")
return j.result
@app.delete("/v1/ocr/jobs/{job_id}")
async def cancel(job_id: str, request: Request):
j = app.state.jobs.get(job_id, owner(request))
app.state.jobs.cancel(j)
return status(j)
@app.post("/v1/ocr/page")
async def page(request: Request, file: UploadFile = File(...), options: str = Form("{}")):
opts = parse_options(options)
opts.max_pages = 1
pages = await app.state.jobs.parse(await upload(file), file.filename or "", opts)
return await app.state.dispatcher.page(pages[0], opts)
@app.get("/metrics")
async def metrics():
import numpy as np
lanes = [app.state.dispatcher.cpu] + (
[app.state.dispatcher.gpu] if app.state.dispatcher.gpu else []
)
return {
lane.backend.name: dict(
queued=lane.queue.qsize(),
active=lane.active,
completed=lane.completed,
errors=lane.errors,
ewma_ms=lane.ewma,
cold_ms=lane.cold_ms,
warm_latency_ms={
str(p): float(np.percentile(lane.samples, p)) if lane.samples else None
for p in (50, 95, 99)
},
)
for lane in lanes
}
from .openai_facade import ChatRequest, completion, image_from_request
@app.post("/v1/chat/completions")
async def chat(body: ChatRequest):
data, filename = image_from_request(body, settings.max_upload)
opts = body.ocr_options.model_copy(deep=True)
opts.sensitivity = sensitivity(opts.sensitivity, settings.min_sensitivity)
opts.mode = {"ocr-auto": "auto", "ocr-cpu": "cpu_only", "ocr-gpu": "gpu_only"}[body.model]
opts.max_pages = 1
pages = await app.state.jobs.parse(data, filename, opts)
result = await app.state.dispatcher.page(pages[0], opts)
return completion(result, body.model)
return app
app = create_app()
|