This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f87aa58e09d Migrate common.ai provider to pydantic-ai 2.x and remove
the <2 cap (#69358)
f87aa58e09d is described below
commit f87aa58e09d18710ef6db57d2f4661a28f6ba6e6
Author: Kaxil Naik <[email protected]>
AuthorDate: Fri Jul 3 22:00:46 2026 +0100
Migrate common.ai provider to pydantic-ai 2.x and remove the <2 cap (#69358)
pydantic-ai 2.x removed the ``instrument`` keyword from the ``Agent()`` and
``Agent.from_file()`` constructors (it is now only the ``agent.instrument``
property), which broke ``PydanticAIHook.create_agent`` when a caller passed
``instrument``. Pop a caller-supplied ``instrument`` out of the constructor
kwargs and apply it through the property instead, so the caller still wins
over
the provider's auto-instrumentation.
- Bump the OTel GenAI semantic-convention format to version 5 (formats 2-4
are
deprecated in 2.x and warn). Agent-run token usage now reports under
``gen_ai.aggregated_usage.*`` while model-request spans keep
``gen_ai.usage.*``;
documented in observability.rst and asserted in the end-to-end span test.
- Update ``Agent[None, ...]`` annotations to ``Agent[object, ...]`` to
match the
2.x generic-deps default (object, not None).
- Raise the floor to ``pydantic-ai-slim>=2.0.0`` and re-lock (slim +
pydantic-graph
1.107.0 -> 2.1.0).
Closes #69122
---
providers/common/ai/README.rst | 2 +-
providers/common/ai/docs/index.rst | 2 +-
providers/common/ai/docs/observability.rst | 8 +++++
providers/common/ai/pyproject.toml | 10 +++----
.../providers/common/ai/hooks/pydantic_ai.py | 34 ++++++++++++++--------
.../airflow/providers/common/ai/observability.py | 11 ++++---
.../airflow/providers/common/ai/operators/agent.py | 2 +-
.../airflow/providers/common/ai/operators/llm.py | 2 +-
.../common/ai/operators/llm_file_analysis.py | 2 +-
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 15 ++++++++--
.../ai/tests/unit/common/ai/test_observability.py | 9 ++++--
uv.lock | 14 ++++-----
12 files changed, 73 insertions(+), 38 deletions(-)
diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst
index 3c5439b4e83..30a30c8e09a 100644
--- a/providers/common/ai/README.rst
+++ b/providers/common/ai/README.rst
@@ -56,7 +56,7 @@ PIP package Version required
``apache-airflow`` ``>=3.0.0``
``apache-airflow-providers-common-compat`` ``>=1.15.0``
``apache-airflow-providers-standard`` ``>=1.12.1``
-``pydantic-ai-slim`` ``>=1.99.0,<2``
+``pydantic-ai-slim`` ``>=2.0.0``
========================================== ==================
Optional cross provider package dependencies
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 0be69bb300b..385ed423861 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -109,7 +109,7 @@ PIP package Version required
``apache-airflow`` ``>=3.0.0``
``apache-airflow-providers-common-compat`` ``>=1.15.0``
``apache-airflow-providers-standard`` ``>=1.12.1``
-``pydantic-ai-slim`` ``>=1.99.0,<2``
+``pydantic-ai-slim`` ``>=2.0.0``
========================================== ==================
Optional cross provider package dependencies
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 6669419c8aa..54d5f5f7e9a 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -52,6 +52,14 @@ How it works
names, and finish reason are recorded. Prompt and completion text is never
emitted unless you opt in (see below).
+.. note::
+
+ On pydantic-ai 2.x the agent-run span reports token usage under
+ ``gen_ai.aggregated_usage.*`` while the per-model-call span keeps
+ ``gen_ai.usage.*``. This avoids double-counting in backends that sum a
+ parent span and its children. Dashboards or alerts that read run-level
token
+ usage from ``gen_ai.usage.*`` should switch to
``gen_ai.aggregated_usage.*``.
+
Enabling it
-----------
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index f4597e957f9..4531d2cfb8c 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -69,9 +69,8 @@ dependencies = [
"apache-airflow>=3.0.0",
"apache-airflow-providers-common-compat>=1.15.0",
"apache-airflow-providers-standard>=1.12.1",
- # Capped to 1.x: pydantic-ai 2.x changed the agent instrumentation API and
breaks the provider.
- # Remove the cap after migrating; tracked at
https://github.com/apache/airflow/issues/69122
- "pydantic-ai-slim>=1.99.0,<2",
+ # Requires the pydantic-ai 2.x agent/instrumentation API (see #69122).
+ "pydantic-ai-slim>=2.0.0",
]
# The optional dependencies should be modified in place in the generated file
@@ -87,9 +86,8 @@ dependencies = [
# Enables AgentOperator(code_mode=True). Monty is pre-1.0; pinned here as an
# opt-in extra so its churn never breaks the base provider install.
"code-mode" = ["pydantic-ai-harness[codemode]>=0.3.0"]
-# Agent Skills (agentskills.io) support. pydantic-ai-skills provides the
toolset
-# (pulls in pydantic-ai-slim>=1.74 transitively; the provider base floor stays
-# 1.71); the git provider supplies GitHook + GitPython for cloning GitSkills
with
+# Agent Skills (agentskills.io) support. pydantic-ai-skills provides the
toolset;
+# the git provider supplies GitHook + GitPython for cloning GitSkills with
# credentials from a `git` connection. Native progressive disclosure is tracked
# upstream in pydantic/pydantic-ai#5230; revisit this extra once that lands.
"skills" = [
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
index eec7c5e7944..c46f974a59c 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
@@ -28,6 +28,11 @@ from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
+# Sentinel distinguishing "caller did not pass ``instrument``" from an explicit
+# ``instrument=None`` / ``instrument=False`` (which mean "do not instrument,
and
+# do not auto-enable it either").
+_UNSET: Any = object()
+
if TYPE_CHECKING:
from pydantic_ai.models import KnownModelName, Model
@@ -196,10 +201,10 @@ class PydanticAIHook(BaseHook):
@overload
def create_agent(
self, output_type: type[OutputT], *, instructions: str, **agent_kwargs
- ) -> Agent[None, OutputT]: ...
+ ) -> Agent[object, OutputT]: ...
@overload
- def create_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[None, str]: ...
+ def create_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[object, str]: ...
@overload
def create_agent(
@@ -209,7 +214,7 @@ class PydanticAIHook(BaseHook):
spec_file: str | Path,
instructions: str | None = ...,
**agent_kwargs,
- ) -> Agent[None, OutputT]: ...
+ ) -> Agent[object, OutputT]: ...
@overload
def create_agent(
@@ -218,7 +223,7 @@ class PydanticAIHook(BaseHook):
spec_file: str | Path,
instructions: str | None = ...,
**agent_kwargs,
- ) -> Agent[None, str]: ...
+ ) -> Agent[object, str]: ...
def create_agent(
self,
@@ -227,7 +232,7 @@ class PydanticAIHook(BaseHook):
instructions: str | None = None,
spec_file: str | Path | None = None,
**agent_kwargs,
- ) -> Agent[None, Any]:
+ ) -> Agent[object, Any]:
"""
Create a pydantic-ai Agent configured with this hook's model.
@@ -247,6 +252,14 @@ class PydanticAIHook(BaseHook):
the spec file's ``model`` is used.
:param agent_kwargs: Additional keyword arguments passed to the Agent
constructor.
"""
+ # ``instrument`` is no longer an ``Agent()`` / ``Agent.from_file()``
+ # constructor argument in pydantic-ai 2.x; it is configured through the
+ # ``agent.instrument`` property (which is unchanged across the 2.x
line).
+ # Pop any caller-supplied value out of the constructor kwargs and apply
+ # it after construction so a caller that passes its own ``instrument``
+ # still wins over the provider's auto-instrumentation.
+ caller_instrument = agent_kwargs.pop("instrument", _UNSET)
+
if spec_file is not None:
from_file_kwargs = dict(agent_kwargs)
model = self._get_conn_if_model_configured()
@@ -264,13 +277,10 @@ class PydanticAIHook(BaseHook):
if instructions is None:
raise ValueError("instructions is required when spec_file is
not provided.")
agent = Agent(self.get_conn(), output_type=output_type,
instructions=instructions, **agent_kwargs)
- if "instrument" not in agent_kwargs:
- # Set the public ``agent.instrument`` surface rather than the
- # ``Agent(instrument=...)`` constructor kwarg, which is deprecated
in
- # current pydantic-ai. Assigning ``agent.instrument`` works across
the
- # provider's ``pydantic-ai-slim>=1.71`` floor (a plain instance
- # attribute on older versions, a property on newer ones). A caller
- # that passed its own ``instrument`` wins.
+
+ if caller_instrument is not _UNSET:
+ agent.instrument = caller_instrument
+ else:
settings = genai_instrumentation_settings()
if settings is not None:
agent.instrument = settings
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/observability.py
b/providers/common/ai/src/airflow/providers/common/ai/observability.py
index e81a6680d04..08a59c40595 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/observability.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/observability.py
@@ -40,10 +40,13 @@ if TYPE_CHECKING:
SECTION = "common.ai"
-# OTel GenAI semantic-convention attribute set. Pinned so the emitted span and
-# ``gen_ai.*`` attribute names stay stable regardless of the pydantic-ai
default
-# (which tracks the latest, still-evolving revision).
-_SEMCONV_VERSION: Literal[4] = 4
+# OTel GenAI semantic-convention format version. Pinned so a change in the
+# pydantic-ai default does not silently shift the emitted span/attribute format
+# between provider releases. Version 5 is the current default in pydantic-ai
2.x;
+# formats 2-4 still work but are deprecated. Note: independent of this version,
+# pydantic-ai 2.x reports agent-run token usage under
``gen_ai.aggregated_usage.*``
+# (model-request spans keep ``gen_ai.usage.*``) -- see docs/observability.rst.
+_SEMCONV_VERSION: Literal[5] = 5
def _otel_export_enabled() -> bool:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 56c9ec5bbb6..2eed3da9435 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -301,7 +301,7 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
}
return PydanticAIHook.get_hook(self.llm_conn_id,
hook_params=hook_params)
- def _build_agent(self) -> Agent[None, Any]:
+ def _build_agent(self) -> Agent[object, Any]:
"""Build and return a pydantic-ai Agent from the operator's config."""
extra_kwargs = dict(self.agent_params)
if self.toolsets:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index c9a22632f2e..95f33bffffa 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -163,7 +163,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
f"str prompt, or disable require_approval."
)
- agent: Agent[None, Any] = self.llm_hook.create_agent(
+ agent: Agent[object, Any] = self.llm_hook.create_agent(
output_type=self.output_type, instructions=self.system_prompt,
**self.agent_params
)
result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
index 0c9d1df1bdd..8e919acf13b 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
@@ -129,7 +129,7 @@ class LLMFileAnalysisOperator(LLMOperator):
self.sample_rows,
)
self.log.debug("Resolved file analysis paths: %s",
request.resolved_paths)
- agent: Agent[None, Any] = self.llm_hook.create_agent(
+ agent: Agent[object, Any] = self.llm_hook.create_agent(
output_type=self.output_type,
instructions=self._build_system_prompt(),
**self.agent_params,
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
index c203a3782a1..4d7baf9b4e7 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@@ -404,12 +404,23 @@ class TestPydanticAIHookCreateAgentInstrumentation:
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
def test_caller_instrument_short_circuits(self, mock_infer_model,
mock_settings, mock_agent_cls):
"""A caller that passes its own ``instrument`` wins; we don't override
it."""
- mock_infer_model.return_value = MagicMock(spec=Model)
+ mock_model = MagicMock(spec=Model)
+ mock_infer_model.return_value = mock_model
hook = self._hook()
conn = Connection(conn_id="test_conn", conn_type="pydanticai")
with patch.object(hook, "get_connection", return_value=conn):
- hook.create_agent(instructions="hi", instrument=False)
+ agent = hook.create_agent(instructions="hi", instrument=False)
+ # ``instrument`` is not an Agent() constructor kwarg in pydantic-ai
2.x:
+ # it must be stripped from the constructor call and applied through the
+ # ``agent.instrument`` property instead, and the provider's own
+ # auto-instrumentation must not override the caller's value.
+ mock_agent_cls.assert_called_once_with(
+ mock_model,
+ output_type=str,
+ instructions="hi",
+ )
+ assert agent.instrument is False
mock_settings.assert_not_called()
diff --git a/providers/common/ai/tests/unit/common/ai/test_observability.py
b/providers/common/ai/tests/unit/common/ai/test_observability.py
index bfb6f0fbe6a..4b0c84f7e50 100644
--- a/providers/common/ai/tests/unit/common/ai/test_observability.py
+++ b/providers/common/ai/tests/unit/common/ai/test_observability.py
@@ -56,7 +56,7 @@ class TestGenaiInstrumentationSettings:
settings = observability.genai_instrumentation_settings()
assert settings is not None
- assert settings.version == observability._SEMCONV_VERSION == 4
+ assert settings.version == observability._SEMCONV_VERSION == 5
# Content is never emitted unless explicitly opted in.
assert settings.include_content is False
assert settings.include_binary_content is False
@@ -123,8 +123,13 @@ class TestEndToEndSpanEmission:
genai, parent_trace_id, attrs_blob = self._run(capture=False)
assert genai, "expected gen_ai spans to be emitted"
- # Token usage is captured even with content off.
+ # Token usage is captured even with content off. In pydantic-ai 2.x the
+ # model-request span keeps ``gen_ai.usage.*`` while the agent-run span
+ # reports ``gen_ai.aggregated_usage.*`` (avoids double-counting when a
+ # backend sums parent and child spans); assert both so a change in that
+ # split is caught here rather than silently shifting users' telemetry.
assert any("gen_ai.usage.input_tokens" in (s.attributes or {}) for s
in genai)
+ assert any("gen_ai.aggregated_usage.input_tokens" in (s.attributes or
{}) for s in genai)
# Parenting is implicit: agent spans share the task span's trace_id.
assert all(s.context.trace_id == parent_trace_id for s in genai)
# The prompt text must not leak when content capture is off.
diff --git a/uv.lock b/uv.lock
index 3df0f85ac23..bf13cb8afb9 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4492,7 +4492,7 @@ requires-dist = [
{ name = "pyarrow", marker = "python_full_version < '3.14' and extra ==
'parquet'", specifier = ">=18.0.0" },
{ name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra ==
'code-mode'", specifier = ">=0.3.0" },
{ name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier =
">=0.11.0" },
- { name = "pydantic-ai-slim", specifier = ">=1.99.0,<2" },
+ { name = "pydantic-ai-slim", specifier = ">=2.0.0" },
{ name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra ==
'anthropic'" },
{ name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra ==
'bedrock'" },
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra ==
'google'" },
@@ -19221,7 +19221,7 @@ wheels = [
[[package]]
name = "pydantic-ai-slim"
-version = "1.107.0"
+version = "2.1.0"
source = { registry = "https://pypi.org/simple"; }
dependencies = [
{ name = "exceptiongroup", marker = "python_full_version < '3.11'" },
@@ -19233,9 +19233,9 @@ dependencies = [
{ name = "pydantic-graph" },
{ name = "typing-inspection" },
]
-sdist = { url =
"https://files.pythonhosted.org/packages/4c/26/ced63dfaabbc77f3beb86d59689cdea748e7ccffb6b419dbaf4780f211e8/pydantic_ai_slim-1.107.0.tar.gz";,
hash =
"sha256:4616f689a92fcfecfecf2a7af27aca22f139a873cf6d7a8929eaeee9c0eedbb4", size
= 779902, upload-time = "2026-06-10T14:53:10.574Z" }
+sdist = { url =
"https://files.pythonhosted.org/packages/ca/d1/78fd15c9c68b95ac0bec2d2afe22feb3b2f46e4b8e3a6dd1cead61cde434/pydantic_ai_slim-2.1.0.tar.gz";,
hash =
"sha256:f79dca2429dbb9d2e32a0e2c613cfb9b9d6f0dc81d42cff8a5c81ede8ca0a6c7", size
= 738698, upload-time = "2026-06-29T09:51:18.634Z" }
wheels = [
- { url =
"https://files.pythonhosted.org/packages/15/57/71044e17f931b08cc3930bc0fe5a1e1fd37fa474ae826be004729ef1cb4a/pydantic_ai_slim-1.107.0-py3-none-any.whl";,
hash =
"sha256:1af49bbae06a6c598f72c54d4734ba377100cac493c9a05fa8e089bebeae0da6", size
= 964046, upload-time = "2026-06-10T14:53:03.333Z" },
+ { url =
"https://files.pythonhosted.org/packages/5e/c7/b531cf65a1f8b221ab39739a16fe08f7b6a56fd24db875b9c33010320c9c/pydantic_ai_slim-2.1.0-py3-none-any.whl";,
hash =
"sha256:2afda56459606226113ab433ee659aac7ff29bb66feefa716c67144bd82e5ce2", size
= 910109, upload-time = "2026-06-29T09:51:11.016Z" },
]
[package.optional-dependencies]
@@ -19387,7 +19387,7 @@ wheels = [
[[package]]
name = "pydantic-graph"
-version = "1.107.0"
+version = "2.1.0"
source = { registry = "https://pypi.org/simple"; }
dependencies = [
{ name = "httpx" },
@@ -19395,9 +19395,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "typing-inspection" },
]
-sdist = { url =
"https://files.pythonhosted.org/packages/dd/c3/6e8c2d13b8701041f1b3eac5deb41f25d4dbfa479a190d5c6becc23f2a49/pydantic_graph-1.107.0.tar.gz";,
hash =
"sha256:278dd89b3e33f3a2963ac949f27a53aef705c5d883a8ce5d06d23e6e3cfbd972", size
= 62564, upload-time = "2026-06-10T14:53:13.366Z" }
+sdist = { url =
"https://files.pythonhosted.org/packages/39/20/018532c826aba3c29ffc37bda46c20f3db4584bb555a8c7dc2769866267d/pydantic_graph-2.1.0.tar.gz";,
hash =
"sha256:36ed6af24543421fb628fee593ccf5553a286c1dd1677018f5064a4e44c4daba", size
= 43052, upload-time = "2026-06-29T09:51:20.832Z" }
wheels = [
- { url =
"https://files.pythonhosted.org/packages/fc/72/621556e3f5068400d43a0375d38e5963de30256eaa5a702aba12e82ed0ff/pydantic_graph-1.107.0-py3-none-any.whl";,
hash =
"sha256:71add94fe7e14c703977a895117c475aae6c0b02a774a036c4d00d9a63c78b00", size
= 80106, upload-time = "2026-06-10T14:53:06.543Z" },
+ { url =
"https://files.pythonhosted.org/packages/42/a8/7f32bdeda6cff28652bba1c21ab2155c81d1733abd4b8bfef46a26e8d3fe/pydantic_graph-2.1.0-py3-none-any.whl";,
hash =
"sha256:bae1e99829abf590a8693442ba5693ace38d2d93a80ae48be55c610a503e3709", size
= 50772, upload-time = "2026-06-29T09:51:14.161Z" },
]
[[package]]
