better

Author: vangheem

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.0.6
 
-* **Support streaming response types for /invoke**
+* **Support streaming response types for /invoke if callable is async generator**
 
 ## 0.0.5
 

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -47,26 +47,15 @@ def log_func_and_body(func: Callable, body: Optional[str] = None) -> None:
         logger.log(level=logger.level, msg=msg)
 
 
-async def run_generator_in_executor(generator_func, **kwargs):
-    loop = asyncio.get_event_loop()
-    # Create a future that will yield the generator values one by one
-    gen = generator_func(**kwargs)
-    while True:
-        result = await loop.run_in_executor(None, next, gen, None)
-        if result is None:
-            break
-        yield result
+async def invoke_async_gen_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> Any:
+    kwargs = kwargs or {}
+    async for val in func(**kwargs):
+        yield val
 
 
 async def invoke_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> Any:
     kwargs = kwargs or {}
-    if inspect.isasyncgenfunction(func):
-        async for val in func(**kwargs):
-            yield val
-    elif inspect.isgeneratorfunction(func):
-        async for val in run_generator_in_executor(func, **kwargs):
-            yield val
-    elif inspect.iscoroutinefunction(func):
+    if inspect.iscoroutinefunction(func):
         yield await func(**kwargs)
     else:
         yield await asyncio.get_event_loop().run_in_executor(None, partial(func, **kwargs))
@@ -136,14 +125,22 @@ async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> In
         else:
             logger.warning("usage data not an expected parameter, omitting")
         try:
-
-            async def _stream_response():
-                async for output in invoke_func(func=func, kwargs=request_dict):
-                    yield InvokeResponse(
-                        usage=usage, status_code=status.HTTP_200_OK, output=output
-                    ).model_dump_json() + "\n"
-
-            return StreamingResponse(_stream_response(), media_type="application/x-ndjson")
+            if inspect.isasyncgenfunction(func):
+                # Stream response if function is an async generator
+
+                async def _stream_response():
+                    async for output in invoke_async_gen_func(func=func, kwargs=request_dict):
+                        yield InvokeResponse(
+                            usage=usage, status_code=status.HTTP_200_OK, output=output
+                        ).model_dump_json() + "\n"
+
+                return StreamingResponse(_stream_response(), media_type="application/x-ndjson")
+            else:
+                return InvokeResponse(
+                    usage=usage,
+                    status_code=status.HTTP_200_OK,
+                    output=await invoke_func(func, request_dict),
+                )
         except Exception as invoke_error:
             logger.error(f"failed to invoke plugin: {invoke_error}", exc_info=True)
             return InvokeResponse(