Skip to content

Validate the connection type in PydanticAIHook.get_hook - #74626

Merged
Lee-W merged 3 commits into
apache:mainfrom
astronomer:validate-connection-type-pydanticai-hook
Oct 14, 2026
Merged

Lee-W merged 3 commits into
apache:mainfrom
astronomer:validate-connection-type-pydanticai-hook

Conversation

@Lee-W

@Lee-W Lee-W commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

A connection of another type, such as openai, made the Pydantic AI operators fail with "BaseHook.__init__() got an unexpected keyword argument 'model_id'", because Connection.get_hook dispatched to a hook from another provider.

PydanticAIHook.get_hook now raises a ValueError naming the connection, its actual type and the registered pydanticai types when the type resolves to a hook that is not a PydanticAIHook. Subclass types (Azure, Bedrock, Vertex, Snowflake) pass, even if their module has not been imported yet. A missing or unregistered type is still left to Connection.get_hook.


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: [Claude] following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@Lee-W
Lee-W force-pushed the validate-connection-type-pydanticai-hook branch from ba01399 to 5245f68 Compare October 13, 2026 09:32
@Lee-W
Lee-W marked this pull request as ready for review October 13, 2026 09:50
Comment thread providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py Outdated
Comment thread providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
Comment thread providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py Outdated
Comment thread providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@Lee-W
Lee-W force-pushed the validate-connection-type-pydanticai-hook branch from 5245f68 to c8f4542 Compare October 13, 2026 18:01
Lee-W added 3 commits October 14, 2026 07:57
A connection of another type, such as openai, made the Pydantic AI
operators fail with "BaseHook.__init__() got an unexpected keyword
argument 'model_id'", because Connection.get_hook dispatched to a hook
from another provider.

PydanticAIHook.get_hook now raises a ValueError naming the connection,
its actual type and the registered pydanticai types when the type
resolves to a hook that is not a PydanticAIHook. Subclass types
(Azure, Bedrock, Vertex, Snowflake) pass, even if their module has not
been imported yet. A missing or unregistered type is still left to
Connection.get_hook.
Generic and email connections are registered with no hook class, so
`PydanticAIHook.get_hook` raised an `AttributeError` from `import_string`
instead of the new "needs one of [...]" message. They now get the same
message as any other non-Pydantic AI connection type.

Tighten the `get_hook` tests: add a `generic` rejection case, run
`_get_registered_hooks` unpatched once, accept a subclass registered
under a type without the `pydanticai` prefix, and assert on the
`Connection.get_hook` call instead of `get_connection`.
@Lee-W
Lee-W force-pushed the validate-connection-type-pydanticai-hook branch from c8f4542 to 8d0e649 Compare October 14, 2026 06:58
@Lee-W
Lee-W merged commit d6bb744 into apache:main Oct 14, 2026
85 checks passed
@Lee-W
Lee-W deleted the validate-connection-type-pydanticai-hook branch October 14, 2026 08:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants