feat(tools/looker): Enable access to the Conversational Analytics API for Looker (#1596)

## Description

This enables the Conversational Analytics API for Looker. The prebuilt
config is separate since it is not a good
idea to use the Looker prebuilt config with CA. Agents get confused as
to whether they should query the data directly
or use the CA tool.
This commit is contained in:
Dr. Strangelove
2025-10-02 13:45:41 -04:00
committed by GitHub
parent 5aed4e136d
commit 2d5a93e312
16 changed files with 999 additions and 9 deletions

View File

@@ -537,6 +537,8 @@ steps:
- "FIRESTORE_PROJECT=$PROJECT_ID"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
- "LOOKER_VERIFY_SSL=$_LOOKER_VERIFY_SSL"
- "LOOKER_PROJECT=$_LOOKER_PROJECT"
- "LOOKER_LOCATION=$_LOOKER_LOCATION"
secretEnv:
[
"CLIENT_ID",
@@ -824,6 +826,8 @@ substitutions:
_DGRAPHURL: "https://play.dgraph.io"
_COUCHBASE_BUCKET: "couchbase-bucket"
_COUCHBASE_SCOPE: "couchbase-scope"
_LOOKER_LOCATION: "us"
_LOOKER_PROJECT: "149671255749"
_LOOKER_VERIFY_SSL: "true"
_TIDB_HOST: 127.0.0.1
_TIDB_PORT: "4000"

View File

@@ -98,6 +98,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorevalidaterules"
_ "github.com/googleapis/genai-toolbox/internal/tools/http"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookeradddashboardelement"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerconversationalanalytics"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdashboards"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdimensions"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetexplores"

View File

@@ -1244,6 +1244,7 @@ func TestPrebuiltTools(t *testing.T) {
mysql_config, _ := prebuiltconfigs.Get("mysql")
mssql_config, _ := prebuiltconfigs.Get("mssql")
looker_config, _ := prebuiltconfigs.Get("looker")
lookerca_config, _ := prebuiltconfigs.Get("looker-conversational-analytics")
postgresconfig, _ := prebuiltconfigs.Get("postgres")
spanner_config, _ := prebuiltconfigs.Get("spanner")
spannerpg_config, _ := prebuiltconfigs.Get("spanner-postgres")
@@ -1327,6 +1328,9 @@ func TestPrebuiltTools(t *testing.T) {
t.Setenv("LOOKER_CLIENT_SECRET", "your_looker_client_secret")
t.Setenv("LOOKER_VERIFY_SSL", "true")
t.Setenv("LOOKER_PROJECT", "your_project_id")
t.Setenv("LOOKER_LOCATION", "us")
t.Setenv("SQLITE_DATABASE", "test.db")
t.Setenv("NEO4J_URI", "bolt://localhost:7687")
@@ -1493,6 +1497,16 @@ func TestPrebuiltTools(t *testing.T) {
},
},
},
{
name: "looker-conversational-analytics prebuilt tools",
in: lookerca_config,
wantToolset: server.ToolsetConfigs{
"looker_conversational_analytics_tools": tools.ToolsetConfig{
Name: "looker_conversational_analytics_tools",
ToolNames: []string{"ask_data_insights", "get_models", "get_explores"},
},
},
},
{
name: "postgres prebuilt tools",
in: postgresconfig,

View File

@@ -358,6 +358,10 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
* `LOOKER_CLIENT_ID`: The client ID for the Looker API.
* `LOOKER_CLIENT_SECRET`: The client secret for the Looker API.
* `LOOKER_VERIFY_SSL`: Whether to verify SSL certificates.
* `LOOKER_USE_CLIENT_OAUTH`: Whether to use OAuth for authentication.
* `LOOKER_SHOW_HIDDEN_MODELS`: Whether to show hidden models.
* `LOOKER_SHOW_HIDDEN_EXPLORES`: Whether to show hidden explores.
* `LOOKER_SHOW_HIDDEN_FIELDS`: Whether to show hidden fields.
* **Permissions:**
* A Looker account with permissions to access the desired models,
explores, and data is required.
@@ -377,6 +381,35 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
* `get_dashboards`: Searches for saved dashboards.
* `make_dashboard`: Creates a new dashboard.
* `add_dashboard_element`: Adds a tile to a dashboard.
* `health_pulse`: Test the health of a Looker instance.
* `health_analyze`: Analyze the LookML usage of a Looker instance.
* `health_vacuum`: Suggest LookML elements that can be removed.
## Looker Conversational Analytics
* `--prebuilt` value: `looker-conversational-analytics`
* **Environment Variables:**
* `LOOKER_BASE_URL`: The URL of your Looker instance.
* `LOOKER_CLIENT_ID`: The client ID for the Looker API.
* `LOOKER_CLIENT_SECRET`: The client secret for the Looker API.
* `LOOKER_VERIFY_SSL`: Whether to verify SSL certificates.
* `LOOKER_USE_CLIENT_OAUTH`: Whether to use OAuth for authentication.
* `LOOKER_PROJECT`: The GCP Project to use for Conversational Analytics.
* `LOOKER_LOCATION`: The GCP Location to use for Conversational Analytics.
* **Permissions:**
* A Looker account with permissions to access the desired models,
explores, and data is required.
* **Looker Instance User** (`roles/looker.instanceUser`): IAM role to
access Looker.
* **Gemini for Google Cloud User** (`roles/cloudaicompanion.user`): IAM
role to access Conversational Analytics.
* **Gemini Data Analytics Stateless Chat User (Beta)**
(`roles/geminidataanalytics.dataAgentStatelessUser`): IAM role to
access Conversational Analytics.
* **Tools:**
* `ask_data_insights`: Ask a question of the data.
* `get_models`: Retrieves the list of LookML models.
* `get_explores`: Retrieves the list of explores in a model.
## Microsoft SQL Server

View File

@@ -16,7 +16,7 @@ in the cloud, on GCP, or on premises.
## Requirements
### Database User
### Looker User
This source only uses API authentication. You will need to
[create an API user][looker-user] to login to Looker.
@@ -24,6 +24,35 @@ This source only uses API authentication. You will need to
[looker-user]:
https://cloud.google.com/looker/docs/api-auth#authentication_with_an_sdk
{{< notice note >}}
To use the Conversational Analytics API, you will need to have the following
Google Cloud Project API enabled and IAM permissions.
{{< /notice >}}
### API Enablement in GCP
Enable the following APIs in your Google Cloud Project:
```
gcloud services enable geminidataanalytics.googleapis.com --project=$PROJECT_ID
gcloud services enable cloudaicompanion.googleapis.com --project=$PROJECT_ID
```
### IAM Permissions in GCP
In addition to [setting the ADC for your server][set-adc], you need to ensure
the IAM identity has been given the following IAM roles (or corresponding
permissions):
- `roles/looker.instanceUser`
- `roles/cloudaicompanion.user`
- `roles/geminidataanalytics.dataAgentStatelessUser`
To initialize the application default credential run `gcloud auth login --update-adc`
in your environment before starting MCP Toolbox.
[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc
## Example
```yaml
@@ -33,6 +62,8 @@ sources:
base_url: http://looker.example.com
client_id: ${LOOKER_CLIENT_ID}
client_secret: ${LOOKER_CLIENT_SECRET}
project: ${LOOKER_PROJECT}
location: ${LOOKER_LOCATION}
verify_ssl: true
timeout: 600s
```
@@ -50,6 +81,8 @@ The client id and client secret are seemingly random character sequences
assigned by the looker server. If you are using Looker OAuth you don't need
these settings
The `project` and `location` fields are utilized **only** when using the conversational analytics tool.
{{< notice tip >}}
Use environment variable replacement with the format ${ENV_NAME}
instead of hardcoding your secrets into the configuration file.
@@ -64,6 +97,8 @@ instead of hardcoding your secrets into the configuration file.
| client_id | string | false | The client id assigned by Looker. |
| client_secret | string | false | The client secret assigned by Looker. |
| verify_ssl | string | false | Whether to check the ssl certificate of the server. |
| project | string | false | The project id to use in Google Cloud. |
| location | string | false | The location to use in Google Cloud. (default: us) |
| timeout | string | false | Maximum time to wait for query execution (e.g. "30s", "2m"). By default, 120s is applied. |
| use_client_oauth | string | false | Use OAuth tokens instead of client_id and client_secret. (default: false) |
| show_hidden_models | string | false | Show or hide hidden models. (default: true) |

View File

@@ -0,0 +1,45 @@
---
title: "looker-conversational-analytics"
type: docs
weight: 1
description: >
The "looker-conversational-analytics" tool will use the Conversational
Analaytics API to analyze data from Looker
aliases:
- /resources/tools/looker-conversational-analytics
---
## About
A `looker-conversational-analytics` tool allows you to ask questions about your Looker data.
It's compatible with the following sources:
- [looker](../../sources/looker.md)
`looker-conversational-analytics` accepts two parameters:
1. `user_query_with_context`: The question asked of the Conversational Analytics system.
2. `explore_references`: A list of one to five explores that can be queried to answer the
question. The form of the entry is `[{"model": "model name", "explore": "explore name"}, ...]`
## Example
```yaml
tools:
ask_data_insights:
kind: looker-conversational-analytics
source: looker-source
description: |
Use this tool to perform data analysis, get insights,
or answer complex questions about the contents of specific
Looker explores.
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:--------:|:------------:|----------------------------------------------------|
| kind | string | true | Must be "lookerca-conversational-analytics". |
| source | string | true | Name of the source the SQL should execute on. |
| description | string | true | Description of the tool that is passed to the LLM. |

1
go.mod
View File

@@ -11,6 +11,7 @@ require (
cloud.google.com/go/cloudsqlconn v1.18.1
cloud.google.com/go/dataplex v1.27.1
cloud.google.com/go/firestore v1.18.0
cloud.google.com/go/geminidataanalytics v0.2.1
cloud.google.com/go/spanner v1.86.0
github.com/ClickHouse/clickhouse-go/v2 v2.40.3
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0

2
go.sum
View File

@@ -309,6 +309,8 @@ cloud.google.com/go/gaming v1.6.0/go.mod h1:YMU1GEvA39Qt3zWGyAVA9bpYz/yAhTvaQ1t2
cloud.google.com/go/gaming v1.7.0/go.mod h1:LrB8U7MHdGgFG851iHAfqUdLcKBdQ55hzXy9xBJz0+w=
cloud.google.com/go/gaming v1.8.0/go.mod h1:xAqjS8b7jAVW0KFYeRUxngo9My3f33kFmua++Pi+ggM=
cloud.google.com/go/gaming v1.9.0/go.mod h1:Fc7kEmCObylSWLO334NcO+O9QMDyz+TKC4v1D7X+Bc0=
cloud.google.com/go/geminidataanalytics v0.2.1 h1:gtG/9VlUJpL67yukFen/twkAEHliYvW7610Rlnn5rpQ=
cloud.google.com/go/geminidataanalytics v0.2.1/go.mod h1:gIsj/ELDCzVbw24185zwjXgbzYiqdGe7TSSK2HrdtA0=
cloud.google.com/go/gkebackup v0.2.0/go.mod h1:XKvv/4LfG829/B8B7xRkk8zRrOEbKtEam6yNfuQNH60=
cloud.google.com/go/gkebackup v0.3.0/go.mod h1:n/E671i1aOQvUxT541aTkCwExO/bTer2HDlj4TsBRAo=
cloud.google.com/go/gkebackup v0.4.0/go.mod h1:byAyBGUwYGEEww7xsbnUTBHIYcOPy/PgUWUtOeRm9Vg=

View File

@@ -37,6 +37,7 @@ var expectedToolSources = []string{
"cloud-sql-postgres",
"dataplex",
"firestore",
"looker-conversational-analytics",
"looker",
"mssql",
"mysql",
@@ -108,6 +109,8 @@ func TestGetPrebuiltTool(t *testing.T) {
cloudsqlmssql_config, _ := Get("cloud-sql-mssql")
dataplex_config, _ := Get("dataplex")
firestoreconfig, _ := Get("firestore")
looker_config, _ := Get("looker")
lookerca_config, _ := Get("looker-conversational-analytics")
mysql_config, _ := Get("mysql")
mssql_config, _ := Get("mssql")
oceanbase_config, _ := Get("oceanbase")
@@ -164,6 +167,12 @@ func TestGetPrebuiltTool(t *testing.T) {
if len(firestoreconfig) <= 0 {
t.Fatalf("unexpected error: could not fetch firestore prebuilt tools yaml")
}
if len(looker_config) <= 0 {
t.Fatalf("unexpected error: could not fetch looker prebuilt tools yaml")
}
if len(lookerca_config) <= 0 {
t.Fatalf("unexpected error: could not fetch looker-conversational-analytics prebuilt tools yaml")
}
if len(mysql_config) <= 0 {
t.Fatalf("unexpected error: could not fetch mysql prebuilt tools yaml")
}

View File

@@ -0,0 +1,57 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
sources:
looker-source:
kind: looker
base_url: ${LOOKER_BASE_URL}
client_id: ${LOOKER_CLIENT_ID:}
client_secret: ${LOOKER_CLIENT_SECRET:}
verify_ssl: ${LOOKER_VERIFY_SSL:true}
timeout: 600s
use_client_oauth: ${LOOKER_USE_CLIENT_OAUTH:false}
project: ${LOOKER_PROJECT:}
location: ${LOOKER_LOCATION:}
tools:
ask_data_insights:
kind: looker-conversational-analytics
source: looker-source
description: |
Use this tool to perform data analysis, get insights,
or answer complex questions about the contents of specific
Looker explores.
get_models:
kind: looker-get-models
source: looker-source
description: |
The get_models tool retrieves the list of LookML models in the Looker system.
It takes no parameters.
get_explores:
kind: looker-get-explores
source: looker-source
description: |
The get_explores tool retrieves the list of explores defined in a LookML model
in the Looker system.
It takes one parameter, the model_name looked up from get_models.
toolsets:
looker_conversational_analytics_tools:
- ask_data_insights
- get_models
- get_explores

View File

@@ -18,10 +18,13 @@ import (
"fmt"
"time"
geminidataanalytics "cloud.google.com/go/geminidataanalytics/apiv1beta"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/util"
"go.opentelemetry.io/otel/trace"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"
"github.com/looker-open-source/sdk-codegen/go/rtl"
v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
@@ -47,6 +50,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
ShowHiddenModels: true,
ShowHiddenExplores: true,
ShowHiddenFields: true,
Location: "us",
} // Default Ssl,timeout, ShowHidden
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
@@ -66,6 +70,8 @@ type Config struct {
ShowHiddenModels bool `yaml:"show_hidden_models"`
ShowHiddenExplores bool `yaml:"show_hidden_explores"`
ShowHiddenFields bool `yaml:"show_hidden_fields"`
Project string `yaml:"project"`
Location string `yaml:"location"`
}
func (r Config) SourceConfigKind() string {
@@ -102,6 +108,9 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
ClientSecret: r.ClientSecret,
}
var tokenSource oauth2.TokenSource
tokenSource, _ = initGoogleCloudConnection(ctx)
s := &Source{
Name: r.Name,
Kind: SourceKind,
@@ -111,6 +120,9 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
ShowHiddenModels: r.ShowHiddenModels,
ShowHiddenExplores: r.ShowHiddenExplores,
ShowHiddenFields: r.ShowHiddenFields,
Project: r.Project,
Location: r.Location,
TokenSource: tokenSource,
}
if !r.UseClientOAuth {
@@ -137,12 +149,48 @@ type Source struct {
Timeout string `yaml:"timeout"`
Client *v4.LookerSDK
ApiSettings *rtl.ApiSettings
UseClientOAuth bool `yaml:"use_client_oauth"`
ShowHiddenModels bool `yaml:"show_hidden_models"`
ShowHiddenExplores bool `yaml:"show_hidden_explores"`
ShowHiddenFields bool `yaml:"show_hidden_fields"`
UseClientOAuth bool `yaml:"use_client_oauth"`
ShowHiddenModels bool `yaml:"show_hidden_models"`
ShowHiddenExplores bool `yaml:"show_hidden_explores"`
ShowHiddenFields bool `yaml:"show_hidden_fields"`
Project string `yaml:"project"`
Location string `yaml:"location"`
TokenSource oauth2.TokenSource
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) GetApiSettings() *rtl.ApiSettings {
return s.ApiSettings
}
func (s *Source) UseClientAuthorization() bool {
return s.UseClientOAuth
}
func (s *Source) GoogleCloudProject() string {
return s.Project
}
func (s *Source) GoogleCloudLocation() string {
return s.Location
}
func (s *Source) GoogleCloudTokenSource() oauth2.TokenSource {
return s.TokenSource
}
func (s *Source) GoogleCloudTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error) {
return google.DefaultTokenSource(ctx, scope)
}
func initGoogleCloudConnection(ctx context.Context) (oauth2.TokenSource, error) {
cred, err := google.FindDefaultCredentials(ctx, geminidataanalytics.DefaultAuthScopes()...)
if err != nil {
return nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", geminidataanalytics.DefaultAuthScopes(), err)
}
return cred.TokenSource, nil
}

View File

@@ -54,6 +54,7 @@ func TestParseFromYamlLooker(t *testing.T) {
ShowHiddenModels: true,
ShowHiddenExplores: true,
ShowHiddenFields: true,
Location: "us",
},
},
},
@@ -90,9 +91,9 @@ func TestFailParseFromYamlLooker(t *testing.T) {
base_url: http://example.looker.com/
client_id: jasdl;k;tjl
client_secret: sdakl;jgflkasdfkfg
project: test-project
schema: test-schema
`,
err: "unable to parse source \"my-looker-instance\" as \"looker\": [5:1] unknown field \"project\"\n 2 | client_id: jasdl;k;tjl\n 3 | client_secret: sdakl;jgflkasdfkfg\n 4 | kind: looker\n> 5 | project: test-project\n ^\n",
err: "unable to parse source \"my-looker-instance\" as \"looker\": [5:1] unknown field \"schema\"\n 2 | client_id: jasdl;k;tjl\n 3 | client_secret: sdakl;jgflkasdfkfg\n 4 | kind: looker\n> 5 | schema: test-schema\n ^\n",
},
{
desc: "missing required field",
@@ -100,6 +101,7 @@ func TestFailParseFromYamlLooker(t *testing.T) {
sources:
my-looker-instance:
kind: looker
client_id: jasdl;k;tjl
`,
err: "unable to parse source \"my-looker-instance\" as \"looker\": Key: 'Config.BaseURL' Error:Field validation for 'BaseURL' failed on the 'required' tag",
},

View File

@@ -0,0 +1,550 @@
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package lookerconversationalanalytics
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
lookerds "github.com/googleapis/genai-toolbox/internal/sources/looker"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/looker-open-source/sdk-codegen/go/rtl"
"golang.org/x/oauth2"
)
const kind string = "looker-conversational-analytics"
const instructions = `**INSTRUCTIONS - FOLLOW THESE RULES:**
1. **CONTENT:** Your answer should present the supporting data and then provide a conclusion based on that data.
2. **OUTPUT FORMAT:** Your entire response MUST be in plain text format ONLY.
3. **NO CHARTS:** You are STRICTLY FORBIDDEN from generating any charts, graphs, images, or any other form of visualization.`
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
}
}
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
actual := Config{Name: name}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type compatibleSource interface {
GetApiSettings() *rtl.ApiSettings
GoogleCloudTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error)
GoogleCloudProject() string
GoogleCloudLocation() string
UseClientAuthorization() bool
}
// Structs for building the JSON payload
type UserMessage struct {
Text string `json:"text"`
}
type Message struct {
UserMessage UserMessage `json:"userMessage"`
}
type LookerExploreReference struct {
LookerInstanceUri string `json:"lookerInstanceUri"`
LookmlModel string `json:"lookmlModel"`
Explore string `json:"explore"`
}
type LookerExploreReferences struct {
ExploreReferences []LookerExploreReference `json:"exploreReferences"`
Credentials Credentials `json:"credentials,omitzero"`
}
type SecretBased struct {
ClientId string `json:"clientId"`
ClientSecret string `json:"clientSecret"`
}
type TokenBased struct {
AccessToken string `json:"accessToken"`
}
type OAuthCredentials struct {
Secret SecretBased `json:"secret,omitzero"`
Token TokenBased `json:"token,omitzero"`
}
type Credentials struct {
OAuth OAuthCredentials `json:"oauth"`
}
type DatasourceReferences struct {
Looker LookerExploreReferences `json:"looker"`
}
type ImageOptions struct {
NoImage map[string]any `json:"noImage"`
}
type ChartOptions struct {
Image ImageOptions `json:"image"`
}
type Python struct {
Enabled bool `json:"enabled"`
}
type AnalysisOptions struct {
Python Python `json:"python"`
}
type ConversationOptions struct {
Chart ChartOptions `json:"chart,omitzero"`
Analysis AnalysisOptions `json:"analysis,omitzero"`
}
type InlineContext struct {
SystemInstruction string `json:"systemInstruction"`
DatasourceReferences DatasourceReferences `json:"datasourceReferences"`
Options ConversationOptions `json:"options"`
}
type CAPayload struct {
Messages []Message `json:"messages"`
InlineContext InlineContext `json:"inlineContext"`
ClientIdEnum string `json:"clientIdEnum"`
}
// validate compatible sources are still compatible
var _ compatibleSource = &lookerds.Source{}
var compatibleSources = [...]string{lookerds.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is compatible
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
if s.GoogleCloudProject() == "" {
return nil, fmt.Errorf("project must be defined for source to use with %q tool", kind)
}
userQueryParameter := tools.NewStringParameter("user_query_with_context", "The user's question, potentially including conversation history and system instructions for context.")
exploreRefsDescription := `An Array of at least one and up to 5 explore references like [{'model': 'MODEL_NAME', 'explore': 'EXPLORE_NAME'}]`
exploreRefsParameter := tools.NewArrayParameter(
"explore_references",
exploreRefsDescription,
tools.NewMapParameter(
"explore_reference",
"An explore reference like {'model': 'MODEL_NAME', 'explore': 'EXPLORE_NAME'}",
"",
),
)
parameters := tools.Parameters{userQueryParameter, exploreRefsParameter}
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
// Get cloud-platform token source for Gemini Data Analytics API during initialization
ctx := context.Background()
ts, err := s.GoogleCloudTokenSourceWithScope(ctx, "https://www.googleapis.com/auth/cloud-platform")
if err != nil {
return nil, fmt.Errorf("failed to get cloud-platform token source: %w", err)
}
// finish tool setup
t := Tool{
Name: cfg.Name,
Kind: kind,
ApiSettings: s.GetApiSettings(),
Project: s.GoogleCloudProject(),
Location: s.GoogleCloudLocation(),
Parameters: parameters,
AuthRequired: cfg.AuthRequired,
UseClientOAuth: s.UseClientAuthorization(),
TokenSource: ts,
manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
// validate interface
var _ tools.Tool = Tool{}
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
ApiSettings *rtl.ApiSettings
AuthRequired []string `yaml:"authRequired"`
UseClientOAuth bool `yaml:"useClientOAuth"`
Parameters tools.Parameters `yaml:"parameters"`
Project string
Location string
TokenSource oauth2.TokenSource
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
var tokenStr string
var err error
// Get credentials for the API call
// Use cloud-platform token source for Gemini Data Analytics API
if t.TokenSource == nil {
return nil, fmt.Errorf("cloud-platform token source is missing")
}
token, err := t.TokenSource.Token()
if err != nil {
return nil, fmt.Errorf("failed to get token from cloud-platform token source: %w", err)
}
tokenStr = token.AccessToken
// Extract parameters from the map
mapParams := params.AsMap()
userQuery, _ := mapParams["user_query_with_context"].(string)
exploreReferences, _ := mapParams["explore_references"].([]any)
ler := make([]LookerExploreReference, 0)
for _, er := range exploreReferences {
ler = append(ler, LookerExploreReference{
LookerInstanceUri: t.ApiSettings.BaseUrl,
LookmlModel: er.(map[string]any)["model"].(string),
Explore: er.(map[string]any)["explore"].(string),
})
}
oauth_creds := OAuthCredentials{}
if t.UseClientOAuth {
oauth_creds.Token = TokenBased{AccessToken: string(accessToken)}
} else {
oauth_creds.Secret = SecretBased{ClientId: t.ApiSettings.ClientId, ClientSecret: t.ApiSettings.ClientSecret}
}
lers := LookerExploreReferences{
ExploreReferences: ler,
Credentials: Credentials{
OAuth: oauth_creds,
},
}
// Construct URL, headers, and payload
projectID := t.Project
location := t.Location
caURL := fmt.Sprintf("https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s:chat", url.PathEscape(projectID), url.PathEscape(location))
headers := map[string]string{
"Authorization": fmt.Sprintf("Bearer %s", tokenStr),
"Content-Type": "application/json",
}
payload := CAPayload{
Messages: []Message{{UserMessage: UserMessage{Text: userQuery}}},
InlineContext: InlineContext{
SystemInstruction: instructions,
DatasourceReferences: DatasourceReferences{
Looker: lers,
},
Options: ConversationOptions{Chart: ChartOptions{Image: ImageOptions{NoImage: map[string]any{}}}},
},
ClientIdEnum: "GENAI_TOOLBOX",
}
// Call the streaming API
response, err := getStream(ctx, caURL, payload, headers)
if err != nil {
return nil, fmt.Errorf("failed to get response from conversational analytics API: %w", err)
}
return response, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
return tools.ParseParams(t.Parameters, data, claims)
}
func (t Tool) Manifest() tools.Manifest {
return t.manifest
}
func (t Tool) McpManifest() tools.McpManifest {
return t.mcpManifest
}
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t Tool) RequiresClientAuthorization() bool {
return t.UseClientOAuth
}
// StreamMessage represents a single message object from the streaming API response.
type StreamMessage struct {
SystemMessage *SystemMessage `json:"systemMessage,omitempty"`
}
// SystemMessage contains different types of system-generated content.
type SystemMessage struct {
Text *TextMessage `json:"text,omitempty"`
Schema *SchemaMessage `json:"schema,omitempty"`
Data *DataMessage `json:"data,omitempty"`
Analysis *AnalysisMessage `json:"analysis,omitempty"`
Error *ErrorMessage `json:"error,omitempty"`
}
// TextMessage contains textual parts of a message.
type TextMessage struct {
Parts []string `json:"parts"`
}
// SchemaMessage contains schema-related information.
type SchemaMessage struct {
Query *SchemaQuery `json:"query,omitempty"`
Result *SchemaResult `json:"result,omitempty"`
}
// SchemaQuery holds the question that prompted a schema lookup.
type SchemaQuery struct {
Question string `json:"question"`
}
// SchemaResult contains the datasources with their schemas.
type SchemaResult struct {
Datasources []Datasource `json:"datasources"`
}
// Datasource represents a data source with its reference and schema.
type Datasource struct {
LookerExploreReference LookerExploreReference `json:"lookerExploreReference"`
}
// DataMessage contains data-related information, like queries and results.
type DataMessage struct {
GeneratedLookerQuery *LookerQuery `json:"generatedLookerQuery,omitempty"`
Result *DataResult `json:"result,omitempty"`
}
type LookerQuery struct {
Model string `json:"model"`
Explore string `json:"explore"`
Fields []string `json:"fields"`
Filters []Filter `json:"filters,omitempty"`
Sorts []string `json:"sorts,omitempty"`
Limit string `json:"limit,omitempty"`
}
type Filter struct {
Field string `json:"field,omitempty"`
Value string `json:"value,omitempty"`
}
// DataResult contains the schema and rows of a query result.
type DataResult struct {
Data []map[string]any `json:"data"`
}
type AnalysisQuery struct {
Question string `json:"question,omitempty"`
DataResultNames []string `json:"dataResultNames,omitempty"`
}
type AnalysisEvent struct {
PlannerReasoning string `json:"plannerReasoning,omitempty"`
CoderInstructions string `json:"coderInstructions,omitempty"`
Code string `json:"code,omitempty"`
ExecutionOutput string `json:"executionOutput,omitempty"`
ExecutionError string `json:"executionError,omitempty"`
ResultVegaChartJson string `json:"resultVegaChartJson,omitempty"`
ResultNaturalLanguage string `json:"resultNaturalLanguage,omitempty"`
ResultCsvData string `json:"resultCsvData,omitempty"`
ResultReferenceData string `json:"resultReferenceData,omitempty"`
Error string `json:"error,omitempty"`
}
type AnalysisMessage struct {
Query AnalysisQuery `json:"query,omitempty"`
ProgressEvent AnalysisEvent `json:"progressEvent,omitempty"`
}
// ErrorResponse represents an error message from the API.
type ErrorMessage struct {
Text string `json:"text"`
}
func getStream(ctx context.Context, url string, payload CAPayload, headers map[string]string) ([]map[string]any, error) {
payloadBytes, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal payload: %w", err)
}
req, err := http.NewRequest("POST", url, bytes.NewBuffer(payloadBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
for k, v := range headers {
req.Header.Set(k, v)
}
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("API returned non-200 status: %d %s", resp.StatusCode, string(body))
}
var messages []map[string]any
decoder := json.NewDecoder(resp.Body)
// The response is a JSON array, so we read the opening bracket.
if _, err := decoder.Token(); err != nil {
if err == io.EOF {
return nil, nil // Empty response is valid
}
return nil, fmt.Errorf("error reading start of json array: %w", err)
}
for decoder.More() {
var msg StreamMessage
if err := decoder.Decode(&msg); err != nil {
if err == io.EOF {
break
}
return nil, fmt.Errorf("error decoding stream message: %w", err)
}
var newMessage map[string]any
if msg.SystemMessage != nil {
if msg.SystemMessage.Text != nil {
newMessage = handleTextResponse(ctx, msg.SystemMessage.Text)
} else if msg.SystemMessage.Schema != nil {
newMessage = handleSchemaResponse(ctx, msg.SystemMessage.Schema)
} else if msg.SystemMessage.Data != nil {
newMessage = handleDataResponse(ctx, msg.SystemMessage.Data)
} else if msg.SystemMessage.Analysis != nil {
newMessage = handleAnalysisResponse(ctx, msg.SystemMessage.Analysis)
} else if msg.SystemMessage.Error != nil {
newMessage = handleError(ctx, msg.SystemMessage.Error)
}
messages = appendMessage(messages, newMessage)
}
}
return messages, nil
}
func formatDatasourceAsDict(ctx context.Context, datasource *Datasource) map[string]any {
logger, _ := util.LoggerFromContext(ctx)
logger.DebugContext(ctx, "Datasource %s", *datasource)
ds := make(map[string]any)
ds["model"] = datasource.LookerExploreReference.LookmlModel
ds["explore"] = datasource.LookerExploreReference.Explore
ds["lookerInstanceUri"] = datasource.LookerExploreReference.LookerInstanceUri
return map[string]any{"Datasource": ds}
}
func handleAnalysisResponse(ctx context.Context, resp *AnalysisMessage) map[string]any {
logger, _ := util.LoggerFromContext(ctx)
jsonData, err := json.Marshal(*resp)
if err != nil {
logger.ErrorContext(ctx, "error marshaling struct: %w", err)
return map[string]any{"Analysis": "error"}
}
return map[string]any{"Analysis": jsonData}
}
func handleTextResponse(ctx context.Context, resp *TextMessage) map[string]any {
logger, _ := util.LoggerFromContext(ctx)
logger.DebugContext(ctx, "Text Response: %s", strings.Join(resp.Parts, ""))
return map[string]any{"Answer": strings.Join(resp.Parts, "")}
}
func handleSchemaResponse(ctx context.Context, resp *SchemaMessage) map[string]any {
if resp.Query != nil {
return map[string]any{"Question": resp.Query.Question}
}
if resp.Result != nil {
var formattedSources []map[string]any
for _, ds := range resp.Result.Datasources {
formattedSources = append(formattedSources, formatDatasourceAsDict(ctx, &ds))
}
return map[string]any{"Schema Resolved": formattedSources}
}
return nil
}
func handleDataResponse(ctx context.Context, resp *DataMessage) map[string]any {
if resp.GeneratedLookerQuery != nil {
logger, _ := util.LoggerFromContext(ctx)
jsonData, err := json.Marshal(resp.GeneratedLookerQuery)
if err != nil {
logger.ErrorContext(ctx, "error marshaling struct: %w", err)
return map[string]any{"Retrieval Query": "error"}
}
return map[string]any{
"Retrieval Query": jsonData,
}
}
if resp.Result != nil {
return map[string]any{
"Data Retrieved": resp.Result.Data,
}
}
return nil
}
func handleError(ctx context.Context, resp *ErrorMessage) map[string]any {
logger, _ := util.LoggerFromContext(ctx)
logger.DebugContext(ctx, "Error Response: %s", resp.Text)
return map[string]any{
"Error": map[string]any{
"Message": resp.Text,
},
}
}
func appendMessage(messages []map[string]any, newMessage map[string]any) []map[string]any {
if newMessage == nil {
return messages
}
if len(messages) > 0 {
if _, ok := messages[len(messages)-1]["Data Retrieved"]; ok {
messages = messages[:len(messages)-1]
}
}
return append(messages, newMessage)
}

View File

@@ -0,0 +1,72 @@
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package lookerconversationalanalytics_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/looker/lookerconversationalanalytics"
)
func TestParseFromYamlLookerConversationalAnalytics(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: looker-conversational-analytics
source: my-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": lookerconversationalanalytics.Config{
Name: "example_tool",
Kind: "looker-conversational-analytics",
Source: "my-instance",
Description: "some description",
AuthRequired: []string{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}

View File

@@ -833,7 +833,7 @@ func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool
}
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", strings.Join(tablesToDrop, ", "))
if _, err := pool.Exec(ctx, dropQuery); err != nil {
t.Fatalf("Failed to drop all tables in 'public' schema: %v", err)
}
@@ -871,7 +871,7 @@ func CleanupMySQLTables(t *testing.T, ctx context.Context, pool *sql.DB) {
}
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s;", strings.Join(tablesToDrop, ", "))
if _, err := pool.ExecContext(ctx, dropQuery); err != nil {
// Try to re-enable checks even if drop fails
if _, err := pool.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 1;"); err != nil {

View File

@@ -15,9 +15,14 @@
package looker
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"regexp"
"strings"
"testing"
"time"
@@ -33,6 +38,8 @@ var (
LookerVerifySsl = os.Getenv("LOOKER_VERIFY_SSL")
LookerClientId = os.Getenv("LOOKER_CLIENT_ID")
LookerClientSecret = os.Getenv("LOOKER_CLIENT_SECRET")
LookerProject = os.Getenv("LOOKER_PROJECT")
LookerLocation = os.Getenv("LOOKER_LOCATION")
)
func getLookerVars(t *testing.T) map[string]any {
@@ -45,6 +52,10 @@ func getLookerVars(t *testing.T) map[string]any {
t.Fatal("'LOOKER_CLIENT_ID' not set")
case LookerClientSecret:
t.Fatal("'LOOKER_CLIENT_SECRET' not set")
case LookerProject:
t.Fatal("'LOOKER_PROJECT' not set")
case LookerLocation:
t.Fatal("'LOOKER_LOCATION' not set")
}
return map[string]any{
@@ -53,6 +64,8 @@ func getLookerVars(t *testing.T) map[string]any {
"verify_ssl": (LookerVerifySsl == "true"),
"client_id": LookerClientId,
"client_secret": LookerClientSecret,
"project": LookerProject,
"location": LookerLocation,
}
}
@@ -131,6 +144,11 @@ func TestLooker(t *testing.T) {
"source": "my-instance",
"description": "Simple tool to test end to end functionality.",
},
"conversational_analytics": map[string]any{
"kind": "looker-conversational-analytics",
"source": "my-instance",
"description": "Simple tool to test end to end functionality.",
},
"health_pulse": map[string]any{
"kind": "looker-health-pulse",
"source": "my-instance",
@@ -633,6 +651,39 @@ func TestLooker(t *testing.T) {
},
},
)
tests.RunToolGetTestByName(t, "conversational_analytics",
map[string]any{
"conversational_analytics": map[string]any{
"description": "Simple tool to test end to end functionality.",
"authRequired": []any{},
"parameters": []any{
map[string]any{
"authSources": []any{},
"description": "The user's question, potentially including conversation history and system instructions for context.",
"name": "user_query_with_context",
"required": true,
"type": "string",
},
map[string]any{
"authSources": []any{},
"description": "An Array of at least one and up to 5 explore references like [{'model': 'MODEL_NAME', 'explore': 'EXPLORE_NAME'}]",
"items": map[string]any{
"additionalProperties": true,
"authSources": []any{},
"name": "explore_reference",
"description": "An explore reference like {'model': 'MODEL_NAME', 'explore': 'EXPLORE_NAME'}",
"required": true,
"type": "object",
},
"name": "explore_references",
"required": true,
"type": "array",
},
},
},
},
)
tests.RunToolGetTestByName(t, "health_pulse",
map[string]any{
"health_pulse": map[string]any{
@@ -789,6 +840,8 @@ func TestLooker(t *testing.T) {
wantResult = "null"
tests.RunToolInvokeParametersTest(t, "get_dashboards", []byte(`{"title": "FOO", "desc": "BAR"}`), wantResult)
runConversationalAnalytics(t, "system__activity", "content_usage")
wantResult = "\"Connection\":\"thelook\""
tests.RunToolInvokeParametersTest(t, "health_pulse", []byte(`{"action": "check_db_connections"}`), wantResult)
@@ -806,4 +859,68 @@ func TestLooker(t *testing.T) {
wantResult = "\"Model\":\"the_look\""
tests.RunToolInvokeParametersTest(t, "health_vacuum", []byte(`{"action": "models"}`), wantResult)
}
func runConversationalAnalytics(t *testing.T, modelName, exploreName string) {
exploreRefsJSON := fmt.Sprintf(`[{"model":"%s","explore":"%s"}]`, modelName, exploreName)
var refs []map[string]any
if err := json.Unmarshal([]byte(exploreRefsJSON), &refs); err != nil {
t.Fatalf("failed to unmarshal explore refs: %v", err)
}
testCases := []struct {
name string
exploreRefs []map[string]any
wantStatusCode int
wantInResult string
wantInError string
}{
{
name: "invoke conversational analytics with explore",
exploreRefs: refs,
wantStatusCode: http.StatusOK,
wantInResult: `Answer`,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
requestBodyMap := map[string]any{
"user_query_with_context": "What is in the explore?",
"explore_references": tc.exploreRefs,
}
bodyBytes, err := json.Marshal(requestBodyMap)
if err != nil {
t.Fatalf("failed to marshal request body: %v", err)
}
url := "http://127.0.0.1:5000/api/tool/conversational_analytics/invoke"
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, url, bytes.NewBuffer(bodyBytes), nil)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("unexpected status code: got %d, want %d. Body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if tc.wantInResult != "" {
var respBody map[string]interface{}
if err := json.Unmarshal(bodyBytes, &respBody); err != nil {
t.Fatalf("error parsing response body: %v", err)
}
got, ok := respBody["result"].(string)
if !ok {
t.Fatalf("unable to find result in response body")
}
if !strings.Contains(got, tc.wantInResult) {
t.Errorf("unexpected result: got %q, want to contain %q", got, tc.wantInResult)
}
}
if tc.wantInError != "" {
if !strings.Contains(string(bodyBytes), tc.wantInError) {
t.Errorf("unexpected error message: got %q, want to contain %q", string(bodyBytes), tc.wantInError)
}
}
})
}
}