Create Genie API wrapper by prithvikannan · Pull Request #2 · databricks/databricks-ai-bridge

prithvikannan · 2024-10-17T06:00:05Z

Create the core Genie API wrapper. The wrapper API follows the Genie convseration APIs for polling.

Unit tested.

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

aravind-segu · 2024-10-23T18:42:33Z

+        }
+
+    def start_conversation(self, content):
+        resp = self.genie._api.do(


question, why are we using the raw _api request here? Can we not use this method directly? https://databricks-sdk-py.readthedocs.io/en/stable/workspace/dashboards/genie.html#databricks.sdk.service.dashboards.GenieAPI.start_conversation

i think the polling mechanisms with the start_conversation API does not work. ill asked the genie team and they recommended to use the _api for now. we can revisit when fixed.

aravind-segu · 2024-10-23T18:48:32Z

+                    headers=self.headers,
+                )
+                if resp["status"] == "EXECUTING_QUERY":
+                    sql = next(r for r in resp["attachments"] if "query" in r)["query"]["query"]


Are we getting the sql statements here only for the debug statement? Do end customers need this information?

eventually these should be part of the trace. currently adding extra traces with autologged traces are not supported, but this is coming soon. we'll use the SQL at that time.

aravind-segu · 2024-10-23T18:50:49Z

+
+        return poll_result()
+
+    def ask_question(self, question):


What do you think of adding a timeout feature here, so users are not stuck waiting forever in the polling loop?

good question. i wonder if genie requests have some RPC level timeout on that side, but i think a client side timeout also makes sense. will update

actually genie has a CANCELED state that happens in the case of timeout, which will hit the else case and break the loop. i dont think we need the client side timeout.

aravind-segu · 2024-10-23T18:52:16Z

+                    logging.debug(f"SQL: {sql}")
+                    return poll_query_results()
+                elif resp["status"] == "COMPLETED":
+                    return next(r for r in resp["attachments"] if "text" in r)["text"]["content"]


For genie, do we only care about the first response in the list?

at this time there's only one text attachment from genie so this is a safe assumption.

aravind-segu · 2024-10-23T18:57:44Z

+        {"status": "COMPLETED", "attachments": [{"text": {"content": "Answer"}}]},
+    ]
+    result = genie.ask_question("What is the meaning of life?")
+    assert result == "Answer"


The answer cant be that simple 🤣

aravind-segu

Looking at the databricks-sdk code, the get api is exactly similar if we can use it directly. But it would be half direct api code and half databricks sdk code. Upto you on if you think that would be cleaner.

prithvikannan · 2024-10-24T00:44:11Z

Looking at the databricks-sdk code, the get api is exactly similar if we can use it directly. But it would be half direct api code and half databricks sdk code. Upto you on if you think that would be cleaner.

Good find. Let's do the refactor in one shot to keep it readable.

…sponse_to_chat_result` (#84) * Enhance ChatResult output with model details and token usage * Update chat_models.py * Dump metadata in llm_output and create new unit test * Fix linting error * revert changes * Fix linting errors (#2) --------- Co-authored-by: Jose Moreno Ortega <j45351@eon.com>

…re catches Review feedback from Ann (comments #3-7): - Remove unused pytestmark from conftest.py (#3) - Add multi-server test: UC + VS in DatabricksMultiServerMCPClient (#4) - Narrow conftest fixture catches: ExceptionGroup with McpError NOT_FOUND check instead of bare Exception — re-raises non-NOT_FOUND errors (#6) - Revert ty/type-ignore changes to non-test files (chat_models.py, lakebase.py, checkpoint.py) — not in scope for this PR (#1, #2) Error path tests across all 3 layers (#7): - Core: bad function (McpError BAD_REQUEST not found), bad tool name (McpError BAD_REQUEST malformed), wrong args (McpError missing parameter) - OpenAI Toolkit: bad function (ValueError wrapping), bad tool (ExceptionGroup), wrong args (ExceptionGroup) - OpenAI Agents: bad function (McpError), bad tool (McpError), wrong args (UserError) - LangChain: bad function (ExceptionGroup > McpError), wrong args (McpError) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Per review on PR #425. Skipping #3 (trim) and #4 (conv_id naming) — addressed separately. #1 Drop owner_pod_id; ownership via heartbeat CAS on attempt_number - Remove owner_pod_id column from Response model - heartbeat_response(response_id, expected_attempt_number) CAS-checks the attempt_number column. If a heartbeat write returns 0 rows, the prior owner has been bumped by another pod's claim and the heartbeat task knows to stop. - _heartbeat() context manager takes attempt_number; passes it through from _run_background_stream / _run_background_invoke. - claim_stale_response() no longer takes a pod parameter. - _POD_LOG_ID retained for log-line identity only (not stored in DB). #2 Simplify prose recovery to json.dumps the events array - _build_prose_recovery_message: was ~110 LOC structural walker (function_call/output pairs, narrative messages, partial-text reassembly). Now ~15 LOC: filter events by prior_attempt_number, json.dumps them, wrap in a directive prompt asking the model to figure out what's done vs interrupted. #5 Drop _inject_conversation_id - The function was defensive injection of response_id into context.conversation_id when no client anchor was supplied. With rotation handling resume, and templates / chatbot consistently setting conv_id, the injection was redundant. Top-level review: proactive stale-scan loop with jitter - New _stale_response_scanner_loop: every ~30s ± 50% jitter, queries responses for in_progress rows with stale heartbeats and tries to claim+resume them. The proactive counterpart to lazy-on-GET claim; ensures crashed responses get recovered even if no client polls. - find_stale_response_ids repository function with LIMIT 50. - Spawned in the FastAPI lifespan alongside init_db; cancelled on shutdown. - Settings: stale_scan_interval_seconds=30.0, stale_scan_jitter_fraction=0.5. #6 Document /_debug/kill_task in AGENTS.md - New §4.4 explaining the test-only debug endpoint, env-var gating, what state it leaves the row in. AGENTS.md updates: - ER diagram: drop owner_pod_id, annotate attempt_number as CAS guard. - New §3.5 documenting the proactive scanner with mermaid flowchart. - §4.3 includes new scanner settings. Tests: 110 pass. Ruff/format/ty all clean. Co-authored-by: Isaac

prithvikannan added 9 commits October 16, 2024 22:59

Create Genie API wrapper

66ce723

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

package setup

ac93eae

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

Add genie tests and resources to run

188ac83

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

update tests

ded9ec2

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

create github workflow

0fa7b70

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

merge

f6d3126

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

remove match

8457766

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

fix test

96ab57a

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

ruff

1b17d84

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

prithvikannan commented Oct 22, 2024

View reviewed changes

Comment thread pyproject.toml

prithvikannan added 3 commits October 22, 2024 13:54

logging

4cadcbc

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

remove version file

0e58f54

Signed-off-by: Prithvi Kannan <prithvi.kannan@databricks.com>

Merge remote-tracking branch 'origin/main' into setup-genie-core

86a74fb

prithvikannan requested review from aravind-segu and smurching October 22, 2024 21:08

prithvikannan mentioned this pull request Oct 22, 2024

Create Langchain Genie #5

Merged

aravind-segu reviewed Oct 23, 2024

View reviewed changes

aravind-segu approved these changes Oct 24, 2024

View reviewed changes

prithvikannan merged commit b314f60 into main Oct 24, 2024

jennsun mentioned this pull request Nov 25, 2025

DatabricksStore: PostgresStore Wrapper SDK #227

Merged

smurching mentioned this pull request Mar 23, 2026

[appkit plugin] [2/x] Add @databricks/appkit-agent package #384

Open

Conversation

prithvikannan commented Oct 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aravind-segu left a comment

Choose a reason for hiding this comment

Uh oh!

prithvikannan commented Oct 24, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

prithvikannan commented Oct 17, 2024 •

edited

Loading