github/genai-toolbox
1
1
Fork 0
mirror of https://github.com/googleapis/genai-toolbox.git synced 2026-04-09 03:02:26 -04:00
Code Issues Packages Projects Releases 14 Wiki Activity

refactor: decouple Source from Tool (#2204)

Browse Source 
This PR update the linking mechanism between Source and Tool.

Tools are directly linked to their Source, either by pointing to the
Source's functions or by assigning values from the source during Tool's
initialization. However, the existing approach means that any
modification to the Source after Tool's initialization might not be
reflected. To address this limitation, each tool should only store a
name reference to the Source, rather than direct link or assigned
values.

Tools will provide interface for `compatibleSource`. This will be used
to determine if a Source is compatible with the Tool.
```
type compatibleSource interface{
    Client() http.Client
    ProjectID() string
}
```

During `Invoke()`, the tool will run the following operations:
* retrieve Source from the `resourceManager` with source's named defined
in Tool's config
* validate Source via `compatibleSource interface{}`
* run the remaining `Invoke()` function. Fields that are needed is
retrieved directly from the source.

With this update, resource manager is also added as input to other
Tool's function that require access to source (e.g.
`RequiresClientAuthorization()`).
This commit is contained in:
Yuan Teoh
2025-12-19 21:27:55 -08:00
committed by GitHub
parent 7daa4111f4
commit 967a72da11
203 changed files with 4242 additions and 6391 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
internal/tools/cloudsql/cloudsqlwaitforoperation/cloudsqlwaitforoperation.go
View File

@@ -25,9 +25,9 @@ import (
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqladmin"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/sqladmin/v1"
)
const kind string = "cloud-sql-wait-for-operation"
@@ -87,6 +87,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
return actual, nil
}
type compatibleSource interface {
GetDefaultProject() string
GetService(context.Context, string) (*sqladmin.Service, error)
UseClientAuthorization() bool
}
// Config defines the configuration for the wait-for-operation tool.
type Config struct {
Name string `yaml:"name" validate:"required"`
@@ -118,12 +124,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
s, ok := rawS.(*cloudsqladmin.Source)
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be `cloud-sql-admin`", kind)
}
project := s.DefaultProject
project := s.GetDefaultProject()
var projectParam parameters.Parameter
if project != "" {
projectParam = parameters.NewStringParameterWithDefault("project", project, "The GCP project ID. This is pre-configured; do not ask for it unless the user explicitly provides a different one.")
@@ -177,7 +183,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
return Tool{
Config: cfg,
Source: s,
AllParams: allParameters,
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
@@ -191,17 +196,15 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
// Tool represents the wait-for-operation tool.
type Tool struct {
Config
Source *cloudsqladmin.Source
AllParams parameters.Parameters `yaml:"allParams"`
AllParams parameters.Parameters `yaml:"allParams"`
manifest tools.Manifest
mcpManifest tools.McpManifest
// Polling configuration
Delay time.Duration
MaxDelay time.Duration
Multiplier float64
MaxRetries int
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) ToConfig() tools.ToolConfig {
@@ -210,6 +213,11 @@ func (t Tool) ToConfig() tools.ToolConfig {
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
if err != nil {
return nil, err
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
@@ -221,7 +229,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, fmt.Errorf("missing 'operation' parameter")
}
service, err := t.Source.GetService(ctx, string(accessToken))
service, err := source.GetService(ctx, string(accessToken))
if err != nil {
return nil, err
}
@@ -267,7 +275,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, fmt.Errorf("could not unmarshal operation: %w", err)
}
if msg, ok := t.generateCloudSQLConnectionMessage(data); ok {
if msg, ok := t.generateCloudSQLConnectionMessage(source, data); ok {
return msg, nil
}
return string(opBytes), nil
@@ -305,11 +313,15 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool {
return true
}
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) bool {
return t.Source.UseClientAuthorization()
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
if err != nil {
return false, err
}
return source.UseClientAuthorization(), nil
}
func (t Tool) generateCloudSQLConnectionMessage(opResponse map[string]any) (string, bool) {
func (t Tool) generateCloudSQLConnectionMessage(source compatibleSource, opResponse map[string]any) (string, bool) {
operationType, ok := opResponse["operationType"].(string)
if !ok || operationType != "CREATE_DATABASE" {
return "", false
@@ -329,7 +341,7 @@ func (t Tool) generateCloudSQLConnectionMessage(opResponse map[string]any) (stri
instance := matches[2]
database := matches[3]
instanceData, err := t.fetchInstanceData(context.Background(), project, instance)
instanceData, err := t.fetchInstanceData(context.Background(), source, project, instance)
if err != nil {
fmt.Printf("error fetching instance data: %v\n", err)
return "", false
@@ -385,8 +397,8 @@ func (t Tool) generateCloudSQLConnectionMessage(opResponse map[string]any) (stri
return b.String(), true
}
func (t Tool) fetchInstanceData(ctx context.Context, project, instance string) (map[string]any, error) {
service, err := t.Source.GetService(ctx, "")
func (t Tool) fetchInstanceData(ctx context.Context, source compatibleSource, project, instance string) (map[string]any, error) {
service, err := source.GetService(ctx, "")
if err != nil {
return nil, err
}
@@ -408,6 +420,6 @@ func (t Tool) fetchInstanceData(ctx context.Context, project, instance string) (
return data, nil
}
func (t Tool) GetAuthTokenHeaderName() string {
return "Authorization"
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
return "Authorization", nil
}