airflow.providers.databricks.hooks.agent

Invoke agents served by Databricks DurableAgentServer.

Classes

DatabricksAgentHook

Invoke a DurableAgentServer deployed on Databricks Apps.

Module Contents

class airflow.providers.databricks.hooks.agent.DatabricksAgentHook(databricks_conn_id='databricks_default', timeout_seconds=180, retry_limit=3, retry_delay=1, caller='Unknown')[source]

Bases: airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook, airflow.providers.common.ai.managed_agents.base.BaseManagedAgentHook

Invoke a DurableAgentServer deployed on Databricks Apps.

With the common.ai extra installed, hook.agent(app_url) exposes the synchronous invocation API through the managed-agent contract. Background invocations remain available through create_invocation and get_invocation without that extra.

Parameters:

  • databricks_conn_id (str) – Databricks workspace connection configured with service principal OAuth credentials or OAuth token federation. Defaults to databricks_default.

  • timeout_seconds (int) – HTTP request timeout in seconds. Defaults to 180.

  • retry_limit (int) – Maximum number of attempts for transient HTTP failures. Defaults to 3.

  • retry_delay (float) – Minimum delay between HTTP retries in seconds. Defaults to 1.

  • caller (str) – Name of the operator or trigger calling the hook. Defaults to Unknown.

agent_platform = 'databricks.agent_runtime'[source]

The platform every ManagedAgentRef from this hook carries.

resolve_agent(agent)[source]

Resolve an app’s HTTPS base URL without contacting Databricks.

Parameters:

agent (str) – HTTPS base URL of the deployed app.

Returns:

Platform-qualified app URL with its trailing slash removed.

Return type:

airflow.providers.common.ai.managed_agents.base.ManagedAgentRef

get_agent_capabilities(agent)[source]

Report session, structured-output and invocation-ID support.

Parameters:

agent (str) – HTTPS base URL of the deployed app.

Returns:

Capabilities of the synchronous invocation API.

Return type:

airflow.providers.common.ai.managed_agents.base.ManagedAgentCapabilities

invoke_agent(agent, request)[source]

Invoke an agent synchronously through the Common AI contract.

Parameters:

  • agent (str) – HTTPS base URL of the deployed app.

  • request (airflow.providers.common.ai.managed_agents.base.ManagedAgentRequest) – Prompt or messages, optional session and request timeout. Vendor options accept input_key to send a prompt under an agent-specific key and invocation_id to reuse a caller-provided UUID. By default, prompts are converted to messages and a new UUID is generated for each call. Without a session, the invocation UUID is also used as a one-shot session ID.

Returns:

Answer text, full invocation response, structured output, session and invocation ID.

Return type:

airflow.providers.common.ai.managed_agents.base.ManagedAgentResponse

create_invocation(invocation_id, input, session_id=None, *, app_url)[source]

Start a background invocation; retrying the same UUID and input is idempotent.

Parameters:

  • invocation_id (str) – UUID identifying the invocation.

  • input (Any) – JSON-serializable input whose schema is defined by the agent.

  • session_id (str | None) – Conversation ID, also sent as the routing key. Defaults to None.

  • app_url (str) – HTTPS base URL of the deployed app.

Returns:

Complete JSON submission response returned by the agent server.

Return type:

dict[str, Any]

get_invocation(invocation_id, session_id=None, *, app_url, timeout_seconds=None)[source]

Get the status and output of an invocation.

Parameters:

  • invocation_id (str) – UUID identifying the invocation.

  • session_id (str | None) – Conversation ID sent as the routing key. Defaults to None.

  • app_url (str) – HTTPS base URL of the deployed app.

  • timeout_seconds (float | None) – Optional polling request/retry budget in seconds. Defaults to None, which uses the hook’s HTTP timeout and retry settings. OAuth refresh uses the base hook’s token timeout; if it exhausts this budget, no invocation request is sent.

Returns:

Complete JSON invocation response, including output when available.

Return type:

dict[str, Any]

async a_get_invocation(invocation_id, session_id=None, *, app_url)[source]

Asynchronously get an invocation inside the hook’s async context manager.

Parameters:

  • invocation_id (str) – UUID identifying the invocation.

  • session_id (str | None) – Conversation ID sent as the routing key. Defaults to None.

  • app_url (str) – HTTPS base URL of the deployed app.

Returns:

Complete JSON invocation response, including output when available.

Return type:

dict[str, Any]

Was this entry helpful?