diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..bf9f951 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,60 @@ +name: Docker CI Pipeline + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + docker-integration-test: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build Docker image + run: | + docker build -t api-ndlocr-lite:latest . + + - name: Start API container + run: | + docker run -d --name ndlocr-api -p 8001:8000 api-ndlocr-lite:latest + + - name: Wait for API to become healthy + run: | + echo "Waiting for API startup (this may take a few seconds)..." + # Poll the healthcheck endpoint up to 12 times with 5s sleep (1 minute total max) + for i in {1..12}; do + if curl -s -f http://localhost:8001/health; then + echo "API is healthy!" + exit 0 + fi + echo "Waiting... (attempt $i/12)" + sleep 5 + done + echo "Error: API did not become healthy in time." + docker logs ndlocr-api + exit 1 + + - name: Run basic API smoke test + run: | + # Query health endpoint + HEALTH_STATUS=$(curl -s http://localhost:8001/health) + echo "Health response: $HEALTH_STATUS" + if [[ "$HEALTH_STATUS" != *"ok"* ]]; then + echo "Unexpected health status response" + exit 1 + fi + + - name: Stop container + if: always() + run: | + docker stop ndlocr-api || true + docker rm ndlocr-api || true diff --git a/Dockerfile b/Dockerfile index 735556c..d53de6e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y \ libxext6 \ libxrender-dev \ git \ + curl \ && rm -rf /var/lib/apt/lists/* # Install uv diff --git a/docker-compose.yml b/docker-compose.yml index e6dd8f7..4b86398 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -14,3 +14,12 @@ services: interval: 30s timeout: 10s retries: 3 + +# Example configuration for integration with LibreChat or other services: +# +# librechat: +# image: ghcr.io/danny-avila/librechat:latest +# ... +# depends_on: +# api: +# condition: service_healthy diff --git a/pyproject.toml b/pyproject.toml index 65f29d4..0635a3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,9 @@ [project] name = "api-ndlocr-lite" version = "0.1.0" -description = "Add your description here" +description = "OpenAI-compatible OCR API server based on NDLOCR-Lite, optimized for LibreChat, Code Interpreter, and secure sandbox environments." readme = "README.md" +keywords = ["librechat", "code-interpreter", "sandbox", "fastapi"] requires-python = ">=3.13" dependencies = [ "dill==0.3.8", diff --git a/src/api/main.py b/src/api/main.py index 00b9a02..4fb72e2 100644 --- a/src/api/main.py +++ b/src/api/main.py @@ -153,11 +153,17 @@ async def ocr_endpoint( if engine is None: raise HTTPException(status_code=503, detail="Engine not initialized") + if await request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + try: # Run CPU-bound OCR processing in a thread pool to avoid blocking the event loop loop = asyncio.get_running_loop() result = await loop.run_in_executor(None, engine.ocr, img, filename) + if await request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + # Convert and return response page = _engine_result_to_ocr_page(result) return OCRResponse(model="ndlocr-lite", pages=[page], usage={"pages": 1}) @@ -182,6 +188,9 @@ async def create_ocr_job( # Extract image first to ensure it's valid before accepting the job img, filename = await _get_image_from_request(request, file) + if await request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + job_id = str(uuid.uuid4()) job_store.set(job_id, OCRJobResult(job_id=job_id, status="pending")) @@ -202,6 +211,7 @@ async def get_ocr_job(job_id: str, job_store: InMemoryJobStore = Depends(get_job @app.post("/v1/chat/completions", response_model=ChatCompletionResponse) async def openai_vision_endpoint( + raw_request: Request, request: ChatCompletionRequest, engine: Optional[NDLOCREngine] = Depends(get_engine), ): @@ -212,6 +222,9 @@ async def openai_vision_endpoint( if engine is None: raise HTTPException(status_code=503, detail="Engine not initialized") + if await raw_request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + img_data = None for message in request.messages: if message.role == "user" and isinstance(message.content, list): @@ -235,10 +248,16 @@ async def openai_vision_endpoint( if img.width * img.height > MAX_PIXELS: raise HTTPException(status_code=400, detail="Image dimensions too large") + if await raw_request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + # Run OCR loop = asyncio.get_running_loop() result = await loop.run_in_executor(None, engine.ocr, img, "openai_image.jpg") + if await raw_request.is_disconnected(): + raise HTTPException(status_code=499, detail="Client closed connection") + # Format response return ChatCompletionResponse( id=f"chatcmpl-{uuid.uuid4()}", @@ -289,12 +308,17 @@ async def _get_image_from_request(request: Request, file: Optional[UploadFile]): Includes security checks for body size, file size, and image dimensions. """ + from starlette.requests import ClientDisconnect img = None filename = "image.jpg" try: if file: # Handle multipart/form-data - contents = await file.read(MAX_IMAGE_SIZE + 1) + try: + contents = await file.read(MAX_IMAGE_SIZE + 1) + except ClientDisconnect: + logger.warning("Client disconnected during file upload") + raise HTTPException(status_code=499, detail="Client closed connection") if len(contents) > MAX_IMAGE_SIZE: raise HTTPException(status_code=413, detail="File too large") img = Image.open(io.BytesIO(contents)) @@ -306,10 +330,14 @@ async def _get_image_from_request(request: Request, file: Optional[UploadFile]): raise HTTPException(status_code=413, detail="Request body too large") body_bytes = b"" - async for chunk in request.stream(): - body_bytes += chunk - if len(body_bytes) > MAX_BODY_SIZE: - raise HTTPException(status_code=413, detail="Request body too large") + try: + async for chunk in request.stream(): + body_bytes += chunk + if len(body_bytes) > MAX_BODY_SIZE: + raise HTTPException(status_code=413, detail="Request body too large") + except ClientDisconnect: + logger.warning("Client disconnected during body streaming") + raise HTTPException(status_code=499, detail="Client closed connection") if not body_bytes: raise HTTPException(status_code=400, detail="Empty request body") diff --git a/src/core/engine.py b/src/core/engine.py index d55b0df..fc6c2c8 100644 --- a/src/core/engine.py +++ b/src/core/engine.py @@ -129,7 +129,7 @@ def _get_recognizer(self, weights_path: str): def shutdown(self): """Shuts down the internal thread pool.""" - self.executor.shutdown() + self.executor.shutdown(wait=False, cancel_futures=True) def _process_cascade(self, alllineobj: List[RecogLine], is_cascade: bool = True) -> List[str]: """ diff --git a/tests/test_boundary_conditions.py b/tests/test_boundary_conditions.py new file mode 100644 index 0000000..f2f10f3 --- /dev/null +++ b/tests/test_boundary_conditions.py @@ -0,0 +1,205 @@ +import asyncio +import io +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from starlette.requests import ClientDisconnect +from starlette.requests import Request +from starlette.datastructures import Headers +from fastapi import HTTPException +from fastapi.testclient import TestClient +from PIL import Image + +from src.api.main import app, _get_image_from_request +from src.core.engine import NDLOCREngine + +@pytest.mark.asyncio +async def test_get_image_from_request_disconnect_file(): + # Mocking file upload exception during read + mock_request = MagicMock(spec=Request) + mock_file = AsyncMock() + mock_file.read.side_effect = ClientDisconnect() + + with pytest.raises(HTTPException) as exc_info: + await _get_image_from_request(mock_request, mock_file) + + assert exc_info.value.status_code == 499 + assert exc_info.value.detail == "Client closed connection" + + +@pytest.mark.asyncio +async def test_get_image_from_request_disconnect_body(): + # Mocking stream request chunk reading exception + mock_request = MagicMock(spec=Request) + mock_request.headers = Headers() + + async def mock_stream(): + raise ClientDisconnect() + yield b"" # unreachable, but defines it as generator/async iterator + + mock_request.stream = mock_stream + + with pytest.raises(HTTPException) as exc_info: + await _get_image_from_request(mock_request, None) + + assert exc_info.value.status_code == 499 + assert exc_info.value.detail == "Client closed connection" + + +def test_ocr_endpoint_early_disconnect(monkeypatch): + # If client is already disconnected when hitting the endpoint + async def mock_is_disconnected(self): + return True + + monkeypatch.setattr(Request, "is_disconnected", mock_is_disconnected) + + img = Image.new('RGB', (100, 100), color=(255, 255, 255)) + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + img_bytes = img_byte_arr.getvalue() + + with TestClient(app) as client: + response = client.post( + "/v1/ocr", + files={"file": ("test.jpg", img_bytes, "image/jpeg")} + ) + assert response.status_code == 499 + assert response.json()["detail"] == "Client closed connection" + + +def test_ocr_endpoint_post_run_disconnect(monkeypatch): + # Client disconnects right after engine.ocr runs + disconnect_call_count = 0 + + async def mock_is_disconnected(self): + nonlocal disconnect_call_count + disconnect_call_count += 1 + # First check (before engine ocr) -> False + # Second check (after engine ocr) -> True + return disconnect_call_count > 1 + + monkeypatch.setattr(Request, "is_disconnected", mock_is_disconnected) + + img = Image.new('RGB', (100, 100), color=(255, 255, 255)) + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + img_bytes = img_byte_arr.getvalue() + + with TestClient(app) as client: + response = client.post( + "/v1/ocr", + files={"file": ("test.jpg", img_bytes, "image/jpeg")} + ) + assert response.status_code == 499 + assert response.json()["detail"] == "Client closed connection" + + +def test_engine_shutdown_with_pending_tasks(): + # Verify that engine shutdown terminates and cancels pending tasks in executor + engine = NDLOCREngine(device="cpu") + + # Mocking actual model execution to avoid overhead or dependencies + engine._load_models = MagicMock() + + # Setup some heavy task in executor + def dummy_task(): + import time + time.sleep(10) + + future = engine.executor.submit(dummy_task) + assert future.running() or future.pending() + + # Shutdown engine + engine.shutdown() + + # Submitting a task post-shutdown should raise a RuntimeError + with pytest.raises(RuntimeError): + engine.executor.submit(dummy_task) + + +def test_openai_vision_endpoint_early_disconnect(monkeypatch): + # Test OpenAI chat completion endpoint when client is disconnected early + async def mock_is_disconnected(self): + return True + + monkeypatch.setattr(Request, "is_disconnected", mock_is_disconnected) + + img = Image.new('RGB', (50, 50), color=(255, 255, 255)) + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + import base64 + b64_str = "data:image/jpeg;base64," + base64.b64encode(img_byte_arr.getvalue()).decode() + + payload = { + "model": "ndlocr-lite", + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": b64_str}} + ] + } + ] + } + + with TestClient(app) as client: + response = client.post("/v1/chat/completions", json=payload) + assert response.status_code == 499 + assert response.json()["detail"] == "Client closed connection" + + +def test_openai_vision_endpoint_post_run_disconnect(monkeypatch): + # Test OpenAI chat completion endpoint when client disconnects right after engine.ocr runs + disconnect_call_count = 0 + + async def mock_is_disconnected(self): + nonlocal disconnect_call_count + disconnect_call_count += 1 + return disconnect_call_count > 1 + + monkeypatch.setattr(Request, "is_disconnected", mock_is_disconnected) + + img = Image.new('RGB', (50, 50), color=(255, 255, 255)) + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + import base64 + b64_str = "data:image/jpeg;base64," + base64.b64encode(img_byte_arr.getvalue()).decode() + + payload = { + "model": "ndlocr-lite", + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": b64_str}} + ] + } + ] + } + + with TestClient(app) as client: + response = client.post("/v1/chat/completions", json=payload) + assert response.status_code == 499 + assert response.json()["detail"] == "Client closed connection" + + +def test_create_ocr_job_early_disconnect(monkeypatch): + # Test /v1/ocr/jobs endpoint when client is disconnected early + async def mock_is_disconnected(self): + return True + + monkeypatch.setattr(Request, "is_disconnected", mock_is_disconnected) + + img = Image.new('RGB', (50, 50), color=(255, 255, 255)) + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + img_bytes = img_byte_arr.getvalue() + + with TestClient(app) as client: + response = client.post( + "/v1/ocr/jobs", + files={"file": ("test.jpg", img_bytes, "image/jpeg")} + ) + assert response.status_code == 499 + assert response.json()["detail"] == "Client closed connection" + +