dhanusharer opened a new pull request, #74419: URL: https://github.com/apache/airflow/pull/74419
<!-- SPDX-License-Identifier: Apache-2.0 https://www.apache.org/licenses/LICENSE-2.0 --> related: #74194 ### Description This PR delivers **Part 1 of 2** for issue #74194 (*"Databricks: support Genie consultations and resumable standalone tasks"*). It introduces `DatabricksGenieHook`, enabling direct conversational consultations with Databricks Genie spaces, and adopts `common.ai`'s managed-agent contract (`BaseManagedAgentHook`). #### Key Features Delivered 1. **`DatabricksGenieHook`**: - Integrates with Databricks Genie endpoints: - `POST /api/2.0/genie/spaces/{space_id}/start-conversation` - `POST /api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages` - `GET /api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}` - `GET /api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}/query-result` - Supports both synchronous and asynchronous operations (`start_conversation`, `create_message`, `get_message`, `get_query_result`, `wait_for_message` + `a_*` equivalents). 2. **Common AI Managed-Agent Contract (`BaseManagedAgentHook`)**: - Implements `resolve_agent`, `get_agent_capabilities`, and `invoke_agent`. - Allows Airflow AI agents in `common.ai` to seamlessly consult Genie spaces via `hook.agent(space_id)`: ```python from airflow.providers.common.ai.toolsets import ManagedAgentToolset from airflow.providers.databricks.hooks.genie import DatabricksGenieHook hook = DatabricksGenieHook(databricks_conn_id="databricks_default") genie_agent = hook.agent("01ef8392-4f3b-1234-9abc-1234567890ab") toolset = ManagedAgentToolset( genie_agent, tool_name="ask_sales_genie", description="Consults Databricks Genie for sales, revenue, and pipeline analytics.", ) ``` - Automatically handles session continuation (`session_id`), extracts responses and markdown-formatted SQL attachments, and translates error states: - Prompt rejections / invalid queries $\rightarrow$ `ManagedAgentRejected` (triggers `pydantic-ai` model retry) - Authentication, authorization (401/403), or space missing (404) $\rightarrow$ `ManagedAgentInvocationError` - Transient network/server errors (5xx, 429) $\rightarrow$ bubbles up for standard Airflow task retry. 3. **Optional Dependency Isolation**: - `common.ai` remains strictly optional. If `apache-airflow-providers-common-ai` is not installed, the hook operates normally for direct API calls, and calling `.agent()` raises `AirflowOptionalProviderFeatureException` instructing the user to install `apache-airflow-providers-databricks[common.ai]`. 4. **Documentation & Unity MCP Distinction**: - Added `docs/operators/genie.rst` detailing when Databricks Unity Gateway MCP (generic discrete tool/function calling) is sufficient versus what native Genie API adds (stateful conversational multi-turn analytics, curated SQL queries, and benchmark context). 5. **Unit Tests**: - Added unit tests in `tests/unit/databricks/hooks/test_genie.py` covering sync/async operations, polling timeout, contract compliance, error translations, and missing extra fallbacks. --- ### Delivery Roadmap - [x] **Part 1 (This PR)**: `DatabricksGenieHook` + `common.ai` `BaseManagedAgentHook` adapter + unit tests + docs. - [ ] **Part 2 (Follow-up PR)**: `DatabricksGenieOperator` + `DatabricksGenieTrigger` with asynchronous deferrable execution, failure window idempotency safeguards, and bounded XCom/result handling. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (Antigravity) <!-- Generated-by: Antigravity following the guidelines (https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) --> -- 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]
