Skip to content

Commit 64cb899

Browse files
Yuan325Yuan325
committed
feat(tools/bigtable): add templateParameters field for bigtable
1 parent b4b986a commit 64cb899

13 files changed

Lines changed: 183 additions & 35 deletions

‎internal/tools/bigtable/bigtable.go‎

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,14 @@ var _ compatibleSource = &bigtabledb.Source{}
5151
var compatibleSources = [...]string{bigtabledb.SourceKind}
5252

5353
type Config struct {
54-
Name string `yaml:"name" validate:"required"`
55-
Kind string `yaml:"kind" validate:"required"`
56-
Source string `yaml:"source" validate:"required"`
57-
Description string `yaml:"description" validate:"required"`
58-
Statement string `yaml:"statement" validate:"required"`
59-
AuthRequired []string `yaml:"authRequired"`
60-
Parameters tools.Parameters `yaml:"parameters"`
54+
Name string `yaml:"name" validate:"required"`
55+
Kind string `yaml:"kind" validate:"required"`
56+
Source string `yaml:"source" validate:"required"`
57+
Description string `yaml:"description" validate:"required"`
58+
Statement string `yaml:"statement" validate:"required"`
59+
AuthRequired []string `yaml:"authRequired"`
60+
Parameters tools.Parameters `yaml:"parameters"`
61+
TemplateParameters tools.Parameters `yaml:"templateParameters"`
6162
}
6263

6364
// validate interface
@@ -80,22 +81,26 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
8081
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
8182
}
8283

84+
allParameters, paramManifest, paramMcpManifest := tools.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)
85+
8386
mcpManifest := tools.McpManifest{
8487
Name: cfg.Name,
8588
Description: cfg.Description,
86-
InputSchema: cfg.Parameters.McpManifest(),
89+
InputSchema: paramMcpManifest,
8790
}
8891

8992
// finish tool setup
9093
t := Tool{
91-
Name: cfg.Name,
92-
Kind: kind,
93-
Parameters: cfg.Parameters,
94-
Statement: cfg.Statement,
95-
AuthRequired: cfg.AuthRequired,
96-
Client: s.BigtableClient(),
97-
manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest(), AuthRequired: cfg.AuthRequired},
98-
mcpManifest: mcpManifest,
94+
Name: cfg.Name,
95+
Kind: kind,
96+
Parameters: cfg.Parameters,
97+
TemplateParameters: cfg.TemplateParameters,
98+
AllParams: allParameters,
99+
Statement: cfg.Statement,
100+
AuthRequired: cfg.AuthRequired,
101+
Client: s.BigtableClient(),
102+
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
103+
mcpManifest: mcpManifest,
99104
}
100105
return t, nil
101106
}
@@ -104,10 +109,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
104109
var _ tools.Tool = Tool{}
105110

106111
type Tool struct {
107-
Name string `yaml:"name"`
108-
Kind string `yaml:"kind"`
109-
AuthRequired []string `yaml:"authRequired"`
110-
Parameters tools.Parameters `yaml:"parameters"`
112+
Name string `yaml:"name"`
113+
Kind string `yaml:"kind"`
114+
AuthRequired []string `yaml:"authRequired"`
115+
Parameters tools.Parameters `yaml:"parameters"`
116+
TemplateParameters tools.Parameters `yaml:"templateParameters"`
117+
AllParams tools.Parameters `yaml:"allParams"`
111118

112119
Client *bigtable.Client
113120
Statement string
@@ -141,21 +148,32 @@ func getMapParamsType(tparams tools.Parameters, params tools.ParamValues) (map[s
141148
}
142149

143150
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
144-
mapParamsType, err := getMapParamsType(t.Parameters, params)
151+
paramsMap := params.AsMap()
152+
newStatement, err := tools.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
153+
if err != nil {
154+
return nil, fmt.Errorf("unable to extract template params %w", err)
155+
}
156+
157+
newParams, err := tools.GetParams(t.Parameters, paramsMap)
158+
if err != nil {
159+
return nil, fmt.Errorf("unable to extract standard params %w", err)
160+
}
161+
162+
mapParamsType, err := getMapParamsType(t.Parameters, newParams)
145163
if err != nil {
146164
return nil, fmt.Errorf("fail to get map params: %w", err)
147165
}
148166

149167
ps, err := t.Client.PrepareStatement(
150168
ctx,
151-
t.Statement,
169+
newStatement,
152170
mapParamsType,
153171
)
154172
if err != nil {
155173
return nil, fmt.Errorf("unable to prepare statement: %w", err)
156174
}
157175

158-
bs, err := ps.Bind(params.AsMap())
176+
bs, err := ps.Bind(newParams.AsMap())
159177
if err != nil {
160178
return nil, fmt.Errorf("unable to bind: %w", err)
161179
}
@@ -183,7 +201,7 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, erro
183201
}
184202

185203
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
186-
return tools.ParseParams(t.Parameters, data, claims)
204+
return tools.ParseParams(t.AllParams, data, claims)
187205
}
188206

189207
func (t Tool) Manifest() tools.Manifest {

‎internal/tools/bigtable/bigtable_test.go‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,3 +82,76 @@ func TestParseFromYamlBigtable(t *testing.T) {
8282
}
8383

8484
}
85+
86+
func TestParseFromYamlWithTemplateBigtable(t *testing.T) {
87+
ctx, err := testutils.ContextWithNewLogger()
88+
if err != nil {
89+
t.Fatalf("unexpected error: %s", err)
90+
}
91+
tcs := []struct {
92+
desc string
93+
in string
94+
want server.ToolConfigs
95+
}{
96+
{
97+
desc: "basic example",
98+
in: `
99+
tools:
100+
example_tool:
101+
kind: bigtable-sql
102+
source: my-pg-instance
103+
description: some description
104+
statement: |
105+
SELECT * FROM SQL_STATEMENT;
106+
parameters:
107+
- name: country
108+
type: string
109+
description: some description
110+
templateParameters:
111+
- name: tableName
112+
type: string
113+
description: The table to select hotels from.
114+
- name: fieldArray
115+
type: array
116+
description: The columns to return for the query.
117+
items:
118+
name: column
119+
type: string
120+
description: A column name that will be returned from the query.
121+
`,
122+
want: server.ToolConfigs{
123+
"example_tool": bigtable.Config{
124+
Name: "example_tool",
125+
Kind: "bigtable-sql",
126+
Source: "my-pg-instance",
127+
Description: "some description",
128+
Statement: "SELECT * FROM SQL_STATEMENT;\n",
129+
AuthRequired: []string{},
130+
Parameters: []tools.Parameter{
131+
tools.NewStringParameter("country", "some description"),
132+
},
133+
TemplateParameters: []tools.Parameter{
134+
tools.NewStringParameter("tableName", "The table to select hotels from."),
135+
tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
136+
},
137+
},
138+
},
139+
},
140+
}
141+
for _, tc := range tcs {
142+
t.Run(tc.desc, func(t *testing.T) {
143+
got := struct {
144+
Tools server.ToolConfigs `yaml:"tools"`
145+
}{}
146+
// Parse contents
147+
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
148+
if err != nil {
149+
t.Fatalf("unable to unmarshal: %s", err)
150+
}
151+
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
152+
t.Fatalf("incorrect parse: diff %v", diff)
153+
}
154+
})
155+
}
156+
157+
}

‎tests/alloydbpg/alloydb_pg_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
170170
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
171171
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
172172
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
173-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
173+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
174174
}
175175

176176
// Test connection with different IP type

‎tests/bigtable/bigtable_integration_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929

3030
"cloud.google.com/go/bigtable"
3131
"github.com/google/uuid"
32+
"github.com/googleapis/genai-toolbox/internal/tools"
3233
"github.com/googleapis/genai-toolbox/tests"
3334
)
3435

@@ -69,6 +70,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
6970

7071
tableName := "param_table" + strings.ReplaceAll(uuid.New().String(), "-", "")
7172
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
73+
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
7274

7375
columnFamilyName := "cf"
7476
muts, rowKeys := getTestData(columnFamilyName)
@@ -87,6 +89,8 @@ func TestBigtableToolEndpoints(t *testing.T) {
8789

8890
// Write config into a file and pass it to command
8991
toolsFile := tests.GetToolsConfig(sourceConfig, BIGTABLE_TOOL_KIND, param_test_statement, auth_tool_statement)
92+
toolsFile = addTemplateParamConfig(t, toolsFile)
93+
9094
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
9195
if err != nil {
9296
t.Fatalf("command initialization returned an error: %s", err)
@@ -109,6 +113,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
109113
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
110114
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
111115
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
116+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, true, true)
112117
}
113118

114119
func getTestData(columnFamilyName string) ([]*bigtable.Mutation, []string) {
@@ -213,3 +218,52 @@ func setupBtTable(t *testing.T, ctx context.Context, projectId string, instance
213218
defer adminClient.Close()
214219
}
215220
}
221+
222+
func addTemplateParamConfig(t *testing.T, config map[string]any) map[string]any {
223+
toolsMap, ok := config["tools"].(map[string]any)
224+
if !ok {
225+
t.Fatalf("unable to get tools from config")
226+
}
227+
toolsMap["select-templateParams-tool"] = map[string]any{
228+
"kind": "bigtable-sql",
229+
"source": "my-instance",
230+
"description": "Create table tool with template parameters",
231+
"statement": "SELECT * FROM {{.tableName}}",
232+
"templateParameters": []tools.Parameter{
233+
tools.NewStringParameter("tableName", "some description"),
234+
},
235+
}
236+
toolsMap["select-templateParams-combined-tool"] = map[string]any{
237+
"kind": "bigtable-sql",
238+
"source": "my-instance",
239+
"description": "Create table tool with template parameters",
240+
"statement": "SELECT * FROM {{.tableName}} WHERE id = @id",
241+
"parameters": []tools.Parameter{tools.NewStringParameter("id", "the id of the user")},
242+
"templateParameters": []tools.Parameter{
243+
tools.NewStringParameter("tableName", "some description"),
244+
},
245+
}
246+
toolsMap["select-fields-templateParams-tool"] = map[string]any{
247+
"kind": "bigtable-sql",
248+
"source": "my-instance",
249+
"description": "Create table tool with template parameters",
250+
"statement": "SELECT {{array .fields}} FROM {{.tableName}}",
251+
"templateParameters": []tools.Parameter{
252+
tools.NewStringParameter("tableName", "some description"),
253+
tools.NewArrayParameter("fields", "The fields to select from", tools.NewStringParameter("field", "A field that will be returned from the query.")),
254+
},
255+
}
256+
toolsMap["select-filter-templateParams-combined-tool"] = map[string]any{
257+
"kind": "bigtable-sql",
258+
"source": "my-instance",
259+
"description": "Create table tool with template parameters",
260+
"statement": "SELECT * FROM {{.tableName}} WHERE {{.columnFilter}} = @name",
261+
"parameters": []tools.Parameter{tools.NewStringParameter("name", "the name of the user")},
262+
"templateParameters": []tools.Parameter{
263+
tools.NewStringParameter("tableName", "some description"),
264+
tools.NewStringParameter("columnFilter", "some description"),
265+
},
266+
}
267+
config["tools"] = toolsMap
268+
return config
269+
}

‎tests/cloudsqlmssql/cloud_sql_mssql_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
157157
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
158158
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
159159
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
160-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
160+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
161161
}
162162

163163
// Test connection with different IP type

‎tests/cloudsqlmysql/cloud_sql_mysql_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ func TestCloudSQLMysqlToolEndpoints(t *testing.T) {
151151
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
152152
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
153153
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
154-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
154+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
155155
}
156156

157157
// Test connection with different IP type

‎tests/cloudsqlpg/cloud_sql_pg_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
155155
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
156156
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
157157
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
158-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
158+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
159159
}
160160

161161
// Test connection with different IP type

‎tests/mssql/mssql_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,5 +129,5 @@ func TestMssqlToolEndpoints(t *testing.T) {
129129
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
130130
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
131131
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
132-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
132+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
133133
}

‎tests/mysql/mysql_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,5 +128,5 @@ func TestMysqlToolEndpoints(t *testing.T) {
128128
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
129129
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
130130
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
131-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
131+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
132132
}

‎tests/postgres/postgres_integration_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,5 +128,5 @@ func TestPostgres(t *testing.T) {
128128
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
129129
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
130130
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
131-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false)
131+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, false, false)
132132
}

0 commit comments

Comments
 (0)