feat!: consolidate "x-postgres-generic" tools to "postgres-sql" tool (#43)

This PR introduces the following breaking change: The
`alloydb-pg-generic`, `cloud-sql-pg-generic`, and
`postgres-generic-tool` have been replaced by the `postgres-sql` tool,
which works with all 3 Postgres sources.

If you were using of the the previous tools, you will need to update it
as follows:
```diff
example_tool:
-    kind: cloud-sql-pg-generic
+    kind: postgres-sql
     source: my-cloud-sql-pg-instance
     description: some description
        statement: |
            SELECT * FROM SQL_STATEMENT;
        parameters:
        - name: country
          type: string
          description: some description
```

I'm proposing this change for the following reasons:
1. It provides greater flexibility between postgres-compatible sources
-- you can change between "postgres" and "alloydb-postgres" without
issue
2. The name "postgres-sql" is more clear that "postgres-generic" -- it
indicates it's a tool that runs SQL on the source
3. It's easier for us to maintain feature compatibility across a single
"postgres-sql" tool
This commit is contained in:
Kurtis Van Gent
2024-11-01 17:17:18 -06:00
committed by GitHub
parent 0a0d206efd
commit f630965937
14 changed files with 81 additions and 355 deletions

View File

@@ -26,7 +26,7 @@ import (
cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
cloudsqlpgtool "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/tools/postgressql"
"github.com/spf13/cobra"
)
@@ -186,7 +186,7 @@ func TestParseToolFile(t *testing.T) {
database: my_db
tools:
example_tool:
kind: cloud-sql-postgres-generic
kind: postgres-sql
source: my-pg-instance
description: some description
statement: |
@@ -210,9 +210,9 @@ func TestParseToolFile(t *testing.T) {
},
},
wantTools: server.ToolConfigs{
"example_tool": cloudsqlpgtool.GenericConfig{
"example_tool": postgressql.Config{
Name: "example_tool",
Kind: cloudsqlpgtool.ToolKind,
Kind: postgressql.ToolKind,
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",

View File

@@ -46,11 +46,7 @@ tools:
We currently support the following types of kinds of tools:
* [alloydb-postgres](./alloydb-pg.md) - Run a SQL statement against an AlloyDB
for PostgreSQL instance.
* [cloud-sql-postgres](./cloud-sql-pg.md) - Run a SQL statement against a Cloud
SQL for PostgreSQL instance.
* [postgres-generic](./postgres.md) - Run a PostgreSQL statement against a
* [postgres-sql](./sql.md) - Run a PostgreSQL statement against a
PostgreSQL-compatible database.

View File

@@ -1,18 +1,25 @@
# PostgreSQL Generic Tool
# Postgres SQL Tool
The "postgres-generic" tool executes a pre-defined SQL statement.
## Requirements
PostgreSQL Generic Tools require one of the following sources:
A "postgres-sql" tool executes a pre-defined SQL statement against a postgres
database. It's compatible with any of the following sources:
- [alloydb-postgres](../sources/alloydb-pg.md)
- [cloud-sql-postgres](../sources/cloud-sql-pg.md)
- [postgres](../sources/postgres.md)
The specified SQL statement is executed as a [prepared statement][pg-prepare],
and specified parameters will inserted according to their position: e.g. "$1"
will be the first parameter specified, "$@" will be the second parameter, and so
on.
[pg-prepare]: https://www.postgresql.org/docs/current/sql-prepare.html
## Example
```yaml
tools:
search_flights_by_number:
kind: postgres-generic
kind: postgres-sql
source: my-pg-instance
statement: |
SELECT * FROM flights

View File

@@ -21,9 +21,7 @@ import (
cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
postgressrc "github.com/googleapis/genai-toolbox/internal/sources/postgres"
"github.com/googleapis/genai-toolbox/internal/tools"
alloydbpgtool "github.com/googleapis/genai-toolbox/internal/tools/alloydbpg"
cloudsqlpgtool "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg"
postgrestool "github.com/googleapis/genai-toolbox/internal/tools/postgres"
"github.com/googleapis/genai-toolbox/internal/tools/postgressql"
"gopkg.in/yaml.v3"
)
@@ -92,7 +90,7 @@ func (c *SourceConfigs) UnmarshalYAML(node *yaml.Node) error {
}
// ToolConfigs is a type used to allow unmarshal of the tool configs
type ToolConfigs map[string]tools.Config
type ToolConfigs map[string]tools.ToolConfig
// validate interface
var _ yaml.Unmarshaler = &ToolConfigs{}
@@ -114,20 +112,8 @@ func (c *ToolConfigs) UnmarshalYAML(node *yaml.Node) error {
return fmt.Errorf("missing 'kind' field for %q", name)
}
switch k.Kind {
case alloydbpgtool.ToolKind:
actual := alloydbpgtool.GenericConfig{Name: name}
if err := n.Decode(&actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
case cloudsqlpgtool.ToolKind:
actual := cloudsqlpgtool.GenericConfig{Name: name}
if err := n.Decode(&actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
(*c)[name] = actual
case postgrestool.ToolKind:
actual := postgrestool.GenericConfig{Name: name}
case postgressql.ToolKind:
actual := postgressql.Config{Name: name}
if err := n.Decode(&actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", k.Kind, err)
}
@@ -140,6 +126,7 @@ func (c *ToolConfigs) UnmarshalYAML(node *yaml.Node) error {
return nil
}
// ToolConfigs is a type used to allow unmarshal of the toolset configs
type ToolsetConfigs map[string]tools.ToolsetConfig
// validate interface

View File

@@ -76,6 +76,10 @@ func (s Source) SourceKind() string {
return SourceKind
}
func (s *Source) PostgresPool() *pgxpool.Pool {
return s.Pool
}
func initAlloyDBPgConnectionPool(project, region, cluster, instance, user, pass, dbname string) (*pgxpool.Pool, error) {
// Configure the driver to connect to the database
dsn := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=disable", user, pass, dbname)

View File

@@ -75,6 +75,10 @@ func (s Source) SourceKind() string {
return SourceKind
}
func (s *Source) PostgresPool() *pgxpool.Pool {
return s.Pool
}
func initCloudSQLPgConnectionPool(project, region, instance, user, pass, dbname string) (*pgxpool.Pool, error) {
// Configure the driver to connect to the database
dsn := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=disable", user, pass, dbname)

View File

@@ -72,6 +72,10 @@ func (s Source) SourceKind() string {
return SourceKind
}
func (s *Source) PostgresPool() *pgxpool.Pool {
return s.Pool
}
func initPostgresConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) {
// urlExample := "postgres:dd//username:password@localhost:5432/database_name"
i := fmt.Sprintf("postgres://%s:%s@%s:%s/%s", user, pass, host, port, dbname)

View File

@@ -1,69 +0,0 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package alloydbpg
import (
"fmt"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/postgres"
)
const ToolKind string = "alloydb-postgres-generic"
// validate interface
var _ tools.Config = GenericConfig{}
type GenericConfig struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Source string `yaml:"source"`
Description string `yaml:"description"`
Statement string `yaml:"statement"`
Parameters tools.Parameters `yaml:"parameters"`
}
func (cfg GenericConfig) ToolKind() string {
return ToolKind
}
func (cfg GenericConfig) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is the right kind
s, ok := rawS.(alloydbpg.Source)
if !ok {
return nil, fmt.Errorf("sources for %q tools must be of kind %q", ToolKind, alloydbpg.SourceKind)
}
// finish tool setup
t := GenericTool{
GenericTool: postgres.NewGenericTool(cfg.Name, cfg.Statement, cfg.Description, s.Pool, cfg.Parameters),
}
return t, nil
}
// validate interface
var _ tools.Tool = GenericTool{}
type GenericTool struct {
postgres.GenericTool
}

View File

@@ -1,79 +0,0 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package alloydbpg_test
import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/alloydbpg"
"gopkg.in/yaml.v3"
)
func TestParseFromYamlAlloyDBPg(t *testing.T) {
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: alloydb-postgres-generic
source: my-pg-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
`,
want: server.ToolConfigs{
"example_tool": alloydbpg.GenericConfig{
Name: "example_tool",
Kind: alloydbpg.ToolKind,
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}

View File

@@ -1,69 +0,0 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package cloudsqlpg
import (
"fmt"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/postgres"
)
const ToolKind string = "cloud-sql-postgres-generic"
// validate interface
var _ tools.Config = GenericConfig{}
type GenericConfig struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Source string `yaml:"source"`
Description string `yaml:"description"`
Statement string `yaml:"statement"`
Parameters tools.Parameters `yaml:"parameters"`
}
func (cfg GenericConfig) ToolKind() string {
return ToolKind
}
func (cfg GenericConfig) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is the right kind
s, ok := rawS.(cloudsqlpg.Source)
if !ok {
return nil, fmt.Errorf("sources for %q tools must be of kind %q", ToolKind, cloudsqlpg.SourceKind)
}
// finish tool setup
t := GenericTool{
GenericTool: postgres.NewGenericTool(cfg.Name, cfg.Statement, cfg.Description, s.Pool, cfg.Parameters),
}
return t, nil
}
// validate interface
var _ tools.Tool = GenericTool{}
type GenericTool struct {
postgres.GenericTool
}

View File

@@ -1,79 +0,0 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package cloudsqlpg_test
import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg"
"gopkg.in/yaml.v3"
)
func TestParseFromYamlCloudSQLPg(t *testing.T) {
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: cloud-sql-postgres-generic
source: my-pg-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
`,
want: server.ToolConfigs{
"example_tool": cloudsqlpg.GenericConfig{
Name: "example_tool",
Kind: cloudsqlpg.ToolKind,
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}

View File

@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package postgres
package postgressql
import (
"context"
@@ -20,17 +20,27 @@ import (
"strings"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/jackc/pgx/v5/pgxpool"
)
const ToolKind string = "postgres-generic"
const ToolKind string = "postgres-sql"
// validate interface
var _ tools.Config = GenericConfig{}
type compatibleSource interface {
PostgresPool() *pgxpool.Pool
}
type GenericConfig struct {
// validate compatible sources are still compatible
var _ compatibleSource = &alloydbpg.Source{}
var _ compatibleSource = &cloudsqlpg.Source{}
var _ compatibleSource = &postgres.Source{}
var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
type Config struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Source string `yaml:"source"`
@@ -39,30 +49,40 @@ type GenericConfig struct {
Parameters tools.Parameters `yaml:"parameters"`
}
func (cfg GenericConfig) ToolKind() string {
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return ToolKind
}
func (cfg GenericConfig) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is the right kind
s, ok := rawS.(postgres.Source)
// verify the source is compatible
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("sources for %q tools must be of kind %q", ToolKind, postgres.SourceKind)
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources)
}
// finish tool setup
t := NewGenericTool(cfg.Name, cfg.Statement, cfg.Description, s.Pool, cfg.Parameters)
t := Tool{
Name: cfg.Name,
Kind: ToolKind,
Parameters: cfg.Parameters,
Statement: cfg.Statement,
Pool: s.PostgresPool(),
manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest()},
}
return t, nil
}
func NewGenericTool(name, stmt, desc string, pool *pgxpool.Pool, parameters tools.Parameters) GenericTool {
return GenericTool{
func NewGenericTool(name, stmt, desc string, pool *pgxpool.Pool, parameters tools.Parameters) Tool {
return Tool{
Name: name,
Kind: ToolKind,
Statement: stmt,
@@ -73,9 +93,9 @@ func NewGenericTool(name, stmt, desc string, pool *pgxpool.Pool, parameters tool
}
// validate interface
var _ tools.Tool = GenericTool{}
var _ tools.Tool = Tool{}
type GenericTool struct {
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Parameters tools.Parameters `yaml:"parameters"`
@@ -85,7 +105,7 @@ type GenericTool struct {
manifest tools.Manifest
}
func (t GenericTool) Invoke(params []any) (string, error) {
func (t Tool) Invoke(params []any) (string, error) {
fmt.Printf("Invoked tool %s\n", t.Name)
results, err := t.Pool.Query(context.Background(), t.Statement, params...)
if err != nil {
@@ -104,10 +124,10 @@ func (t GenericTool) Invoke(params []any) (string, error) {
return fmt.Sprintf("Stub tool call for %q! Parameters parsed: %q \n Output: %s", t.Name, params, out.String()), nil
}
func (t GenericTool) ParseParams(data map[string]any) ([]any, error) {
func (t Tool) ParseParams(data map[string]any) ([]any, error) {
return tools.ParseParams(t.Parameters, data)
}
func (t GenericTool) Manifest() tools.Manifest {
func (t Tool) Manifest() tools.Manifest {
return t.manifest
}

View File

@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package postgres_test
package postgressql_test
import (
"testing"
@@ -21,7 +21,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/postgres"
"github.com/googleapis/genai-toolbox/internal/tools/postgressql"
"gopkg.in/yaml.v3"
)
@@ -36,7 +36,7 @@ func TestParseFromYamlPostgres(t *testing.T) {
in: `
tools:
example_tool:
kind: postgres-generic
kind: postgres-sql
source: my-pg-instance
description: some description
statement: |
@@ -47,9 +47,9 @@ func TestParseFromYamlPostgres(t *testing.T) {
description: some description
`,
want: server.ToolConfigs{
"example_tool": postgres.GenericConfig{
"example_tool": postgressql.Config{
Name: "example_tool",
Kind: postgres.ToolKind,
Kind: postgressql.ToolKind,
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",

View File

@@ -18,8 +18,8 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
)
type Config interface {
ToolKind() string
type ToolConfig interface {
ToolConfigKind() string
Initialize(map[string]sources.Source) (Tool, error)
}