Introduce BaseAIHook to common-ai provider#67373
Draft
gopidesupavan wants to merge 1 commit into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Introduces
BaseAIHook, an abstract hook that defines the contract everyagent-framework backend must implement.
AgentOperator(and the@task.agentdecorator) now selects the backend at runtime from the Airflow connection's
conn_type— no new operator class is needed when adding another framework.What changed
New:
BaseAIHookcontract (hooks/base_ai.py)BaseAIHook(BaseHook)— abstract base with three abstract methods:get_conn()— return the backend model/clientcreate_agent(output_type, instructions, **kwargs)— build an agentrun_agent(agent, *, prompt, usage_limits, message_history) → AgentRunResultAgentRunResultdataclass — backend-neutral result:output,message_history,model_name,usage,tool_namesAgentUsagedataclass — normalized token/request counters:requests,tool_calls,input_tokens,output_tokens,total_tokenssupports_toolsets,supports_durable,supports_usage_limitsBaseAIHook.get_agent_hook(conn_id)— resolves the registered hook fromconn_typeand raisesTypeErrorif it is not aBaseAIHookRefactor:
PydanticAIHook(BaseAIHook)PydanticAIHook(and its subclassesPydanticAIAzureHook,PydanticAIBedrockHook,PydanticAIVertexHook) now subclassBaseAIHook.run_agent()wraps the pydantic-aiRunResultintoAgentRunResult, so allcallers receive a backend-neutral object.
AgentOperatordispatch viaget_agent_hook()AgentOperator.execute()callsBaseAIHook.get_agent_hook(conn_id)instead ofimporting
PydanticAIHookdirectly. This is the main extensibility seam:a future hook registered under a new
conn_typeis picked up with no operatorchanges.
LLMOperator/LLMFileAnalysisOperatoralignedBoth operators previously called
agent.run_sync()directly, bypassing the hookabstraction and receiving a pydantic-ai
RunResultrather than anAgentRunResult. They now callself.llm_hook.run_agent(agent, prompt=…),which means:
log_run_summary()receives a consistentAgentRunResultfrom all operatorsAgentOperatorlogging.pybackend-agnosticlog_run_summary()now readsresult.model_name,result.usage.*, andresult.tool_namesfromAgentRunResultdirectly — no more pydantic-airesult.response.model_nameorresult.usage()callable.Tests
tests/unit/common/ai/hooks/test_base_ai.pycovering the contract,get_agent_hook()dispatch, andTypeErroron non-agent hooktest_pydantic_ai.py,test_agent.py,test_llm.py,test_llm_file_analysis.pyto mockhook.run_agent()returningAgentRunResultinstead of a raw pydantic-aiRunResultWas generative AI tooling used to co-author this PR?
{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.