mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-09 15:38:08 -05:00
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:
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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) |
|
||||
|
||||
@@ -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
1
go.mod
@@ -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
2
go.sum
@@ -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=
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user