From b966694041fc615afb361b06d371a16568ae6090 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Fri, 3 Jul 2026 23:16:28 +0530 Subject: [PATCH] Migrate common.ai provider to pydantic-ai 2.x and remove the <2 cap 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 ++++++++++++------- .../providers/common/ai/observability.py | 11 +++--- .../providers/common/ai/operators/agent.py | 2 +- .../providers/common/ai/operators/llm.py | 2 +- .../common/ai/operators/llm_file_analysis.py | 2 +- .../unit/common/ai/hooks/test_pydantic_ai.py | 15 ++++++-- .../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 3c5439b4e8398..30a30c8e09a54 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 0be69bb300b9a..385ed423861d4 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 6669419c8aac7..54d5f5f7e9a36 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 f4597e957f9c5..4531d2cfb8c5d 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 eec7c5e7944cc..c46f974a59c2e 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 @@ 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 @@ def _get_conn_if_model_configured(self) -> Model | None: @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 @@ def create_agent( spec_file: str | Path, instructions: str | None = ..., **agent_kwargs, - ) -> Agent[None, OutputT]: ... + ) -> Agent[object, OutputT]: ... @overload def create_agent( @@ -218,7 +223,7 @@ def create_agent( spec_file: str | Path, instructions: str | None = ..., **agent_kwargs, - ) -> Agent[None, str]: ... + ) -> Agent[object, str]: ... def create_agent( self, @@ -227,7 +232,7 @@ def create_agent( 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 @@ def create_agent( 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 @@ def create_agent( 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 e81a6680d04ae..08a59c40595f8 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 @@ 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 56c9ec5bbb65a..2eed3da943516 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 @@ def llm_hook(self) -> PydanticAIHook: } 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 c9a22632f2e3d..95f33bffffaa7 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 @@ def execute(self, context: Context) -> Any: 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 0c9d1df1bdd20..8e919acf13b27 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 @@ def execute(self, context: Context) -> Any: 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 c203a3782a1ba..4d7baf9b4e7c6 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 @@ def test_no_instrument_when_settings_none(self, mock_settings): @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 bfb6f0fbe6aa3..4b0c84f7e508f 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 @@ def test_returns_settings_when_enabled_with_provider(self): 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 @@ def test_spans_emitted_and_nested_without_content_by_default(self): 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 3df0f85ac235c..bf13cb8afb9b1 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]]