Invoke a Databricks agent¶
Use 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.
Alternatively, use OAuth token federation through federated_token_provider,
federated_aws or federated_k8s, as described in
Databricks Connection. These modes reuse the base Databricks hook’s
token exchange and refresh without requiring a client secret. Configure the
corresponding Databricks federation policy and identity provider for your workload.
The service principal needs CAN USE permission on the deployed app. The hook
requires Databricks workspace OAuth; personal access tokens and Azure AD tokens
are not supported. Do not combine OAuth settings with PAT or Azure AD
authentication settings, which take precedence in the base hook.
Supply the app’s HTTPS base URL as app_url.
Invoke an agent¶
invoke = DatabricksAgentInvokeOperator(
task_id="invoke_agent",
app_url=APP_URL,
databricks_conn_id=CONN_ID,
session_id=SESSION_ID,
input=INPUT,
deferrable=False,
timeout=3600,
)
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 uses the remaining wait to set request timeouts and limit retries,
with each GET timeout capped by the hook’s timeout_seconds. It
raises DatabricksAgentInvocationTimeout
when it detects that the budget has expired. Requests’ HTTP timeout limits network
inactivity, not total response duration: a server that keeps sending data can delay
that check. OAuth refresh also has a separate HTTP timeout and can delay reporting
that the wait expired. In deferrable mode, Airflow enforces the timeout on the deferred task.
invoke_deferrable = DatabricksAgentInvokeOperator(
task_id="invoke_agent_deferrable",
app_url=APP_URL,
databricks_conn_id=CONN_ID,
session_id=SESSION_ID,
input=INPUT,
deferrable=True,
timeout=3600,
)
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
DatabricksAgentHook and its
create_invocation and get_invocation methods. Configure the hook with
databricks_conn_id and pass app_url to each call, including
a_get_invocation for asynchronous polling. The same hook can address different
apps using the same connection:
hook = DatabricksAgentHook(databricks_conn_id="databricks_oauth")
submitted = hook.create_invocation(invocation_id, input_data, app_url=app_url)
status = hook.get_invocation(invocation_id, app_url=app_url)
See Query agents deployed on Databricks 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.
To let an agent running under Common AI’s AgentOperator consult the deployed
Databricks agent as a tool, pass hook.agent(app_url) to ManagedAgentToolset:
from airflow.providers.common.ai.toolsets import ManagedAgentToolset
from airflow.providers.databricks.hooks.agent import DatabricksAgentHook
hook = DatabricksAgentHook(databricks_conn_id="databricks_oauth")
toolset = ManagedAgentToolset(
hook.agent("https://my-agent.databricksapps.com"),
tool_name="ask_analyst",
description="Answers revenue questions.",
timeout=300,
)
The app must serve DurableAgentServer’s invocation API. The toolset sends each
prompt through that API and returns the normalized answer text to the calling model.
For a direct invocation, import ManagedAgentRequest from airflow.providers.common.ai.managed_agents
and DatabricksAgentHook from airflow.providers.databricks.hooks.agent:
@task
def invoke_managed_agent(app_url: str, conn_id: str, session_id: str) -> dict:
agent = DatabricksAgentHook(databricks_conn_id=conn_id).agent(app_url)
result = agent.invoke(
ManagedAgentRequest(
messages=INPUT["messages"],
session_id=session_id,
timeout=60,
)
)
return result.raw
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. A caller-provided UUID must be unique
per Dag run and task instance, while remaining unchanged across retries of that
invocation. Reusing it in a later Dag run returns the stored result for the same
input or an HTTP 409 conflict if the input changed.
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.
For CLI template outputs containing a list of messages, text uses the last
assistant message’s content when it is a string or a list of text blocks.
Non-text blocks are ignored when extracting text from a list; if no text blocks
are found, the output is returned as JSON text.
Both OpenAI role="assistant" and LangGraph type="ai" messages are recognized.
Unrecognized 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 supplies a shared budget for the HTTP call, its retries and
the status lookup. It is checked between attempts and after a successful HTTP response;
it does not interrupt an HTTP response that keeps sending data. OAuth refresh
retains its separate timeout.
ManagedAgentToolset defaults to timeout=None, so each synchronous HTTP
attempt uses the hook’s timeout_seconds (180 seconds by default). After the
hook’s retries are exhausted (three attempts by default), the call fails without
cancelling the remote invocation. A task retry generates a new invocation UUID
and can start another run while the previous one is still running. For long-running
agents, set ManagedAgentToolset(..., timeout=...) to the expected request budget.
The operator and background hook methods do not require the Common AI extra.