feat(source/bigquery): add semantic search integration tests

Author: duwenxin99

internal/tools/bigquery/bigquerysql/bigquerysql.go

@@ -219,7 +219,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 }
 
 func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
-	return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
+	return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, embeddingmodels.FormatVectorForPgvector)
 }
 
 func (t Tool) Manifest() tools.Manifest {

tests/bigquery/bigquery_integration_test.go

@@ -156,6 +156,14 @@ func TestBigQueryToolEndpoints(t *testing.T) {
 	tmplSelectCombined, tmplSelectFilterCombined := getBigQueryTmplToolStatement()
 	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, BigqueryToolType, tmplSelectCombined, tmplSelectFilterCombined, "")
 
+	// Set up table for semantic search
+	vectorTableName, teardownVectorTable := setupBigQueryVectorTable(t, ctx, client, datasetName)
+	defer teardownVectorTable(t)
+
+	// Add semantic search tool config
+	insertStmt, searchStmt := getBigQueryVectorSearchStmts(vectorTableName)
+	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, BigqueryToolType, insertStmt, searchStmt)
+
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
 		t.Fatalf("command initialization returned an error: %s", err)
@@ -205,6 +213,7 @@ func TestBigQueryToolEndpoints(t *testing.T) {
 	runBigQueryGetTableInfoToolInvokeTest(t, datasetName, tableName, tableInfoWant)
 	runBigQueryConversationalAnalyticsInvokeTest(t, datasetName, tableName, dataInsightsWant)
 	runBigQuerySearchCatalogToolInvokeTest(t, datasetName, tableName)
+	tests.RunSemanticSearchToolInvokeTest(t, ddlWant, "", "The quick brown fox")
 }
 
 func TestBigQueryToolWithDatasetRestriction(t *testing.T) {
@@ -3259,3 +3268,24 @@ func runAnalyzeContributionWithRestriction(t *testing.T, allowedTableFullName, d
 		})
 	}
 }
+
+// setupBigQueryVectorTable creates a vector table in BigQuery for semantic search testing
+func setupBigQueryVectorTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string) (string, func(*testing.T)) {
+	tableName := fmt.Sprintf("vector_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
+	fullTableName := fmt.Sprintf("`%s.%s.%s`", BigqueryProject, datasetName, tableName)
+	createStatement := fmt.Sprintf(`CREATE TABLE %s (
+		id INT64,
+		content STRING,
+		embedding ARRAY<FLOAT64>
+	)`, fullTableName)
+
+	teardownTable := setupBigQueryTable(t, ctx, client, createStatement, "", datasetName, fullTableName, nil)
+	return fullTableName, teardownTable
+}
+
+// getBigQueryVectorSearchStmts returns statements for bigquery semantic search
+func getBigQueryVectorSearchStmts(vectorTableName string) (string, string) {
+	insertStmt := fmt.Sprintf("INSERT INTO %s (id, content, embedding) VALUES (1, @content, (SELECT ARRAY_AGG(CAST(x AS FLOAT64)) FROM UNNEST(JSON_VALUE_ARRAY(PARSE_JSON(@text_to_embed))) AS x))", vectorTableName)
+	searchStmt := fmt.Sprintf("SELECT id, content, ML.DISTANCE(embedding, (SELECT ARRAY_AGG(CAST(x AS FLOAT64)) FROM UNNEST(JSON_VALUE_ARRAY(PARSE_JSON(@query))) AS x), 'COSINE') AS distance FROM %s ORDER BY distance LIMIT 1", vectorTableName)
+	return insertStmt, searchStmt
+}