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>
97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
import base64
|
|
import io
|
|
from unittest.mock import patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
from PIL import Image
|
|
|
|
|
|
def _png_data_url() -> str:
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (8, 8), "white").save(buf, format="PNG")
|
|
b64 = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
return f"data:image/png;base64,{b64}"
|
|
|
|
|
|
def _client() -> TestClient:
|
|
from surya.endpoint.app import app
|
|
|
|
return TestClient(app)
|
|
|
|
|
|
def _image_message():
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image_url", "image_url": {"url": _png_data_url()}},
|
|
{"type": "text", "text": "extract the text"},
|
|
],
|
|
}
|
|
]
|
|
|
|
|
|
def test_returns_openai_chat_completion_shape():
|
|
body = {"model": "surya-ocr", "messages": _image_message()}
|
|
fake = {"text_lines": "hello\nworld", "ocr_text_json": {"blocks": []}, "elapsed_seconds": 1.23}
|
|
with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 200
|
|
d = r.json()
|
|
assert d["object"] == "chat.completion"
|
|
assert d["model"] == "surya-ocr"
|
|
assert d["choices"][0]["message"]["content"] == "hello\nworld"
|
|
assert d["choices"][0]["finish_reason"] == "stop"
|
|
assert d["surya"]["ocr_text_json"] == {"blocks": []}
|
|
assert d["surya"]["elapsed_seconds"] == 1.23
|
|
assert "usage" in d
|
|
|
|
|
|
def test_mode_full_page_skips_text_detection():
|
|
body = {"model": "m", "mode": "full_page", "messages": _image_message()}
|
|
fake = {"text_lines": "x", "ocr_text_json": {}, "elapsed_seconds": None}
|
|
with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake) as m:
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 200
|
|
assert m.call_args.kwargs["skip_text_detection"] is True
|
|
|
|
|
|
def test_mode_table_routes_to_table_image():
|
|
body = {"model": "m", "mode": "table", "messages": _image_message()}
|
|
fake = {"text_lines": "tbl", "ocr_text_json": {"x": 1}, "elapsed_seconds": None}
|
|
with patch("surya.endpoint.service.table_image", return_value=fake) as m:
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 200
|
|
assert m.called
|
|
assert r.json()["choices"][0]["message"]["content"] == "tbl"
|
|
|
|
|
|
def test_ocr_with_boxes_false_omits_structured_json():
|
|
body = {"model": "m", "ocr_with_boxes": False, "messages": _image_message()}
|
|
fake = {"text_lines": "x", "ocr_text_json": {"a": 1}, "elapsed_seconds": None}
|
|
with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.json()["surya"]["ocr_text_json"] is None
|
|
|
|
|
|
def test_stream_true_returns_400():
|
|
body = {"model": "m", "stream": True, "messages": _image_message()}
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 400
|
|
assert r.json()["error"]["type"] == "invalid_request_error"
|
|
|
|
|
|
def test_missing_image_returns_400():
|
|
body = {"model": "m", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 400
|
|
assert r.json()["error"]["type"] == "invalid_request_error"
|
|
|
|
|
|
def test_pipeline_error_returns_500():
|
|
body = {"model": "m", "messages": _image_message()}
|
|
with patch("surya.endpoint.service.ocr_via_batcher", side_effect=RuntimeError("boom")):
|
|
r = _client().post("/v1/chat/completions", json=body)
|
|
assert r.status_code == 500
|
|
assert r.json()["error"]["type"] == "internal_error"
|