diff --git a/internal/tools/alloydbainl/alloydbainl.go b/internal/tools/alloydbainl/alloydbainl.go index 4a0b8b9ba8..98cf20870b 100644 --- a/internal/tools/alloydbainl/alloydbainl.go +++ b/internal/tools/alloydbainl/alloydbainl.go @@ -17,12 +17,14 @@ package alloydbainl import ( "context" "fmt" + "net/http" "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" "github.com/googleapis/genai-toolbox/internal/util/parameters" "github.com/jackc/pgx/v5/pgxpool" ) @@ -127,10 +129,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) } sliceParams := params.AsSlice() @@ -143,7 +145,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para resp, err := source.RunSQL(ctx, t.Statement, allParamValues) if err != nil { - return nil, fmt.Errorf("%w. Query: %v , Values: %v. Toolbox v0.19.0+ is only compatible with AlloyDB AI NL v1.0.3+. Please ensure that you are using the latest AlloyDB AI NL extension", err, t.Statement, allParamValues) + return nil, util.NewClientServerError(fmt.Sprintf("error running SQL query: %v. Query: %v , Values: %v. Toolbox v0.19.0+ is only compatible with AlloyDB AI NL v1.0.3+. Please ensure that you are using the latest AlloyDB AI NL extension", err, t.Statement, allParamValues), http.StatusBadRequest, err) } return resp, nil } diff --git a/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go b/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go index e14cfea511..157740c1bb 100644 --- a/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go +++ b/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go @@ -18,6 +18,7 @@ import ( "context" "encoding/json" "fmt" + "net/http" "strings" bigqueryapi "cloud.google.com/go/bigquery" @@ -152,25 +153,25 @@ 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() sql, ok := paramsMap["sql"].(string) if !ok { - return nil, fmt.Errorf("unable to cast sql parameter %s", paramsMap["sql"]) + return nil, util.NewAgentError(fmt.Sprintf("unable to cast sql parameter %s", paramsMap["sql"]), nil) } dryRun, ok := paramsMap["dry_run"].(bool) if !ok { - return nil, fmt.Errorf("unable to cast dry_run parameter %s", paramsMap["dry_run"]) + return nil, util.NewAgentError(fmt.Sprintf("unable to cast dry_run parameter %s", paramsMap["dry_run"]), nil) } bqClient, restService, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } var connProps []*bigqueryapi.ConnectionProperty @@ -178,7 +179,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if source.BigQueryWriteMode() == bigqueryds.WriteModeProtected { session, err = source.BigQuerySession()(ctx) if err != nil { - return nil, fmt.Errorf("failed to get BigQuery session for protected mode: %w", err) + return nil, util.NewClientServerError("failed to get BigQuery session for protected mode", http.StatusInternalServerError, err) } connProps = []*bigqueryapi.ConnectionProperty{ {Key: "session_id", Value: session.ID}, @@ -187,7 +188,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para dryRunJob, err := bqutil.DryRunQuery(ctx, restService, bqClient.Project(), bqClient.Location, sql, nil, connProps) if err != nil { - return nil, fmt.Errorf("query validation failed: %w", err) + return nil, util.NewClientServerError("query validation failed", http.StatusInternalServerError, err) } statementType := dryRunJob.Statistics.Query.StatementType @@ -195,13 +196,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para switch source.BigQueryWriteMode() { case bigqueryds.WriteModeBlocked: if statementType != "SELECT" { - return nil, fmt.Errorf("write mode is 'blocked', only SELECT statements are allowed") + return nil, util.NewAgentError("write mode is 'blocked', only SELECT statements are allowed", nil) } case bigqueryds.WriteModeProtected: if dryRunJob.Configuration != nil && dryRunJob.Configuration.Query != nil { if dest := dryRunJob.Configuration.Query.DestinationTable; dest != nil && dest.DatasetId != session.DatasetID { - return nil, fmt.Errorf("protected write mode only supports SELECT statements, or write operations in the anonymous "+ - "dataset of a BigQuery session, but destination was %q", dest.DatasetId) + return nil, util.NewAgentError(fmt.Sprintf("protected write mode only supports SELECT statements, or write operations in the anonymous "+ + "dataset of a BigQuery session, but destination was %q", dest.DatasetId), nil) } } } @@ -209,11 +210,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if len(source.BigQueryAllowedDatasets()) > 0 { switch statementType { case "CREATE_SCHEMA", "DROP_SCHEMA", "ALTER_SCHEMA": - return nil, fmt.Errorf("dataset-level operations like '%s' are not allowed when dataset restrictions are in place", statementType) + return nil, util.NewAgentError(fmt.Sprintf("dataset-level operations like '%s' are not allowed when dataset restrictions are in place", statementType), nil) case "CREATE_FUNCTION", "CREATE_TABLE_FUNCTION", "CREATE_PROCEDURE": - return nil, fmt.Errorf("creating stored routines ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType) + return nil, util.NewAgentError(fmt.Sprintf("creating stored routines ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType), nil) case "CALL": - return nil, fmt.Errorf("calling stored procedures ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType) + return nil, util.NewAgentError(fmt.Sprintf("calling stored procedures ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType), nil) } // Use a map to avoid duplicate table names. @@ -244,7 +245,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para parsedTables, parseErr := bqutil.TableParser(sql, source.BigQueryClient().Project()) if parseErr != nil { // If parsing fails (e.g., EXECUTE IMMEDIATE), we cannot guarantee safety, so we must fail. - return nil, fmt.Errorf("could not parse tables from query to validate against allowed datasets: %w", parseErr) + return nil, util.NewAgentError("could not parse tables from query to validate against allowed datasets", parseErr) } tableNames = parsedTables } @@ -254,7 +255,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if len(parts) == 3 { projectID, datasetID := parts[0], parts[1] if !source.IsDatasetAllowed(projectID, datasetID) { - return nil, fmt.Errorf("query accesses dataset '%s.%s', which is not in the allowed list", projectID, datasetID) + return nil, util.NewAgentError(fmt.Sprintf("query accesses dataset '%s.%s', which is not in the allowed list", projectID, datasetID), nil) } } } @@ -264,7 +265,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if dryRunJob != nil { jobJSON, err := json.MarshalIndent(dryRunJob, "", " ") if err != nil { - return nil, fmt.Errorf("failed to marshal dry run job to JSON: %w", err) + return nil, util.NewClientServerError("failed to marshal dry run job to JSON", http.StatusInternalServerError, err) } return string(jobJSON), nil } @@ -275,10 +276,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para // Log the query executed for debugging. logger, err := util.LoggerFromContext(ctx) if err != nil { - return nil, fmt.Errorf("error getting logger: %s", err) + return nil, util.NewClientServerError("error getting logger", http.StatusInternalServerError, err) } logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", resourceType, sql)) - return source.RunSQL(ctx, bqClient, sql, statementType, nil, connProps) + resp, err := source.RunSQL(ctx, bqClient, sql, statementType, nil, connProps) + if err != nil { + return nil, util.NewClientServerError("error running sql", http.StatusInternalServerError, 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/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index b3844d20cd..ab9577276e 100644 --- a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -17,6 +17,7 @@ package bigquerygetdatasetinfo import ( "context" "fmt" + "net/http" bigqueryapi "cloud.google.com/go/bigquery" yaml "github.com/goccy/go-yaml" @@ -24,6 +25,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/googleapis/genai-toolbox/internal/util/parameters" bigqueryrestapi "google.golang.org/api/bigquery/v2" ) @@ -120,38 +122,38 @@ type Tool struct { 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() projectId, ok := mapParams[projectKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey) + // Updated: Use fmt.Sprintf for formatting, pass nil as cause + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil) } datasetId, ok := mapParams[datasetKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil) } bqClient, _, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } if !source.IsDatasetAllowed(projectId, datasetId) { - return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) + return nil, util.NewClientServerError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), http.StatusInternalServerError, nil) } dsHandle := bqClient.DatasetInProject(projectId, datasetId) metadata, err := dsHandle.Metadata(ctx) if err != nil { - return nil, fmt.Errorf("failed to get metadata for dataset %s (in project %s): %w", datasetId, projectId, err) + return nil, util.ProecessGcpError(err) } return metadata, nil diff --git a/internal/tools/bigquery/bigquerygettableinfo/bigquerygettableinfo.go b/internal/tools/bigquery/bigquerygettableinfo/bigquerygettableinfo.go index b7131df89f..61a31bd137 100644 --- a/internal/tools/bigquery/bigquerygettableinfo/bigquerygettableinfo.go +++ b/internal/tools/bigquery/bigquerygettableinfo/bigquerygettableinfo.go @@ -17,6 +17,7 @@ package bigquerygettableinfo import ( "context" "fmt" + "net/http" bigqueryapi "cloud.google.com/go/bigquery" yaml "github.com/goccy/go-yaml" @@ -24,6 +25,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/googleapis/genai-toolbox/internal/util/parameters" bigqueryrestapi "google.golang.org/api/bigquery/v2" ) @@ -125,35 +127,35 @@ 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() projectId, ok := mapParams[projectKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil) } datasetId, ok := mapParams[datasetKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil) } tableId, ok := mapParams[tableKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", tableKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", tableKey), nil) } if !source.IsDatasetAllowed(projectId, datasetId) { - return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) + return nil, util.NewClientServerError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), http.StatusInternalServerError, nil) } bqClient, _, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } dsHandle := bqClient.DatasetInProject(projectId, datasetId) @@ -161,7 +163,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para metadata, err := tableHandle.Metadata(ctx) if err != nil { - return nil, fmt.Errorf("failed to get metadata for table %s.%s.%s: %w", projectId, datasetId, tableId, err) + return nil, util.ProecessGcpError(err) } return metadata, nil diff --git a/internal/tools/bigquery/bigquerylistdatasetids/bigquerylistdatasetids.go b/internal/tools/bigquery/bigquerylistdatasetids/bigquerylistdatasetids.go index 186ad7be54..0ff199137d 100644 --- a/internal/tools/bigquery/bigquerylistdatasetids/bigquerylistdatasetids.go +++ b/internal/tools/bigquery/bigquerylistdatasetids/bigquerylistdatasetids.go @@ -17,12 +17,14 @@ package bigquerylistdatasetids import ( "context" "fmt" + "net/http" bigqueryapi "cloud.google.com/go/bigquery" 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" "github.com/googleapis/genai-toolbox/internal/util/parameters" bigqueryrestapi "google.golang.org/api/bigquery/v2" "google.golang.org/api/iterator" @@ -120,10 +122,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) } if len(source.BigQueryAllowedDatasets()) > 0 { @@ -132,12 +134,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para mapParams := params.AsMap() projectId, ok := mapParams[projectKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil) } bqClient, _, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } datasetIterator := bqClient.Datasets(ctx) datasetIterator.ProjectID = projectId @@ -149,7 +151,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para break } if err != nil { - return nil, fmt.Errorf("unable to iterate through datasets: %w", err) + return nil, util.ProecessGcpError(err) } // Remove leading and trailing quotes diff --git a/internal/tools/bigquery/bigquerylisttableids/bigquerylisttableids.go b/internal/tools/bigquery/bigquerylisttableids/bigquerylisttableids.go index 4390a89961..a336ae86e9 100644 --- a/internal/tools/bigquery/bigquerylisttableids/bigquerylisttableids.go +++ b/internal/tools/bigquery/bigquerylisttableids/bigquerylisttableids.go @@ -17,6 +17,7 @@ package bigquerylisttableids import ( "context" "fmt" + "net/http" bigqueryapi "cloud.google.com/go/bigquery" yaml "github.com/goccy/go-yaml" @@ -24,6 +25,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/googleapis/genai-toolbox/internal/util/parameters" bigqueryrestapi "google.golang.org/api/bigquery/v2" "google.golang.org/api/iterator" @@ -123,31 +125,30 @@ type Tool struct { 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() projectId, ok := mapParams[projectKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil) } datasetId, ok := mapParams[datasetKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil) } if !source.IsDatasetAllowed(projectId, datasetId) { - return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) + return nil, util.NewClientServerError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), http.StatusInternalServerError, nil) } bqClient, _, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } dsHandle := bqClient.DatasetInProject(projectId, datasetId) @@ -160,7 +161,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para break } if err != nil { - return nil, fmt.Errorf("failed to iterate through tables in dataset %s.%s: %w", projectId, datasetId, err) + return nil, util.ProecessGcpError(err) } // Remove leading and trailing quotes diff --git a/internal/tools/bigquery/bigquerysql/bigquerysql.go b/internal/tools/bigquery/bigquerysql/bigquerysql.go index 78685deaa3..a86f9e37bf 100644 --- a/internal/tools/bigquery/bigquerysql/bigquerysql.go +++ b/internal/tools/bigquery/bigquerysql/bigquerysql.go @@ -17,6 +17,7 @@ package bigquerysql import ( "context" "fmt" + "net/http" "reflect" "strings" @@ -27,6 +28,7 @@ import ( bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/googleapis/genai-toolbox/internal/util/parameters" bigqueryrestapi "google.golang.org/api/bigquery/v2" ) @@ -103,11 +105,10 @@ type Tool struct { 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) } highLevelParams := make([]bigqueryapi.QueryParameter, 0, len(t.Parameters)) @@ -116,7 +117,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para paramsMap := params.AsMap() newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap) if err != nil { - return nil, fmt.Errorf("unable to extract template params %w", err) + return nil, util.NewAgentError("unable to extract template params", err) } for _, p := range t.Parameters { @@ -127,13 +128,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if arrayParam, ok := p.(*parameters.ArrayParameter); ok { arrayParamValue, ok := value.([]any) if !ok { - return nil, fmt.Errorf("unable to convert parameter `%s` to []any", name) + return nil, util.NewAgentError(fmt.Sprintf("unable to convert parameter `%s` to []any", name), nil) } itemType := arrayParam.GetItems().GetType() var err error value, err = parameters.ConvertAnySliceToTyped(arrayParamValue, itemType) if err != nil { - return nil, fmt.Errorf("unable to convert parameter `%s` from []any to typed slice: %w", name, err) + return nil, util.NewAgentError(fmt.Sprintf("unable to convert parameter `%s` from []any to typed slice", name), err) } } @@ -161,7 +162,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para lowLevelParam.ParameterType.Type = "ARRAY" itemType, err := bqutil.BQTypeStringFromToolType(arrayParam.GetItems().GetType()) if err != nil { - return nil, err + return nil, util.NewAgentError("unable to get BigQuery type from tool parameter type", err) } lowLevelParam.ParameterType.ArrayType = &bigqueryrestapi.QueryParameterType{Type: itemType} @@ -178,7 +179,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para // Handle scalar types based on their defined type. bqType, err := bqutil.BQTypeStringFromToolType(p.GetType()) if err != nil { - return nil, err + return nil, util.NewAgentError("unable to get BigQuery type from tool parameter type", err) } lowLevelParam.ParameterType.Type = bqType lowLevelParam.ParameterValue.Value = fmt.Sprintf("%v", value) @@ -190,7 +191,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if source.BigQuerySession() != nil { session, err := source.BigQuerySession()(ctx) if err != nil { - return nil, fmt.Errorf("failed to get BigQuery session: %w", err) + return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err) } if session != nil { // Add session ID to the connection properties for subsequent calls. @@ -200,17 +201,20 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para bqClient, restService, err := source.RetrieveClientAndService(accessToken) if err != nil { - return nil, err + return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err) } dryRunJob, err := bqutil.DryRunQuery(ctx, restService, bqClient.Project(), bqClient.Location, newStatement, lowLevelParams, connProps) if err != nil { - return nil, fmt.Errorf("query validation failed: %w", err) + return nil, util.ProecessGcpError(err) } statementType := dryRunJob.Statistics.Query.StatementType - - return source.RunSQL(ctx, bqClient, newStatement, statementType, highLevelParams, connProps) + resp, err := source.RunSQL(ctx, bqClient, newStatement, statementType, highLevelParams, connProps) + 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/bigtable/bigtable.go b/internal/tools/bigtable/bigtable.go index 4c47ca945e..222cb6021e 100644 --- a/internal/tools/bigtable/bigtable.go +++ b/internal/tools/bigtable/bigtable.go @@ -17,12 +17,14 @@ package bigtable import ( "context" "fmt" + "net/http" "cloud.google.com/go/bigtable" 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" "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -96,24 +98,28 @@ type Tool struct { 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() newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap) if err != nil { - return nil, fmt.Errorf("unable to extract template params %w", err) + return nil, util.NewAgentError("unable to extract template params", err) } newParams, err := parameters.GetParams(t.Parameters, paramsMap) if err != nil { - return nil, fmt.Errorf("unable to extract standard params %w", err) + return nil, util.NewAgentError("unable to extract standard params", err) } - return source.RunSQL(ctx, newStatement, t.Parameters, newParams) + + resp, err := source.RunSQL(ctx, newStatement, t.Parameters, newParams) + 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/cloudsql/cloudsqlcreateusers/cloudsqlcreateusers.go b/internal/tools/cloudsql/cloudsqlcreateusers/cloudsqlcreateusers.go index 1594b81dd3..811418e500 100644 --- a/internal/tools/cloudsql/cloudsqlcreateusers/cloudsqlcreateusers.go +++ b/internal/tools/cloudsql/cloudsqlcreateusers/cloudsqlcreateusers.go @@ -17,11 +17,13 @@ package cloudsqlcreateusers import ( "context" "fmt" + "net/http" "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" "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -119,30 +121,34 @@ 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) } instance, ok := paramsMap["instance"].(string) if !ok { - return nil, fmt.Errorf("missing 'instance' parameter") + return nil, util.NewAgentError("missing 'instance' parameter", nil) } name, ok := paramsMap["name"].(string) if !ok { - return nil, fmt.Errorf("missing 'name' parameter") + return nil, util.NewAgentError("missing 'name' parameter", nil) } iamUser, _ := paramsMap["iamUser"].(bool) password, _ := paramsMap["password"].(string) - return source.CreateUsers(ctx, project, instance, name, password, iamUser, string(accessToken)) + resp, err := source.CreateUsers(ctx, project, instance, name, password, iamUser, 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) {