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:
You're right about #70075 -- it 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']`. My reasoning for
`list[str]`/`bool`/`dict` was wrong.
The old `isinstance(self.output_type, type)` check was also guarding the
case where `output_type` isn't a plain type at all. `AgentOperator` passes it
straight through to `Agent(...)` (agent.py:342 -> hooks/pydantic_ai.py:279), so
pydantic-ai's multi-output and marker forms land here, and `TypeAdapter` fails
on them outside the helper's `except (ValidationError, ValueError, TypeError)`:
```
output_type=[Foo, Bar] -> PydanticSchemaGenerationError
(PydanticUserError -> RuntimeError, so not caught)
output_type=ToolOutput(Foo) -> AttributeError: 'ToolOutput' object has no
attribute '__mro__'
```
On pydantic-ai 2.13 / pydantic 2.13.4, `Agent(TestModel(), output_type=[Foo,
Bar])` returns `Foo(a=0)` and `_to_string` gives `'{"a":0}'`. Before this PR
that fell through to `json.loads` and came back as `{'a': 0}`; now it raises,
after the model call and after the reviewer has already approved, so the
approved output is lost.
Widening the helper's `except` doesn't get you back to the old result -- it
returns `raw`, so callers get `'{"a":0}'` instead of the dict. Separating
schema-build failure from validation does:
```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
```
Putting it in `utils/output_type.py` rather than here also covers
`LLMOperator`'s `require_approval` path, which 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]
