From bf94b0f201b1c4f98e7b4019f39a7bd9c1834d54 Mon Sep 17 00:00:00 2001
From: Nathan Van Gheem <vangheem@gmail.com>
Date: Fri, 31 Jan 2025 17:07:06 -0500
Subject: [PATCH] handle errors in streaming (#43)

---
 CHANGELOG.md                                  |  4 ++++
 unstructured_platform_plugins/__version__.py  |  2 +-
 .../etl_uvicorn/api_generator.py              | 24 ++++++++++++-------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2047e57..888c75e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.0.23
+
+* **Handle errors in streaming responses**
+
 ## 0.0.22
 
 * **Bump `unstructured-ingest` to 0.4.0**
diff --git a/unstructured_platform_plugins/__version__.py b/unstructured_platform_plugins/__version__.py
index 8e89aab..7ec71da 100644
--- a/unstructured_platform_plugins/__version__.py
+++ b/unstructured_platform_plugins/__version__.py
@@ -1 +1 @@
-__version__ = "0.0.22"  # pragma: no cover
+__version__ = "0.0.23"  # pragma: no cover
diff --git a/unstructured_platform_plugins/etl_uvicorn/api_generator.py b/unstructured_platform_plugins/etl_uvicorn/api_generator.py
index 0920a62..132146b 100644
--- a/unstructured_platform_plugins/etl_uvicorn/api_generator.py
+++ b/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")