From 4ca592c3665a26b502c7b552a16ff74376796cd5 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Tue, 3 Feb 2026 01:37:42 -0500 Subject: [PATCH] add dataplex firestore --- .../cloudsqlmssqlcreateinstance.go | 24 ++++++++----- .../cloudsqlmysqlcreateinstance.go | 24 ++++++++----- .../cloudsqlpgcreateinstances.go | 24 ++++++++----- .../cloudsqlpgupgradeprecheck.go | 22 ++++++------ .../dataplexlookupentry.go | 14 +++++--- .../dataplexsearchaspecttypes.go | 27 +++++++++++---- .../dataplexsearchentries.go | 27 +++++++++++---- .../firestoreadddocuments.go | 26 ++++++++------ .../firestoredeletedocuments.go | 26 ++++++++------ .../firestoregetdocuments.go | 26 ++++++++------ .../firestoregetrules/firestoregetrules.go | 12 +++++-- internal/util/errors.go | 34 +++++++++++++++++++ 12 files changed, 200 insertions(+), 86 deletions(-) diff --git a/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance/cloudsqlmssqlcreateinstance.go b/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance/cloudsqlmssqlcreateinstance.go index 23d4d2d2e4..7927c50edd 100644 --- a/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance/cloudsqlmssqlcreateinstance.go +++ b/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance/cloudsqlmssqlcreateinstance.go @@ -17,12 +17,14 @@ package cloudsqlmssqlcreateinstance import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "strings" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" "google.golang.org/api/sqladmin/v1" ) @@ -121,33 +123,33 @@ 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) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() project, ok := paramsMap["project"].(string) if !ok { - return nil, fmt.Errorf("error casting 'project' parameter: %s", paramsMap["project"]) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'project' parameter: %s", paramsMap["project"]), nil) } name, ok := paramsMap["name"].(string) if !ok { - return nil, fmt.Errorf("error casting 'name' parameter: %s", paramsMap["name"]) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %s", paramsMap["name"]), nil) } dbVersion, ok := paramsMap["databaseVersion"].(string) if !ok { - return nil, fmt.Errorf("error casting 'databaseVersion' parameter: %s", paramsMap["databaseVersion"]) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'databaseVersion' parameter: %s", paramsMap["databaseVersion"]), nil) } rootPassword, ok := paramsMap["rootPassword"].(string) if !ok { - return nil, fmt.Errorf("error casting 'rootPassword' parameter: %s", paramsMap["rootPassword"]) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'rootPassword' parameter: %s", paramsMap["rootPassword"]), nil) } editionPreset, ok := paramsMap["editionPreset"].(string) if !ok { - return nil, fmt.Errorf("error casting 'editionPreset' parameter: %s", paramsMap["editionPreset"]) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'editionPreset' parameter: %s", paramsMap["editionPreset"]), nil) } settings := sqladmin.Settings{} switch strings.ToLower(editionPreset) { @@ -164,9 +166,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para settings.DataDiskSizeGb = 100 settings.DataDiskType = "PD_SSD" default: - return nil, fmt.Errorf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset) + return nil, util.NewAgentError(fmt.Sprintf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset), nil) } - return source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + resp, err := source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance/cloudsqlmysqlcreateinstance.go b/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance/cloudsqlmysqlcreateinstance.go index ab78fdc6b7..800439acbf 100644 --- a/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance/cloudsqlmysqlcreateinstance.go +++ b/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance/cloudsqlmysqlcreateinstance.go @@ -17,12 +17,14 @@ package cloudsqlmysqlcreateinstance import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "strings" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" sqladmin "google.golang.org/api/sqladmin/v1" ) @@ -121,33 +123,33 @@ 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) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() project, ok := paramsMap["project"].(string) if !ok { - return nil, fmt.Errorf("missing 'project' parameter") + return nil, util.NewAgentError("missing 'project' parameter", nil) } name, ok := paramsMap["name"].(string) if !ok { - return nil, fmt.Errorf("missing 'name' parameter") + return nil, util.NewAgentError("missing 'name' parameter", nil) } dbVersion, ok := paramsMap["databaseVersion"].(string) if !ok { - return nil, fmt.Errorf("missing 'databaseVersion' parameter") + return nil, util.NewAgentError("missing 'databaseVersion' parameter", nil) } rootPassword, ok := paramsMap["rootPassword"].(string) if !ok { - return nil, fmt.Errorf("missing 'rootPassword' parameter") + return nil, util.NewAgentError("missing 'rootPassword' parameter", nil) } editionPreset, ok := paramsMap["editionPreset"].(string) if !ok { - return nil, fmt.Errorf("missing 'editionPreset' parameter") + return nil, util.NewAgentError("missing 'editionPreset' parameter", nil) } settings := sqladmin.Settings{} @@ -165,10 +167,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para settings.DataDiskSizeGb = 100 settings.DataDiskType = "PD_SSD" default: - return nil, fmt.Errorf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset) + return nil, util.NewAgentError(fmt.Sprintf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset), nil) } - return source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + resp, err := source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances/cloudsqlpgcreateinstances.go b/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances/cloudsqlpgcreateinstances.go index 93639c84d5..03093ca01d 100644 --- a/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances/cloudsqlpgcreateinstances.go +++ b/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances/cloudsqlpgcreateinstances.go @@ -17,12 +17,14 @@ package cloudsqlpgcreateinstances import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "strings" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" sqladmin "google.golang.org/api/sqladmin/v1" ) @@ -121,33 +123,33 @@ 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) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() project, ok := paramsMap["project"].(string) if !ok { - return nil, fmt.Errorf("missing 'project' parameter") + return nil, util.NewAgentError("missing 'project' parameter", nil) } name, ok := paramsMap["name"].(string) if !ok { - return nil, fmt.Errorf("missing 'name' parameter") + return nil, util.NewAgentError("missing 'name' parameter", nil) } dbVersion, ok := paramsMap["databaseVersion"].(string) if !ok { - return nil, fmt.Errorf("missing 'databaseVersion' parameter") + return nil, util.NewAgentError("missing 'databaseVersion' parameter", nil) } rootPassword, ok := paramsMap["rootPassword"].(string) if !ok { - return nil, fmt.Errorf("missing 'rootPassword' parameter") + return nil, util.NewAgentError("missing 'rootPassword' parameter", nil) } editionPreset, ok := paramsMap["editionPreset"].(string) if !ok { - return nil, fmt.Errorf("missing 'editionPreset' parameter") + return nil, util.NewAgentError("missing 'editionPreset' parameter", nil) } settings := sqladmin.Settings{} @@ -165,9 +167,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para settings.DataDiskSizeGb = 100 settings.DataDiskType = "PD_SSD" default: - return nil, fmt.Errorf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset) + return nil, util.NewAgentError(fmt.Sprintf("invalid 'editionPreset': %q. Must be either 'Production' or 'Development'", editionPreset), nil) } - return source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + resp, err := source.CreateInstance(ctx, project, name, dbVersion, rootPassword, settings, string(accessToken)) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck/cloudsqlpgupgradeprecheck.go b/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck/cloudsqlpgupgradeprecheck.go index f5d57750e6..eb2f89674b 100644 --- a/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck/cloudsqlpgupgradeprecheck.go +++ b/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck/cloudsqlpgupgradeprecheck.go @@ -17,12 +17,14 @@ package cloudsqlpgupgradeprecheck import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "time" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" sqladmin "google.golang.org/api/sqladmin/v1" ) @@ -132,31 +134,31 @@ 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) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() project, ok := paramsMap["project"].(string) if !ok || project == "" { - return nil, fmt.Errorf("missing or empty 'project' parameter") + return nil, util.NewAgentError("missing or empty 'project' parameter", nil) } instanceName, ok := paramsMap["instance"].(string) if !ok || instanceName == "" { - return nil, fmt.Errorf("missing or empty 'instance' parameter") + return nil, util.NewAgentError("missing or empty 'instance' parameter", nil) } targetVersion, ok := paramsMap["targetDatabaseVersion"].(string) if !ok || targetVersion == "" { // This should not happen due to the default value - return nil, fmt.Errorf("missing or empty 'targetDatabaseVersion' parameter") + return nil, util.NewAgentError("missing or empty 'targetDatabaseVersion' parameter", nil) } service, err := source.GetService(ctx, string(accessToken)) if err != nil { - return nil, fmt.Errorf("failed to get HTTP client from source: %w", err) + return nil, util.ProecessGcpError(err) } reqBody := &sqladmin.InstancesPreCheckMajorVersionUpgradeRequest{ @@ -168,7 +170,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para call := service.Instances.PreCheckMajorVersionUpgrade(project, instanceName, reqBody).Context(ctx) op, err := call.Do() if err != nil { - return nil, fmt.Errorf("failed to start pre-check operation: %w", err) + return nil, util.ProecessGcpError(err) } const pollTimeout = 20 * time.Second @@ -177,7 +179,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para for time.Now().Before(cutoffTime) { currentOp, err := service.Operations.Get(project, op.Name).Context(ctx).Do() if err != nil { - return nil, fmt.Errorf("failed to get operation status: %w", err) + return nil, util.ProecessGcpError(err) } if currentOp.Status == "DONE" { @@ -186,7 +188,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if currentOp.Error.Errors[0].Code != "" { errMsg = fmt.Sprintf("%s (Code: %s)", errMsg, currentOp.Error.Errors[0].Code) } - return nil, fmt.Errorf("%s", errMsg) + return nil, util.NewClientServerError(errMsg, http.StatusInternalServerError, fmt.Errorf("pre-check operation failed with error: %s", errMsg)) } var preCheckItems []*sqladmin.PreCheckResponse @@ -199,7 +201,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para select { case <-ctx.Done(): - return nil, ctx.Err() + return nil, util.NewClientServerError("timed out waiting for operation", http.StatusRequestTimeout, ctx.Err()) case <-time.After(5 * time.Second): } } diff --git a/internal/tools/dataplex/dataplexlookupentry/dataplexlookupentry.go b/internal/tools/dataplex/dataplexlookupentry/dataplexlookupentry.go index fdfe656eb8..726cbaad85 100644 --- a/internal/tools/dataplex/dataplexlookupentry/dataplexlookupentry.go +++ b/internal/tools/dataplex/dataplexlookupentry/dataplexlookupentry.go @@ -17,12 +17,14 @@ package dataplexlookupentry import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb" "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -110,10 +112,10 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() @@ -122,10 +124,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para view, _ := paramsMap["view"].(int) aspectTypeSlice, err := parameters.ConvertAnySliceToTyped(paramsMap["aspectTypes"].([]any), "string") if err != nil { - return nil, fmt.Errorf("can't convert aspectTypes to array of strings: %s", err) + return nil, util.NewAgentError(fmt.Sprintf("can't convert aspectTypes to array of strings: %s", err), err) } aspectTypes := aspectTypeSlice.([]string) - return source.LookupEntry(ctx, name, view, aspectTypes, entry) + resp, err := source.LookupEntry(ctx, name, view, aspectTypes, entry) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/dataplex/dataplexsearchaspecttypes/dataplexsearchaspecttypes.go b/internal/tools/dataplex/dataplexsearchaspecttypes/dataplexsearchaspecttypes.go index b57b598fca..a542e954e8 100644 --- a/internal/tools/dataplex/dataplexsearchaspecttypes/dataplexsearchaspecttypes.go +++ b/internal/tools/dataplex/dataplexsearchaspecttypes/dataplexsearchaspecttypes.go @@ -17,12 +17,14 @@ package dataplexsearchaspecttypes import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "cloud.google.com/go/dataplex/apiv1/dataplexpb" "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -93,16 +95,29 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() - query, _ := paramsMap["query"].(string) - pageSize, _ := paramsMap["pageSize"].(int) - orderBy, _ := paramsMap["orderBy"].(string) - return source.SearchAspectTypes(ctx, query, pageSize, orderBy) + query, ok := paramsMap["query"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'query' parameter: %v", paramsMap["query"]), nil) + } + pageSize, ok := paramsMap["pageSize"].(int) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'pageSize' parameter: %v", paramsMap["pageSize"]), nil) + } + orderBy, ok := paramsMap["orderBy"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'orderBy' parameter: %v", paramsMap["orderBy"]), nil) + } + resp, err := source.SearchAspectTypes(ctx, query, pageSize, orderBy) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/dataplex/dataplexsearchentries/dataplexsearchentries.go b/internal/tools/dataplex/dataplexsearchentries/dataplexsearchentries.go index b3dafbff98..8bb1bf4e93 100644 --- a/internal/tools/dataplex/dataplexsearchentries/dataplexsearchentries.go +++ b/internal/tools/dataplex/dataplexsearchentries/dataplexsearchentries.go @@ -17,12 +17,14 @@ package dataplexsearchentries import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError "cloud.google.com/go/dataplex/apiv1/dataplexpb" "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -93,16 +95,29 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } paramsMap := params.AsMap() - query, _ := paramsMap["query"].(string) - pageSize, _ := paramsMap["pageSize"].(int) - orderBy, _ := paramsMap["orderBy"].(string) - return source.SearchEntries(ctx, query, pageSize, orderBy) + query, ok := paramsMap["query"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'query' parameter: %v", paramsMap["query"]), nil) + } + pageSize, ok := paramsMap["pageSize"].(int) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'pageSize' parameter: %v", paramsMap["pageSize"]), nil) + } + orderBy, ok := paramsMap["orderBy"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'orderBy' parameter: %v", paramsMap["orderBy"]), nil) + } + resp, err := source.SearchEntries(ctx, query, pageSize, orderBy) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments.go b/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments.go index 20c6163335..ddc66ff9b1 100644 --- a/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments.go +++ b/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments.go @@ -17,13 +17,15 @@ package firestoreadddocuments import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError firestoreapi "cloud.google.com/go/firestore" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + fsUtil "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -128,32 +130,32 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } mapParams := params.AsMap() // Get collection path collectionPath, ok := mapParams[collectionPathKey].(string) if !ok || collectionPath == "" { - return nil, fmt.Errorf("invalid or missing '%s' parameter", collectionPathKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter", collectionPathKey), nil) } // Validate collection path - if err := util.ValidateCollectionPath(collectionPath); err != nil { - return nil, fmt.Errorf("invalid collection path: %w", err) + if err := fsUtil.ValidateCollectionPath(collectionPath); err != nil { + return nil, util.NewAgentError(fmt.Sprintf("invalid collection path: %v", err), err) } // Get document data documentDataRaw, ok := mapParams[documentDataKey] if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter", documentDataKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter", documentDataKey), nil) } // Convert the document data from JSON format to Firestore format // The client is passed to handle referenceValue types - documentData, err := util.JSONToFirestoreValue(documentDataRaw, source.FirestoreClient()) + documentData, err := fsUtil.JSONToFirestoreValue(documentDataRaw, source.FirestoreClient()) if err != nil { - return nil, fmt.Errorf("failed to convert document data: %w", err) + return nil, util.NewAgentError(fmt.Sprintf("failed to convert document data: %v", err), err) } // Get return document data flag @@ -161,7 +163,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if val, ok := mapParams[returnDocumentDataKey].(bool); ok { returnData = val } - return source.AddDocuments(ctx, collectionPath, documentData, returnData) + resp, err := source.AddDocuments(ctx, collectionPath, documentData, returnData) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments.go b/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments.go index 1610c6a038..1eefeadda7 100644 --- a/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments.go +++ b/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments.go @@ -17,13 +17,15 @@ package firestoredeletedocuments import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError firestoreapi "cloud.google.com/go/firestore" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + fsUtil "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -94,39 +96,43 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } mapParams := params.AsMap() documentPathsRaw, ok := mapParams[documentPathsKey].([]any) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected an array", documentPathsKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected an array", documentPathsKey), nil) } if len(documentPathsRaw) == 0 { - return nil, fmt.Errorf("'%s' parameter cannot be empty", documentPathsKey) + return nil, util.NewAgentError(fmt.Sprintf("'%s' parameter cannot be empty", documentPathsKey), nil) } // Use ConvertAnySliceToTyped to convert the slice typedSlice, err := parameters.ConvertAnySliceToTyped(documentPathsRaw, "string") if err != nil { - return nil, fmt.Errorf("failed to convert document paths: %w", err) + return nil, util.NewAgentError(fmt.Sprintf("failed to convert document paths: %v", err), err) } documentPaths, ok := typedSlice.([]string) if !ok { - return nil, fmt.Errorf("unexpected type conversion error for document paths") + return nil, util.NewAgentError("unexpected type conversion error for document paths", nil) } // Validate each document path for i, path := range documentPaths { - if err := util.ValidateDocumentPath(path); err != nil { - return nil, fmt.Errorf("invalid document path at index %d: %w", i, err) + if err := fsUtil.ValidateDocumentPath(path); err != nil { + return nil, util.NewAgentError(fmt.Sprintf("invalid document path at index %d: %v", i, err), err) } } - return source.DeleteDocuments(ctx, documentPaths) + resp, err := source.DeleteDocuments(ctx, documentPaths) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments.go b/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments.go index 5ccc68ef9b..4535299142 100644 --- a/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments.go +++ b/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments.go @@ -17,13 +17,15 @@ package firestoregetdocuments import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError firestoreapi "cloud.google.com/go/firestore" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + fsUtil "github.com/googleapis/genai-toolbox/internal/tools/firestore/util" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -94,40 +96,44 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } mapParams := params.AsMap() documentPathsRaw, ok := mapParams[documentPathsKey].([]any) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected an array", documentPathsKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected an array", documentPathsKey), nil) } if len(documentPathsRaw) == 0 { - return nil, fmt.Errorf("'%s' parameter cannot be empty", documentPathsKey) + return nil, util.NewAgentError(fmt.Sprintf("'%s' parameter cannot be empty", documentPathsKey), nil) } // Use ConvertAnySliceToTyped to convert the slice typedSlice, err := parameters.ConvertAnySliceToTyped(documentPathsRaw, "string") if err != nil { - return nil, fmt.Errorf("failed to convert document paths: %w", err) + return nil, util.NewAgentError(fmt.Sprintf("failed to convert document paths: %v", err), err) } documentPaths, ok := typedSlice.([]string) if !ok { - return nil, fmt.Errorf("unexpected type conversion error for document paths") + return nil, util.NewAgentError("unexpected type conversion error for document paths", nil) } // Validate each document path for i, path := range documentPaths { - if err := util.ValidateDocumentPath(path); err != nil { - return nil, fmt.Errorf("invalid document path at index %d: %w", i, err) + if err := fsUtil.ValidateDocumentPath(path); err != nil { + return nil, util.NewAgentError(fmt.Sprintf("invalid document path at index %d: %v", i, err), err) } } - return source.GetDocuments(ctx, documentPaths) + resp, err := source.GetDocuments(ctx, documentPaths) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/tools/firestore/firestoregetrules/firestoregetrules.go b/internal/tools/firestore/firestoregetrules/firestoregetrules.go index 13453c4e30..56ecba6316 100644 --- a/internal/tools/firestore/firestoregetrules/firestoregetrules.go +++ b/internal/tools/firestore/firestoregetrules/firestoregetrules.go @@ -17,11 +17,13 @@ package firestoregetrules import ( "context" "fmt" + "net/http" // Added for http.StatusInternalServerError yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" // Added for util.ToolboxError "github.com/googleapis/genai-toolbox/internal/util/parameters" "google.golang.org/api/firebaserules/v1" ) @@ -92,12 +94,16 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { - return nil, err + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) } - return source.GetRules(ctx) + resp, err := source.GetRules(ctx) + if err != nil { + return nil, util.ProecessGcpError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { diff --git a/internal/util/errors.go b/internal/util/errors.go index 625ccbb0c0..57445cc587 100644 --- a/internal/util/errors.go +++ b/internal/util/errors.go @@ -16,6 +16,7 @@ import ( "errors" "fmt" "net/http" + "strings" "google.golang.org/api/googleapi" ) @@ -66,6 +67,8 @@ func NewClientServerError(msg string, code int, cause error) *ClientServerError return &ClientServerError{Msg: msg, Code: code, Cause: cause} } +// ProecessGcpError catches auth related errors and return 401/403 error codes +// Returns AgentError for all other errors func ProecessGcpError(err error) ToolboxError { var gErr *googleapi.Error if errors.As(err, &gErr) { @@ -86,3 +89,34 @@ func ProecessGcpError(err error) ToolboxError { } return NewAgentError("error processing GCP request", err) } + +// ProcessGeneralError handles generic errors by inspecting the error string +// for common status code patterns. +func ProcessGeneralError(err error) ToolboxError { + if err == nil { + return nil + } + + errStr := err.Error() + + // Check for Unauthorized + if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "status 401") { + return NewClientServerError( + "failed to access resource", + http.StatusUnauthorized, + err, + ) + } + + // Check for Forbidden + if strings.Contains(errStr, "Error 403") || strings.Contains(errStr, "status 403") { + return NewClientServerError( + "failed to access resource", + http.StatusForbidden, + err, + ) + } + + // Default to AgentError for logical failures (task execution failed) + return NewAgentError("error processing request", err) +}