fix: streaming payload

Author: guillaq

workflowai/core/client/api.py

@@ -1,4 +1,3 @@
-import json
 from typing import Any, AsyncIterator, Literal, Optional, TypeVar, overload
 
 import httpx
@@ -71,5 +70,4 @@ async def stream(
             ) as response:
                 async for chunk in response.aiter_bytes():
                     stripped = chunk.removeprefix(b"data: ").removesuffix(b"\n\n")
-                    parsed = json.loads(stripped)
-                    yield returns.model_construct(None, **parsed)
+                    yield returns.model_validate_json(stripped)

workflowai/core/client/client.py

@@ -11,6 +11,7 @@
     ImportExampleRequest,
     ImportRunRequest,
     RunRequest,
+    RunTaskStreamChunk,
     TaskRunResponse,
 )
 from workflowai.core.domain.cache_usage import CacheUsage
@@ -113,9 +114,13 @@ async def run(
 
             return res.to_domain(task)
 
-        return self.api.stream(
-            method="POST", path=route, data=request, returns=task.output_class
-        )
+        async def _stream():
+            async for chunk in self.api.stream(
+                method="POST", path=route, data=request, returns=RunTaskStreamChunk
+            ):
+                yield task.output_class.model_construct(None, **chunk.task_output)
+
+        return _stream()
 
     async def import_run(
         self, run: TaskRun[TaskInput, TaskOutput]

workflowai/core/client/client_test.py

@@ -65,9 +65,9 @@ async def test_stream(self, httpx_mock: HTTPXMock, client: Client):
         httpx_mock.add_response(
             stream=IteratorStream(
                 [
-                    b'data: {"message": ""}',
-                    b'data: {"message": "hel"}',
-                    b'data: {"message": "hello"}',
+                    b'data: {"run_id":"1","task_output":{"message":""}}',
+                    b'data: {"run_id":"1","task_output":{"message":"hel"}}',
+                    b'data: {"run_id":"1","task_output":{"message":"hello"}}',
                 ]
             )
         )

workflowai/core/client/models.py

@@ -80,6 +80,11 @@ def from_domain(cls, task_run: TaskRun[TaskInput, TaskOutput]):
         )
 
 
+class RunTaskStreamChunk(BaseModel):
+    run_id: str
+    task_output: dict[str, Any]
+
+
 class TaskRunResponse(BaseModel):
     id: str
     task_id: str