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 47cabf7ce6ccc..52dcf1d9ea068 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 @@ 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 5a0f141a76154..e01654f907548 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 25a176e3297a6..1a9d82bd065bb 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 67ea067d160c0..9a1c389fef513 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 023af790d36da..37238bff00b62 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 41fc750d5051c..cac472b765763 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 @@ 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 df3c3f571c2e7..0664b6c52b18c 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 5b2e4b6e6e3bf..d44b5d016f56f 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 4a0b08cf1483e..57eaa552238db 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 2a707752fdfe3..d5ef8228d3528 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 82f710cc10a0b..390137e9fc9be 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 6c970e4326383..7c955a160b4f4 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 d719162bc00f7..c7a23de583847 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 e8e31c6f5de88..a994ae3d1cd2f 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 230335a0e027a..bbc91df0850d8 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 = []