perf(api): reduce response serialization copies

prepline_general/api/general.py

@@ -13,6 +13,7 @@
 from typing import IO, Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast
 
 import backoff
+import orjson
 import pandas as pd
 import psutil
 import requests
@@ -25,12 +26,13 @@
     UploadFile,
     status,
 )
-from fastapi.responses import PlainTextResponse, StreamingResponse
+from fastapi.responses import PlainTextResponse, Response, StreamingResponse
 from pypdf import PageObject, PdfReader, PdfWriter
 from pypdf.errors import FileNotDecryptedError, PdfReadError
 from starlette.datastructures import Headers
 from starlette.types import Send
 
+from prepline_general.api import __version__ as api_version
 from prepline_general.api.filetypes import get_validated_mimetype
 from prepline_general.api.models.form_params import GeneralFormParams
 from unstructured.documents.elements import Element
@@ -41,7 +43,6 @@
     elements_from_json,
 )
 from unstructured_inference.models.base import UnknownModelException
-from prepline_general.api import __version__ as api_version
 
 app = FastAPI()
 router = APIRouter()
@@ -574,14 +575,16 @@ def _build_part_headers(self, headers: Dict[str, Any]) -> bytes:
             header_bytes += f"{header}: {value}".encode() + self.CRLF
         return header_bytes
 
-    def build_part(self, chunk: bytes) -> bytes:
+    def _build_part_prefix(self, content_length: int) -> bytes:
         part = self.boundary + self.CRLF
-        part_headers = {"Content-Length": len(chunk), "Content-Transfer-Encoding": "base64"}
+        part_headers = {
+            "Content-Length": content_length,
+            "Content-Transfer-Encoding": "base64",
+        }
         if self.content_type is not None:
             part_headers["Content-Type"] = self.content_type
         part += self._build_part_headers(part_headers)
-        part += self.CRLF + chunk + self.CRLF
-        return part
+        return part + self.CRLF
 
     async def stream_response(self, send: Send) -> None:
         await send(
@@ -594,10 +597,16 @@ async def stream_response(self, send: Send) -> None:
         async for chunk in self.body_iterator:
             if not isinstance(chunk, bytes):
                 chunk = chunk.encode(self.charset)  # type: ignore
-                chunk = b64encode(chunk)
+            chunk = b64encode(chunk)
             await send(
-                {"type": "http.response.body", "body": self.build_part(chunk), "more_body": True}
+                {
+                    "type": "http.response.body",
+                    "body": self._build_part_prefix(len(chunk)),
+                    "more_body": True,
+                }
             )
+            await send({"type": "http.response.body", "body": chunk, "more_body": True})
+            await send({"type": "http.response.body", "body": self.CRLF, "more_body": True})
 
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
@@ -726,7 +735,7 @@ def response_generator(is_multipart: bool):
             )
 
             yield (
-                json.dumps(response)
+                orjson.dumps(response)
                 if is_multipart and type(response) not in [str, bytes]
                 else (
                     PlainTextResponse(response)
@@ -755,17 +764,17 @@ def join_responses(
                 )
         return PlainTextResponse(data.to_csv())
 
-    return (
-        MultipartMixedResponse(
+    if accept_type == "multipart/mixed":
+        return MultipartMixedResponse(
             response_generator(is_multipart=True), content_type=form_params.output_format
         )
-        if accept_type == "multipart/mixed"
-        else (
-            list(response_generator(is_multipart=False))[0]
-            if len(files) == 1
-            else join_responses(list(response_generator(is_multipart=False)))
-        )
-    )
+
+    responses = list(response_generator(is_multipart=False))
+    data = responses[0] if len(files) == 1 else join_responses(responses)
+    if isinstance(data, PlainTextResponse):
+        return data
+
+    return Response(content=orjson.dumps(data), media_type="application/json")
 
 
 app.include_router(router)

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "pandas >=3.0.0, <4.0.0",
     "psutil >=7.0.0, <8.0.0",
     "pypdf >=6.0.0, <7.0.0",
+    "orjson >=3.11.7, <4.0.0",
     "requests >=2.32.0, <3.0.0",
 ]
 

test_general/api/test_response_serialization.py

@@ -0,0 +1,55 @@
+import asyncio
+import io
+from base64 import b64encode
+
+import orjson
+import pytest
+from fastapi.testclient import TestClient
+
+from prepline_general.api import general
+from prepline_general.api.app import app
+from prepline_general.api.general import MultipartMixedResponse
+
+
+def test_multipart_response_sends_framing_and_payload_separately():
+    response = MultipartMixedResponse(
+        iter([b"serialized response"]), content_type="application/json"
+    )
+    events = []
+
+    async def send(event):
+        events.append(event)
+
+    asyncio.run(response.stream_response(send))
+
+    body_events = [event for event in events if event["type"] == "http.response.body"]
+    encoded_payload = b64encode(b"serialized response")
+
+    assert body_events[0]["body"].startswith(response.boundary)
+    assert f"Content-Length: {len(encoded_payload)}".encode() in body_events[0]["body"]
+    assert body_events[1]["body"] == encoded_payload
+    assert body_events[2]["body"] == response.CRLF
+    assert body_events[3] == {
+        "type": "http.response.body",
+        "body": b"",
+        "more_body": False,
+    }
+
+
+@pytest.mark.parametrize("file_count", [1, 2])
+def test_json_endpoint_serializes_response_directly_to_bytes(monkeypatch, file_count):
+    file_response = [{"text": "partitioned text", "metadata": {}}]
+    monkeypatch.delenv("UNSTRUCTURED_API_KEY", raising=False)
+    monkeypatch.setattr(general, "get_validated_mimetype", lambda *args, **kwargs: "text/plain")
+    monkeypatch.setattr(general, "pipeline_api", lambda *args, **kwargs: file_response)
+    files = [
+        ("files", (f"sample-{idx}.txt", io.BytesIO(b"sample"), "text/plain"))
+        for idx in range(file_count)
+    ]
+
+    response = TestClient(app).post("/general/v0/general", files=files)
+
+    expected = file_response if file_count == 1 else [file_response] * file_count
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/json"
+    assert response.content == orjson.dumps(expected)