diff --git a/cmd/root.go b/cmd/root.go index 3d96987327..0af1f5b37f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -32,6 +32,30 @@ import ( "github.com/googleapis/genai-toolbox/internal/server" "github.com/googleapis/genai-toolbox/internal/telemetry" "github.com/googleapis/genai-toolbox/internal/util" + + // Import tool packages for side effect of registration + _ "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigquery" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigqueryexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigquerygetdatasetinfo" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigquerygettableinfo" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigquerylistdatasetids" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigquerylisttableids" + _ "github.com/googleapis/genai-toolbox/internal/tools/bigtable" + _ "github.com/googleapis/genai-toolbox/internal/tools/couchbase" + _ "github.com/googleapis/genai-toolbox/internal/tools/dgraph" + _ "github.com/googleapis/genai-toolbox/internal/tools/http" + _ "github.com/googleapis/genai-toolbox/internal/tools/mssqlexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" + _ "github.com/googleapis/genai-toolbox/internal/tools/mysqlexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/mysqlsql" + _ "github.com/googleapis/genai-toolbox/internal/tools/neo4j" + _ "github.com/googleapis/genai-toolbox/internal/tools/postgresexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/postgressql" + _ "github.com/googleapis/genai-toolbox/internal/tools/spanner" + _ "github.com/googleapis/genai-toolbox/internal/tools/spannerexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql" + "github.com/spf13/cobra" ) diff --git a/cmd/root_test.go b/cmd/root_test.go index 04f58a4caf..e741a09cac 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -358,7 +358,7 @@ func TestParseToolFile(t *testing.T) { Tools: server.ToolConfigs{ "example_tool": postgressql.Config{ Name: "example_tool", - Kind: postgressql.ToolKind, + Kind: "postgres-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", @@ -489,7 +489,7 @@ func TestParseToolFileWithAuth(t *testing.T) { Tools: server.ToolConfigs{ "example_tool": postgressql.Config{ Name: "example_tool", - Kind: postgressql.ToolKind, + Kind: "postgres-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", @@ -588,7 +588,7 @@ func TestParseToolFileWithAuth(t *testing.T) { Tools: server.ToolConfigs{ "example_tool": postgressql.Config{ Name: "example_tool", - Kind: postgressql.ToolKind, + Kind: "postgres-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", @@ -689,7 +689,7 @@ func TestParseToolFileWithAuth(t *testing.T) { Tools: server.ToolConfigs{ "example_tool": postgressql.Config{ Name: "example_tool", - Kind: postgressql.ToolKind, + Kind: "postgres-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", @@ -842,7 +842,7 @@ func TestEnvVarReplacement(t *testing.T) { Tools: server.ToolConfigs{ "example_tool": http.Config{ Name: "example_tool", - Kind: http.ToolKind, + Kind: "http", Source: "my-instance", Method: "GET", Path: "search?name=alice&pet=cat", diff --git a/internal/server/config.go b/internal/server/config.go index 5385b91c32..07c732c25e 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -38,27 +38,6 @@ import ( spannersrc "github.com/googleapis/genai-toolbox/internal/sources/spanner" sqlitesrc "github.com/googleapis/genai-toolbox/internal/sources/sqlite" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl" - "github.com/googleapis/genai-toolbox/internal/tools/bigquery" - "github.com/googleapis/genai-toolbox/internal/tools/bigqueryexecutesql" - "github.com/googleapis/genai-toolbox/internal/tools/bigquerygetdatasetinfo" - "github.com/googleapis/genai-toolbox/internal/tools/bigquerygettableinfo" - "github.com/googleapis/genai-toolbox/internal/tools/bigquerylistdatasetids" - "github.com/googleapis/genai-toolbox/internal/tools/bigquerylisttableids" - "github.com/googleapis/genai-toolbox/internal/tools/bigtable" - couchbasetool "github.com/googleapis/genai-toolbox/internal/tools/couchbase" - "github.com/googleapis/genai-toolbox/internal/tools/dgraph" - httptool "github.com/googleapis/genai-toolbox/internal/tools/http" - "github.com/googleapis/genai-toolbox/internal/tools/mssqlexecutesql" - "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" - "github.com/googleapis/genai-toolbox/internal/tools/mysqlexecutesql" - "github.com/googleapis/genai-toolbox/internal/tools/mysqlsql" - neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j" - "github.com/googleapis/genai-toolbox/internal/tools/postgresexecutesql" - "github.com/googleapis/genai-toolbox/internal/tools/postgressql" - "github.com/googleapis/genai-toolbox/internal/tools/spanner" - "github.com/googleapis/genai-toolbox/internal/tools/spannerexecutesql" - "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql" "github.com/googleapis/genai-toolbox/internal/util" ) @@ -340,146 +319,25 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac v["authRequired"] = []string{} } - kind, ok := v["kind"] + kindVal, ok := v["kind"] if !ok { - return fmt.Errorf("missing 'kind' field for %q", name) + return fmt.Errorf("missing 'kind' field for tool %q", name) + } + kindStr, ok := kindVal.(string) + if !ok { + return fmt.Errorf("invalid 'kind' field for tool %q (must be a string)", name) } - dec, err := util.NewStrictDecoder(v) + yamlDecoder, err := util.NewStrictDecoder(v) if err != nil { - return fmt.Errorf("error creating decoder: %w", err) - } - switch kind { - case bigtable.ToolKind: - actual := bigtable.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case postgressql.ToolKind: - actual := postgressql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case alloydbainl.ToolKind: - actual := alloydbainl.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case mysqlsql.ToolKind: - actual := mysqlsql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case spanner.ToolKind: - actual := spanner.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case neo4jtool.ToolKind: - actual := neo4jtool.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case mssqlsql.ToolKind: - actual := mssqlsql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case dgraph.ToolKind: - actual := dgraph.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case httptool.ToolKind: - actual := httptool.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigquery.ToolKind: - actual := bigquery.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case sqlitesql.ToolKind: - actual := sqlitesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case postgresexecutesql.ToolKind: - actual := postgresexecutesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case mysqlexecutesql.ToolKind: - actual := mysqlexecutesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case spannerexecutesql.ToolKind: - actual := spannerexecutesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case mssqlexecutesql.ToolKind: - actual := mssqlexecutesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case couchbasetool.ToolKind: - actual := couchbasetool.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigqueryexecutesql.ToolKind: - actual := bigqueryexecutesql.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigquerylistdatasetids.ToolKind: - actual := bigquerylistdatasetids.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigquerygetdatasetinfo.ToolKind: - actual := bigquerygetdatasetinfo.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigquerylisttableids.ToolKind: - actual := bigquerylisttableids.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - case bigquerygettableinfo.ToolKind: - actual := bigquerygettableinfo.Config{Name: name} - if err := dec.DecodeContext(ctx, &actual); err != nil { - return fmt.Errorf("unable to parse as %q: %w", kind, err) - } - (*c)[name] = actual - default: - return fmt.Errorf("%q is not a valid kind of tool", kind) + return fmt.Errorf("error creating YAML decoder for tool %q: %w", name, err) } + toolCfg, err := tools.DecodeConfig(ctx, kindStr, name, yamlDecoder) + if err != nil { + return err + } + (*c)[name] = toolCfg } return nil } diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index 049727609a..765b886d2b 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -19,13 +19,28 @@ import ( "fmt" "strings" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" "github.com/googleapis/genai-toolbox/internal/tools" "github.com/jackc/pgx/v5/pgxpool" ) -const ToolKind string = "alloydb-ai-nl" +const kind string = "alloydb-ai-nl" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { PostgresPool() *pgxpool.Pool @@ -50,7 +65,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -63,7 +78,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } numParams := len(cfg.NLConfigParameters) @@ -112,7 +127,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.NLConfigParameters, Statement: stmt, NLConfig: cfg.NLConfig, diff --git a/internal/tools/alloydbainl/alloydbainl_test.go b/internal/tools/alloydbainl/alloydbainl_test.go index 50ebe6e443..2061450c0a 100644 --- a/internal/tools/alloydbainl/alloydbainl_test.go +++ b/internal/tools/alloydbainl/alloydbainl_test.go @@ -57,7 +57,7 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { want: server.ToolConfigs{ "example_tool": alloydbainl.Config{ Name: "example_tool", - Kind: alloydbainl.ToolKind, + Kind: "alloydb-ai-nl", Source: "my-alloydb-instance", Description: "AlloyDB natural language query tool", NLConfig: "my_nl_config", @@ -98,7 +98,7 @@ func TestParseFromYamlAlloyDBNLA(t *testing.T) { want: server.ToolConfigs{ "complex_tool": alloydbainl.Config{ Name: "complex_tool", - Kind: alloydbainl.ToolKind, + Kind: "alloydb-ai-nl", Source: "my-alloydb-instance", Description: "AlloyDB natural language query tool with multiple parameters", NLConfig: "complex_nl_config", diff --git a/internal/tools/bigquery/bigquery.go b/internal/tools/bigquery/bigquery.go index 636cdb4b5c..2ea5621f2a 100644 --- a/internal/tools/bigquery/bigquery.go +++ b/internal/tools/bigquery/bigquery.go @@ -20,13 +20,28 @@ import ( "strings" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "bigquery-sql" +const kind string = "bigquery-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/bigquery/bigquery_test.go b/internal/tools/bigquery/bigquery_test.go index 00e5b39325..5732172eab 100644 --- a/internal/tools/bigquery/bigquery_test.go +++ b/internal/tools/bigquery/bigquery_test.go @@ -53,7 +53,7 @@ func TestParseFromYamlBigQuery(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigquery.Config{ Name: "example_tool", - Kind: bigquery.ToolKind, + Kind: "bigquery-sql", Source: "my-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/bigqueryexecutesql/bigqueryexecutesql.go b/internal/tools/bigqueryexecutesql/bigqueryexecutesql.go index fd8b77292e..7b3b42675b 100644 --- a/internal/tools/bigqueryexecutesql/bigqueryexecutesql.go +++ b/internal/tools/bigqueryexecutesql/bigqueryexecutesql.go @@ -19,13 +19,28 @@ import ( "fmt" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "bigquery-execute-sql" +const kind string = "bigquery-execute-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -48,7 +63,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -61,7 +76,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } sqlParameter := tools.NewStringParameter("sql", "The sql to execute.") @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Client: s.BigQueryClient(), diff --git a/internal/tools/bigqueryexecutesql/bigqueryexecutesql_test.go b/internal/tools/bigqueryexecutesql/bigqueryexecutesql_test.go index 120ad2339a..934a7e28b9 100644 --- a/internal/tools/bigqueryexecutesql/bigqueryexecutesql_test.go +++ b/internal/tools/bigqueryexecutesql/bigqueryexecutesql_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlBigQueryExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigqueryexecutesql.Config{ Name: "example_tool", - Kind: bigqueryexecutesql.ToolKind, + Kind: "bigquery-execute-sql", Source: "my-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index e74d9b8671..ab8bf5b421 100644 --- a/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -19,12 +19,27 @@ import ( "fmt" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "bigquery-get-dataset-info" +const kind string = "bigquery-get-dataset-info" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -47,7 +62,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -60,7 +75,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } datasetParameter := tools.NewStringParameter("dataset", "The dataset to get metadata information.") @@ -75,7 +90,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Client: s.BigQueryClient(), diff --git a/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo_test.go b/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo_test.go index bc27372bb7..403ced65f8 100644 --- a/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo_test.go +++ b/internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlBigQueryGetDatasetInfo(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigquerygetdatasetinfo.Config{ Name: "example_tool", - Kind: bigquerygetdatasetinfo.ToolKind, + Kind: "bigquery-get-dataset-info", Source: "my-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/bigquerygettableinfo/bigquerygettableinfo.go b/internal/tools/bigquerygettableinfo/bigquerygettableinfo.go index dde0155b15..550e9f8bcc 100644 --- a/internal/tools/bigquerygettableinfo/bigquerygettableinfo.go +++ b/internal/tools/bigquerygettableinfo/bigquerygettableinfo.go @@ -19,12 +19,27 @@ import ( "fmt" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "bigquery-get-table-info" +const kind string = "bigquery-get-table-info" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -47,7 +62,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -60,7 +75,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } datasetParameter := tools.NewStringParameter("dataset", "The table's parent dataset.") @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Client: s.BigQueryClient(), diff --git a/internal/tools/bigquerygettableinfo/bigquerygettableinfo_test.go b/internal/tools/bigquerygettableinfo/bigquerygettableinfo_test.go index 1a9884b92c..3553c76a64 100644 --- a/internal/tools/bigquerygettableinfo/bigquerygettableinfo_test.go +++ b/internal/tools/bigquerygettableinfo/bigquerygettableinfo_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlBigQueryGetTableInfo(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigquerygettableinfo.Config{ Name: "example_tool", - Kind: bigquerygettableinfo.ToolKind, + Kind: "bigquery-get-table-info", Source: "my-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids.go b/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids.go index f33250f2b7..8f01894c0d 100644 --- a/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids.go +++ b/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids.go @@ -19,13 +19,28 @@ import ( "fmt" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "bigquery-list-dataset-ids" +const kind string = "bigquery-list-dataset-ids" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -48,7 +63,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -61,7 +76,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } parameters := tools.Parameters{} @@ -75,7 +90,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Client: s.BigQueryClient(), diff --git a/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids_test.go b/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids_test.go index 8fcca753ba..bf33a2c851 100644 --- a/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids_test.go +++ b/internal/tools/bigquerylistdatasetids/bigquerylistdatasetids_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlBigQueryListDatasetIds(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigquerylistdatasetids.Config{ Name: "example_tool", - Kind: bigquerylistdatasetids.ToolKind, + Kind: "bigquery-list-dataset-ids", Source: "my-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/bigquerylisttableids/bigquerylisttableids.go b/internal/tools/bigquerylisttableids/bigquerylisttableids.go index 3cf3db43c4..3165de45fd 100644 --- a/internal/tools/bigquerylisttableids/bigquerylisttableids.go +++ b/internal/tools/bigquerylisttableids/bigquerylisttableids.go @@ -19,13 +19,28 @@ import ( "fmt" bigqueryapi "cloud.google.com/go/bigquery" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "bigquery-list-table-ids" +const kind string = "bigquery-list-table-ids" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigQueryClient() *bigqueryapi.Client @@ -48,7 +63,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -61,7 +76,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } datasetParameter := tools.NewStringParameter("dataset", "The dataset to list table ids.") @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Client: s.BigQueryClient(), diff --git a/internal/tools/bigquerylisttableids/bigquerylisttableids_test.go b/internal/tools/bigquerylisttableids/bigquerylisttableids_test.go index 62cbf81615..76ec8d445e 100644 --- a/internal/tools/bigquerylisttableids/bigquerylisttableids_test.go +++ b/internal/tools/bigquerylisttableids/bigquerylisttableids_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlBigQueryListTableIds(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigquerylisttableids.Config{ Name: "example_tool", - Kind: bigquerylisttableids.ToolKind, + Kind: "bigquery-list-table-ids", Source: "my-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/bigtable/bigtable.go b/internal/tools/bigtable/bigtable.go index 1eef123694..7e9ebb9432 100644 --- a/internal/tools/bigtable/bigtable.go +++ b/internal/tools/bigtable/bigtable.go @@ -19,12 +19,27 @@ import ( "fmt" "cloud.google.com/go/bigtable" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" bigtabledb "github.com/googleapis/genai-toolbox/internal/sources/bigtable" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "bigtable-sql" +const kind string = "bigtable-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { BigtableClient() *bigtable.Client @@ -49,7 +64,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -62,7 +77,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -74,7 +89,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/bigtable/bigtable_test.go b/internal/tools/bigtable/bigtable_test.go index fc41f0f675..ec0b5e49c8 100644 --- a/internal/tools/bigtable/bigtable_test.go +++ b/internal/tools/bigtable/bigtable_test.go @@ -53,7 +53,7 @@ func TestParseFromYamlBigtable(t *testing.T) { want: server.ToolConfigs{ "example_tool": bigtable.Config{ Name: "example_tool", - Kind: bigtable.ToolKind, + Kind: "bigtable-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/couchbase/couchbase.go b/internal/tools/couchbase/couchbase.go index e7437498a8..fe2319180a 100644 --- a/internal/tools/couchbase/couchbase.go +++ b/internal/tools/couchbase/couchbase.go @@ -20,12 +20,27 @@ import ( "fmt" "github.com/couchbase/gocb/v2" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/couchbase" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "couchbase-sql" +const kind string = "couchbase-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { CouchbaseScope() *gocb.Scope @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -75,7 +90,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, Scope: s.CouchbaseScope(), diff --git a/internal/tools/couchbase/couchbase_test.go b/internal/tools/couchbase/couchbase_test.go index 3a719fb3c3..9d79a4d823 100644 --- a/internal/tools/couchbase/couchbase_test.go +++ b/internal/tools/couchbase/couchbase_test.go @@ -50,7 +50,7 @@ func TestParseFromYamlCouchbase(t *testing.T) { want: server.ToolConfigs{ "example_tool": couchbase.Config{ Name: "example_tool", - Kind: couchbase.ToolKind, + Kind: "couchbase-sql", AuthRequired: []string{}, Source: "my-couchbase-instance", Description: "some tool description", diff --git a/internal/tools/dgraph/dgraph.go b/internal/tools/dgraph/dgraph.go index 782418ce6b..1e0c12260f 100644 --- a/internal/tools/dgraph/dgraph.go +++ b/internal/tools/dgraph/dgraph.go @@ -19,12 +19,27 @@ import ( "encoding/json" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/dgraph" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "dgraph-dql" +const kind string = "dgraph-dql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { DgraphClient() *dgraph.DgraphClient @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/dgraph/dgraph_test.go b/internal/tools/dgraph/dgraph_test.go index 7c12e3baca..cb4f672ff3 100644 --- a/internal/tools/dgraph/dgraph_test.go +++ b/internal/tools/dgraph/dgraph_test.go @@ -50,7 +50,7 @@ func TestParseFromYamlDgraph(t *testing.T) { want: server.ToolConfigs{ "example_tool": dgraph.Config{ Name: "example_tool", - Kind: dgraph.ToolKind, + Kind: "dgraph-dql", Source: "my-dgraph-instance", AuthRequired: []string{}, Description: "some tool description", @@ -74,7 +74,7 @@ func TestParseFromYamlDgraph(t *testing.T) { want: server.ToolConfigs{ "example_tool": dgraph.Config{ Name: "example_tool", - Kind: dgraph.ToolKind, + Kind: "dgraph-dql", Source: "my-dgraph-instance", Description: "some tool description", AuthRequired: []string{}, diff --git a/internal/tools/http/http.go b/internal/tools/http/http.go index 541a9a2287..31f759fd38 100644 --- a/internal/tools/http/http.go +++ b/internal/tools/http/http.go @@ -27,12 +27,27 @@ import ( "maps" "text/template" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" httpsrc "github.com/googleapis/genai-toolbox/internal/sources/http" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "http" +const kind string = "http" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type Config struct { Name string `yaml:"name" validate:"required"` @@ -53,7 +68,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -66,7 +81,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(*httpsrc.Source) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be `http`", ToolKind) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `http`", kind) } // Create URL based on BaseURL and Path @@ -153,7 +168,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup return Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, URL: u, Method: cfg.Method, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/http/http_test.go b/internal/tools/http/http_test.go index 4be67b1fa0..3a862aa38a 100644 --- a/internal/tools/http/http_test.go +++ b/internal/tools/http/http_test.go @@ -82,7 +82,7 @@ func TestParseFromYamlHTTP(t *testing.T) { want: server.ToolConfigs{ "example_tool": http.Config{ Name: "example_tool", - Kind: http.ToolKind, + Kind: "http", Source: "my-instance", Method: "GET", Path: "search?name=alice&pet=cat", diff --git a/internal/tools/mssqlexecutesql/mssqlexecutesql.go b/internal/tools/mssqlexecutesql/mssqlexecutesql.go index b7aca85cf3..1aaaae950d 100644 --- a/internal/tools/mssqlexecutesql/mssqlexecutesql.go +++ b/internal/tools/mssqlexecutesql/mssqlexecutesql.go @@ -19,13 +19,28 @@ import ( "database/sql" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql" "github.com/googleapis/genai-toolbox/internal/sources/mssql" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "mssql-execute-sql" +const kind string = "mssql-execute-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { MSSQLDB() *sql.DB @@ -49,7 +64,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -62,7 +77,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } sqlParameter := tools.NewStringParameter("sql", "The sql to execute.") @@ -77,7 +92,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Pool: s.MSSQLDB(), diff --git a/internal/tools/mssqlexecutesql/mssqlexecutesql_test.go b/internal/tools/mssqlexecutesql/mssqlexecutesql_test.go index 84bc98dfa2..825d7d4594 100644 --- a/internal/tools/mssqlexecutesql/mssqlexecutesql_test.go +++ b/internal/tools/mssqlexecutesql/mssqlexecutesql_test.go @@ -49,7 +49,7 @@ func TestParseFromYamlExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": mssqlexecutesql.Config{ Name: "example_tool", - Kind: mssqlexecutesql.ToolKind, + Kind: "mssql-execute-sql", Source: "my-instance", Description: "some description", AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, diff --git a/internal/tools/mssqlsql/mssqlsql.go b/internal/tools/mssqlsql/mssqlsql.go index 131f78bacc..a55e63449e 100644 --- a/internal/tools/mssqlsql/mssqlsql.go +++ b/internal/tools/mssqlsql/mssqlsql.go @@ -20,13 +20,28 @@ import ( "fmt" "strings" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql" "github.com/googleapis/genai-toolbox/internal/sources/mssql" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "mssql-sql" +const kind string = "mssql-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { MSSQLDB() *sql.DB @@ -52,7 +67,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -65,7 +80,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -77,7 +92,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/mssqlsql/mssqlsql_test.go b/internal/tools/mssqlsql/mssqlsql_test.go index 69bc734c6c..01f4af67a0 100644 --- a/internal/tools/mssqlsql/mssqlsql_test.go +++ b/internal/tools/mssqlsql/mssqlsql_test.go @@ -61,7 +61,7 @@ func TestParseFromYamlMssql(t *testing.T) { want: server.ToolConfigs{ "example_tool": mssqlsql.Config{ Name: "example_tool", - Kind: mssqlsql.ToolKind, + Kind: "mssql-sql", Source: "my-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/mysqlexecutesql/mysqlexecutesql.go b/internal/tools/mysqlexecutesql/mysqlexecutesql.go index 7a492db47e..4eb813bf7d 100644 --- a/internal/tools/mysqlexecutesql/mysqlexecutesql.go +++ b/internal/tools/mysqlexecutesql/mysqlexecutesql.go @@ -19,13 +19,28 @@ import ( "database/sql" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql" "github.com/googleapis/genai-toolbox/internal/sources/mysql" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "mysql-execute-sql" +const kind string = "mysql-execute-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { MySQLPool() *sql.DB @@ -49,7 +64,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -62,7 +77,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } sqlParameter := tools.NewStringParameter("sql", "The sql to execute.") @@ -77,7 +92,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Pool: s.MySQLPool(), diff --git a/internal/tools/mysqlexecutesql/mysqlexecutesql_test.go b/internal/tools/mysqlexecutesql/mysqlexecutesql_test.go index e552fdfada..9c07753f27 100644 --- a/internal/tools/mysqlexecutesql/mysqlexecutesql_test.go +++ b/internal/tools/mysqlexecutesql/mysqlexecutesql_test.go @@ -49,7 +49,7 @@ func TestParseFromYamlExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": mysqlexecutesql.Config{ Name: "example_tool", - Kind: mysqlexecutesql.ToolKind, + Kind: "mysql-execute-sql", Source: "my-instance", Description: "some description", AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, diff --git a/internal/tools/mysqlsql/mysqlsql.go b/internal/tools/mysqlsql/mysqlsql.go index ed306f37f3..56f26e3731 100644 --- a/internal/tools/mysqlsql/mysqlsql.go +++ b/internal/tools/mysqlsql/mysqlsql.go @@ -19,13 +19,28 @@ import ( "database/sql" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql" "github.com/googleapis/genai-toolbox/internal/sources/mysql" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "mysql-sql" +const kind string = "mysql-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { MySQLPool() *sql.DB @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/mysqlsql/mysqlsql_test.go b/internal/tools/mysqlsql/mysqlsql_test.go index e67602b027..e17d932287 100644 --- a/internal/tools/mysqlsql/mysqlsql_test.go +++ b/internal/tools/mysqlsql/mysqlsql_test.go @@ -61,7 +61,7 @@ func TestParseFromYamlMySQL(t *testing.T) { want: server.ToolConfigs{ "example_tool": mysqlsql.Config{ Name: "example_tool", - Kind: mysqlsql.ToolKind, + Kind: "mysql-sql", Source: "my-mysql-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/neo4j/neo4j.go b/internal/tools/neo4j/neo4j.go index 0256457eb3..c967e62d54 100644 --- a/internal/tools/neo4j/neo4j.go +++ b/internal/tools/neo4j/neo4j.go @@ -18,6 +18,7 @@ import ( "context" "fmt" + yaml "github.com/goccy/go-yaml" neo4jsc "github.com/googleapis/genai-toolbox/internal/sources/neo4j" "github.com/neo4j/neo4j-go-driver/v5/neo4j" @@ -25,7 +26,21 @@ import ( "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "neo4j-cypher" +const kind string = "neo4j-cypher" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { Neo4jDriver() neo4j.DriverWithContext @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -76,7 +91,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/neo4j/neo4j_test.go b/internal/tools/neo4j/neo4j_test.go index 5162308c4c..20a272fcde 100644 --- a/internal/tools/neo4j/neo4j_test.go +++ b/internal/tools/neo4j/neo4j_test.go @@ -56,7 +56,7 @@ func TestParseFromYamlNeo4j(t *testing.T) { want: server.ToolConfigs{ "example_tool": neo4j.Config{ Name: "example_tool", - Kind: neo4j.ToolKind, + Kind: "neo4j-cypher", Source: "my-neo4j-instance", Description: "some tool description", AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, diff --git a/internal/tools/postgresexecutesql/postgresexecutesql.go b/internal/tools/postgresexecutesql/postgresexecutesql.go index a8f2f206a6..b0f32e953e 100644 --- a/internal/tools/postgresexecutesql/postgresexecutesql.go +++ b/internal/tools/postgresexecutesql/postgresexecutesql.go @@ -18,6 +18,7 @@ import ( "context" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg" @@ -26,7 +27,21 @@ import ( "github.com/jackc/pgx/v5/pgxpool" ) -const ToolKind string = "postgres-execute-sql" +const kind string = "postgres-execute-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { PostgresPool() *pgxpool.Pool @@ -51,7 +66,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -64,7 +79,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } sqlParameter := tools.NewStringParameter("sql", "The sql to execute.") @@ -79,7 +94,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, Pool: s.PostgresPool(), diff --git a/internal/tools/postgresexecutesql/postgresexecutesql_test.go b/internal/tools/postgresexecutesql/postgresexecutesql_test.go index 6c7229eb32..09d5c72e9d 100644 --- a/internal/tools/postgresexecutesql/postgresexecutesql_test.go +++ b/internal/tools/postgresexecutesql/postgresexecutesql_test.go @@ -49,7 +49,7 @@ func TestParseFromYamlExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": postgresexecutesql.Config{ Name: "example_tool", - Kind: postgresexecutesql.ToolKind, + Kind: "postgres-execute-sql", Source: "my-instance", Description: "some description", AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, diff --git a/internal/tools/postgressql/postgressql.go b/internal/tools/postgressql/postgressql.go index 33f0d59cd4..756faf4a2c 100644 --- a/internal/tools/postgressql/postgressql.go +++ b/internal/tools/postgressql/postgressql.go @@ -18,6 +18,7 @@ import ( "context" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg" @@ -26,7 +27,21 @@ import ( "github.com/jackc/pgx/v5/pgxpool" ) -const ToolKind string = "postgres-sql" +const kind string = "postgres-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { PostgresPool() *pgxpool.Pool @@ -53,7 +68,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -66,7 +81,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -78,7 +93,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/postgressql/postgressql_test.go b/internal/tools/postgressql/postgressql_test.go index 2786c3d6c2..997860e90b 100644 --- a/internal/tools/postgressql/postgressql_test.go +++ b/internal/tools/postgressql/postgressql_test.go @@ -61,7 +61,7 @@ func TestParseFromYamlPostgres(t *testing.T) { want: server.ToolConfigs{ "example_tool": postgressql.Config{ Name: "example_tool", - Kind: postgressql.ToolKind, + Kind: "postgres-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/spanner/spanner.go b/internal/tools/spanner/spanner.go index dfd34acb0f..0901d875cb 100644 --- a/internal/tools/spanner/spanner.go +++ b/internal/tools/spanner/spanner.go @@ -20,13 +20,28 @@ import ( "strings" "cloud.google.com/go/spanner" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" spannerdb "github.com/googleapis/genai-toolbox/internal/sources/spanner" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "spanner-sql" +const kind string = "spanner-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { SpannerClient() *spanner.Client @@ -53,7 +68,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -66,7 +81,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -78,7 +93,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/spanner/spanner_test.go b/internal/tools/spanner/spanner_test.go index 8ba07c7ce3..19de8545ec 100644 --- a/internal/tools/spanner/spanner_test.go +++ b/internal/tools/spanner/spanner_test.go @@ -53,7 +53,7 @@ func TestParseFromYamlSpanner(t *testing.T) { want: server.ToolConfigs{ "example_tool": spanner.Config{ Name: "example_tool", - Kind: spanner.ToolKind, + Kind: "spanner-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", @@ -83,7 +83,7 @@ func TestParseFromYamlSpanner(t *testing.T) { want: server.ToolConfigs{ "example_tool": spanner.Config{ Name: "example_tool", - Kind: spanner.ToolKind, + Kind: "spanner-sql", Source: "my-pg-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/spannerexecutesql/spannerexecutesql.go b/internal/tools/spannerexecutesql/spannerexecutesql.go index dbadf25eba..9ecea6aac4 100644 --- a/internal/tools/spannerexecutesql/spannerexecutesql.go +++ b/internal/tools/spannerexecutesql/spannerexecutesql.go @@ -19,13 +19,28 @@ import ( "fmt" "cloud.google.com/go/spanner" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" spannerdb "github.com/googleapis/genai-toolbox/internal/sources/spanner" "github.com/googleapis/genai-toolbox/internal/tools" "google.golang.org/api/iterator" ) -const ToolKind string = "spanner-execute-sql" +const kind string = "spanner-execute-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { SpannerClient() *spanner.Client @@ -50,7 +65,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -63,7 +78,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } sqlParameter := tools.NewStringParameter("sql", "The sql to execute.") @@ -78,7 +93,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: parameters, AuthRequired: cfg.AuthRequired, ReadOnly: cfg.ReadOnly, diff --git a/internal/tools/spannerexecutesql/spannerexecutesql_test.go b/internal/tools/spannerexecutesql/spannerexecutesql_test.go index 989d86cf66..abc2bb8f49 100644 --- a/internal/tools/spannerexecutesql/spannerexecutesql_test.go +++ b/internal/tools/spannerexecutesql/spannerexecutesql_test.go @@ -46,7 +46,7 @@ func TestParseFromYamlExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": spannerexecutesql.Config{ Name: "example_tool", - Kind: spannerexecutesql.ToolKind, + Kind: "spanner-execute-sql", Source: "my-spanner-instance", Description: "some description", AuthRequired: []string{}, @@ -67,7 +67,7 @@ func TestParseFromYamlExecuteSql(t *testing.T) { want: server.ToolConfigs{ "example_tool": spannerexecutesql.Config{ Name: "example_tool", - Kind: spannerexecutesql.ToolKind, + Kind: "spanner-execute-sql", Source: "my-spanner-instance", Description: "some description", AuthRequired: []string{}, diff --git a/internal/tools/sqlitesql/sqlitesql.go b/internal/tools/sqlitesql/sqlitesql.go index 6ecbc9c43c..254b303ea5 100644 --- a/internal/tools/sqlitesql/sqlitesql.go +++ b/internal/tools/sqlitesql/sqlitesql.go @@ -19,12 +19,27 @@ import ( "database/sql" "fmt" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/sources/sqlite" "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "sqlite-sql" +const kind string = "sqlite-sql" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} type compatibleSource interface { SQLiteDB() *sql.DB @@ -49,7 +64,7 @@ type Config struct { var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigKind() string { - return ToolKind + return kind } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { @@ -62,7 +77,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // verify the source is compatible s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } mcpManifest := tools.McpManifest{ @@ -74,7 +89,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // finish tool setup t := Tool{ Name: cfg.Name, - Kind: ToolKind, + Kind: kind, Parameters: cfg.Parameters, Statement: cfg.Statement, AuthRequired: cfg.AuthRequired, diff --git a/internal/tools/sqlitesql/sqlitesql_test.go b/internal/tools/sqlitesql/sqlitesql_test.go index 00d19eb455..7da958e1a1 100644 --- a/internal/tools/sqlitesql/sqlitesql_test.go +++ b/internal/tools/sqlitesql/sqlitesql_test.go @@ -61,7 +61,7 @@ func TestParseFromYamlSQLite(t *testing.T) { want: server.ToolConfigs{ "example_tool": sqlitesql.Config{ Name: "example_tool", - Kind: sqlitesql.ToolKind, + Kind: "sqlite-sql", Source: "my-sqlite-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/internal/tools/tools.go b/internal/tools/tools.go index 8735ed4f36..9a2129ad83 100644 --- a/internal/tools/tools.go +++ b/internal/tools/tools.go @@ -16,11 +16,48 @@ package tools import ( "context" + "fmt" "slices" + yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" ) +// ToolConfigFactory defines the signature for a function that creates and +// decodes a specific tool's configuration. It takes the context, the tool's +// name, and a YAML decoder to parse the config. +type ToolConfigFactory func(ctx context.Context, name string, decoder *yaml.Decoder) (ToolConfig, error) + +var toolRegistry = make(map[string]ToolConfigFactory) + +// Register allows individual tool packages to register their configuration +// factory function. This is typically called from an init() function in the +// tool's package. It associates a 'kind' string with a function that can +// produce the specific ToolConfig type. It returns true if the registration was +// successful, and false if a tool with the same kind was already registered. +func Register(kind string, factory ToolConfigFactory) bool { + if _, exists := toolRegistry[kind]; exists { + // Tool with this kind already exists, do not overwrite. + return false + } + toolRegistry[kind] = factory + return true +} + +// DecodeConfig looks up the registered factory for the given kind and uses it +// to decode the tool configuration. +func DecodeConfig(ctx context.Context, kind string, name string, decoder *yaml.Decoder) (ToolConfig, error) { + factory, found := toolRegistry[kind] + if !found { + return nil, fmt.Errorf("unknown tool kind: %q", kind) + } + toolConfig, err := factory(ctx, name, decoder) + if err != nil { + return nil, fmt.Errorf("unable to parse tool %q as kind %q: %w", name, kind, err) + } + return toolConfig, nil +} + type ToolConfig interface { ToolConfigKind() string Initialize(map[string]sources.Source) (Tool, error) diff --git a/main.go b/main.go index 2688aaf50f..4ad18d9d77 100644 --- a/main.go +++ b/main.go @@ -14,7 +14,9 @@ package main -import "github.com/googleapis/genai-toolbox/cmd" +import ( + "github.com/googleapis/genai-toolbox/cmd" +) func main() { cmd.Execute()