kaxil commented on code in PR #75208:
URL: https://github.com/apache/airflow/pull/75208#discussion_r4352314851
##########
providers/common/ai/docs/observability.rst:
##########
@@ -92,31 +92,26 @@ How it works
``usage`` above, so a downstream task, a Jinja template, or a REST client
can read which
model responded without parsing the decision record or a span
(``ti.xcom_pull(task_ids="my_agent",
key="__AIRFLOW__COMMON_AI_MODEL_NAME__")``).
- Only ``AgentOperator`` /
- ``@task.agent`` and ``LLMOperator`` / ``@task.llm`` push it today; like
``run_id`` and
- ``usage``, with ``enable_hitl_review`` it is the initial run's model, not a
- human-feedback regeneration's, since ``regenerate_with_feedback`` doesn't
re-push it.
- ``LLMOperator``'s subclasses that override ``execute()`` instead of calling
- ``super().execute()`` -- ``LLMBranchOperator``, ``LLMSQLQueryOperator``,
- ``LLMSchemaCompareOperator``, ``LLMFileAnalysisOperator`` -- don't publish
it yet;
- this is upcoming work.
-* **Scope.** The ``run_id`` / ``usage`` XComs, the model name, and the
``airflow.*``
- identity attributes come from ``AgentOperator`` / ``@task.agent`` and
``LLMOperator`` /
- ``@task.llm``, apart from a Strands, ADK or OpenAI Agents SDK agent run
inside
- ``agent_framework_tracing`` (see below). The other LLM operators still emit
GenAI spans
- correlated to the task span by nesting, but without the identity attributes,
the run
- join key, the model name, or a ``usage`` XCom.
+ ``AgentOperator`` and ``LLMOperator`` and its subclasses (and their
decorator forms)
+ push it; like ``run_id`` and ``usage``, with ``enable_hitl_review`` it is
the initial
+ run's model, not a human-feedback regeneration's, since
``regenerate_with_feedback``
+ doesn't re-push it.
+* **Scope.** The ``run_id`` / ``usage`` XComs and the model name come from
Review Comment:
This page is right now, but the four subclass pages still don't say they
push `usage`, `run_id` and the model name: the Logging sections in
`llm_sql.rst`, `llm_branch.rst` and `llm_schema_compare.rst` stop at the log
summary, and `llm_file_analysis.rst` has none. `llm_branch.rst`'s "pushes a
`decision` XCom next to its return value" also reads like the full list. A
one-line pointer to this section from each page would close it.
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py:
##########
@@ -252,9 +252,23 @@ def test_execute_with_schema_context(self, mock_hook_cls,
make_mock_run_result):
assert result == "SELECT id, name FROM users WHERE active = true"
mock_agent.run_sync.assert_called_once_with(
- "Get active users", usage_limits=None, cancellation_token=ANY
+ "Get active users", usage=ANY, usage_limits=None, run_id=ANY,
cancellation_token=ANY
)
+ @pytest.mark.parametrize("model_name", ["gpt-5", None])
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_run_usage_run_id_and_model_name_pushed_to_xcom(
+ self, mock_hook_cls, make_reporting_agent, model_name
+ ):
+ mock_agent, expected_pushes = make_reporting_agent("SELECT 1",
model_name=model_name)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMSQLQueryOperator(task_id="test", prompt="Get one",
llm_conn_id="my_llm")
+ ti = MagicMock(id="ti-1")
Review Comment:
Nit: this `MagicMock(id="ti-1")` (and the same one in the file-analysis and
schema-compare tests, plus `op.llm_hook = MagicMock()` there) has no spec, so
`build_run_identity_attributes` reads whatever it asks for off it. A specced
task-instance fixture next to `make_reporting_agent` would cover all three.
--
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]
