fix: issue with streaming partial nested objects

Author: guillaq

workflowai/core/_common_types.py

@@ -19,7 +19,18 @@
 
 
 class OutputValidator(Protocol, Generic[AgentOutputCov]):
-    def __call__(self, data: dict[str, Any], has_tool_call_requests: bool) -> AgentOutputCov: ...
+    def __call__(self, data: dict[str, Any], partial: bool) -> AgentOutputCov:
+        """A way to convert a json object into an AgentOutput
+
+        Args:
+            data (dict[str, Any]): The json object to convert
+            partial (bool): Whether the json is partial, meaning that
+            it may not contain all the fields required by the AgentOutput model.
+
+        Returns:
+            AgentOutputCov: The converted AgentOutput
+        """
+        ...
 
 
 class VersionRunParams(TypedDict):

workflowai/core/client/_fn_utils.py

@@ -23,7 +23,7 @@
     RunParams,
     RunTemplate,
 )
-from workflowai.core.client._utils import intolerant_validator
+from workflowai.core.client._utils import default_validator
 from workflowai.core.client.agent import Agent
 from workflowai.core.domain.errors import InvalidGenerationError
 from workflowai.core.domain.model import ModelOrStr
@@ -144,14 +144,15 @@ async def __call__(self, input: AgentInput, **kwargs: Unpack[RunParams[AgentOutp
         except InvalidGenerationError as e:
             if e.partial_output and e.run_id:
                 try:
-                    validator, _ = self._sanitize_validator(kwargs, intolerant_validator(self.output_cls))
+                    validator, _ = self._sanitize_validator(kwargs, default_validator(self.output_cls))
                     run = self._build_run_no_tools(
                         chunk=RunResponse(
                             id=e.run_id,
                             task_output=e.partial_output,
                         ),
                         schema_id=self.schema_id or 0,
                         validator=validator,
+                        partial=False,
                     )
                     run.error = e.error
                     return run

workflowai/core/client/_models.py

@@ -134,12 +134,18 @@ def to_domain(
         task_id: str,
         task_schema_id: int,
         validator: OutputValidator[AgentOutput],
+        partial: Optional[bool] = None,
     ) -> Run[AgentOutput]:
+        # We do partial validation if either:
+        # - there are tool call requests, which means that the output can be empty
+        # - the run has not yet finished, for exmaple when streaming, in which case the duration_seconds is None
+        if partial is None:
+            partial = bool(self.tool_call_requests) or self.duration_seconds is None
         return Run(
             id=self.id,
             agent_id=task_id,
             schema_id=task_schema_id,
-            output=validator(self.task_output, self.tool_call_requests is not None),
+            output=validator(self.task_output, partial),
             version=self.version and self.version.to_domain(),
             duration_seconds=self.duration_seconds,
             cost_usd=self.cost_usd,

workflowai/core/client/_models_test.py

@@ -5,7 +5,7 @@
 
 from tests.utils import fixture_text
 from workflowai.core.client._models import RunResponse
-from workflowai.core.client._utils import intolerant_validator, tolerant_validator
+from workflowai.core.client._utils import default_validator
 from workflowai.core.domain.run import Run
 from workflowai.core.domain.tool_call import ToolCallRequest
 
@@ -41,7 +41,7 @@ def test_no_version_not_optional(self):
         with pytest.raises(ValidationError):  # sanity
             _TaskOutput.model_validate({"a": 1})
 
-        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=tolerant_validator(_TaskOutput))
+        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=default_validator(_TaskOutput))
         assert isinstance(parsed, Run)
         assert parsed.output.a == 1
         # b is not defined
@@ -52,18 +52,19 @@ def test_no_version_optional(self):
         chunk = RunResponse.model_validate_json('{"id": "1", "task_output": {"a": 1}}')
         assert chunk
 
-        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=tolerant_validator(_TaskOutputOpt))
+        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=default_validator(_TaskOutputOpt))
         assert isinstance(parsed, Run)
         assert parsed.output.a == 1
         assert parsed.output.b is None
 
     def test_with_version(self):
+        """Full output is validated since the duration is passed and there are no tool calls"""
         chunk = RunResponse.model_validate_json(
             '{"id": "1", "task_output": {"a": 1, "b": "test"}, "cost_usd": 0.1, "duration_seconds": 1, "version": {"properties": {"a": 1, "b": "test"}}}',  # noqa: E501
         )
         assert chunk
 
-        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=tolerant_validator(_TaskOutput))
+        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=default_validator(_TaskOutput))
         assert isinstance(parsed, Run)
         assert parsed.output.a == 1
         assert parsed.output.b == "test"
@@ -73,17 +74,19 @@ def test_with_version(self):
 
     def test_with_version_validation_fails(self):
         chunk = RunResponse.model_validate_json(
-            '{"id": "1", "task_output": {"a": 1}, "version": {"properties": {"a": 1, "b": "test"}}}',
+            """{"id": "1", "task_output": {"a": 1},
+            "version": {"properties": {"a": 1, "b": "test"}}, "duration_seconds": 1}""",
         )
         with pytest.raises(ValidationError):
-            chunk.to_domain(task_id="1", task_schema_id=1, validator=intolerant_validator(_TaskOutput))
+            chunk.to_domain(task_id="1", task_schema_id=1, validator=default_validator(_TaskOutput))
 
     def test_with_tool_calls(self):
         chunk = RunResponse.model_validate_json(
-            '{"id": "1", "task_output": {}, "tool_call_requests": [{"id": "1", "name": "test", "input": {"a": 1}}]}',
+            """{"id": "1", "task_output": {},
+            "tool_call_requests": [{"id": "1", "name": "test", "input": {"a": 1}}], "duration_seconds": 1}""",
         )
         assert chunk
 
-        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=tolerant_validator(_TaskOutput))
+        parsed = chunk.to_domain(task_id="1", task_schema_id=1, validator=default_validator(_TaskOutput))
         assert isinstance(parsed, Run)
         assert parsed.tool_call_requests == [ToolCallRequest(id="1", name="test", input={"a": 1})]

workflowai/core/client/_utils.py

@@ -87,18 +87,15 @@ async def _wait_for_exception(e: WorkflowAIError):
     return _should_retry, _wait_for_exception
 
 
-def tolerant_validator(m: type[AgentOutput]) -> OutputValidator[AgentOutput]:
-    def _validator(data: dict[str, Any], has_tool_call_requests: bool) -> AgentOutput:  # noqa: ARG001
-        return construct_model_recursive(m, data)
-
-    return _validator
-
-
-def intolerant_validator(m: type[AgentOutput]) -> OutputValidator[AgentOutput]:
-    def _validator(data: dict[str, Any], has_tool_call_requests: bool) -> AgentOutput:
+def default_validator(m: type[AgentOutput]) -> OutputValidator[AgentOutput]:
+    def _validator(data: dict[str, Any], partial: bool) -> AgentOutput:
         # When we have tool call requests, the output can be empty
-        if has_tool_call_requests:
-            return tolerant_validator(m)(data, has_tool_call_requests)
+        if partial:
+            try:
+                return construct_model_recursive(m, data)
+            except Exception:  # noqa: BLE001
+                logger.warning("Failed to validate partial data: %s", data)
+                return m.model_construct(None, **data)
 
         return m.model_validate(data)
 

workflowai/core/client/_utils_test.py

@@ -6,9 +6,9 @@
 
 from workflowai.core.client._utils import (
     build_retryable_wait,
+    default_validator,
     global_default_version_reference,
     split_chunks,
-    tolerant_validator,
 )
 from workflowai.core.domain.errors import BaseError, WorkflowAIError
 
@@ -59,13 +59,13 @@ class Ingredient(BaseModel):
     ingredients: list[Ingredient]
 
 
-class TestTolerantValidator:
+class TestValidator:
     def test_tolerant_validator_nested_object(self):
-        validated = tolerant_validator(Recipe)(
+        validated = default_validator(Recipe)(
             {
                 "ingredients": [{"name": "salt"}],
             },
-            has_tool_call_requests=False,
+            partial=True,
         )
         for ingredient in validated.ingredients:
             assert isinstance(ingredient, Recipe.Ingredient)

workflowai/core/client/agent.py

@@ -22,9 +22,8 @@
 from workflowai.core.client._types import RunParams
 from workflowai.core.client._utils import (
     build_retryable_wait,
+    default_validator,
     global_default_version_reference,
-    intolerant_validator,
-    tolerant_validator,
 )
 from workflowai.core.domain.completion import Completion
 from workflowai.core.domain.errors import BaseError, WorkflowAIError
@@ -295,8 +294,9 @@ def _build_run_no_tools(
         chunk: RunResponse,
         schema_id: int,
         validator: OutputValidator[AgentOutput],
+        partial: Optional[bool] = None,
     ) -> Run[AgentOutput]:
-        run = chunk.to_domain(self.agent_id, schema_id, validator)
+        run = chunk.to_domain(self.agent_id, schema_id, validator, partial)
         run._agent = self  # pyright: ignore [reportPrivateUsage]
         return run
 
@@ -362,7 +362,7 @@ async def run(
             Run[AgentOutput]: The task run object.
         """
         prepared_run = await self._prepare_run(agent_input, stream=False, **kwargs)
-        validator, new_kwargs = self._sanitize_validator(kwargs, intolerant_validator(self.output_cls))
+        validator, new_kwargs = self._sanitize_validator(kwargs, default_validator(self.output_cls))
 
         last_error = None
         while prepared_run.should_retry():
@@ -374,7 +374,6 @@ async def run(
                     validator,
                     current_iteration=0,
                     # TODO[test]: add test with custom validator
-                    # We popped validator above
                     **new_kwargs,
                 )
             except WorkflowAIError as e:  # noqa: PERF203
@@ -419,10 +418,11 @@ async def stream(
             AsyncIterator[Run[AgentOutput]]: An async iterator yielding task run objects.
         """
         prepared_run = await self._prepare_run(agent_input, stream=True, **kwargs)
-        validator, new_kwargs = self._sanitize_validator(kwargs, tolerant_validator(self.output_cls))
+        validator, new_kwargs = self._sanitize_validator(kwargs, default_validator(self.output_cls))
 
         while prepared_run.should_retry():
             try:
+                chunk: Optional[RunResponse] = None
                 async for chunk in self.api.stream(
                     method="POST",
                     path=prepared_run.route,
@@ -462,7 +462,7 @@ async def reply(
         """
 
         prepared_run = await self._prepare_reply(run_id, user_message, tool_results, stream=False, **kwargs)
-        validator, new_kwargs = self._sanitize_validator(kwargs, intolerant_validator(self.output_cls))
+        validator, new_kwargs = self._sanitize_validator(kwargs, default_validator(self.output_cls))
 
         res = await self.api.post(prepared_run.route, prepared_run.request, returns=RunResponse, run=True)
         return await self._build_run(