chore: import upstream snapshot with attribution
CF: Deploy Dev Docs / deploy (push) Has been cancelled
Sync Labels / build (push) Has been cancelled
tests / unit tests (macos-latest) (push) Has been cancelled
tests / unit tests (windows-latest) (push) Has been cancelled
tests / unit tests (ubuntu-latest) (push) Has been cancelled
CF: Deploy Dev Docs / deploy (push) Has been cancelled
Sync Labels / build (push) Has been cancelled
tests / unit tests (macos-latest) (push) Has been cancelled
tests / unit tests (windows-latest) (push) Has been cancelled
tests / unit tests (ubuntu-latest) (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,199 @@
|
||||
// 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 cloudsqlpgcreateinstances
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/sources"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
sqladmin "google.golang.org/api/sqladmin/v1"
|
||||
)
|
||||
|
||||
const resourceType string = "cloud-sql-postgres-create-instance"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
GetDefaultProject() string
|
||||
UseClientAuthorization() bool
|
||||
CreateInstance(context.Context, string, string, string, string, sqladmin.Settings, string) (any, error)
|
||||
}
|
||||
|
||||
// Config defines the configuration for the create-instances tool.
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigType returns the type of the tool.
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
// Initialize initializes the tool from the configuration.
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "Creates a Postgres instance using `Production` and `Development` presets. For the `Development` template, it chooses a 2 vCPU, 16 GiB RAM, 100 GiB SSD configuration with Non-HA/zonal availability. For the `Production` template, it chooses an 8 vCPU, 64 GiB RAM, 250 GiB SSD configuration with HA/regional availability. The Enterprise Plus edition is used in both cases. The default database version is `POSTGRES_17`. The agent should ask the user if they want to use a different version."
|
||||
}
|
||||
params := buildParams("")
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: params.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
params,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Tool represents the create-instances tool.
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
// Invoke executes the tool's logic.
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
project, ok := paramsMap["project"].(string)
|
||||
if !ok {
|
||||
return nil, util.NewAgentError("missing 'project' parameter", nil)
|
||||
}
|
||||
name, ok := paramsMap["name"].(string)
|
||||
if !ok {
|
||||
return nil, util.NewAgentError("missing 'name' parameter", nil)
|
||||
}
|
||||
dbVersion, ok := paramsMap["databaseVersion"].(string)
|
||||
if !ok {
|
||||
return nil, util.NewAgentError("missing 'databaseVersion' parameter", nil)
|
||||
}
|
||||
rootPassword, ok := paramsMap["rootPassword"].(string)
|
||||
if !ok {
|
||||
return nil, util.NewAgentError("missing 'rootPassword' parameter", nil)
|
||||
}
|
||||
editionPreset, ok := paramsMap["editionPreset"].(string)
|
||||
if !ok {
|
||||
return nil, util.NewAgentError("missing 'editionPreset' parameter", nil)
|
||||
}
|
||||
|
||||
settings := sqladmin.Settings{}
|
||||
switch strings.ToLower(editionPreset) {
|
||||
case "production":
|
||||
settings.AvailabilityType = "REGIONAL"
|
||||
settings.Edition = "ENTERPRISE_PLUS"
|
||||
settings.Tier = "db-perf-optimized-N-8"
|
||||
settings.DataDiskSizeGb = 250
|
||||
settings.DataDiskType = "PD_SSD"
|
||||
case "development":
|
||||
settings.AvailabilityType = "ZONAL"
|
||||
settings.Edition = "ENTERPRISE_PLUS"
|
||||
settings.Tier = "db-perf-optimized-N-2"
|
||||
settings.DataDiskSizeGb = 100
|
||||
settings.DataDiskType = "PD_SSD"
|
||||
default:
|
||||
return nil, util.NewAgentError(fmt.Sprintf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset), nil)
|
||||
}
|
||||
resp, err := source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken))
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// Authorized checks if the tool is authorized.
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return source.UseClientAuthorization(), nil
|
||||
}
|
||||
|
||||
// buildParams builds the tool's parameters. A non-empty project means the source has a
|
||||
// configured default project, which is baked into the project param; otherwise the plain form is used.
|
||||
func buildParams(project string) parameters.Parameters {
|
||||
projectParam := parameters.NewStringParameter("project", "The project ID")
|
||||
if project != "" {
|
||||
projectParam = parameters.NewStringParameter("project", "The GCP project ID. This is pre-configured; do not ask for it unless the user explicitly provides a different one.", parameters.WithStringDefault(project))
|
||||
}
|
||||
return parameters.Parameters{
|
||||
projectParam,
|
||||
parameters.NewStringParameter("name", "The name of the instance"),
|
||||
parameters.NewStringParameter("databaseVersion", "The database version for Postgres. If not specified, defaults to the latest available version (e.g., POSTGRES_17).", parameters.WithStringDefault("POSTGRES_17")),
|
||||
parameters.NewStringParameter("rootPassword", "The root password for the instance"),
|
||||
parameters.NewStringParameter("editionPreset", "The edition of the instance. Can be `Production` or `Development`. This determines the default machine type and availability. Defaults to `Development`.", parameters.WithStringDefault("Development")),
|
||||
}
|
||||
}
|
||||
|
||||
// resolveParams builds the tool's parameters using the source's configured default GCP project.
|
||||
func (t Tool) resolveParams(srcs map[string]sources.Source) (parameters.Parameters, error) {
|
||||
s, err := tools.GetCompatibleSourceFromMap[compatibleSource](srcs, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return buildParams(s.GetDefaultProject()), nil
|
||||
}
|
||||
|
||||
// GetParameters returns the tool's parameters, resolved against the source.
|
||||
func (t Tool) GetParameters(srcs map[string]sources.Source) (parameters.Parameters, error) {
|
||||
return t.resolveParams(srcs)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest, resolved against the source.
|
||||
func (t Tool) Manifest(srcs map[string]sources.Source) (tools.Manifest, error) {
|
||||
allParameters, err := t.resolveParams(srcs)
|
||||
if err != nil {
|
||||
return tools.Manifest{}, err
|
||||
}
|
||||
return tools.Manifest{Description: t.Cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: t.Cfg.AuthRequired}, nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// 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 cloudsqlpgcreateinstances_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: create-instance-tool
|
||||
type: cloud-sql-postgres-create-instance
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"create-instance-tool": cloudsqlpgcreateinstances.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "create-instance-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "cloud-sql-postgres-create-instance",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
// 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 cloudsqlpgupgradeprecheck
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
sqladmin "google.golang.org/api/sqladmin/v1"
|
||||
)
|
||||
|
||||
const resourceType string = "postgres-upgrade-precheck"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
GetService(context.Context, string) (*sqladmin.Service, error)
|
||||
UseClientAuthorization() bool
|
||||
}
|
||||
|
||||
// Config defines the configuration for the precheck-upgrade tool.
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigType returns the type of the tool.
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
// Initialize initializes the tool from the configuration.
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("project", "The project ID"),
|
||||
parameters.NewStringParameter("instance", "The name of the instance to check"),
|
||||
parameters.NewStringParameter("targetDatabaseVersion", "The target PostgreSQL version for the upgrade (e.g., POSTGRES_18). If not specified, defaults to the PostgreSQL 18.", parameters.WithStringDefault("POSTGRES_18")),
|
||||
}
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "Analyzes a Cloud SQL PostgreSQL instance for major version upgrade readiness. Results are provided to guide customer actions:\n" +
|
||||
"ERROR: Action Required. These are critical issues blocking the upgrade. Customers must resolve these using the provided actions_required steps before attempting the upgrade.\n" +
|
||||
"WARNING: Review Recommended. These are potential issues. Customers should review the message and actions_required. While not blocking, addressing these is advised to prevent future problems or unexpected behavior post-upgrade.\n" +
|
||||
"INFO: No Action Needed. Informational messages only. This pre-check helps customers proactively fix problems, preventing upgrade failures and ensuring a smoother transition."
|
||||
}
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewReadOnlyAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Tool represents the precheck-upgrade tool.
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
// PreCheckResultItem holds the details of a single check result.
|
||||
type PreCheckResultItem struct {
|
||||
Message string `json:"message"`
|
||||
MessageType string `json:"messageType"` // INFO, WARNING, ERROR
|
||||
ActionsRequired []string `json:"actionsRequired"`
|
||||
}
|
||||
|
||||
// PreCheckAPIResponse holds the array of pre-check results.
|
||||
type PreCheckAPIResponse struct {
|
||||
Items []PreCheckResultItem `json:"preCheckResponse"`
|
||||
}
|
||||
|
||||
// Helper function to convert from []*sqladmin.PreCheckResponse to []PreCheckResultItem
|
||||
func convertResults(items []*sqladmin.PreCheckResponse) []PreCheckResultItem {
|
||||
if len(items) == 0 { // Handle nil or empty slice
|
||||
return []PreCheckResultItem{}
|
||||
}
|
||||
results := make([]PreCheckResultItem, len(items))
|
||||
for i, item := range items {
|
||||
results[i] = PreCheckResultItem{
|
||||
Message: item.Message,
|
||||
MessageType: item.MessageType,
|
||||
ActionsRequired: item.ActionsRequired,
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
// Invoke executes the tool's logic.
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
project, ok := paramsMap["project"].(string)
|
||||
if !ok || project == "" {
|
||||
return nil, util.NewAgentError("missing or empty 'project' parameter", nil)
|
||||
}
|
||||
instanceName, ok := paramsMap["instance"].(string)
|
||||
if !ok || instanceName == "" {
|
||||
return nil, util.NewAgentError("missing or empty 'instance' parameter", nil)
|
||||
}
|
||||
targetVersion, ok := paramsMap["targetDatabaseVersion"].(string)
|
||||
if !ok || targetVersion == "" {
|
||||
// This should not happen due to the default value
|
||||
return nil, util.NewAgentError("missing or empty 'targetDatabaseVersion' parameter", nil)
|
||||
}
|
||||
|
||||
service, err := source.GetService(ctx, string(accessToken))
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
|
||||
reqBody := &sqladmin.InstancesPreCheckMajorVersionUpgradeRequest{
|
||||
PreCheckMajorVersionUpgradeContext: &sqladmin.PreCheckMajorVersionUpgradeContext{
|
||||
TargetDatabaseVersion: targetVersion,
|
||||
},
|
||||
}
|
||||
|
||||
call := service.Instances.PreCheckMajorVersionUpgrade(project, instanceName, reqBody).Context(ctx)
|
||||
op, err := call.Do()
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
|
||||
const pollTimeout = 20 * time.Second
|
||||
cutoffTime := time.Now().Add(pollTimeout)
|
||||
|
||||
for time.Now().Before(cutoffTime) {
|
||||
currentOp, err := service.Operations.Get(project, op.Name).Context(ctx).Do()
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
|
||||
if currentOp.Status == "DONE" {
|
||||
if currentOp.Error != nil && len(currentOp.Error.Errors) > 0 {
|
||||
errMsg := fmt.Sprintf("pre-check operation LRO failed: %s", currentOp.Error.Errors[0].Message)
|
||||
if currentOp.Error.Errors[0].Code != "" {
|
||||
errMsg = fmt.Sprintf("%s (Code: %s)", errMsg, currentOp.Error.Errors[0].Code)
|
||||
}
|
||||
return nil, util.NewClientServerError(errMsg, http.StatusInternalServerError, fmt.Errorf("pre-check operation failed with error: %s", errMsg))
|
||||
}
|
||||
|
||||
var preCheckItems []*sqladmin.PreCheckResponse
|
||||
if currentOp.PreCheckMajorVersionUpgradeContext != nil {
|
||||
preCheckItems = currentOp.PreCheckMajorVersionUpgradeContext.PreCheckResponse
|
||||
}
|
||||
// convertResults handles nil or empty preCheckItems
|
||||
return PreCheckAPIResponse{Items: convertResults(preCheckItems)}, nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, util.NewClientServerError("timed out waiting for operation", http.StatusRequestTimeout, ctx.Err())
|
||||
case <-time.After(5 * time.Second):
|
||||
}
|
||||
}
|
||||
return op, nil
|
||||
}
|
||||
|
||||
// Authorized checks if the tool is authorized.
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return source.UseClientAuthorization(), nil
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
// 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 cloudsqlpgupgradeprecheck_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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 precheck example",
|
||||
in: `
|
||||
kind: tool
|
||||
name: precheck-upgrade-tool
|
||||
type: postgres-upgrade-precheck
|
||||
description: a precheck test description
|
||||
source: some-admin-source
|
||||
authRequired:
|
||||
- https://www.googleapis.com/auth/cloud-platform
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"precheck-upgrade-tool": cloudsqlpgupgradeprecheck.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "precheck-upgrade-tool",
|
||||
Description: "a precheck test description",
|
||||
AuthRequired: []string{"https://www.googleapis.com/auth/cloud-platform"},
|
||||
},
|
||||
Type: "postgres-upgrade-precheck",
|
||||
Source: "some-admin-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "precheck example with no auth",
|
||||
in: `
|
||||
kind: tool
|
||||
name: precheck-upgrade-tool-no-auth
|
||||
type: postgres-upgrade-precheck
|
||||
description: a precheck test description no auth
|
||||
source: other-admin-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"precheck-upgrade-tool-no-auth": cloudsqlpgupgradeprecheck.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "precheck-upgrade-tool-no-auth",
|
||||
Description: "a precheck test description no auth",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "postgres-upgrade-precheck",
|
||||
Source: "other-admin-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
// Copyright 2026 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 vectorassistapplyspec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-apply-spec"
|
||||
|
||||
const applySpecQuery = `
|
||||
SELECT * FROM vector_assist.apply_spec(spec_id => @spec_id::TEXT, table_name => @table_name::TEXT,
|
||||
column_name => @column_name::TEXT, schema_name => @schema_name::TEXT);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
// parameters are marked required/ optional based on the vector assist function defintions
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("spec_id", "The unique ID of the vector specification to apply.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("table_name", "The name of the table to apply the vector specification to (in case of a single spec defined on the table).", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("column_name", "The text_column_name or vector_column_name of the spec to identify the exact spec in case there are multiple specs defined on a table.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("schema_name", "The schema name for the table.", parameters.WithStringRequired(false)),
|
||||
}
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool automatically executes all the SQL recommendations associated with a specific vector specification (spec_id) or table. It runs the necessary commands in the correct sequence to provision the workload, marking each step as applied once successful. Use this tool when the user has reviewed the generated recommendations from a defined (or modified) spec and is ready to apply the changes directly to their database instance to finalize the vector search setup. This tool can be used as a follow-up action after invoking the 'define_spec' or 'modify_spec' tool."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
nil,
|
||||
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
// Convert our parsed parameters directly into pgx.NamedArgs
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
// As long as source.RunSQL unwraps args into pgx.Query(ctx, sql, args...), pgx handles the mapping of @param to the named parameter.
|
||||
resp, err := source.RunSQL(ctx, applySpecQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistapplyspec_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistapplyspec"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: apply-spec-tool
|
||||
type: vector-assist-apply-spec
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"apply-spec-tool": vectorassistapplyspec.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "apply-spec-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-apply-spec",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
// Copyright 2026 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 vectorassistdefinespec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-define-spec"
|
||||
|
||||
const defineSpecQuery = `
|
||||
SELECT recommendation_id, vector_spec_id, table_name, schema_name, query, recommendation, applied, modified, created_at
|
||||
FROM vector_assist.define_spec(table_name => @table_name::TEXT, schema_name => @schema_name::TEXT, spec_id => @spec_id::TEXT,
|
||||
vector_column_name => @vector_column_name::TEXT, text_column_name => @text_column_name::TEXT,
|
||||
vector_index_type => @vector_index_type::TEXT, embeddings_available => @embeddings_available::BOOLEAN,
|
||||
num_vectors => @num_vectors::INTEGER, dimensionality => @dimensionality::INTEGER,
|
||||
embedding_model => @embedding_model::TEXT, prefilter_column_names => @prefilter_column_names,
|
||||
distance_func => @distance_func::TEXT, quantization => @quantization::TEXT,
|
||||
memory_budget_kb => @memory_budget_kb::INTEGER, target_recall => @target_recall::FLOAT,
|
||||
target_top_k => @target_top_k::INTEGER, tune_vector_index => @tune_vector_index::BOOLEAN);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
// parameters are marked required/ optional based on the vector assist function defintions
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("table_name", "Table name on which vector workload needs to be set up.", parameters.WithStringRequired(true)),
|
||||
parameters.NewStringParameter("schema_name", "Schema containing the given table.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("spec_id", "Unique ID for the vector spec. Auto-generated, if not specified.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("vector_column_name", "Column name for the column with vector embeddings.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("text_column_name", "Column name for the column with text on which vector search needs to be set up.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("vector_index_type", "Type of the vector index to be created (Allowed inputs: 'hnsw', 'ivfflat', 'scann').", parameters.WithStringRequired(false)),
|
||||
parameters.NewBooleanParameter("embeddings_available", "Boolean parameter to know if vector embeddings are already available in the table.", parameters.WithBooleanRequired(false)),
|
||||
parameters.NewIntParameter("num_vectors", "Number of vectors expected in the dataset.", parameters.WithIntRequired(false)),
|
||||
parameters.NewIntParameter("dimensionality", "If vectors are already generated, set to dimension of vectors. If not, set to dimensionality of the embedding_model.", parameters.WithIntRequired(false)),
|
||||
parameters.NewStringParameter("embedding_model", "Optional parameter: Model to be used for generating embeddings. If not provided, it has an internally selected default value.", parameters.WithStringRequired(false)),
|
||||
parameters.NewArrayParameter("prefilter_column_names", "Columns based on which prefiltering will happen in vector search queries.", parameters.NewStringParameter("prefilter_column_name", "Pre filter column name"), parameters.WithArrayRequired(false)),
|
||||
parameters.NewStringParameter("distance_func", "Distance function to be used for comparing vectors (Allowed inputs: 'cosine', 'ip', 'l2', 'l1').", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("quantization", "Quantization to be used for creating the vector indexes (Allowed inputs: 'none', 'halfvec', 'bit').", parameters.WithStringRequired(false)),
|
||||
parameters.NewIntParameter("memory_budget_kb", "Maximum size in KB that the index can consume in memory while building.", parameters.WithIntRequired(false)),
|
||||
parameters.NewFloatParameter("target_recall", "The recall that the user would like to target with the given index for standard vector queries.", parameters.WithFloatRequired(false)),
|
||||
parameters.NewIntParameter("target_top_k", "The top-K values that need to be retrieved for the given query.", parameters.WithIntRequired(false)),
|
||||
parameters.NewBooleanParameter("tune_vector_index", "Boolean parameter to specify if the auto tuning is required for the index.", parameters.WithBooleanRequired(false)),
|
||||
}
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool defines a new vector specification by capturing the user's intent and requirements for a vector search workload. This generates a complete, ordered set of SQL recommendations required to set up the database, embeddings, and vector indexes. While highly customizable, any optional parameters left unspecified will use internally determined defaults optimized for the specific workload. Use this tool at the very beginning of the vector setup process when a user first wants to configure a table for vector search, generate embeddings, or create a new vector index."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
nil,
|
||||
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
// Convert our parsed parameters directly into pgx.NamedArgs
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
// As long as source.RunSQL unwraps args into pgx.Query(ctx, sql, args...), pgx handles the mapping of @param to the named parameter.
|
||||
resp, err := source.RunSQL(ctx, defineSpecQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistdefinespec_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistdefinespec"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: define-spec-tool
|
||||
type: vector-assist-define-spec
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"define-spec-tool": vectorassistdefinespec.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "define-spec-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-define-spec",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
// Copyright 2026 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 vectorassistdeletespec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-delete-spec"
|
||||
|
||||
const deleteSpecQuery = `
|
||||
SELECT vector_assist.delete_spec(spec_id => @spec_id::TEXT);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("spec_id", "Unique ID for the vector spec to delete.", parameters.WithStringRequired(true)),
|
||||
}
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool deletes an existing vector specification using its spec_id. Use this tool when a user explicitly requests to delete, remove, or clean up an existing vector specification which was created in the context of the vector assist tools."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
resp, err := source.RunSQL(ctx, deleteSpecQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistdeletespec_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistdeletespec"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: delete-spec-tool
|
||||
type: vector-assist-delete-spec
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"delete-spec-tool": vectorassistdeletespec.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "delete-spec-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-delete-spec",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
// Copyright 2026 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 vectorassistgeneratequery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-generate-query"
|
||||
|
||||
const generateQueryStatement = `
|
||||
SELECT vector_assist.generate_query(
|
||||
spec_id => @spec_id::TEXT, table_name => @table_name::TEXT,
|
||||
schema_name => @schema_name::TEXT, column_name => @column_name::TEXT,
|
||||
search_text => @search_text::TEXT, search_vector => @search_vector::vector,
|
||||
output_column_names => @output_column_names,
|
||||
top_k => @top_k::INTEGER,
|
||||
filter_expressions => @filter_expressions,
|
||||
target_recall => @target_recall::FLOAT,
|
||||
iterative_index_search => @iterative_index_search::BOOLEAN
|
||||
);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
// parameters are marked required/ optional based on the vector assist function defintions
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("spec_id", "Generate the vector query corresponding to this vector spec.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("table_name", "Generate the vector query corresponding to this table (in case of a single spec defined on the table).", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("schema_name", "Schema name for the table related to the vector query generation.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("column_name", "text_column_name or vector_column_name of the spec to identify the exact spec in case there are multiple specs defined on a table.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("search_text", "Text search for which query needs to be generated. Embeddings are generated using the model defined in the vector spec.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("search_vector", "Vector for which query needs to be generated. Only one of search_text or search_vector must be populated.", parameters.WithStringRequired(false)),
|
||||
parameters.NewArrayParameter("output_column_names", "Column names to retrieve in the output search query. Defaults to retrieving all columns.", parameters.NewStringParameter("output_column_name", "Output column name"), parameters.WithArrayRequired(false)),
|
||||
parameters.NewIntParameter("top_k", "Number of nearest neighbors to be returned in the vector search query. Defaults to 10.", parameters.WithIntRequired(false)),
|
||||
parameters.NewArrayParameter("filter_expressions", "Any filter expressions to be applied on the vector search query.", parameters.NewStringParameter("filter_expression", "Filter expression"), parameters.WithArrayRequired(false)),
|
||||
parameters.NewFloatParameter("target_recall", "The recall that the user would like to target with the given query. Overrides the spec-level target_recall.", parameters.WithFloatRequired(false)),
|
||||
parameters.NewBooleanParameter("iterative_index_search", "Perform iterative index search for filtered queries to ensure enough results are returned.", parameters.WithBooleanRequired(false)),
|
||||
}
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool generates optimized SQL queries for vector search by leveraging the metadata and vector specifications defined in a specific spec_id. It may return a single query or a sequence of multiple SQL queries that can be executed sequentially. Use this tool when a user wants to perform semantic or similarity searches on their data. It serves as the primary actionable tool to invoke for generating the executable SQL required to retrieve relevant results based on vector similarity. The 'execute_sql' tool can be used as a follow-up action after invoking this tool."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
nil,
|
||||
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
// Convert our parsed parameters directly into pgx.NamedArgs
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
// As long as source.RunSQL unwraps args into pgx.Query(ctx, sql, args...), pgx handles the mapping of @param to the named parameter.
|
||||
resp, err := source.RunSQL(ctx, generateQueryStatement, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistgeneratequery_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistgeneratequery"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: generate-query-tool
|
||||
type: vector-assist-generate-query
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"generate-query-tool": vectorassistgeneratequery.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "generate-query-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-generate-query",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright 2026 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 vectorassistgetspec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-get-spec"
|
||||
|
||||
const getSpecQuery = `
|
||||
SELECT *
|
||||
FROM vector_assist.get_spec(spec_id => @spec_id::TEXT);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("spec_id", "Unique ID for the vector spec.", parameters.WithStringRequired(true)),
|
||||
}
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool retrieves the details of an existing vector specification using its unique 'spec_id'. Use this tool to retrieve a vector specification which was created in the context of the vector assist tools."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
resp, err := source.RunSQL(ctx, getSpecQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistgetspec_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistgetspec"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: get-spec-tool
|
||||
type: vector-assist-get-spec
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"get-spec-tool": vectorassistgetspec.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "get-spec-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-get-spec",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
// Copyright 2026 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 vectorassistimprovequeryrecall
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-improve-query-recall"
|
||||
|
||||
// Query to check if the index exists and if it is an HNSW index.
|
||||
const checkIndexQuery = `
|
||||
SELECT
|
||||
(COUNT(1) > 0) AS index_present,
|
||||
COALESCE(BOOL_OR(indexdef ILIKE '%USING hnsw%'), false) AS is_hnsw
|
||||
FROM pg_indexes
|
||||
WHERE schemaname = @schema_name::TEXT
|
||||
AND tablename = @table_name::TEXT
|
||||
AND indexname = @index_name::TEXT
|
||||
AND indexdef ILIKE '%' || @vector_column_name::TEXT || '%';
|
||||
`
|
||||
|
||||
// Query to find the optimal index parameters
|
||||
const improveRecallQuery = `
|
||||
SELECT output_ef_search AS ef_search
|
||||
FROM vector_assist.find_ef_search_for_target_recall(
|
||||
table_name => @table_name::TEXT,
|
||||
schema_name => @schema_name::TEXT,
|
||||
column_name => @vector_column_name::TEXT,
|
||||
top_k => @top_k::INT,
|
||||
target_recall => @target_recall::FLOAT,
|
||||
distance_func => @distance_func::TEXT
|
||||
);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("schema_name", "Optional parameter: Schema name of the table.", parameters.WithStringDefault("public")),
|
||||
parameters.NewStringParameter("table_name", "Table name experiencing degraded vector search recall.", parameters.WithStringRequired(true)),
|
||||
parameters.NewStringParameter("vector_column_name", "Column name containing the vector embeddings.", parameters.WithStringRequired(true)),
|
||||
parameters.NewStringParameter("index_name", "Name of the vector index to tune.", parameters.WithStringRequired(true)),
|
||||
parameters.NewIntParameter("top_k", "Optional parameter: Top k value for the vector search.", parameters.WithIntDefault(10)),
|
||||
parameters.NewFloatParameter("target_recall", "Optional parameter: Target recall value for search results.", parameters.WithFloatDefault(0.95)),
|
||||
parameters.NewStringParameter("distance_func", "Optional parameter: Distance function used for the vector search similarity.", parameters.WithStringDefault("cosine")),
|
||||
}
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "Use this tool to troubleshoot and optimize existing vector search workloads when a user reports irrelevant results, poor accuracy, or degraded recall. It determines the optimal tuning parameter (such as ef_search) for an active vector index to improve the search results. The tool outputs an actionable SQL query recommendation to be executed as a next action using the 'execute_sql' tool."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
// Check if the index exists and if it is an HNSW index.
|
||||
checkResp, err := source.RunSQL(ctx, checkIndexQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
|
||||
checkBytes, marshalErr := json.Marshal(checkResp)
|
||||
if marshalErr != nil {
|
||||
return nil, util.NewClientServerError("failed to process index check response", http.StatusInternalServerError, marshalErr)
|
||||
}
|
||||
|
||||
var checkRows []map[string]interface{}
|
||||
if unmarshalErr := json.Unmarshal(checkBytes, &checkRows); unmarshalErr != nil || len(checkRows) == 0 {
|
||||
return nil, util.NewClientServerError("unexpected empty response from database", http.StatusInternalServerError, unmarshalErr)
|
||||
}
|
||||
|
||||
row := checkRows[0]
|
||||
indexPresent, ok := row["index_present"].(bool)
|
||||
if !ok {
|
||||
// If the key is missing or isn't a boolean, it's likely a server-side/query issue.
|
||||
return nil, util.NewClientServerError("Internal error: 'index_present' is missing or has an invalid type.", http.StatusInternalServerError, nil)
|
||||
}
|
||||
if !indexPresent {
|
||||
return nil, util.NewClientServerError("Index not found for the given table and vector column. If the table lacks an existing vector setup, use the 'define_spec' tool to configure the database.", http.StatusBadRequest, nil)
|
||||
}
|
||||
|
||||
isHnsw, ok := row["is_hnsw"].(bool)
|
||||
if !ok {
|
||||
return nil, util.NewClientServerError("Internal error: 'is_hnsw' is missing or has an invalid type.", http.StatusInternalServerError, nil)
|
||||
}
|
||||
if !isHnsw {
|
||||
return nil, util.NewClientServerError("Unsupported index type for recall optimization. Only HNSW index is supported.", http.StatusBadRequest, nil)
|
||||
}
|
||||
|
||||
// Calculate the optimal index parameters to achieve the target recall.
|
||||
tuningResp, err := source.RunSQL(ctx, improveRecallQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
|
||||
tuningBytes, marshalErr := json.Marshal(tuningResp)
|
||||
if marshalErr != nil {
|
||||
return nil, util.NewClientServerError("failed to process tuning response", http.StatusInternalServerError, marshalErr)
|
||||
}
|
||||
|
||||
var tuningRows []map[string]interface{}
|
||||
if unmarshalErr := json.Unmarshal(tuningBytes, &tuningRows); unmarshalErr != nil || len(tuningRows) == 0 {
|
||||
return nil, util.NewClientServerError("unexpected empty tuning response from database", http.StatusInternalServerError, unmarshalErr)
|
||||
}
|
||||
|
||||
// Extract ef_search (JSON decoder defaults numbers to float64)
|
||||
efSearchVal, ok := tuningRows[0]["ef_search"].(float64)
|
||||
if !ok {
|
||||
return nil, util.NewClientServerError("Failed to calculate appropriate efSearch value", http.StatusInternalServerError, nil)
|
||||
}
|
||||
|
||||
queryRecommendation := fmt.Sprintf("SET hnsw.ef_search = %d;", int(efSearchVal))
|
||||
return queryRecommendation, nil
|
||||
}
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistimprovequeryrecall_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistimprovequeryrecall"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: improve-query-recall-tool
|
||||
type: vector-assist-improve-query-recall
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"improve-query-recall-tool": vectorassistimprovequeryrecall.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "improve-query-recall-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-improve-query-recall",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
// Copyright 2026 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 vectorassistlistspecs
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-list-specs"
|
||||
|
||||
const listSpecsQuery = `
|
||||
SELECT *
|
||||
FROM vector_assist.list_specs(table_name => @table_name::TEXT, column_name => @column_name::TEXT);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("table_name", "Table name to list vector specifications for.", parameters.WithStringRequired(true)),
|
||||
parameters.NewStringParameter("column_name", "Column name to list vector specifications for.", parameters.WithStringRequired(false)),
|
||||
}
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool lists all defined vector specifications for a given table and column name. Use this tool to list vector specifications which were created in the context of the vector assist tools."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
|
||||
tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
resp, err := source.RunSQL(ctx, listSpecsQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistlistspecs_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistlistspecs"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: list-specs-tool
|
||||
type: vector-assist-list-specs
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"list-specs-tool": vectorassistlistspecs.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "list-specs-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-list-specs",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
// Copyright 2026 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 vectorassistmodifyspec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util"
|
||||
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const resourceType string = "vector-assist-modify-spec"
|
||||
|
||||
const modifySpecQuery = `
|
||||
SELECT recommendation_id, vector_spec_id, table_name, schema_name, query, recommendation, applied, modified, created_at
|
||||
FROM vector_assist.modify_spec(spec_id => @spec_id::TEXT, table_name => @table_name::TEXT, schema_name => @schema_name::TEXT,
|
||||
vector_column_name => @vector_column_name::TEXT, text_column_name => @text_column_name::TEXT,
|
||||
vector_index_type => @vector_index_type::TEXT, embeddings_available => @embeddings_available::BOOLEAN,
|
||||
num_vectors => @num_vectors::INTEGER, dimensionality => @dimensionality::INTEGER,
|
||||
embedding_model => @embedding_model::TEXT, prefilter_column_names => @prefilter_column_names,
|
||||
distance_func => @distance_func::TEXT, quantization => @quantization::TEXT,
|
||||
memory_budget_kb => @memory_budget_kb::INTEGER, target_recall => @target_recall::FLOAT,
|
||||
target_top_k => @target_top_k::INTEGER, tune_vector_index => @tune_vector_index::BOOLEAN);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(resourceType, newConfig) {
|
||||
panic(fmt.Sprintf("tool type %q already registered", resourceType))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
RunSQL(context.Context, string, []any) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
tools.ConfigBase `yaml:",inline"`
|
||||
Type string `yaml:"type" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
}
|
||||
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigType() string {
|
||||
return resourceType
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
|
||||
// parameters are marked required/ optional based on the vector assist function defintions
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameter("spec_id", "Unique ID for the vector spec you want to modify.", parameters.WithStringRequired(true)),
|
||||
parameters.NewStringParameter("table_name", "Modify the table name on which vector workload needs to be set up.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("schema_name", "Modify the schema containing the given table.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("vector_column_name", "Modify the column name for the column with vector embeddings.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("text_column_name", "Modify the column name for the column with text on which vector search needs to be set up.", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("vector_index_type", "Modify the type of the vector index to be created (Allowed inputs: 'hnsw', 'ivfflat', 'scann').", parameters.WithStringRequired(false)),
|
||||
parameters.NewBooleanParameter("embeddings_available", "Modify whether vector embeddings are already available in the table.", parameters.WithBooleanRequired(false)),
|
||||
parameters.NewIntParameter("num_vectors", "Modify the number of vectors expected in the dataset.", parameters.WithIntRequired(false)),
|
||||
parameters.NewIntParameter("dimensionality", "Modify the dimensionality of the vectors or embedding model.", parameters.WithIntRequired(false)),
|
||||
parameters.NewStringParameter("embedding_model", "Modify the model used for generating embeddings.", parameters.WithStringRequired(false)),
|
||||
parameters.NewArrayParameter("prefilter_column_names", "Modify the column(s) based on which prefiltering will happen in vector search queries.", parameters.NewStringParameter("prefilter_column_name", "Pre filter column name"), parameters.WithArrayRequired(false)),
|
||||
parameters.NewStringParameter("distance_func", "Modify the distance function to be used for comparing vectors (Allowed inputs: 'cosine', 'ip', 'l2', 'l1').", parameters.WithStringRequired(false)),
|
||||
parameters.NewStringParameter("quantization", "Modify the quantization to be used for creating the vector indexes (Allowed inputs: 'none', 'halfvec', 'bit').", parameters.WithStringRequired(false)),
|
||||
parameters.NewIntParameter("memory_budget_kb", "Modify the maximum size that the index can consume in memory while building.", parameters.WithIntRequired(false)),
|
||||
parameters.NewFloatParameter("target_recall", "Modify the recall that the user would like to target with the given index.", parameters.WithFloatRequired(false)),
|
||||
parameters.NewIntParameter("target_top_k", "Modify the Top-K matching values that need to be retrieved for the given query.", parameters.WithIntRequired(false)),
|
||||
parameters.NewBooleanParameter("tune_vector_index", "Modify whether to tune vector index build and search parameters.", parameters.WithBooleanRequired(false)),
|
||||
}
|
||||
|
||||
if cfg.Description == "" {
|
||||
cfg.Description = "This tool modifies an existing vector specification (identified by a required spec_id) with new parameters or overrides. Upon modification, it automatically recalculates and refreshes the list of generated SQL recommendations to match the updated requirements. This tool provides a way to modify column(s) in the vector spec before applying and taking action on the recommendations. While highly customizable, any optional parameters left unspecified will use internally determined defaults optimized for the specific workload. Use this tool to modify configurations established via 'define_spec' tool such as adjusting target recall, embedding models, or quantization settings, etc."
|
||||
}
|
||||
|
||||
return Tool{
|
||||
BaseTool: tools.NewBaseTool(
|
||||
cfg,
|
||||
nil,
|
||||
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
allParameters,
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
tools.BaseTool[Config]
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
|
||||
if err != nil {
|
||||
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
|
||||
}
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
// Convert our parsed parameters directly into pgx.NamedArgs
|
||||
namedArgs := pgx.NamedArgs{}
|
||||
for key, value := range paramsMap {
|
||||
namedArgs[key] = value
|
||||
}
|
||||
|
||||
// As long as source.RunSQL unwraps args into pgx.Query(ctx, sql, args...), pgx handles the mapping of @param to the named parameter.
|
||||
resp, err := source.RunSQL(ctx, modifySpecQuery, []any{namedArgs})
|
||||
if err != nil {
|
||||
return nil, util.ProcessGeneralError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Cfg
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright 2026 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 vectorassistmodifyspec_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/mcp-toolbox/internal/server"
|
||||
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools"
|
||||
"github.com/googleapis/mcp-toolbox/internal/tools/cloudsqlpg/vectorassistmodifyspec"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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: `
|
||||
kind: tool
|
||||
name: modify-spec-tool
|
||||
type: vector-assist-modify-spec
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"modify-spec-tool": vectorassistmodifyspec.Config{
|
||||
ConfigBase: tools.ConfigBase{
|
||||
Name: "modify-spec-tool",
|
||||
Description: "a test description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
Type: "vector-assist-modify-spec",
|
||||
Source: "a-source",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user