From 1ad6517f4cb4d4f098d4b2ecfcf5656032cfb96b Mon Sep 17 00:00:00 2001 From: totoleon Date: Mon, 31 Mar 2025 22:23:18 +0000 Subject: [PATCH 01/14] feat: Added AlloyDB NLA tool --- docs/en/resources/tools/alloydb-nla.md | 50 ++++++ internal/server/config.go | 7 + internal/tools/alloydbnla/alloydbnla.go | 169 +++++++++++++++++++ internal/tools/alloydbnla/alloydbnla_test.go | 131 ++++++++++++++ 4 files changed, 357 insertions(+) create mode 100644 docs/en/resources/tools/alloydb-nla.md create mode 100644 internal/tools/alloydbnla/alloydbnla.go create mode 100644 internal/tools/alloydbnla/alloydbnla_test.go diff --git a/docs/en/resources/tools/alloydb-nla.md b/docs/en/resources/tools/alloydb-nla.md new file mode 100644 index 000000000000..beeb6de73520 --- /dev/null +++ b/docs/en/resources/tools/alloydb-nla.md @@ -0,0 +1,50 @@ +--- +title: "alloydb-nla" +type: docs +weight: 1 +description: > + A "alloydb-nla" tool leverages AlloyDB's AI functions to execute natural language questions against the database. +--- + +## About + +A `alloydb-nla` tool leverages AlloyDB's AI functions to execute natural language questions against the database. It allows users to query database information using natural language instead of SQL. It's compatible with the following sources: +- [alloydb-postgres](../sources/alloydb-pg.md) + +The tool uses AlloyDB's natural language processing capabilities to interpret questions and convert them into appropriate SQL queries, which are then executed against the database. TODO: link to AlloyDB's documentation. + +## Fields + +`nlConfig` is the name of the `nl_config` created in AlloyDB. + +`nlConfigParameters` are the list of the parameters and values for the AlloyDB [PSV (parameterized secure views)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#sanitize_queries_with_parameterized_secure_views). + +When using this tool, all the PSV parameters should be from filled with values from an auth service or a bounded param. These parameters should not be visible to the LLM agent. Instead, the LLM will only see one argument when using this tool - `question`, with the description being "The natural language question to ask." + +## Example + +```yaml +tools: + ask_questions: + kind: alloydb-nla + source: my-alloydb-source + description: "Ask questions to check information about flights" + nlConfig: "cymbal_air_nl_config" + nlConfigParameters: + - name: user_email + type: string + description: User ID of the logged in user. + authServices: + - name: my_google_service + field: email +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| +| kind | string | true | Must be "alloydb-nla". | +| source | string | true | Name of the AlloyDB source the natural language query should execute on. | +| description | string | true | Description of the tool that is passed to the LLM. | +| nlConfig | string | true | The name of the `nl_config` in AlloyDB | +| nlConfigParameters | [parameters](_index#specifying-parameters) | true | List of PSV parameters defined in the `nl_config` | diff --git a/internal/server/config.go b/internal/server/config.go index 58facfea885e..b78b11d749fc 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -40,6 +40,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/tools/mysqlsql" neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j" "github.com/googleapis/genai-toolbox/internal/tools/postgressql" + "github.com/googleapis/genai-toolbox/internal/tools/alloydbnla" "github.com/googleapis/genai-toolbox/internal/tools/spanner" "github.com/googleapis/genai-toolbox/internal/util" ) @@ -307,6 +308,12 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac return fmt.Errorf("unable to parse as %q: %w", kind, err) } (*c)[name] = actual + case alloydbnla.ToolKind: + actual := alloydbnla.Config{Name: name} + if err := dec.DecodeContext(ctx, &actual); err != nil { + return fmt.Errorf("unable to parse as %q: %w", kind, err) + } + (*c)[name] = actual case mysqlsql.ToolKind: actual := mysqlsql.Config{Name: name} if err := dec.DecodeContext(ctx, &actual); err != nil { diff --git a/internal/tools/alloydbnla/alloydbnla.go b/internal/tools/alloydbnla/alloydbnla.go new file mode 100644 index 000000000000..b7986f264027 --- /dev/null +++ b/internal/tools/alloydbnla/alloydbnla.go @@ -0,0 +1,169 @@ +// Copyright 2024 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 alloydbnla + +import ( + "context" + "fmt" + "strings" + + "github.com/googleapis/genai-toolbox/internal/sources" + "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/jackc/pgx/v5/pgxpool" +) + +const ToolKind string = "alloydb-nla" + +type compatibleSource interface { + PostgresPool() *pgxpool.Pool +} + +// validate compatible sources are still compatible +var _ compatibleSource = &alloydbpg.Source{} + +var compatibleSources = [...]string{alloydbpg.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"` + NLConfig string `yaml:"nlConfig" validate:"required"` + AuthRequired []string `yaml:"authRequired"` + NLConfigParameters tools.Parameters `yaml:"nlConfigParameters"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return ToolKind +} + +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", ToolKind, compatibleSources) + } + + paramNames := make([]string, 0, len(cfg.NLConfigParameters)) + for _, paramDef := range cfg.NLConfigParameters { + paramNames = append(paramNames, paramDef.GetName()) + } + quotedParamNames := make([]string, len(paramNames)) + for i, name := range paramNames { + // Basic escaping for single quotes within the name itself + escapedName := strings.ReplaceAll(name, "'", "''") + quotedParamNames[i] = fmt.Sprintf("'%s'", escapedName) + } + paramNamesSQL := "ARRAY []" // Default for no parameters + if len(quotedParamNames) > 0 { + paramNamesSQL = fmt.Sprintf("ARRAY [%s]", strings.Join(quotedParamNames, ", ")) + } + paramValuePlaceholders := make([]string, len(paramNames)) + for i := 0; i < len(paramNames); i++ { + // Placeholders start from $2 ($1 is reserved for the natural language query) + paramValuePlaceholders[i] = fmt.Sprintf("$%d", i+2) + } + paramValuesSQL := "ARRAY []" // Default for no parameters + if len(paramValuePlaceholders) > 0 { + paramValuesSQL = fmt.Sprintf("ARRAY [%s]", strings.Join(paramValuePlaceholders, ", ")) + } + + // execute_nl_query is the AlloyDB AI function that executes the natural language query + // The first parameter is the natural language query, which is passed as $1 + // The second parameter is the NLConfig, which is passed as a string + // The third and fourth parameters are the list of nl_config parameter names and values, respectively + stmtFormat := "SELECT alloydb_ai_nl.execute_nl_query($1, '%s', param_names => %s, param_values => %s);" + stmt := fmt.Sprintf(stmtFormat, cfg.NLConfig, paramNamesSQL, paramValuesSQL) + + newQuestionParam := tools.NewStringParameter( + "question", // name + "The natural language question to ask.", // description + ) + + cfg.NLConfigParameters = append([]tools.Parameter{newQuestionParam}, cfg.NLConfigParameters...) + + t := Tool{ + Name: cfg.Name, + Kind: ToolKind, + Parameters: cfg.NLConfigParameters, + Statement: stmt, + AuthRequired: cfg.AuthRequired, + Pool: s.PostgresPool(), + manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.NLConfigParameters.Manifest()}, + } + + return t, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + + Pool *pgxpool.Pool + Statement string + manifest tools.Manifest +} + +func (t Tool) Invoke(params tools.ParamValues) ([]any, error) { + sliceParams := params.AsSlice() + results, err := t.Pool.Query(context.Background(), t.Statement, sliceParams...) + if err != nil { + return nil, fmt.Errorf("unable to execute query: %w. Query: %v , Values: %v", err, t.Statement, sliceParams) + } + + fields := results.FieldDescriptions() + + var out []any + for results.Next() { + v, err := results.Values() + if err != nil { + return nil, fmt.Errorf("unable to parse row: %w", err) + } + vMap := make(map[string]any) + for i, f := range fields { + vMap[f.Name] = v[i] + } + out = append(out, vMap) + } + + return out, 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) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} diff --git a/internal/tools/alloydbnla/alloydbnla_test.go b/internal/tools/alloydbnla/alloydbnla_test.go new file mode 100644 index 000000000000..538947b1981d --- /dev/null +++ b/internal/tools/alloydbnla/alloydbnla_test.go @@ -0,0 +1,131 @@ +// 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 alloydbnla_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" + "github.com/googleapis/genai-toolbox/internal/tools/alloydbnla" +) + +func TestParseFromYamlAlloyDBNLA(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: alloydb-nla + source: my-alloydb-instance + description: AlloyDB natural language query tool + nlConfig: 'my_nl_config' + authRequired: + - my-google-auth-service + nlConfigParameters: + - name: user_id + type: string + description: user_id to use + authServices: + - name: my-google-auth-service + field: sub + `, + want: server.ToolConfigs{ + "example_tool": alloydbnla.Config{ + Name: "example_tool", + Kind: alloydbnla.ToolKind, + Source: "my-alloydb-instance", + Description: "AlloyDB natural language query tool", + NLConfig: "my_nl_config", + AuthRequired: []string{"my-google-auth-service"}, + NLConfigParameters: []tools.Parameter{ + tools.NewStringParameterWithAuth("user_id", "user_id to use", + []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "sub"}}), + }, + }, + }, + }, + { + desc: "with multiple parameters", + in: ` + tools: + complex_tool: + kind: alloydb-nla + source: my-alloydb-instance + description: AlloyDB natural language query tool with multiple parameters + nlConfig: 'complex_nl_config' + authRequired: + - my-google-auth-service + - other-auth-service + nlConfigParameters: + - name: user_id + type: string + description: user_id to use + authServices: + - name: my-google-auth-service + field: sub + - name: user_email + type: string + description: user_email to use + authServices: + - name: my-google-auth-service + field: user_email + `, + want: server.ToolConfigs{ + "complex_tool": alloydbnla.Config{ + Name: "complex_tool", + Kind: alloydbnla.ToolKind, + Source: "my-alloydb-instance", + Description: "AlloyDB natural language query tool with multiple parameters", + NLConfig: "complex_nl_config", + AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, + NLConfigParameters: []tools.Parameter{ + tools.NewStringParameterWithAuth("user_id", "user_id to use", + []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "sub"}}), + tools.NewStringParameterWithAuth("user_email", "user_email to use", + []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_email"}}), + }, + }, + }, + }, + } + 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) + } + }) + } +} From fe39446ba56a84c0f5475be31baa96432b78a52b Mon Sep 17 00:00:00 2001 From: totoleon Date: Wed, 2 Apr 2025 21:21:11 +0000 Subject: [PATCH 02/14] Changed AlloyDB NLA tool name to be alloydb-ai-nl --- docs/en/resources/tools/alloydb-nla.md | 10 +++++----- internal/tools/alloydbnla/alloydbnla.go | 2 +- internal/tools/alloydbnla/alloydbnla_test.go | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/en/resources/tools/alloydb-nla.md b/docs/en/resources/tools/alloydb-nla.md index beeb6de73520..7f6bccfef5dd 100644 --- a/docs/en/resources/tools/alloydb-nla.md +++ b/docs/en/resources/tools/alloydb-nla.md @@ -1,14 +1,14 @@ --- -title: "alloydb-nla" +title: "alloydb-ai-nl" type: docs weight: 1 description: > - A "alloydb-nla" tool leverages AlloyDB's AI functions to execute natural language questions against the database. + A "alloydb-ai-nl" tool leverages AlloyDB's AI functions to execute natural language questions against the database. --- ## About -A `alloydb-nla` tool leverages AlloyDB's AI functions to execute natural language questions against the database. It allows users to query database information using natural language instead of SQL. It's compatible with the following sources: +A `alloydb-ai-nl` tool leverages AlloyDB's AI functions to execute natural language questions against the database. It allows users to query database information using natural language instead of SQL. It's compatible with the following sources: - [alloydb-postgres](../sources/alloydb-pg.md) The tool uses AlloyDB's natural language processing capabilities to interpret questions and convert them into appropriate SQL queries, which are then executed against the database. TODO: link to AlloyDB's documentation. @@ -26,7 +26,7 @@ When using this tool, all the PSV parameters should be from filled with values f ```yaml tools: ask_questions: - kind: alloydb-nla + kind: alloydb-ai-nl source: my-alloydb-source description: "Ask questions to check information about flights" nlConfig: "cymbal_air_nl_config" @@ -43,7 +43,7 @@ tools: | **field** | **type** | **required** | **description** | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| -| kind | string | true | Must be "alloydb-nla". | +| kind | string | true | Must be "alloydb-ai-nl". | | source | string | true | Name of the AlloyDB source the natural language query should execute on. | | description | string | true | Description of the tool that is passed to the LLM. | | nlConfig | string | true | The name of the `nl_config` in AlloyDB | diff --git a/internal/tools/alloydbnla/alloydbnla.go b/internal/tools/alloydbnla/alloydbnla.go index b7986f264027..70d3b3682969 100644 --- a/internal/tools/alloydbnla/alloydbnla.go +++ b/internal/tools/alloydbnla/alloydbnla.go @@ -25,7 +25,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" ) -const ToolKind string = "alloydb-nla" +const ToolKind string = "alloydb-ai-nl" type compatibleSource interface { PostgresPool() *pgxpool.Pool diff --git a/internal/tools/alloydbnla/alloydbnla_test.go b/internal/tools/alloydbnla/alloydbnla_test.go index 538947b1981d..3382fdd58211 100644 --- a/internal/tools/alloydbnla/alloydbnla_test.go +++ b/internal/tools/alloydbnla/alloydbnla_test.go @@ -40,7 +40,7 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { in: ` tools: example_tool: - kind: alloydb-nla + kind: alloydb-ai-nl source: my-alloydb-instance description: AlloyDB natural language query tool nlConfig: 'my_nl_config' @@ -74,7 +74,7 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { in: ` tools: complex_tool: - kind: alloydb-nla + kind: alloydb-ai-nl source: my-alloydb-instance description: AlloyDB natural language query tool with multiple parameters nlConfig: 'complex_nl_config' From 2ff0bf99dc4b56220e90be882a3dac90ca7493ae Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 3 Apr 2025 00:34:51 +0000 Subject: [PATCH 03/14] Addressing comments --- .../{alloydb-nla.md => alloydb-ai-nl.md} | 27 ++++++++-- internal/server/config.go | 6 +-- .../alloydbainl.go} | 51 ++++++++++--------- .../alloydbainl_test.go} | 12 ++--- 4 files changed, 57 insertions(+), 39 deletions(-) rename docs/en/resources/tools/{alloydb-nla.md => alloydb-ai-nl.md} (58%) rename internal/tools/{alloydbnla/alloydbnla.go => alloydbainl/alloydbainl.go} (78%) rename internal/tools/{alloydbnla/alloydbnla_test.go => alloydbainl/alloydbainl_test.go} (93%) diff --git a/docs/en/resources/tools/alloydb-nla.md b/docs/en/resources/tools/alloydb-ai-nl.md similarity index 58% rename from docs/en/resources/tools/alloydb-nla.md rename to docs/en/resources/tools/alloydb-ai-nl.md index 7f6bccfef5dd..a621caa2ac24 100644 --- a/docs/en/resources/tools/alloydb-nla.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -3,23 +3,40 @@ title: "alloydb-ai-nl" type: docs weight: 1 description: > - A "alloydb-ai-nl" tool leverages AlloyDB's AI functions to execute natural language questions against the database. + The "alloydb-ai-nl" tool leverages AlloyDB's AI next-generation + [AI natural language]([alloydb-ai-nl-overview] support to provide the + ability to query the database directly using natural language. --- ## About -A `alloydb-ai-nl` tool leverages AlloyDB's AI functions to execute natural language questions against the database. It allows users to query database information using natural language instead of SQL. It's compatible with the following sources: +The "alloydb-ai-nl" tool leverages AlloyDB's next-generation AI natural +language feature to allow an Agent the ability to query the database directly +using natural language. Natural language streamlines the development of +generative AI applications by transferring the complexity of converting +natural language to SQL from the application layer to the database layer. + +This tool is compatible with the following sources: - [alloydb-postgres](../sources/alloydb-pg.md) -The tool uses AlloyDB's natural language processing capabilities to interpret questions and convert them into appropriate SQL queries, which are then executed against the database. TODO: link to AlloyDB's documentation. +AlloyDB AI natural language delivers secure and accurate responses for +application end user natural language questions. Natural language streamlines +the development of generative AI applications by transferring the complexity +of converting natural language to SQL from the application layer to the +database layer. ## Fields `nlConfig` is the name of the `nl_config` created in AlloyDB. -`nlConfigParameters` are the list of the parameters and values for the AlloyDB [PSV (parameterized secure views)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#sanitize_queries_with_parameterized_secure_views). +`nlConfigParameters` are the list of the parameters and values for the AlloyDB +[PSV (parameterized secure views)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#sanitize_queries_with_parameterized_secure_views). -When using this tool, all the PSV parameters should be from filled with values from an auth service or a bounded param. These parameters should not be visible to the LLM agent. Instead, the LLM will only see one argument when using this tool - `question`, with the description being "The natural language question to ask." +When using this tool, all the PSV parameters should be from filled with values +from an auth service or a bounded param. These parameters should not be +visible to the LLM agent. Instead, the LLM will only see one argument when +using this tool - `question`, with the description being "The natural +language question to ask." ## Example diff --git a/internal/server/config.go b/internal/server/config.go index b78b11d749fc..6c35f1abef49 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -40,7 +40,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/tools/mysqlsql" neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j" "github.com/googleapis/genai-toolbox/internal/tools/postgressql" - "github.com/googleapis/genai-toolbox/internal/tools/alloydbnla" + "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl" "github.com/googleapis/genai-toolbox/internal/tools/spanner" "github.com/googleapis/genai-toolbox/internal/util" ) @@ -308,8 +308,8 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac return fmt.Errorf("unable to parse as %q: %w", kind, err) } (*c)[name] = actual - case alloydbnla.ToolKind: - actual := alloydbnla.Config{Name: name} + case alloydbainl.ToolKind: + actual := alloydbainl.Config{Name: name} if err := dec.DecodeContext(ctx, &actual); err != nil { return fmt.Errorf("unable to parse as %q: %w", kind, err) } diff --git a/internal/tools/alloydbnla/alloydbnla.go b/internal/tools/alloydbainl/alloydbainl.go similarity index 78% rename from internal/tools/alloydbnla/alloydbnla.go rename to internal/tools/alloydbainl/alloydbainl.go index 70d3b3682969..367172997f92 100644 --- a/internal/tools/alloydbnla/alloydbnla.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// 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. @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package alloydbnla +package alloydbainl import ( "context" @@ -66,39 +66,40 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) } - paramNames := make([]string, 0, len(cfg.NLConfigParameters)) - for _, paramDef := range cfg.NLConfigParameters { - paramNames = append(paramNames, paramDef.GetName()) - } - quotedParamNames := make([]string, len(paramNames)) - for i, name := range paramNames { - // Basic escaping for single quotes within the name itself - escapedName := strings.ReplaceAll(name, "'", "''") - quotedParamNames[i] = fmt.Sprintf("'%s'", escapedName) - } - paramNamesSQL := "ARRAY []" // Default for no parameters - if len(quotedParamNames) > 0 { - paramNamesSQL = fmt.Sprintf("ARRAY [%s]", strings.Join(quotedParamNames, ", ")) - } - paramValuePlaceholders := make([]string, len(paramNames)) - for i := 0; i < len(paramNames); i++ { - // Placeholders start from $2 ($1 is reserved for the natural language query) - paramValuePlaceholders[i] = fmt.Sprintf("$%d", i+2) + numParams := len(cfg.NLConfigParameters) + quotedNameParts := make([]string, 0, numParams) + placeholderParts := make([]string, 0, numParams) + + for i, paramDef := range cfg.NLConfigParameters { + name := paramDef.GetName() + escapedName := strings.ReplaceAll(name, "'", "''") // Escape for SQL literal + quotedNameParts = append(quotedNameParts, fmt.Sprintf("'%s'", escapedName)) + placeholderParts = append(placeholderParts, fmt.Sprintf("$%d", i+2)) // $1 reserved } - paramValuesSQL := "ARRAY []" // Default for no parameters - if len(paramValuePlaceholders) > 0 { - paramValuesSQL = fmt.Sprintf("ARRAY [%s]", strings.Join(paramValuePlaceholders, ", ")) + + var paramNamesSQL string + var paramValuesSQL string + + if numParams > 0 { + paramNamesSQL = fmt.Sprintf("ARRAY[%s]", strings.Join(quotedNameParts, ", ")) + paramValuesSQL = fmt.Sprintf("ARRAY[%s]", strings.Join(placeholderParts, ", ")) + } else { + paramNamesSQL = "ARRAY[]::TEXT[]" + paramValuesSQL = "ARRAY[]::TEXT[]" } // execute_nl_query is the AlloyDB AI function that executes the natural language query // The first parameter is the natural language query, which is passed as $1 // The second parameter is the NLConfig, which is passed as a string - // The third and fourth parameters are the list of nl_config parameter names and values, respectively + // The following params are the list of nl_config parameter names and values, respectively + // Example SQL statement being executed: + // SELECT alloydb_ai_nl.execute_nl_query('How many tickets do I have?', 'cymbal_air_nl_config', param_names => ARRAY ['user_email'], param_values => ARRAY ['hailongli@google.com']); stmtFormat := "SELECT alloydb_ai_nl.execute_nl_query($1, '%s', param_names => %s, param_values => %s);" stmt := fmt.Sprintf(stmtFormat, cfg.NLConfig, paramNamesSQL, paramValuesSQL) + newQuestionParam := tools.NewStringParameter( - "question", // name + "question", // name "The natural language question to ask.", // description ) diff --git a/internal/tools/alloydbnla/alloydbnla_test.go b/internal/tools/alloydbainl/alloydbainl_test.go similarity index 93% rename from internal/tools/alloydbnla/alloydbnla_test.go rename to internal/tools/alloydbainl/alloydbainl_test.go index 3382fdd58211..50ebe6e443ee 100644 --- a/internal/tools/alloydbnla/alloydbnla_test.go +++ b/internal/tools/alloydbainl/alloydbainl_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package alloydbnla_test +package alloydbainl_test import ( "testing" @@ -22,7 +22,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/server" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/alloydbnla" + "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl" ) func TestParseFromYamlAlloyDBNLA(t *testing.T) { @@ -55,9 +55,9 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { field: sub `, want: server.ToolConfigs{ - "example_tool": alloydbnla.Config{ + "example_tool": alloydbainl.Config{ Name: "example_tool", - Kind: alloydbnla.ToolKind, + Kind: alloydbainl.ToolKind, Source: "my-alloydb-instance", Description: "AlloyDB natural language query tool", NLConfig: "my_nl_config", @@ -96,9 +96,9 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { field: user_email `, want: server.ToolConfigs{ - "complex_tool": alloydbnla.Config{ + "complex_tool": alloydbainl.Config{ Name: "complex_tool", - Kind: alloydbnla.ToolKind, + Kind: alloydbainl.ToolKind, Source: "my-alloydb-instance", Description: "AlloyDB natural language query tool with multiple parameters", NLConfig: "complex_nl_config", From 657e65bc4ca4c6fb05f11034f5a8b77dbff91f47 Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 3 Apr 2025 21:08:54 +0000 Subject: [PATCH 04/14] Addressing comments --- docs/en/resources/tools/alloydb-ai-nl.md | 43 +++++++++++++++++++----- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index a621caa2ac24..c30009a3c50e 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -25,18 +25,42 @@ the development of generative AI applications by transferring the complexity of converting natural language to SQL from the application layer to the database layer. -## Fields +## Requirements +AlloyDB AI natural language is currently in gated public preview. For more +information on availability and limitations, please see +[AlloyDB AI natural language overview](!https://cloud.google.com/alloydb/docs/natural-language-questions-overview). + +To enable AlloyDB AI natural language for your AlloyDB cluster, please follow +the steps listed in the [Generate SQL queries that answer natural language questions](!https://cloud.google.com/alloydb/docs/alloydb/docs/ai/generate-queries-natural-language), including +enabling the extension and configuring context for your application. + + +## Configuration + +### Configuration ID +`nlConfig` is the name of the `nl_config` created in AlloyDB. A `nl_config` +is a configuration associates an application to schema objects, examples and +other contexts that can be used. A large application can also use different +configurations for different parts of the app, as long as the right +configuration can be specified when a question is sent from that part of +the application. -`nlConfig` is the name of the `nl_config` created in AlloyDB. `nlConfigParameters` are the list of the parameters and values for the AlloyDB [PSV (parameterized secure views)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#sanitize_queries_with_parameterized_secure_views). -When using this tool, all the PSV parameters should be from filled with values -from an auth service or a bounded param. These parameters should not be -visible to the LLM agent. Instead, the LLM will only see one argument when -using this tool - `question`, with the description being "The natural -language question to ask." +When using this tool, we strongly recommend all the PSV parameters should be +from filled with values from an auth service or a bounded param. These +parameters should not be visible to the LLM agent. + +[Parameterized Secure Views (PSVs)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#parameterized_secure_views) +are a feature unique to AlloyDB that allow you allow you to require one or +more named parameter values passed to the view when querying it, somewhat +like bind variables with ordinary database queries. You **must** supply +all parameters required for all PSVs in the context. These parameters can be +used with features like [Authenticated Parameters](../tools/#array-parameters) +to provide secure access to queries generated using natural language. + ## Example @@ -51,6 +75,9 @@ tools: - name: user_email type: string description: User ID of the logged in user. + # note: we strongly recommend using features like Authenticated or + # Bound parameters to prevent the LLM from seeing these params and + # specifying values it shouldn't in the tool input authServices: - name: my_google_service field: email @@ -63,5 +90,5 @@ tools: | kind | string | true | Must be "alloydb-ai-nl". | | source | string | true | Name of the AlloyDB source the natural language query should execute on. | | description | string | true | Description of the tool that is passed to the LLM. | -| nlConfig | string | true | The name of the `nl_config` in AlloyDB | +| nlConfig | string | true | The name of the `nl_config` in AlloyDB | | nlConfigParameters | [parameters](_index#specifying-parameters) | true | List of PSV parameters defined in the `nl_config` | From 0940c76095de478d92d9c60536328783adb27e09 Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 3 Apr 2025 21:45:00 +0000 Subject: [PATCH 05/14] Tweaking doc --- docs/en/resources/tools/alloydb-ai-nl.md | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index c30009a3c50e..2f1e64fc8e56 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -5,7 +5,7 @@ weight: 1 description: > The "alloydb-ai-nl" tool leverages AlloyDB's AI next-generation [AI natural language]([alloydb-ai-nl-overview] support to provide the - ability to query the database directly using natural language. + ability to query the database directly using natural language. --- ## About @@ -26,17 +26,19 @@ of converting natural language to SQL from the application layer to the database layer. ## Requirements +{{< notice tip >}} AlloyDB AI natural language is currently in gated public preview. For more information on availability and limitations, please see -[AlloyDB AI natural language overview](!https://cloud.google.com/alloydb/docs/natural-language-questions-overview). +[AlloyDB AI natural language overview][https://cloud.google.com/alloydb/docs/natural-language-questions-overview]. +{{< /notice >}} To enable AlloyDB AI natural language for your AlloyDB cluster, please follow -the steps listed in the [Generate SQL queries that answer natural language questions](!https://cloud.google.com/alloydb/docs/alloydb/docs/ai/generate-queries-natural-language), including -enabling the extension and configuring context for your application. +the steps listed in the [Generate SQL queries that answer natural language questions][alloydb-ai-gen-nl], including enabling the extension and configuring context for your application. +[alloydb-ai-gen-nl]: https://cloud.google.com/alloydb/docs/alloydb/docs/ai/generate-queries-natural-language -## Configuration +## Configuration ### Configuration ID `nlConfig` is the name of the `nl_config` created in AlloyDB. A `nl_config` is a configuration associates an application to schema objects, examples and @@ -47,13 +49,13 @@ the application. `nlConfigParameters` are the list of the parameters and values for the AlloyDB -[PSV (parameterized secure views)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#sanitize_queries_with_parameterized_secure_views). +[PSV][alloydb-psv]. When using this tool, we strongly recommend all the PSV parameters should be from filled with values from an auth service or a bounded param. These parameters should not be visible to the LLM agent. -[Parameterized Secure Views (PSVs)](!https://cloud.google.com/alloydb/docs/ai/use-psvs#parameterized_secure_views) +[Parameterized Secure Views (PSVs)][alloydb-psv] are a feature unique to AlloyDB that allow you allow you to require one or more named parameter values passed to the view when querying it, somewhat like bind variables with ordinary database queries. You **must** supply @@ -61,6 +63,7 @@ all parameters required for all PSVs in the context. These parameters can be used with features like [Authenticated Parameters](../tools/#array-parameters) to provide secure access to queries generated using natural language. +[alloydb-psv]: https://cloud.google.com/alloydb/docs/ai/use-psvs#parameterized_secure_views ## Example From 3faecb91d53c37c0cd62d3a5c691687b1823913f Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 3 Apr 2025 22:11:17 +0000 Subject: [PATCH 06/14] Tweaking doc --- docs/en/resources/tools/alloydb-ai-nl.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index 2f1e64fc8e56..0f6e6b1a5d1c 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -10,7 +10,7 @@ description: > ## About -The "alloydb-ai-nl" tool leverages AlloyDB's next-generation AI natural +The `alloydb-ai-nl` tool leverages AlloyDB's next-generation AI natural language feature to allow an Agent the ability to query the database directly using natural language. Natural language streamlines the development of generative AI applications by transferring the complexity of converting @@ -49,13 +49,13 @@ the application. `nlConfigParameters` are the list of the parameters and values for the AlloyDB -[PSV][alloydb-psv]. +[Parameterized Secure Views (PSVs)][alloydb-psv]. When using this tool, we strongly recommend all the PSV parameters should be from filled with values from an auth service or a bounded param. These parameters should not be visible to the LLM agent. -[Parameterized Secure Views (PSVs)][alloydb-psv] +[PSVs][alloydb-psv] are a feature unique to AlloyDB that allow you allow you to require one or more named parameter values passed to the view when querying it, somewhat like bind variables with ordinary database queries. You **must** supply From e6be266e8beb4bd4dce744049a98bf32e2cec38e Mon Sep 17 00:00:00 2001 From: totoleon Date: Fri, 4 Apr 2025 05:27:10 +0000 Subject: [PATCH 07/14] Pass nl_config to the query template --- internal/tools/alloydbainl/alloydbainl.go | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index 367172997f92..e941fddc9489 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -74,7 +74,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) name := paramDef.GetName() escapedName := strings.ReplaceAll(name, "'", "''") // Escape for SQL literal quotedNameParts = append(quotedNameParts, fmt.Sprintf("'%s'", escapedName)) - placeholderParts = append(placeholderParts, fmt.Sprintf("$%d", i+2)) // $1 reserved + placeholderParts = append(placeholderParts, fmt.Sprintf("$%d", i + 3)) // $1, $2 reserved } var paramNamesSQL string @@ -90,12 +90,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // execute_nl_query is the AlloyDB AI function that executes the natural language query // The first parameter is the natural language query, which is passed as $1 - // The second parameter is the NLConfig, which is passed as a string - // The following params are the list of nl_config parameter names and values, respectively + // The second parameter is the NLConfig, which is passed as a $2 + // The following params are the list of PSV values passed to the NLConfig // Example SQL statement being executed: // SELECT alloydb_ai_nl.execute_nl_query('How many tickets do I have?', 'cymbal_air_nl_config', param_names => ARRAY ['user_email'], param_values => ARRAY ['hailongli@google.com']); - stmtFormat := "SELECT alloydb_ai_nl.execute_nl_query($1, '%s', param_names => %s, param_values => %s);" - stmt := fmt.Sprintf(stmtFormat, cfg.NLConfig, paramNamesSQL, paramValuesSQL) + stmtFormat := "SELECT alloydb_ai_nl.execute_nl_query($1, $2, param_names => %s, param_values => %s);" + stmt := fmt.Sprintf(stmtFormat, paramNamesSQL, paramValuesSQL) newQuestionParam := tools.NewStringParameter( @@ -110,6 +110,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) Kind: ToolKind, Parameters: cfg.NLConfigParameters, Statement: stmt, + NLConfig: cfg.NLConfig, AuthRequired: cfg.AuthRequired, Pool: s.PostgresPool(), manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.NLConfigParameters.Manifest()}, @@ -129,14 +130,22 @@ type Tool struct { Pool *pgxpool.Pool Statement string + NLConfig string manifest tools.Manifest } func (t Tool) Invoke(params tools.ParamValues) ([]any, error) { sliceParams := params.AsSlice() - results, err := t.Pool.Query(context.Background(), t.Statement, sliceParams...) + allParamValues := make([]any, len(sliceParams)+1) + allParamValues[0] = fmt.Sprintf("%s", sliceParams[0]) // nl_question + allParamValues[1] = fmt.Sprintf("%s", t.NLConfig) // nl_config + for i, param := range sliceParams[1:] { + allParamValues[i+2] = fmt.Sprintf("%s", param) + } + + results, err := t.Pool.Query(context.Background(), t.Statement, allParamValues...) if err != nil { - return nil, fmt.Errorf("unable to execute query: %w. Query: %v , Values: %v", err, t.Statement, sliceParams) + return nil, fmt.Errorf("unable to execute query: %w. Query: %v , Values: %v", err, t.Statement, allParamValues) } fields := results.FieldDescriptions() From 07318b6f9e442e16d7d66a8b9f1d4578b173fd53 Mon Sep 17 00:00:00 2001 From: totoleon Date: Fri, 4 Apr 2025 05:45:09 +0000 Subject: [PATCH 08/14] Addressing comments --- docs/en/resources/tools/alloydb-ai-nl.md | 55 ++++++++++++----------- internal/tools/alloydbainl/alloydbainl.go | 4 +- 2 files changed, 30 insertions(+), 29 deletions(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index 0f6e6b1a5d1c..d7af0aeaf9b6 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -4,14 +4,14 @@ type: docs weight: 1 description: > The "alloydb-ai-nl" tool leverages AlloyDB's AI next-generation - [AI natural language]([alloydb-ai-nl-overview] support to provide the + AI natural language support to provide the ability to query the database directly using natural language. --- ## About -The `alloydb-ai-nl` tool leverages AlloyDB's next-generation AI natural -language feature to allow an Agent the ability to query the database directly +The `alloydb-ai-nl` tool leverages [AlloyDB's next-generation AI natural language][alloydb-ai-nl-overview] +to allow an Agent the ability to query the database directly using natural language. Natural language streamlines the development of generative AI applications by transferring the complexity of converting natural language to SQL from the application layer to the database layer. @@ -29,39 +29,40 @@ database layer. {{< notice tip >}} AlloyDB AI natural language is currently in gated public preview. For more information on availability and limitations, please see -[AlloyDB AI natural language overview][https://cloud.google.com/alloydb/docs/natural-language-questions-overview]. +[AlloyDB AI naturaloverview language ][alloydb-ai-nl-overview] {{< /notice >}} To enable AlloyDB AI natural language for your AlloyDB cluster, please follow the steps listed in the [Generate SQL queries that answer natural language questions][alloydb-ai-gen-nl], including enabling the extension and configuring context for your application. +[alloydb-ai-nl-overview]: https://cloud.google.com/alloydb/docs/natural-language-questions-overview [alloydb-ai-gen-nl]: https://cloud.google.com/alloydb/docs/alloydb/docs/ai/generate-queries-natural-language ## Configuration -### Configuration ID -`nlConfig` is the name of the `nl_config` created in AlloyDB. A `nl_config` -is a configuration associates an application to schema objects, examples and -other contexts that can be used. A large application can also use different -configurations for different parts of the app, as long as the right -configuration can be specified when a question is sent from that part of -the application. - - -`nlConfigParameters` are the list of the parameters and values for the AlloyDB -[Parameterized Secure Views (PSVs)][alloydb-psv]. - -When using this tool, we strongly recommend all the PSV parameters should be -from filled with values from an auth service or a bounded param. These -parameters should not be visible to the LLM agent. - -[PSVs][alloydb-psv] -are a feature unique to AlloyDB that allow you allow you to require one or -more named parameter values passed to the view when querying it, somewhat -like bind variables with ordinary database queries. You **must** supply -all parameters required for all PSVs in the context. These parameters can be -used with features like [Authenticated Parameters](../tools/#array-parameters) -to provide secure access to queries generated using natural language. + +### Specifying an `nl_config` +A `nl_config` is a configuration associates an application to schema objects, +examples and other contexts that can be used. A large application can also +use different configurations for different parts of the app, as long as the +right configuration can be specified when a question is sent from that part +of the application. + +Once you've follow the steps for configuring context, you can use the `context` +field when configuring a `alloydb-ai-nl` tool. When this tool is invoked, the +SQL will be generated and executed using this context. + +### Specifying Parameters to PSV's + +[Parameterized Secure Views (PSVs)][alloydb-psv] are a feature unique to AlloyDB +that allow you allow you to require one or more named parameter values passed +to the view when querying it, somewhat like bind variables with ordinary database queries. + +You can use the `nlConfigParameters` to list the parameters required for your +`nl_config`. You **must** supply all parameters required for all PSVs in the context. +It's strongly recommended to use features like [Authenticated Parameters](../tools/#array-parameters) +or Bound Parameters to provide secure access to queries generated using natural language, as these +parameters are not visible to the LLM. [alloydb-psv]: https://cloud.google.com/alloydb/docs/ai/use-psvs#parameterized_secure_views diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index e941fddc9489..46a7bf4462b4 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -99,8 +99,8 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) newQuestionParam := tools.NewStringParameter( - "question", // name - "The natural language question to ask.", // description + "question", // name + "The natural language question to ask.", // description ) cfg.NLConfigParameters = append([]tools.Parameter{newQuestionParam}, cfg.NLConfigParameters...) From d32f7c3cc973ee3332b8e28ce6c5ccc415f688c8 Mon Sep 17 00:00:00 2001 From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> Date: Fri, 4 Apr 2025 09:58:04 -0600 Subject: [PATCH 09/14] Update docs/en/resources/tools/alloydb-ai-nl.md Co-authored-by: Averi Kitsch --- docs/en/resources/tools/alloydb-ai-nl.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index d7af0aeaf9b6..c8613b510609 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -3,7 +3,7 @@ title: "alloydb-ai-nl" type: docs weight: 1 description: > - The "alloydb-ai-nl" tool leverages AlloyDB's AI next-generation + The "alloydb-ai-nl" tool leverages [AlloyDB AI](https://cloud.google.com/alloydb/ai) next-generation AI natural language support to provide the ability to query the database directly using natural language. --- From 2b41553872f5e75aa28ab8c866053891785e0863 Mon Sep 17 00:00:00 2001 From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> Date: Fri, 4 Apr 2025 09:58:24 -0600 Subject: [PATCH 10/14] Update docs/en/resources/tools/alloydb-ai-nl.md Co-authored-by: Averi Kitsch --- docs/en/resources/tools/alloydb-ai-nl.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index c8613b510609..05be6ccf608e 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -10,7 +10,7 @@ description: > ## About -The `alloydb-ai-nl` tool leverages [AlloyDB's next-generation AI natural language][alloydb-ai-nl-overview] +The `alloydb-ai-nl` tool leverages [AlloyDB AI next-generation natural language][alloydb-ai-nl-overview] to allow an Agent the ability to query the database directly using natural language. Natural language streamlines the development of generative AI applications by transferring the complexity of converting From 31f6681272073863176a68a5aac11e7b71d70cc8 Mon Sep 17 00:00:00 2001 From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> Date: Fri, 4 Apr 2025 09:58:32 -0600 Subject: [PATCH 11/14] Update docs/en/resources/tools/alloydb-ai-nl.md Co-authored-by: Averi Kitsch --- docs/en/resources/tools/alloydb-ai-nl.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index 05be6ccf608e..94eb3e8337b9 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -4,7 +4,7 @@ type: docs weight: 1 description: > The "alloydb-ai-nl" tool leverages [AlloyDB AI](https://cloud.google.com/alloydb/ai) next-generation - AI natural language support to provide the + natural language support to provide the ability to query the database directly using natural language. --- From c315991dccb87926241532796dda0c56e952b752 Mon Sep 17 00:00:00 2001 From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> Date: Fri, 4 Apr 2025 16:17:21 +0000 Subject: [PATCH 12/14] chore: address feedback --- docs/en/resources/tools/alloydb-ai-nl.md | 82 +++++++++++++----------- 1 file changed, 44 insertions(+), 38 deletions(-) diff --git a/docs/en/resources/tools/alloydb-ai-nl.md b/docs/en/resources/tools/alloydb-ai-nl.md index 94eb3e8337b9..15ad0ae4f9da 100644 --- a/docs/en/resources/tools/alloydb-ai-nl.md +++ b/docs/en/resources/tools/alloydb-ai-nl.md @@ -3,37 +3,41 @@ title: "alloydb-ai-nl" type: docs weight: 1 description: > - The "alloydb-ai-nl" tool leverages [AlloyDB AI](https://cloud.google.com/alloydb/ai) next-generation - natural language support to provide the - ability to query the database directly using natural language. + The "alloydb-ai-nl" tool leverages + [AlloyDB AI](https://cloud.google.com/alloydb/ai) next-generation Natural + Language support to provide the ability to query the database directly using + natural language. --- ## About -The `alloydb-ai-nl` tool leverages [AlloyDB AI next-generation natural language][alloydb-ai-nl-overview] -to allow an Agent the ability to query the database directly -using natural language. Natural language streamlines the development of -generative AI applications by transferring the complexity of converting -natural language to SQL from the application layer to the database layer. +The `alloydb-ai-nl` tool leverages [AlloyDB AI next-generation natural +Language][alloydb-ai-nl-overview] support to allow an Agent the ability to query +the database directly using natural language. Natural language streamlines the +development of generative AI applications by transferring the complexity of +converting natural language to SQL from the application layer to the database +layer. This tool is compatible with the following sources: - [alloydb-postgres](../sources/alloydb-pg.md) -AlloyDB AI natural language delivers secure and accurate responses for +AlloyDB AI Natural Language delivers secure and accurate responses for application end user natural language questions. Natural language streamlines the development of generative AI applications by transferring the complexity of converting natural language to SQL from the application layer to the database layer. ## Requirements -{{< notice tip >}} -AlloyDB AI natural language is currently in gated public preview. For more -information on availability and limitations, please see -[AlloyDB AI naturaloverview language ][alloydb-ai-nl-overview] +{{< notice tip >}} AlloyDB AI natural language is currently in gated public +preview. For more information on availability and limitations, please see +[AlloyDB AI natural language +overview](https://cloud.google.com/alloydb/docs/natural-language-questions-overview) {{< /notice >}} -To enable AlloyDB AI natural language for your AlloyDB cluster, please follow -the steps listed in the [Generate SQL queries that answer natural language questions][alloydb-ai-gen-nl], including enabling the extension and configuring context for your application. +To enable AlloyDB AI natural language for your AlloyDB cluster, please follow +the steps listed in the [Generate SQL queries that answer natural language +questions][alloydb-ai-gen-nl], including enabling the extension and configuring +context for your application. [alloydb-ai-nl-overview]: https://cloud.google.com/alloydb/docs/natural-language-questions-overview [alloydb-ai-gen-nl]: https://cloud.google.com/alloydb/docs/alloydb/docs/ai/generate-queries-natural-language @@ -42,27 +46,29 @@ the steps listed in the [Generate SQL queries that answer natural language quest ## Configuration ### Specifying an `nl_config` -A `nl_config` is a configuration associates an application to schema objects, -examples and other contexts that can be used. A large application can also -use different configurations for different parts of the app, as long as the -right configuration can be specified when a question is sent from that part -of the application. +A `nl_config` is a configuration that associates an application to schema +objects, examples and other contexts that can be used. A large application can +also use different configurations for different parts of the app, as long as the +correct configuration can be specified when a question is sent from that part of +the application. -Once you've follow the steps for configuring context, you can use the `context` -field when configuring a `alloydb-ai-nl` tool. When this tool is invoked, the -SQL will be generated and executed using this context. +Once you've followed the steps for configuring context, you can use the +`context` field when configuring a `alloydb-ai-nl` tool. When this tool is +invoked, the SQL will be generated and executed using this context. ### Specifying Parameters to PSV's -[Parameterized Secure Views (PSVs)][alloydb-psv] are a feature unique to AlloyDB -that allow you allow you to require one or more named parameter values passed -to the view when querying it, somewhat like bind variables with ordinary database queries. +[Parameterized Secure Views (PSVs)][alloydb-psv] are a feature unique to AlloyDB +that allows you allow you to require one or more named parameter values passed +to the view when querying it, somewhat like bind variables with ordinary +database queries. -You can use the `nlConfigParameters` to list the parameters required for your -`nl_config`. You **must** supply all parameters required for all PSVs in the context. -It's strongly recommended to use features like [Authenticated Parameters](../tools/#array-parameters) -or Bound Parameters to provide secure access to queries generated using natural language, as these -parameters are not visible to the LLM. +You can use the `nlConfigParameters` to list the parameters required for your +`nl_config`. You **must** supply all parameters required for all PSVs in the +context. It's strongly recommended to use features like [Authenticated +Parameters](../tools/#array-parameters) or Bound Parameters to provide secure +access to queries generated using natural language, as these parameters are not +visible to the LLM. [alloydb-psv]: https://cloud.google.com/alloydb/docs/ai/use-psvs#parameterized_secure_views @@ -89,10 +95,10 @@ tools: ## Reference -| **field** | **type** | **required** | **description** | -|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| -| kind | string | true | Must be "alloydb-ai-nl". | -| source | string | true | Name of the AlloyDB source the natural language query should execute on. | -| description | string | true | Description of the tool that is passed to the LLM. | -| nlConfig | string | true | The name of the `nl_config` in AlloyDB | -| nlConfigParameters | [parameters](_index#specifying-parameters) | true | List of PSV parameters defined in the `nl_config` | +| **field** | **type** | **required** | **description** | +|--------------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------| +| kind | string | true | Must be "alloydb-ai-nl". | +| source | string | true | Name of the AlloyDB source the natural language query should execute on. | +| description | string | true | Description of the tool that is passed to the LLM. | +| nlConfig | string | true | The name of the `nl_config` in AlloyDB | +| nlConfigParameters | [parameters](_index#specifying-parameters) | true | List of PSV parameters defined in the `nl_config` | From 68f9ebe41b653c410410898c641dfda0f2bf91b7 Mon Sep 17 00:00:00 2001 From: Yuan <45984206+Yuan325@users.noreply.github.com> Date: Fri, 4 Apr 2025 10:51:44 -0700 Subject: [PATCH 13/14] ci: add integration test for alloydb ai nl (#387) --- .ci/integration.cloudbuild.yaml | 27 ++ .golangci.yaml | 1 + internal/tools/alloydbainl/alloydbainl.go | 10 +- tests/alloydb_ai_nl_integration_test.go | 326 ++++++++++++++++++++++ 4 files changed, 359 insertions(+), 5 deletions(-) create mode 100644 tests/alloydb_ai_nl_integration_test.go diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index d976a77ba8e2..92558347c18f 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.ci/integration.cloudbuild.yaml @@ -66,6 +66,27 @@ steps: - | go test -race -v -tags=integration,alloydb ./tests + - id: "alloydb-ai-nl" + name: golang:1 + waitFor: ["install-dependencies"] + entrypoint: /bin/bash + env: + - "GOPATH=/gopath" + - "ALLOYDB_AI_NL_PROJECT=$PROJECT_ID" + - "ALLOYDB_AI_NL_CLUSTER=$_ALLOYDB_AI_NL_CLUSTER" + - "ALLOYDB_AI_NL_INSTANCE=$_ALLOYDB_AI_NL_INSTANCE" + - "ALLOYDB_AI_NL_DATABASE=$_DATABASE_NAME" + - "ALLOYDB_AI_NL_REGION=$_REGION" + - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL" + secretEnv: ["ALLOYDB_AI_NL_USER", "ALLOYDB_AI_NL_PASS", "CLIENT_ID"] + volumes: + - name: "go" + path: "/gopath" + args: + - -c + - | + go test -race -v -tags=integration,alloydb_ai_nl ./tests + - id: "postgres" name: golang:1 waitFor: ["install-dependencies"] @@ -239,6 +260,10 @@ availableSecrets: env: ALLOYDB_POSTGRES_USER - versionName: projects/$PROJECT_ID/secrets/alloydb_pg_pass/versions/latest env: ALLOYDB_POSTGRES_PASS + - versionName: projects/$PROJECT_ID/secrets/alloydb_ai_nl_user/versions/latest + env: ALLOYDB_AI_NL_USER + - versionName: projects/$PROJECT_ID/secrets/alloydb_ai_nl_pass/versions/latest + env: ALLOYDB_AI_NL_PASS - versionName: projects/$PROJECT_ID/secrets/postgres_user/versions/latest env: POSTGRES_USER - versionName: projects/$PROJECT_ID/secrets/postgres_pass/versions/latest @@ -281,6 +306,8 @@ substitutions: _CLOUD_SQL_POSTGRES_INSTANCE: "cloud-sql-pg-testing" _ALLOYDB_POSTGRES_CLUSTER: "alloydb-pg-testing" _ALLOYDB_POSTGRES_INSTANCE: "alloydb-pg-testing-instance" + _ALLOYDB_AI_NL_CLUSTER: "alloydb-ai-nl-testing" + _ALLOYDB_AI_NL_INSTANCE: "alloydb-ai-nl-testing-instance" _POSTGRES_HOST: 127.0.0.1 _POSTGRES_PORT: "5432" _SPANNER_INSTANCE: "spanner-testing" diff --git a/.golangci.yaml b/.golangci.yaml index a384e6435a4f..6dbb7a296cd9 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -44,3 +44,4 @@ run: - mssql - mysql - http + - alloydb_ai_nl diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index 46a7bf4462b4..e94e38ec8387 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -37,11 +37,11 @@ var _ compatibleSource = &alloydbpg.Source{} var compatibleSources = [...]string{alloydbpg.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"` - NLConfig string `yaml:"nlConfig" validate:"required"` + Name string `yaml:"name" validate:"required"` + Kind string `yaml:"kind" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description" validate:"required"` + NLConfig string `yaml:"nlConfig" validate:"required"` AuthRequired []string `yaml:"authRequired"` NLConfigParameters tools.Parameters `yaml:"nlConfigParameters"` } diff --git a/tests/alloydb_ai_nl_integration_test.go b/tests/alloydb_ai_nl_integration_test.go new file mode 100644 index 000000000000..53c081fc4ea6 --- /dev/null +++ b/tests/alloydb_ai_nl_integration_test.go @@ -0,0 +1,326 @@ +//go:build integration && alloydb_ai_nl + +// 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 tests + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "os" + "reflect" + "regexp" + "testing" + "time" +) + +var ( + ALLOYDB_AI_NL_SOURCE_KIND = "alloydb-postgres" + ALLOYDB_AI_NL_TOOL_KIND = "alloydb-ai-nl" + ALLOYDB_AI_NL_PROJECT = os.Getenv("ALLOYDB_AI_NL_PROJECT") + ALLOYDB_AI_NL_REGION = os.Getenv("ALLOYDB_AI_NL_REGION") + ALLOYDB_AI_NL_CLUSTER = os.Getenv("ALLOYDB_AI_NL_CLUSTER") + ALLOYDB_AI_NL_INSTANCE = os.Getenv("ALLOYDB_AI_NL_INSTANCE") + ALLOYDB_AI_NL_DATABASE = os.Getenv("ALLOYDB_AI_NL_DATABASE") + ALLOYDB_AI_NL_USER = os.Getenv("ALLOYDB_AI_NL_USER") + ALLOYDB_AI_NL_PASS = os.Getenv("ALLOYDB_AI_NL_PASS") +) + +func getAlloyDBAiNlVars(t *testing.T) map[string]any { + switch "" { + case ALLOYDB_AI_NL_PROJECT: + t.Fatal("'ALLOYDB_AI_NL_PROJECT' not set") + case ALLOYDB_AI_NL_REGION: + t.Fatal("'ALLOYDB_AI_NL_REGION' not set") + case ALLOYDB_AI_NL_CLUSTER: + t.Fatal("'ALLOYDB_AI_NL_CLUSTER' not set") + case ALLOYDB_AI_NL_INSTANCE: + t.Fatal("'ALLOYDB_AI_NL_INSTANCE' not set") + case ALLOYDB_AI_NL_DATABASE: + t.Fatal("'ALLOYDB_AI_NL_DATABASE' not set") + case ALLOYDB_AI_NL_USER: + t.Fatal("'ALLOYDB_AI_NL_USER' not set") + case ALLOYDB_AI_NL_PASS: + t.Fatal("'ALLOYDB_AI_NL_PASS' not set") + } + return map[string]any{ + "kind": ALLOYDB_AI_NL_SOURCE_KIND, + "project": ALLOYDB_AI_NL_PROJECT, + "cluster": ALLOYDB_AI_NL_CLUSTER, + "instance": ALLOYDB_AI_NL_INSTANCE, + "region": ALLOYDB_AI_NL_REGION, + "database": ALLOYDB_AI_NL_DATABASE, + "user": ALLOYDB_AI_NL_USER, + "password": ALLOYDB_AI_NL_PASS, + } +} + +func TestAlloyDBAiNlToolEndpoints(t *testing.T) { + sourceConfig := getAlloyDBAiNlVars(t) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + var args []string + + // Write config into a file and pass it to command + toolsFile := getAiNlToolsConfig(sourceConfig) + + cmd, cleanup, err := StartCmd(ctx, toolsFile, args...) + if err != nil { + t.Fatalf("command initialization returned an error: %s", err) + } + defer cleanup() + + waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`)) + if err != nil { + t.Logf("toolbox command logs: \n%s", out) + t.Fatalf("toolbox didn't start successfully: %s", err) + } + + runAiNlToolGetTest(t) + + runAiNlToolInvokeTest(t) +} + +func runAiNlToolGetTest(t *testing.T) { + // Test tool get endpoint + tcs := []struct { + name string + api string + want map[string]any + }{ + { + name: "get my-simple-tool", + api: "http://127.0.0.1:5000/api/tool/my-simple-tool/", + want: map[string]any{ + "my-simple-tool": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "parameters": []any{ + map[string]any{ + "name": "question", + "type": "string", + "description": "The natural language question to ask.", + "authSources": []any{}, + }, + }, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.name, func(t *testing.T) { + resp, err := http.Get(tc.api) + if err != nil { + t.Fatalf("error when sending a request: %s", err) + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + t.Fatalf("response status code is not 200") + } + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + if err != nil { + t.Fatalf("error parsing response body") + } + + got, ok := body["tools"] + if !ok { + t.Fatalf("unable to find tools in response body") + } + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("got %q, want %q", got, tc.want) + } + }) + } +} + +func runAiNlToolInvokeTest(t *testing.T) { + // Get ID token + idToken, err := GetGoogleIdToken(ClientId) + if err != nil { + t.Fatalf("error getting Google ID token: %s", err) + } + + // Test tool invoke endpoint + invokeTcs := []struct { + name string + api string + requestHeader map[string]string + requestBody io.Reader + want string + isErr bool + }{ + { + name: "invoke my-simple-tool", + api: "http://127.0.0.1:5000/api/tool/my-simple-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + want: "[{\"execute_nl_query\":{\"?column?\":1}}]", + isErr: false, + }, + { + name: "Invoke my-tool without parameters", + api: "http://127.0.0.1:5000/api/tool/my-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(`{}`)), + isErr: true, + }, + { + name: "Invoke my-auth-tool with auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(`{"question": "can you show me the name of this user?"}`)), + want: "[{\"execute_nl_query\":{\"name\":\"Alice\"}}]", + isErr: false, + }, + { + name: "Invoke my-auth-tool with invalid auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + isErr: true, + }, + { + name: "Invoke my-auth-tool without auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + isErr: true, + }, + { + name: "Invoke my-auth-required-tool with auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + isErr: false, + want: "[{\"execute_nl_query\":{\"?column?\":1}}]", + }, + { + name: "Invoke my-auth-required-tool with invalid auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + isErr: true, + }, + { + name: "Invoke my-auth-required-tool without auth token", + api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)), + isErr: true, + }, + } + for _, tc := range invokeTcs { + t.Run(tc.name, func(t *testing.T) { + // Send Tool invocation request + req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + for k, v := range tc.requestHeader { + req.Header.Add(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("unable to send request: %s", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + if tc.isErr == true { + return + } + bodyBytes, _ := io.ReadAll(resp.Body) + t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) + } + + // Check response body + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + if err != nil { + t.Fatalf("error parsing response body") + } + got, ok := body["result"].(string) + if !ok { + t.Fatalf("unable to find result in response body") + } + + if got != tc.want { + t.Fatalf("unexpected value: got %q, want %q", got, tc.want) + } + }) + } + +} + +func getAiNlToolsConfig(sourceConfig map[string]any) map[string]any { + // Write config into a file and pass it to command + toolsFile := map[string]any{ + "sources": map[string]any{ + "my-instance": sourceConfig, + }, + "authServices": map[string]any{ + "my-google-auth": map[string]any{ + "kind": "google", + "clientId": ClientId, + }, + }, + "tools": map[string]any{ + "my-simple-tool": map[string]any{ + "kind": ALLOYDB_AI_NL_TOOL_KIND, + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + "nlConfig": "my_nl_config", + }, + "my-auth-tool": map[string]any{ + "kind": ALLOYDB_AI_NL_TOOL_KIND, + "source": "my-instance", + "description": "Tool to test authenticated parameters.", + "nlConfig": "my_nl_config", + "nlConfigParameters": []map[string]any{ + { + "name": "email", + "type": "string", + "description": "user email", + "authServices": []map[string]string{ + { + "name": "my-google-auth", + "field": "email", + }, + }, + }, + }, + }, + "my-auth-required-tool": map[string]any{ + "kind": ALLOYDB_AI_NL_TOOL_KIND, + "source": "my-instance", + "description": "Tool to test auth required invocation.", + "nlConfig": "my_nl_config", + "authRequired": []string{ + "my-google-auth", + }, + }, + }, + } + + return toolsFile +} From 6df808a59a0bb6b4da82985205e17cbbc6c397c6 Mon Sep 17 00:00:00 2001 From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> Date: Fri, 4 Apr 2025 17:59:46 +0000 Subject: [PATCH 14/14] feat: add mcp manifest for ai-nl tool --- internal/tools/alloydbainl/alloydbainl.go | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index e94e38ec8387..6877ba74f6cc 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -74,7 +74,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) name := paramDef.GetName() escapedName := strings.ReplaceAll(name, "'", "''") // Escape for SQL literal quotedNameParts = append(quotedNameParts, fmt.Sprintf("'%s'", escapedName)) - placeholderParts = append(placeholderParts, fmt.Sprintf("$%d", i + 3)) // $1, $2 reserved + placeholderParts = append(placeholderParts, fmt.Sprintf("$%d", i+3)) // $1, $2 reserved } var paramNamesSQL string @@ -97,7 +97,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) stmtFormat := "SELECT alloydb_ai_nl.execute_nl_query($1, $2, param_names => %s, param_values => %s);" stmt := fmt.Sprintf(stmtFormat, paramNamesSQL, paramValuesSQL) - newQuestionParam := tools.NewStringParameter( "question", // name "The natural language question to ask.", // description @@ -105,6 +104,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) cfg.NLConfigParameters = append([]tools.Parameter{newQuestionParam}, cfg.NLConfigParameters...) + mcpManifest := tools.McpManifest{ + Name: cfg.Name, + Description: cfg.Description, + InputSchema: cfg.NLConfigParameters.McpManifest(), + } + t := Tool{ Name: cfg.Name, Kind: ToolKind, @@ -114,6 +119,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) AuthRequired: cfg.AuthRequired, Pool: s.PostgresPool(), manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.NLConfigParameters.Manifest()}, + mcpManifest: mcpManifest, } return t, nil @@ -128,10 +134,11 @@ type Tool struct { AuthRequired []string `yaml:"authRequired"` Parameters tools.Parameters `yaml:"parameters"` - Pool *pgxpool.Pool - Statement string - NLConfig string - manifest tools.Manifest + Pool *pgxpool.Pool + Statement string + NLConfig string + manifest tools.Manifest + mcpManifest tools.McpManifest } func (t Tool) Invoke(params tools.ParamValues) ([]any, error) { @@ -174,6 +181,10 @@ 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) }