kaxil commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3680650264
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -493,16 +492,11 @@ def execute(self, context: Context) -> Any:
output,
message_history=result.all_messages(),
)
- if isinstance(self.output_type, type) and
issubclass(self.output_type, BaseModel):
- return rehydrate_pydantic_output(
- self.output_type,
- result_str,
- serialize_output=self._serialize_model_output,
- )
- try:
- return json.loads(result_str)
- except (ValueError, TypeError):
- return result_str
+ return rehydrate_pydantic_output(
Review Comment:
Half of this is mine to withdraw: you're right that #70075 is an ancestor of
this branch, and `rehydrate_pydantic_output(list[str], '["tag-a","tag-b"]',
serialize_output=False)` does return `['tag-a', 'tag-b']`. For `list[str]`,
`bool` and `dict` my reasoning was wrong, and the `json.loads` fallback is
genuinely redundant for those.
The objection survives for the `output_type` forms that aren't a plain
`type`, which is what the old `isinstance(self.output_type, type)` branch was
routing around. `AgentOperator` forwards `output_type` verbatim to `Agent(...)`
(agent.py:342 -> hooks/pydantic_ai.py:279), so pydantic-ai's multi-output and
marker forms reach this line, and `TypeAdapter` fails on them outside the
helper's `except (ValidationError, ValueError, TypeError)`:
```
output_type=[Foo, Bar] -> PydanticSchemaGenerationError (MRO is
PydanticUserError -> RuntimeError, not TypeError)
output_type=ToolOutput(Foo) -> AttributeError: 'ToolOutput' object has no
attribute '__mro__'
```
Verified on pydantic-ai 2.13 / pydantic 2.13.4: `Agent(TestModel(),
output_type=[Foo, Bar])` runs and returns `Foo(a=0)`, `_to_string` gives
`'{"a":0}'`, the old branch fell through to `json.loads` and returned `{'a':
0}`, and the new unconditional call raises. Same for `ToolOutput`. So those
configs go from working to a hard task failure -- and it raises after the LLM
spend and after the reviewer already clicked approve, so the approved output is
lost.
Restoring `json.loads` at the call site isn't the fix I'd suggest, but
widening the helper's `except` isn't equivalent either: it falls through to
`return raw` and hands downstream `'{"a":0}'` instead of the dict it used to
get. Splitting schema-build failure from validation failure keeps the old
result:
```python
try:
adapter = TypeAdapter(output_type)
except Exception:
# No pydantic schema for this output_type:
ToolOutput/NativeOutput/PromptedOutput
# markers, output functions, or a `[A, B]` union list. Fall back to
plain JSON.
try:
return json.loads(raw)
except (ValueError, TypeError):
return raw
try:
return adapter.validate_json(raw)
except (ValidationError, ValueError, TypeError):
return raw
```
Doing it in `utils/output_type.py` rather than here also closes the same
hole on `LLMOperator`'s `require_approval` path, which already calls the helper
at llm.py:187.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
