handle errors in streaming (#43)

Author: vangheem

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.23
+
+* **Handle errors in streaming responses**
+
 ## 0.0.22
 
 * **Bump `unstructured-ingest` to 0.4.0**

unstructured_platform_plugins/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.22"  # pragma: no cover
+__version__ = "0.0.23"  # pragma: no cover

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -132,7 +132,7 @@ def _wrap_in_fastapi(
     class InvokeResponse(BaseModel):
         usage: list[UsageData]
         status_code: int
-        filedata_meta: filedata_meta_model
+        filedata_meta: Optional[filedata_meta_model]
         status_code_text: Optional[str] = None
         output: Optional[response_type] = None
 
@@ -156,16 +156,24 @@ async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> Re
         try:
             if inspect.isasyncgenfunction(func):
                 # Stream response if function is an async generator
-
                 async def _stream_response():
-                    async for output in func(**(request_dict or {})):
+                    try:
+                        async for output in func(**(request_dict or {})):
+                            yield InvokeResponse(
+                                usage=usage,
+                                filedata_meta=filedata_meta_model.model_validate(
+                                    filedata_meta.model_dump()
+                                ),
+                                status_code=status.HTTP_200_OK,
+                                output=output,
+                            ).model_dump_json() + "\n"
+                    except Exception as e:
+                        logger.error(f"Failure streaming response: {e}", exc_info=True)
                         yield InvokeResponse(
                             usage=usage,
-                            filedata_meta=filedata_meta_model.model_validate(
-                                filedata_meta.model_dump()
-                            ),
-                            status_code=status.HTTP_200_OK,
-                            output=output,
+                            filedata_meta=None,
+                            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                            status_code_text=f"[{e.__class__.__name__}] {e}",
                         ).model_dump_json() + "\n"
 
                 return StreamingResponse(_stream_response(), media_type="application/x-ndjson")