kaxil commented on code in PR #75208:
URL: https://github.com/apache/airflow/pull/75208#discussion_r4336828647
##########
providers/common/ai/tests/unit/common/ai/conftest.py:
##########
@@ -104,6 +106,46 @@ def _make(output, *, cost=None):
return _make
[email protected]
+def make_reporting_agent(make_mock_run_result):
+ """Factory fixture creating a mock agent whose run spends a known usage,
as ``(agent, expected_pushes)``.
+
+ ``expected_pushes`` maps each XCom key the run must report (``usage``,
``run_id``, and the
+ model name unless ``model_name`` is ``None``) to its value, for a task
instance with id ``ti-1``.
+ """
+
+ def _make(output, *, model_name):
+ result = make_mock_run_result(output)
+ result.run_id = "ti-1"
Review Comment:
Since `result.run_id` is fixed to `"ti-1"` here and `_run_sync` ignores the
`run_id=` it's called with, the four new subclass XCom tests still pass if
`_run_llm` stops forwarding `run_id` (or forwards the wrong key). Same for the
identity stamp: `instrument` is a bare MagicMock, so
`stamp_identity_on_agent_spans` is a no-op, and moving the stamp back into
`LLMOperator.execute` would keep every subclass test green. Could `_run_sync`
set `result.run_id = kwargs["run_id"]`, and one subclass test patch
`stamp_identity_on_agent_spans` (autospec) and assert it got the agent?
##########
providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py:
##########
@@ -70,8 +78,16 @@ def
test_applies_to_scopes_to_operators_that_publish_model_name(self):
operator_names = app["applies_to"]["operator_name"]
assert set(operator_names) == {
Review Comment:
This catches renames. A new `LLMOperator` subclass, though, would publish
the model name through `_run_llm` and still not get the tab, and this test
wouldn't notice. Could it build the expected set from the in-tree `LLMOperator`
subclasses (recursively, plus `AgentOperator`) instead of a hand list?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -497,6 +462,60 @@ def _run_llm(
user_prompt: Any,
*,
usage_limits: UsageLimits | None,
+ ) -> AgentRunResult[Any]:
+ """
+ Run the agent and publish the run's ``usage``, ``run_id`` and model
name on XCom.
+
+ Every ``execute`` of this operator and its subclasses runs the agent
through here, so they
+ all report the same run metadata.
+ """
+ # A per-attempt key (the task-instance id on Airflow 3, which is
regenerated on each
+ # retry; dag/run/task/map/try on Airflow 2) ties this run to its GenAI
trace, the same
+ # join key AgentOperator uses. A hand-built context with no task
instance runs untracked.
+ ti = self._resolve_task_instance(context)
+ run_kwargs: dict[str, Any] = {}
+ if ti is not None:
+ stamp_identity_on_agent_spans(agent,
build_run_identity_attributes(ti))
+ run_kwargs["run_id"] = make_task_instance_run_key(ti)
+
+ attempt_usage = RunUsage()
+ try:
+ result = self._run_agent_within_budget(
+ context, agent, user_prompt, usage_limits=usage_limits,
run_usage=attempt_usage, **run_kwargs
+ )
+ except BaseException:
+ # A run that raises partway through (e.g. a cost_limit-triggered
UsageLimitExceeded)
+ # still reports what it spent: the capped run is the one cost
tracking most needs.
+ try:
+ log_run_usage(self.log, attempt_usage, outcome="failed")
+ except Exception:
+ self.log.warning("Failed to log partial usage for the failed
run", exc_info=True)
+ try:
+ self._push_xcom(context, "usage",
format_usage_for_xcom(attempt_usage))
+ except Exception:
+ self.log.warning("Failed to push usage XCom for the failed
run", exc_info=True)
+ if ti is not None:
+ try:
+ self._push_xcom(context, "run_id", run_kwargs["run_id"])
+ except Exception:
+ self.log.warning("Failed to push run_id XCom for the
failed run", exc_info=True)
+ raise
+
+ self._push_xcom(context, "usage", format_usage_for_xcom(attempt_usage))
+ if ti is not None:
+ self._push_xcom(context, "run_id", result.run_id)
+ model_name = ModelConfidence.from_result(result).model
+ if model_name is not None:
+ self._push_xcom(context, MODEL_NAME_XCOM_KEY, model_name)
+ return result
+
+ def _run_agent_within_budget(
+ self,
+ context: Context,
+ agent: Agent[Any, Any],
+ user_prompt: Any,
+ *,
+ usage_limits: UsageLimits | None,
run_usage: RunUsage | None = None,
Review Comment:
`_run_llm` is the only caller now and always passes `run_usage`, so this
could be a required `run_usage: RunUsage` and the two `if run_usage is not
None` guards (and the "when given" in the docstring) can go.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
