kaxil commented on code in PR #74386: URL: https://github.com/apache/airflow/pull/74386#discussion_r4300827651
########## providers/databricks/docs/operators/agent.rst: ########## @@ -0,0 +1,151 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + + +.. _howto/operator:DatabricksAgentInvokeOperator: + +.. spelling:word-list:: + + MLflow + +Invoke a Databricks agent +============================ + +Use :class:`~airflow.providers.databricks.operators.agent.DatabricksAgentInvokeOperator` +to invoke an agent deployed on Databricks Apps using ``DurableAgentServer``. +The operator submits a background request to the invocation API and returns the +invocation response, including its status and output, through XCom. + +Authentication +-------------- + +Configure a Databricks connection with the workspace URL in ``host``, the service +principal's client ID in ``login``, its client secret in ``password``, and +``{"service_principal_oauth": true}`` in ``extra``. The service principal needs +permission to use the deployed app. Databricks Apps requires OAuth; personal +access tokens are not supported. Supply the app's HTTPS base URL as ``app_url``. + +Invoke an agent +--------------- + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py + :language: python + :start-after: [START howto_operator_databricks_agent_invoke] + :end-before: [END howto_operator_databricks_agent_invoke] + +The agent defines the schema of ``input``. Agents generated by the Agent Bricks +CLI templates require a ``session_id``. Reuse the session ID across invocations +to continue a conversation. The hook sends it as ``X-Routing-Key`` when supplied. + +By default, the operator waits for the result. Set ``deferrable=True`` to release +the worker while waiting, or ``wait_for_termination=False`` to return the initial +submission response immediately. ``polling_period_seconds`` controls polling. +The timeout limits the wait after submission; it does not cancel the remote run. +Synchronous polling bounds each request and its retries by the remaining wait +and raises :class:`~airflow.providers.databricks.exceptions.DatabricksAgentInvocationTimeout` +when that budget expires. OAuth refresh has a separate HTTP timeout and can delay +reporting that the wait expired. In deferrable mode, Airflow enforces the timeout on the deferred task. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py + :language: python + :start-after: [START howto_operator_databricks_agent_invoke_deferrable] + :end-before: [END howto_operator_databricks_agent_invoke_deferrable] + +Retries and results +------------------- + +An invocation requires a UUID. Unless ``invocation_id`` is provided, the operator +generates a stable UUID from the app URL, Dag ID, task ID, run ID, map index, +rendered input and session ID. App URLs with and without a trailing slash produce +the same UUID. Retries and cleared tasks with unchanged input and session in the +same Dag run reuse the invocation. Changing the rendered input or session generates +a new UUID. To deliberately repeat the same request in that run, supply a new UUID. +Databricks retains idempotency while it retains the invocation record. Reusing +an explicit ID with different input or session returns an HTTP 409 error. + +A stored ``failed`` invocation raises ``AirflowFailException``, preventing automatic +task retries: submitting the same UUID returns that failure without running the +agent again. Clearing the task with unchanged input and session also returns the +stored failure. Supply a new ``invocation_id`` to run it again. + +The runtime states are ``queued``, ``active``, ``completed`` and ``failed``. +CLI template agents report a pause as ``status="completed"`` with +``output.status="interrupted"``. The operator returns this response successfully +so downstream tasks can handle human input. To resume it, submit a new invocation +with the same session ID and the appropriate ``resume`` input. +The returned response preserves the agent's output without assuming its schema. + +The operator targets ``DurableAgentServer``'s ``/api/invocations`` API. Legacy +MLflow agent servers using ``/responses`` and Model Serving endpoints are not +supported by this operator. + +For direct calls, use +:class:`~airflow.providers.databricks.hooks.agent.DatabricksAgentHook` and its +``create_invocation`` and ``get_invocation`` methods. +See `Query agents deployed on Databricks +<https://docs.databricks.com/aws/en/agents/query-llms>`_ for the API contract. + +Common AI managed-agent interface +--------------------------------- + +Install ``apache-airflow-providers-databricks[common.ai]`` to use the hook with +Common AI consumers such as ``ManagedAgentToolset``. Bind the deployed app's URL +to a hook configured with the same OAuth connection: Review Comment: Could this section also show the toolset path, like the Vertex AI page does (`ManagedAgentToolset(hook.agent(app_url), tool_name=..., description=...)`)? The intro names the toolset, but the example only calls `agent.invoke` directly. Also, this line ends with a colon that leads into another sentence instead of the example. ########## providers/databricks/docs/operators/agent.rst: ########## @@ -0,0 +1,151 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + + +.. _howto/operator:DatabricksAgentInvokeOperator: + +.. spelling:word-list:: + + MLflow + +Invoke a Databricks agent +============================ + +Use :class:`~airflow.providers.databricks.operators.agent.DatabricksAgentInvokeOperator` +to invoke an agent deployed on Databricks Apps using ``DurableAgentServer``. +The operator submits a background request to the invocation API and returns the +invocation response, including its status and output, through XCom. + +Authentication +-------------- + +Configure a Databricks connection with the workspace URL in ``host``, the service +principal's client ID in ``login``, its client secret in ``password``, and +``{"service_principal_oauth": true}`` in ``extra``. The service principal needs +permission to use the deployed app. Databricks Apps requires OAuth; personal +access tokens are not supported. Supply the app's HTTPS base URL as ``app_url``. + +Invoke an agent +--------------- + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py + :language: python + :start-after: [START howto_operator_databricks_agent_invoke] + :end-before: [END howto_operator_databricks_agent_invoke] + +The agent defines the schema of ``input``. Agents generated by the Agent Bricks +CLI templates require a ``session_id``. Reuse the session ID across invocations +to continue a conversation. The hook sends it as ``X-Routing-Key`` when supplied. + +By default, the operator waits for the result. Set ``deferrable=True`` to release +the worker while waiting, or ``wait_for_termination=False`` to return the initial +submission response immediately. ``polling_period_seconds`` controls polling. +The timeout limits the wait after submission; it does not cancel the remote run. +Synchronous polling bounds each request and its retries by the remaining wait +and raises :class:`~airflow.providers.databricks.exceptions.DatabricksAgentInvocationTimeout` +when that budget expires. OAuth refresh has a separate HTTP timeout and can delay +reporting that the wait expired. In deferrable mode, Airflow enforces the timeout on the deferred task. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py + :language: python + :start-after: [START howto_operator_databricks_agent_invoke_deferrable] + :end-before: [END howto_operator_databricks_agent_invoke_deferrable] + +Retries and results +------------------- + +An invocation requires a UUID. Unless ``invocation_id`` is provided, the operator +generates a stable UUID from the app URL, Dag ID, task ID, run ID, map index, +rendered input and session ID. App URLs with and without a trailing slash produce +the same UUID. Retries and cleared tasks with unchanged input and session in the +same Dag run reuse the invocation. Changing the rendered input or session generates +a new UUID. To deliberately repeat the same request in that run, supply a new UUID. +Databricks retains idempotency while it retains the invocation record. Reusing +an explicit ID with different input or session returns an HTTP 409 error. + +A stored ``failed`` invocation raises ``AirflowFailException``, preventing automatic +task retries: submitting the same UUID returns that failure without running the +agent again. Clearing the task with unchanged input and session also returns the +stored failure. Supply a new ``invocation_id`` to run it again. + +The runtime states are ``queued``, ``active``, ``completed`` and ``failed``. +CLI template agents report a pause as ``status="completed"`` with +``output.status="interrupted"``. The operator returns this response successfully +so downstream tasks can handle human input. To resume it, submit a new invocation +with the same session ID and the appropriate ``resume`` input. +The returned response preserves the agent's output without assuming its schema. + +The operator targets ``DurableAgentServer``'s ``/api/invocations`` API. Legacy +MLflow agent servers using ``/responses`` and Model Serving endpoints are not +supported by this operator. + +For direct calls, use +:class:`~airflow.providers.databricks.hooks.agent.DatabricksAgentHook` and its +``create_invocation`` and ``get_invocation`` methods. +See `Query agents deployed on Databricks +<https://docs.databricks.com/aws/en/agents/query-llms>`_ for the API contract. + +Common AI managed-agent interface +--------------------------------- + +Install ``apache-airflow-providers-databricks[common.ai]`` to use the hook with +Common AI consumers such as ``ManagedAgentToolset``. Bind the deployed app's URL +to a hook configured with the same OAuth connection: + +Import ``ManagedAgentRequest`` from ``airflow.providers.common.ai.managed_agents`` +and ``DatabricksAgentHook`` from ``airflow.providers.databricks.hooks.agent``: + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py + :language: python + :start-after: [START howto_databricks_managed_agent] + :end-before: [END howto_databricks_managed_agent] + +This interface uses the synchronous invocation API. It converts a prompt to a +user message, or passes supplied messages under ``input.messages``. For an agent +that expects a prompt under another key, set ``vendor_options={"input_key": "question"}``. +Other input schemas remain available through ``create_invocation`` and the operator. + +Each call generates a new invocation UUID. For one direct ``agent.invoke(...)`` +per task, a stable UUID in ``vendor_options={"invocation_id": "YOUR_UUID"}`` can +reuse that invocation across task retries. Do not set a fixed invocation ID in +``ManagedAgentToolset`` vendor options: the toolset forwards those options on every +call, so a different prompt would conflict and a repeated prompt would return the +stored result. + +HTTP retries within a call reuse the same UUID. ``session_id`` continues the +conversation and is sent as the routing key. When omitted, the hook uses the +invocation UUID as a one-shot session ID, including for toolset calls. A hook can bind multiple app URLs; +vendor options cannot override the app URL or connection. + +``ManagedAgentResponse.raw`` preserves the full invocation response. String outputs +become ``text``; an object with a string ``output`` field uses that field as text; +other outputs become JSON text. Non-string outputs are also available in +``structured``. The invocation ID is returned as ``trace_ref``. + +The managed-agent interface requires a ``completed`` invocation whose output is +not marked ``interrupted`` by the CLI template. Stored failures, interruptions, +missing or unexpected statuses raise ``ManagedAgentInvocationError``; use the +operator to handle an interrupted invocation's response directly. +Terminal HTTP errors raise the same exception, while transient HTTP and connection +errors propagate after the hook's configured retries. A synchronous agent failure +returns HTTP 500. After HTTP retries, the hook queries the known invocation ID: +a stored ``failed`` status becomes ``ManagedAgentInvocationError``. If the lookup +cannot confirm a stored failure, the original HTTP error propagates as transient. +The request timeout bounds the HTTP call, its retries and the status lookup; Review Comment: Worth a sentence on `timeout=None`, which is what `ManagedAgentToolset` sends by default. Each sync attempt is then capped at the hook's `timeout_seconds` (180s), so an agent that runs longer is given up on after three attempts, and the task retry starts it again under a new UUID while the first run keeps going. Pointing users at the toolset's `timeout=` for long-running agents would cover it. -- 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]
