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 1c40bfb27ec Access AgentRunResult.usage as a property in common.ai
logging (#68405)
1c40bfb27ec is described below
commit 1c40bfb27eca87e4a9b939db5de3c254dd32ec95
Author: Kaxil Naik <[email protected]>
AuthorDate: Fri Jun 12 20:45:19 2026 +0100
Access AgentRunResult.usage as a property in common.ai logging (#68405)
pydantic-ai deprecated calling AgentRunResult.usage() as a method
(PydanticAIDeprecationWarning); it is a property since the version range
already covered by the provider floor (pydantic-ai-slim>=1.99.0 exposes
both forms). Drop the parentheses at the call site and update the test
mocks to set .usage as an attribute instead of a return_value.
---
.../common/ai/src/airflow/providers/common/ai/utils/logging.py | 2 +-
providers/common/ai/tests/unit/common/ai/conftest.py | 6 ++----
providers/common/ai/tests/unit/common/ai/decorators/test_agent.py | 4 +---
providers/common/ai/tests/unit/common/ai/decorators/test_llm.py | 4 +---
.../common/ai/tests/unit/common/ai/decorators/test_llm_branch.py | 4 +---
.../ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py | 2 +-
.../ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py | 4 +---
providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py | 4 +---
providers/common/ai/tests/unit/common/ai/operators/test_agent.py | 4 +---
providers/common/ai/tests/unit/common/ai/operators/test_llm.py | 4 +---
.../common/ai/tests/unit/common/ai/operators/test_llm_branch.py | 4 +---
.../ai/tests/unit/common/ai/operators/test_llm_file_analysis.py | 2 +-
.../ai/tests/unit/common/ai/operators/test_llm_schema_compare.py | 4 +---
providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py | 4 +---
providers/common/ai/tests/unit/common/ai/utils/test_logging.py | 2 +-
15 files changed, 16 insertions(+), 38 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
index 47cabf7ce6c..52dcf1d9ea0 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
@@ -36,7 +36,7 @@ _MAX_OUTPUT_LEN = 500
def log_run_summary(logger: Logger | logging.Logger, result:
AgentRunResult[Any]) -> None:
"""Log model name, token usage, and tool call sequence from an agent
run."""
- usage = result.usage()
+ usage = result.usage
model_name = getattr(result.response, "model_name", "unknown")
logger.info(
"::group::LLM run complete: model=%s, requests=%s, tool_calls=%s, "
diff --git a/providers/common/ai/tests/unit/common/ai/conftest.py
b/providers/common/ai/tests/unit/common/ai/conftest.py
index 5a0f141a761..e01654f9075 100644
--- a/providers/common/ai/tests/unit/common/ai/conftest.py
+++ b/providers/common/ai/tests/unit/common/ai/conftest.py
@@ -34,14 +34,12 @@ def isolate_hook_lineage_collector(hook_lineage_collector):
def make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary.
- Returns a MagicMock with .output, .usage(), .response, and .all_messages()
+ Returns a MagicMock with .output, .usage, .response, and .all_messages()
configured so that log_run_summary can read them without error.
"""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
index 25a176e3297..1a9d82bd065 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
@@ -44,9 +44,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
index 67ea067d160..9a1c389fef5 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
@@ -28,9 +28,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
index 023af790d36..37238bff00b 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
@@ -30,9 +30,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
index 41fc750d505..cac472b7657 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
@@ -27,7 +27,7 @@ from airflow.providers.common.ai.utils.file_analysis import
FileAnalysisRequest
def _make_mock_run_result(output):
mock_result = MagicMock(spec=["output", "usage", "response",
"all_messages"])
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
+ mock_result.usage = MagicMock(
spec=["requests", "tool_calls", "input_tokens", "output_tokens",
"total_tokens"],
requests=1,
tool_calls=0,
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
index df3c3f571c2..0664b6c52b1 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
@@ -32,9 +32,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
index 5b2e4b6e6e3..d44b5d016f5 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
@@ -28,9 +28,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 4a0b08cf148..57eaa552238 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -48,9 +48,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 2a707752fdf..d5ef8228d35 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -57,9 +57,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index 82f710cc10a..390137e9fc9 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -29,9 +29,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
index 6c970e43263..7c955a160b4 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
@@ -46,7 +46,7 @@ class Summary(BaseModel):
def _make_mock_run_result(output):
mock_result = MagicMock(spec=["output", "usage", "response",
"all_messages"])
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
+ mock_result.usage = MagicMock(
spec=["requests", "tool_calls", "input_tokens", "output_tokens",
"total_tokens"],
requests=1,
tool_calls=0,
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
index d719162bc00..c7a23de5838 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
@@ -35,9 +35,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
index e8e31c6f5de..a994ae3d1cd 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
@@ -37,9 +37,7 @@ def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
mock_result = MagicMock()
mock_result.output = output
- mock_result.usage.return_value = MagicMock(
- requests=1, tool_calls=0, input_tokens=0, output_tokens=0,
total_tokens=0
- )
+ mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0,
output_tokens=0, total_tokens=0)
mock_result.response = MagicMock(model_name="test-model")
mock_result.all_messages.return_value = []
return mock_result
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
index 230335a0e02..bbc91df0850 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
@@ -44,7 +44,7 @@ def _make_mock_result(model_name="gpt-5", tool_names=None,
usage_kwargs=None):
"total_tokens": 3359,
}
result = MagicMock()
- result.usage.return_value = MagicMock(**usage_kwargs)
+ result.usage = MagicMock(**usage_kwargs)
result.response = MagicMock(model_name=model_name)
messages: list = []
