Skip to content

Commit 9a63018

Browse files
Myst9rahulpinto19
authored andcommitted
feat(tools/alloydb-list-instances): Add custom tool kind for AlloyDB list_instances (#1357)
## Description --- This pull request introduces a new custom tool kind `alloydb-list-instances` that allows users to list the AlloyDB instances in a given project, cluster and location. ### Example Configuration ```yaml tools: list_instances: kind: alloydb-list-instances source: alloydb-admin-source description: Use this tool to list all AlloyDB instances for a given project, cluster and location. ``` ### Example Request ``` curl -X POST http://127.0.0.1:5000/api/tool/list_instances/invoke \ -H "Content-Type: application/json" \ -d '{ "projectId": "example-project", "locationId": "us-central1", "clusterId": "example-cluster", }' ``` ## PR Checklist --- > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<issue_number_goes_here>
1 parent f148cf1 commit 9a63018

File tree

8 files changed

+448
-19
lines changed

8 files changed

+448
-19
lines changed

.ci/integration.cloudbuild.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ steps:
7070
- "GOPATH=/gopath"
7171
- "ALLOYDB_PROJECT=$PROJECT_ID"
7272
- "ALLOYDB_CLUSTER=$_ALLOYDB_POSTGRES_CLUSTER"
73+
- "ALLOYDB_INSTANCE=$_ALLOYDB_POSTGRES_INSTANCE"
7374
- "ALLOYDB_REGION=$_REGION"
7475
secretEnv: ["ALLOYDB_POSTGRES_USER"]
7576
volumes:

cmd/root.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
// Import tool packages for side effect of registration
4545
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl"
4646
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistclusters"
47+
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistinstances"
4748
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistusers"
4849
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryanalyzecontribution"
4950
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryconversationalanalytics"

cmd/root_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1339,7 +1339,7 @@ func TestPrebuiltTools(t *testing.T) {
13391339
wantToolset: server.ToolsetConfigs{
13401340
"alloydb-postgres-admin-tools": tools.ToolsetConfig{
13411341
Name: "alloydb-postgres-admin-tools",
1342-
ToolNames: []string{"alloydb-create-cluster", "alloydb-operations-get", "alloydb-create-instance", "list_clusters", "alloydb-list-instances", "list_users", "alloydb-create-user"},
1342+
ToolNames: []string{"alloydb-create-cluster", "alloydb-operations-get", "alloydb-create-instance", "list_clusters", "list_instances", "list_users", "alloydb-create-user"},
13431343
},
13441344
},
13451345
},
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
---
2+
title: "alloydb-list-instances"
3+
type: docs
4+
weight: 1
5+
description: >
6+
The "alloydb-list-instances" tool lists the AlloyDB instances for a given project, cluster and location.
7+
aliases:
8+
- /resources/tools/alloydb-list-instances
9+
---
10+
11+
## About
12+
13+
The `alloydb-list-instances` tool retrieves AlloyDB instance information for all or specified clusters and locations in a given project. It is compatible with [alloydb-admin](../../sources/alloydb-admin.md) source.
14+
15+
`alloydb-list-instances` tool lists the detailed information of AlloyDB instances (instance name, type, IP address, state, configuration, etc) for a given project, cluster and location. The tool takes the following input parameters:
16+
17+
| Parameter | Type | Description | Required |
18+
| :--------- | :----- | :--------------------------------------------------------------------------------------- | :------- |
19+
| `projectId` | string | The GCP project ID to list instances for. | Yes |
20+
| `clusterId` | string | The ID of the cluster to list instances from. Use '-' to get results for all clusters. Default: `-`.| No |
21+
| `locationId` | string | The location of the cluster (e.g., 'us-central1'). Use '-' to get results for all locations. Default: `-`.| No |
22+
> **Note**
23+
> This tool authenticates using the credentials configured in its [alloydb-admin](../../sources/alloydb-admin.md) source which can be either [Application Default Credentials](https://cloud.google.com/docs/authentication/application-default-credentials) or client-side OAuth.
24+
25+
## Example
26+
27+
```yaml
28+
tools:
29+
list_instances:
30+
kind: alloydb-list-instances
31+
source: alloydb-admin-source
32+
description: Use this tool to list all AlloyDB instances for a given project, cluster and location.
33+
```
34+
## Reference
35+
| **field** | **type** | **required** | **description** |
36+
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
37+
| kind | string | true | Must be alloydb-list-instances. | |
38+
| source | string | true | The name of an alloydb-admin source. |
39+
| description | string | true | Description of the tool that is passed to the agent. |

internal/prebuiltconfigs/tools/alloydb-postgres-admin.yaml

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -134,24 +134,10 @@ tools:
134134
kind: alloydb-list-clusters
135135
source: alloydb-admin-source
136136
description: "Lists all AlloyDB clusters in a given project and location."
137-
alloydb-list-instances:
138-
kind: http
139-
source: alloydb-api-source
140-
method: GET
141-
path: /v1/projects/{{.projectId}}/locations/{{.locationId}}/clusters/{{.clusterId}}/instances
137+
list_instances:
138+
kind: alloydb-list-instances
139+
source: alloydb-admin-source
142140
description: "Lists all AlloyDB instances within a specific cluster."
143-
pathParams:
144-
- name: projectId
145-
type: string
146-
description: "The GCP project ID."
147-
- name: locationId
148-
type: string
149-
description: "The location of the cluster (e.g., 'us-central1'). Use '-' to get results for all regions."
150-
default: "-"
151-
- name: clusterId
152-
type: string
153-
description: "The ID of the cluster to list instances from. Use '-' to get results for all clusters."
154-
default: "-"
155141
list_users:
156142
kind: alloydb-list-users
157143
source: alloydb-admin-source
@@ -210,6 +196,6 @@ toolsets:
210196
- alloydb-operations-get
211197
- alloydb-create-instance
212198
- list_clusters
213-
- alloydb-list-instances
199+
- list_instances
214200
- list_users
215201
- alloydb-create-user
Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package alloydblistinstances
16+
17+
import (
18+
"context"
19+
"fmt"
20+
21+
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/sources"
23+
alloydbadmin "github.com/googleapis/genai-toolbox/internal/sources/alloydbadmin"
24+
"github.com/googleapis/genai-toolbox/internal/tools"
25+
"google.golang.org/api/alloydb/v1"
26+
"google.golang.org/api/option"
27+
)
28+
29+
const kind string = "alloydb-list-instances"
30+
31+
func init() {
32+
if !tools.Register(kind, newConfig) {
33+
panic(fmt.Sprintf("tool kind %q already registered", kind))
34+
}
35+
}
36+
37+
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
38+
actual := Config{Name: name}
39+
if err := decoder.DecodeContext(ctx, &actual); err != nil {
40+
return nil, err
41+
}
42+
return actual, nil
43+
}
44+
45+
// Configuration for the list-instances tool.
46+
type Config struct {
47+
Name string `yaml:"name" validate:"required"`
48+
Kind string `yaml:"kind" validate:"required"`
49+
Source string `yaml:"source" validate:"required"`
50+
Description string `yaml:"description" validate:"required"`
51+
AuthRequired []string `yaml:"authRequired"`
52+
BaseURL string `yaml:"baseURL"`
53+
}
54+
55+
// validate interface
56+
var _ tools.ToolConfig = Config{}
57+
58+
// ToolConfigKind returns the kind of the tool.
59+
func (cfg Config) ToolConfigKind() string {
60+
return kind
61+
}
62+
63+
// Initialize initializes the tool from the configuration.
64+
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
65+
rawS, ok := srcs[cfg.Source]
66+
if !ok {
67+
return nil, fmt.Errorf("source %q not found", cfg.Source)
68+
}
69+
70+
s, ok := rawS.(*alloydbadmin.Source)
71+
if !ok {
72+
return nil, fmt.Errorf("invalid source for %q tool: source kind must be `%s`", kind, alloydbadmin.SourceKind)
73+
}
74+
75+
allParameters := tools.Parameters{
76+
tools.NewStringParameter("projectId", "The GCP project ID to list instances for."),
77+
tools.NewStringParameterWithDefault("locationId", "-", "Optional: The location of the cluster (e.g., 'us-central1'). Use '-' to get results for all regions.(Default: '-')"),
78+
tools.NewStringParameterWithDefault("clusterId", "-", "Optional: The ID of the cluster to list instances from. Use '-' to get results for all clusters.(Default: '-')"),
79+
}
80+
paramManifest := allParameters.Manifest()
81+
82+
inputSchema := allParameters.McpManifest()
83+
inputSchema.Required = []string{"projectId"}
84+
85+
mcpManifest := tools.McpManifest{
86+
Name: cfg.Name,
87+
Description: cfg.Description,
88+
InputSchema: inputSchema,
89+
}
90+
91+
return Tool{
92+
Name: cfg.Name,
93+
Kind: kind,
94+
Source: s,
95+
AllParams: allParameters,
96+
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest},
97+
mcpManifest: mcpManifest,
98+
}, nil
99+
}
100+
101+
// Tool represents the list-instances tool.
102+
type Tool struct {
103+
Name string `yaml:"name"`
104+
Kind string `yaml:"kind"`
105+
Description string `yaml:"description"`
106+
107+
Source *alloydbadmin.Source
108+
AllParams tools.Parameters `yaml:"allParams"`
109+
110+
manifest tools.Manifest
111+
mcpManifest tools.McpManifest
112+
}
113+
114+
// Invoke executes the tool's logic.
115+
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
116+
paramsMap := params.AsMap()
117+
118+
projectId, ok := paramsMap["projectId"].(string)
119+
if !ok {
120+
return nil, fmt.Errorf("invalid or missing 'projectId' parameter; expected a string")
121+
}
122+
locationId, ok := paramsMap["locationId"].(string)
123+
if !ok {
124+
return nil, fmt.Errorf("invalid 'locationId' parameter; expected a string")
125+
}
126+
clusterId, ok := paramsMap["clusterId"].(string)
127+
if !ok {
128+
return nil, fmt.Errorf("invalid 'clusterId' parameter; expected a string")
129+
}
130+
131+
// Get an authenticated HTTP client from the source
132+
client, err := t.Source.GetClient(ctx, string(accessToken))
133+
if err != nil {
134+
return nil, fmt.Errorf("error getting authorized client: %w", err)
135+
}
136+
137+
// Create a new AlloyDB service client using the authorized client
138+
alloydbService, err := alloydb.NewService(ctx, option.WithHTTPClient(client))
139+
if err != nil {
140+
return nil, fmt.Errorf("error creating AlloyDB service: %w", err)
141+
}
142+
143+
urlString := fmt.Sprintf("projects/%s/locations/%s/clusters/%s", projectId, locationId, clusterId)
144+
145+
resp, err := alloydbService.Projects.Locations.Clusters.Instances.List(urlString).Do()
146+
if err != nil {
147+
return nil, fmt.Errorf("error listing AlloyDB instances: %w", err)
148+
}
149+
150+
return resp, nil
151+
}
152+
153+
// ParseParams parses the parameters for the tool.
154+
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
155+
return tools.ParseParams(t.AllParams, data, claims)
156+
}
157+
158+
// Manifest returns the tool's manifest.
159+
func (t Tool) Manifest() tools.Manifest {
160+
return t.manifest
161+
}
162+
163+
// McpManifest returns the tool's MCP manifest.
164+
func (t Tool) McpManifest() tools.McpManifest {
165+
return t.mcpManifest
166+
}
167+
168+
// Authorized checks if the tool is authorized.
169+
func (t Tool) Authorized(verifiedAuthServices []string) bool {
170+
return true
171+
}
172+
173+
func (t Tool) RequiresClientAuthorization() bool {
174+
return t.Source.UseClientAuthorization()
175+
}
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package alloydblistinstances_test
16+
17+
import (
18+
"testing"
19+
20+
yaml "github.com/goccy/go-yaml"
21+
"github.com/google/go-cmp/cmp"
22+
"github.com/googleapis/genai-toolbox/internal/server"
23+
"github.com/googleapis/genai-toolbox/internal/testutils"
24+
alloydblistinstances "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistinstances"
25+
)
26+
27+
func TestParseFromYaml(t *testing.T) {
28+
ctx, err := testutils.ContextWithNewLogger()
29+
if err != nil {
30+
t.Fatalf("unexpected error: %s", err)
31+
}
32+
tcs := []struct {
33+
desc string
34+
in string
35+
want server.ToolConfigs
36+
}{
37+
{
38+
desc: "basic example",
39+
in: `
40+
tools:
41+
list-my-instances:
42+
kind: alloydb-list-instances
43+
source: my-alloydb-admin-source
44+
description: some description
45+
`,
46+
want: server.ToolConfigs{
47+
"list-my-instances": alloydblistinstances.Config{
48+
Name: "list-my-instances",
49+
Kind: "alloydb-list-instances",
50+
Source: "my-alloydb-admin-source",
51+
Description: "some description",
52+
AuthRequired: []string{},
53+
},
54+
},
55+
},
56+
{
57+
desc: "with auth required",
58+
in: `
59+
tools:
60+
list-my-instances-auth:
61+
kind: alloydb-list-instances
62+
source: my-alloydb-admin-source
63+
description: some description
64+
authRequired:
65+
- my-google-auth-service
66+
- other-auth-service
67+
`,
68+
want: server.ToolConfigs{
69+
"list-my-instances-auth": alloydblistinstances.Config{
70+
Name: "list-my-instances-auth",
71+
Kind: "alloydb-list-instances",
72+
Source: "my-alloydb-admin-source",
73+
Description: "some description",
74+
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
75+
},
76+
},
77+
},
78+
}
79+
for _, tc := range tcs {
80+
t.Run(tc.desc, func(t *testing.T) {
81+
got := struct {
82+
Tools server.ToolConfigs `yaml:"tools"`
83+
}{}
84+
// Parse contents
85+
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86+
if err != nil {
87+
t.Fatalf("unable to unmarshal: %s", err)
88+
}
89+
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90+
t.Fatalf("incorrect parse: diff %v", diff)
91+
}
92+
})
93+
}
94+
}

0 commit comments

Comments
 (0)