diff --git a/libs/genie/client.go b/libs/genie/client.go index 0a9f4b53b6..73522afde5 100644 --- a/libs/genie/client.go +++ b/libs/genie/client.go @@ -18,6 +18,11 @@ import ( // "onechat" even though the CLI command is "genie", so the path keeps that name. const genieResponsesPath = "/api/2.0/data-rooms/tools/onechat/responses" +// SourceDatabricksCLI is the source tag sent on every request so the OneChat +// backend attributes CLI traffic distinctly (maps to OneChatEventSource. +// DATABRICKS_CLI / ConversationSource.CONVERSATION_SOURCE_DATABRICKS_CLI). +const SourceDatabricksCLI = "databricks_cli" + // StreamingTimeoutSeconds is how long the SDK waits between body reads // before canceling the stream. The Genie agent can take minutes between SSE // events when executing multi-step tool calls (search, SQL, etc.), so this @@ -29,6 +34,10 @@ const StreamingTimeoutSeconds = 600 // Genie picks a default warehouse itself. A non-empty conversationID continues // that conversation instead of starting a new one. func BuildRequest(question, warehouseID, conversationID string) GenieRequest { + // The CLI runs in a terminal that cannot render visualizations, so viz is + // always disabled. Bound as a local because GenieRequest.EnableViz is a + // *bool (see the type's comment for why it can't be a plain bool). + enableViz := false return GenieRequest{ Input: []InputItem{ { @@ -41,6 +50,8 @@ func BuildRequest(question, warehouseID, conversationID string) GenieRequest { }, WarehouseID: warehouseID, ConversationID: conversationID, + Source: SourceDatabricksCLI, + EnableViz: &enableViz, } } diff --git a/libs/genie/client_test.go b/libs/genie/client_test.go index a8ed0c06c6..28cecffc00 100644 --- a/libs/genie/client_test.go +++ b/libs/genie/client_test.go @@ -17,7 +17,9 @@ import ( func TestBuildRequest_WireFormat(t *testing.T) { // The backend expects camelCase warehouseId; this pins the wire format, - // not just the Go field values. + // not just the Go field values. source and enable_viz are always sent: + // the CLI tags its traffic (databricks_cli) and disables viz because a + // terminal cannot render charts. b, err := json.Marshal(BuildRequest("What tables exist?", "abc123", "")) require.NoError(t, err) assert.JSONEq(t, `{ @@ -26,10 +28,22 @@ func TestBuildRequest_WireFormat(t *testing.T) { "role": "user", "content": [{"type": "input_text", "text": "What tables exist?"}] }], - "warehouseId": "abc123" + "warehouseId": "abc123", + "source": "databricks_cli", + "enable_viz": false }`, string(b)) } +func TestBuildRequest_TagsSourceAndDisablesViz(t *testing.T) { + // enable_viz:false must always be on the wire (a plain false would be + // dropped by omitempty, letting the backend default viz on), and source + // tags CLI traffic for backend attribution. + b, err := json.Marshal(BuildRequest("q", "", "")) + require.NoError(t, err) + assert.Contains(t, string(b), `"source":"databricks_cli"`) + assert.Contains(t, string(b), `"enable_viz":false`) +} + func TestBuildRequest_OmitsEmptyWarehouse(t *testing.T) { b, err := json.Marshal(BuildRequest("q", "", "")) require.NoError(t, err) diff --git a/libs/genie/types.go b/libs/genie/types.go index fb7afa3c37..b63e6777e5 100644 --- a/libs/genie/types.go +++ b/libs/genie/types.go @@ -7,6 +7,15 @@ type GenieRequest struct { // Tag is camelCase (conversationId), unlike the response's snake_case // conversation_id; the server does not read the snake_case form on input. ConversationID string `json:"conversationId,omitempty"` + // Source identifies this client to the OneChat backend so CLI traffic is + // attributed distinctly from webapp/Slack/MCP. Always SourceDatabricksCLI; + // the server matches it case-insensitively against a fixed allowlist. + Source string `json:"source,omitempty"` + // EnableViz is a pointer so it always serializes (a plain false would be + // dropped by omitempty). The CLI sends false: a terminal cannot render the + // chart/plot items, so we tell the agent to skip registering visualization + // tools. Nil would let the backend default (viz on) apply. + EnableViz *bool `json:"enable_viz,omitempty"` } // InputItem is a message in the input array.