AlejandroMorgante commented on code in PR #74386: URL: https://github.com/apache/airflow/pull/74386#discussion_r4238758989
########## providers/databricks/docs/operators/agent.rst: ########## @@ -0,0 +1,131 @@ + .. 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 and map index. +Retries and cleared tasks in the same Dag run therefore reuse the invocation. +To deliberately start a new invocation in that run, supply a new UUID. +Databricks retains idempotency while it retains the invocation record. Reusing +an ID with different input returns an HTTP 409 error. + +A ``failed`` invocation fails the task. An ``interrupted`` invocation is returned +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. To reuse an invocation across task +retries, supply a stable UUID with ``vendor_options={"invocation_id": "YOUR_UUID"}``. Review Comment: Clarified that a fixed invocation ID is for one direct agent.invoke(...) per task and should not be used in `ManagedAgentToolset` vendor options -- 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]
