Fu DaiandClaude Opus 4.8
1a585693be
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
import time
|
|
import uuid
|
|
from typing import Any, Dict, List, Literal, Optional, Union
|
|
|
|
from fastapi import APIRouter, Request
|
|
from fastapi.responses import JSONResponse
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
from surya.endpoint import service
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class _ImageUrl(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
url: str
|
|
|
|
|
|
class _ContentPart(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
type: str
|
|
image_url: Optional[_ImageUrl] = None
|
|
text: Optional[str] = None
|
|
|
|
|
|
class _ChatMessage(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
role: str
|
|
content: Union[str, List[_ContentPart]]
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
# The OpenAI SDK merges `extra_body={...}` into the TOP LEVEL of the request
|
|
# JSON, so the OCR control flags live here rather than nested under a key.
|
|
model_config = ConfigDict(extra="allow")
|
|
model: str = "surya-ocr"
|
|
messages: List[_ChatMessage]
|
|
stream: Optional[bool] = False
|
|
skip_text_detection: bool = False
|
|
recognize_math: bool = False
|
|
skip_table_detection: bool = False
|
|
ocr_with_boxes: bool = True
|
|
mode: Literal["block", "full_page", "table"] = "block"
|
|
|
|
|
|
def _openai_error(message: str, err_type: str, status_code: int) -> JSONResponse:
|
|
return JSONResponse(
|
|
status_code=status_code,
|
|
content={"error": {"message": message, "type": err_type, "code": None}},
|
|
)
|
|
|
|
|
|
def _build_chat_completion(
|
|
content: str, model: str, ocr_json: Any, elapsed: Optional[float]
|
|
) -> Dict[str, Any]:
|
|
# Token accounting happens inside vLLM and is not surfaced to this layer, so
|
|
# usage counts are reported as 0 (prompt_tokens intentionally 0).
|
|
return {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
"message": {"role": "assistant", "content": content},
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
|
"surya": {"ocr_text_json": ocr_json, "elapsed_seconds": elapsed},
|
|
}
|
|
|
|
|
|
@router.post("/v1/chat/completions")
|
|
def chat_completions(req: ChatCompletionRequest, request: Request):
|
|
request_id = getattr(request.state, "request_id", "-")
|
|
if req.stream:
|
|
return _openai_error(
|
|
"streaming is not supported by this OCR endpoint",
|
|
"invalid_request_error",
|
|
400,
|
|
)
|
|
try:
|
|
image = service.extract_last_image(req.messages)
|
|
except service.OcrRequestError as exc:
|
|
return _openai_error(str(exc), "invalid_request_error", 400)
|
|
|
|
try:
|
|
if req.mode == "table":
|
|
data = service.table_image(image, req.skip_table_detection)
|
|
else:
|
|
skip = req.skip_text_detection or req.mode == "full_page"
|
|
data = service.ocr_via_batcher(
|
|
image,
|
|
skip_text_detection=skip,
|
|
recognize_math=req.recognize_math,
|
|
request_id=request_id,
|
|
)
|
|
except service.OcrRequestError as exc:
|
|
return _openai_error(str(exc), "invalid_request_error", 400)
|
|
except Exception as exc:
|
|
service._log_exception("chat_completions", request, exc, {})
|
|
return _openai_error(str(exc), "internal_error", 500)
|
|
|
|
content = data.get("text_lines", "") or ""
|
|
ocr_json = data.get("ocr_text_json") if req.ocr_with_boxes else None
|
|
return _build_chat_completion(content, req.model, ocr_json, data.get("elapsed_seconds"))
|