AlejandroMorgante commented on code in PR #74386: URL: https://github.com/apache/airflow/pull/74386#discussion_r4320401632
########## 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: Done, added the `ManagedAgentToolset` example following the Vertex AI page and fixed the introductory sentence -- 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]
