Skip to content

Commit 89e7faf

Browse files
Genesis929averikitsch
authored andcommitted
chore(bigquery): refactor conversational analytics handlers and fix tests (googleapis#2567)
## Description BigQuery Conversational Analytics tests have been failing, and it appears the cause is a recent change in the API's response format and streaming behavior. This PR refactors the handlers to aggregate response fields correctly. ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<issue_number_goes_here> Co-authored-by: Averi Kitsch <akitsch@google.com>
1 parent 07b29fa commit 89e7faf

2 files changed

Lines changed: 33 additions & 21 deletions

File tree

internal/tools/bigquery/bigqueryconversationalanalytics/bigqueryconversationalanalytics.go

Lines changed: 31 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -482,30 +482,33 @@ func handleTextResponse(resp *TextResponse) map[string]any {
482482
}
483483

484484
func handleSchemaResponse(resp *SchemaResponse) map[string]any {
485+
res := make(map[string]any)
485486
if resp.Query != nil {
486-
return map[string]any{"Question": resp.Query.Question}
487+
res["Question"] = resp.Query.Question
487488
}
488489
if resp.Result != nil {
489490
var formattedSources []map[string]any
490491
for _, ds := range resp.Result.Datasources {
491492
formattedSources = append(formattedSources, formatDatasourceAsDict(&ds))
492493
}
493-
return map[string]any{"Schema Resolved": formattedSources}
494+
res["Schema Resolved"] = formattedSources
494495
}
495-
return nil
496+
if len(res) == 0 {
497+
return nil
498+
}
499+
return res
496500
}
497501

498502
func handleDataResponse(resp *DataResponse, maxRows int) map[string]any {
503+
res := make(map[string]any)
499504
if resp.Query != nil {
500-
return map[string]any{
501-
"Retrieval Query": map[string]any{
502-
"Query Name": resp.Query.Name,
503-
"Question": resp.Query.Question,
504-
},
505+
res["Retrieval Query"] = map[string]any{
506+
"Query Name": resp.Query.Name,
507+
"Question": resp.Query.Question,
505508
}
506509
}
507510
if resp.GeneratedSQL != "" {
508-
return map[string]any{"SQL Generated": resp.GeneratedSQL}
511+
res["SQL Generated"] = resp.GeneratedSQL
509512
}
510513
if resp.Result != nil {
511514
var headers []string
@@ -533,15 +536,16 @@ func handleDataResponse(resp *DataResponse, maxRows int) map[string]any {
533536
summary = fmt.Sprintf("Showing the first %d of %d total rows.", numRowsToDisplay, totalRows)
534537
}
535538

536-
return map[string]any{
537-
"Data Retrieved": map[string]any{
538-
"headers": headers,
539-
"rows": compactRows,
540-
"summary": summary,
541-
},
539+
res["Data Retrieved"] = map[string]any{
540+
"headers": headers,
541+
"rows": compactRows,
542+
"summary": summary,
542543
}
543544
}
544-
return nil
545+
if len(res) == 0 {
546+
return nil
547+
}
548+
return res
545549
}
546550

547551
func handleError(resp *ErrorResponse) map[string]any {
@@ -557,9 +561,17 @@ func appendMessage(messages []map[string]any, newMessage map[string]any) []map[s
557561
if newMessage == nil {
558562
return messages
559563
}
560-
if len(messages) > 0 {
561-
if _, ok := messages[len(messages)-1]["Data Retrieved"]; ok {
562-
messages = messages[:len(messages)-1]
564+
565+
if _, hasData := newMessage["Data Retrieved"]; hasData {
566+
// Only keep the last data result while preserving SQL and other metadata.
567+
for i := len(messages) - 1; i >= 0; i-- {
568+
if _, ok := messages[i]["Data Retrieved"]; ok {
569+
delete(messages[i], "Data Retrieved")
570+
if len(messages[i]) == 0 {
571+
messages = append(messages[:i], messages[i+1:]...)
572+
}
573+
break
574+
}
563575
}
564576
}
565577
return append(messages, newMessage)

tests/bigquery/bigquery_integration_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ func TestBigQueryToolEndpoints(t *testing.T) {
173173
datasetInfoWant := "\"Location\":\"US\",\"DefaultTableExpiration\":0,\"Labels\":null,\"Access\":"
174174
tableInfoWant := "{\"Name\":\"\",\"Location\":\"US\",\"Description\":\"\",\"Schema\":[{\"Name\":\"id\""
175175
ddlWant := `"Query executed successfully and returned no content."`
176-
dataInsightsWant := `(?s)Schema Resolved.*Retrieval Query.*SQL Generated.*Answer`
176+
dataInsightsWant := `(?s)(Schema Resolved.*)?(Retrieval Query.*)?SQL Generated.*Data Retrieved.*Answer`
177177
// Partial message; the full error message is too long.
178178
mcpMyFailToolWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"error processing GCP request: failed to insert dry run job: googleapi: Error 400: Syntax error: Unexpected identifier \"SELEC\" at [1:1]`
179179
mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"f0_\":1}"}]}}`
@@ -2393,7 +2393,7 @@ func runBigQueryConversationalAnalyticsInvokeTest(t *testing.T, datasetName, tab
23932393
`{"user_query_with_context": "What are the names in the table?", "table_references": %q}`,
23942394
tableRefsJSON,
23952395
))),
2396-
want: "[{\"f0_\":1}]",
2396+
want: dataInsightsWant,
23972397
isErr: false,
23982398
},
23992399
{

0 commit comments

Comments
 (0)