diff --git a/internal/tools/cassandra/cassandracql/cassandracql.go b/internal/tools/cassandra/cassandracql/cassandracql.go index 6dcd2a013a..2cdcd92e57 100644 --- a/internal/tools/cassandra/cassandracql/cassandracql.go +++ b/internal/tools/cassandra/cassandracql/cassandracql.go @@ -4,7 +4,7 @@ // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // -// http://www.apache.org/licenses/LICENSE-2.0 +// http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, @@ -17,12 +17,14 @@ package cassandracql import ( "context" "fmt" + "net/http" gocql "github.com/apache/cassandra-gocql-driver/v2" 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" ) @@ -107,23 +109,27 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool { } // Invoke implements tools.Tool. -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, newParams) + resp, err := source.RunSQL(ctx, newStatement, newParams) + if err != nil { + return nil, util.ProcessGeneralError(err) + } + return resp, nil } // Manifest implements tools.Tool. diff --git a/internal/tools/clickhouse/clickhouseexecutesql/clickhouseexecutesql.go b/internal/tools/clickhouse/clickhouseexecutesql/clickhouseexecutesql.go index eefa02c6fa..8b69d71b60 100644 --- a/internal/tools/clickhouse/clickhouseexecutesql/clickhouseexecutesql.go +++ b/internal/tools/clickhouse/clickhouseexecutesql/clickhouseexecutesql.go @@ -17,11 +17,13 @@ package clickhouse import ( "context" "fmt" + "net/http" 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" ) @@ -87,18 +89,22 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token 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) } - return source.RunSQL(ctx, sql, nil) + resp, err := source.RunSQL(ctx, sql, nil) + if err != nil { + return nil, util.ProcessGeneralError(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/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go b/internal/tools/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go index 317c462935..1a710550fc 100644 --- a/internal/tools/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go +++ b/internal/tools/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go @@ -17,11 +17,13 @@ package clickhouse import ( "context" "fmt" + "net/http" 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" ) @@ -86,10 +88,10 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token 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) } // Query to list all databases @@ -97,7 +99,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para out, err := source.RunSQL(ctx, query, nil) if err != nil { - return nil, err + return nil, util.ProcessGeneralError(err) } return out, nil @@ -129,4 +131,4 @@ func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, func (t Tool) GetParameters() parameters.Parameters { return t.Parameters -} +} \ No newline at end of file diff --git a/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go index 492bc281ad..10fb432d55 100644 --- a/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go +++ b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go @@ -17,11 +17,13 @@ package clickhouse import ( "context" "fmt" + "net/http" 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" ) @@ -90,34 +92,37 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token 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() database, ok := mapParams[databaseKey].(string) if !ok { - return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", databaseKey) + return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", databaseKey), nil) } + // Query to list all tables in the specified database + // Note: formatting identifier directly is risky if input is untrusted, but standard for this tool structure. query := fmt.Sprintf("SHOW TABLES FROM %s", database) out, err := source.RunSQL(ctx, query, nil) if err != nil { - return nil, err + return nil, util.ProcessGeneralError(err) } res, ok := out.([]any) if !ok { - return nil, fmt.Errorf("unable to convert result to list") + return nil, util.NewClientServerError("unable to convert result to list", http.StatusInternalServerError, nil) } + var tables []map[string]any for _, item := range res { tableMap, ok := item.(map[string]any) if !ok { - return nil, fmt.Errorf("unexpected type in result: got %T, want map[string]any", item) + return nil, util.NewClientServerError(fmt.Sprintf("unexpected type in result: got %T, want map[string]any", item), http.StatusInternalServerError, nil) } tableMap["database"] = database tables = append(tables, tableMap) diff --git a/internal/tools/clickhouse/clickhousesql/clickhousesql.go b/internal/tools/clickhouse/clickhousesql/clickhousesql.go index 10645d309a..aafd98b2e0 100644 --- a/internal/tools/clickhouse/clickhousesql/clickhousesql.go +++ b/internal/tools/clickhouse/clickhousesql/clickhousesql.go @@ -17,11 +17,13 @@ package clickhouse import ( "context" "fmt" + "net/http" 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" ) @@ -88,24 +90,28 @@ func (t Tool) ToConfig() tools.ToolConfig { return t.Config } -func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) { +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token 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, newParams) + resp, err := source.RunSQL(ctx, newStatement, newParams) + if err != nil { + return nil, util.ProcessGeneralError(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/couchbase/couchbase.go b/internal/tools/couchbase/couchbase.go index b15515d623..439d7cb053 100644 --- a/internal/tools/couchbase/couchbase.go +++ b/internal/tools/couchbase/couchbase.go @@ -17,12 +17,14 @@ package couchbase import ( "context" "fmt" + "net/http" "github.com/couchbase/gocb/v2" 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" ) @@ -58,7 +60,6 @@ type Config struct { TemplateParameters parameters.Parameters `yaml:"templateParameters"` } -// validate interface var _ tools.ToolConfig = Config{} func (cfg Config) ToolConfigType() string { @@ -72,7 +73,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) } mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil) - // finish tool setup t := Tool{ Config: cfg, AllParams: allParameters, @@ -82,7 +82,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return t, nil } -// validate interface var _ tools.Tool = Tool{} type Tool struct { @@ -96,23 +95,28 @@ 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) } namedParamsMap := params.AsMap() newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, namedParamsMap) 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, namedParamsMap) 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(newStatement, newParams) + + resp, err := source.RunSQL(newStatement, newParams) + if err != nil { + return nil, util.ProcessGeneralError(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/dataform/dataformcompilelocal/dataformcompilelocal.go b/internal/tools/dataform/dataformcompilelocal/dataformcompilelocal.go index 61b77e79cf..6f8bc383d1 100644 --- a/internal/tools/dataform/dataformcompilelocal/dataformcompilelocal.go +++ b/internal/tools/dataform/dataformcompilelocal/dataformcompilelocal.go @@ -24,6 +24,7 @@ import ( "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" ) @@ -86,18 +87,19 @@ 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) { paramsMap := params.AsMap() projectDir, ok := paramsMap["project_dir"].(string) if !ok || projectDir == "" { - return nil, fmt.Errorf("error casting 'project_dir' to string or invalid value") + return nil, util.NewAgentError("error casting 'project_dir' to string or invalid value", nil) } cmd := exec.CommandContext(ctx, "dataform", "compile", projectDir, "--json") output, err := cmd.CombinedOutput() if err != nil { - return nil, fmt.Errorf("error executing dataform compile: %w\nOutput: %s", err, string(output)) + // Compilation failures are considered AgentErrors (invalid user code/project) + return nil, util.NewAgentError(fmt.Sprintf("error executing dataform compile: %v\nOutput: %s", err, string(output)), err) } return strings.TrimSpace(string(output)), nil diff --git a/internal/tools/dgraph/dgraph.go b/internal/tools/dgraph/dgraph.go index d5e4cb72bf..40998e9655 100644 --- a/internal/tools/dgraph/dgraph.go +++ b/internal/tools/dgraph/dgraph.go @@ -17,12 +17,14 @@ package dgraph import ( "context" "fmt" + "net/http" 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/sources/dgraph" "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/googleapis/genai-toolbox/internal/util/parameters" ) @@ -91,12 +93,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.RunSQL(t.Statement, params, t.IsQuery, t.Timeout) + resp, err := source.RunSQL(t.Statement, params, t.IsQuery, t.Timeout) + if err != nil { + return nil, util.ProcessGeneralError(err) + } + return resp, nil } func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { @@ -125,4 +131,4 @@ func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, func (t Tool) GetParameters() parameters.Parameters { return t.Parameters -} +} \ No newline at end of file diff --git a/internal/tools/firebird/firebirdexecutesql/firebirdexecutesql.go b/internal/tools/firebird/firebirdexecutesql/firebirdexecutesql.go index 40e9195ee7..b1d97f1235 100644 --- a/internal/tools/firebird/firebirdexecutesql/firebirdexecutesql.go +++ b/internal/tools/firebird/firebirdexecutesql/firebirdexecutesql.go @@ -18,6 +18,7 @@ import ( "context" "database/sql" "fmt" + "net/http" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" @@ -90,25 +91,30 @@ 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) + sqlStr, ok := paramsMap["sql"].(string) if !ok { - return nil, fmt.Errorf("unable to get cast %s", paramsMap["sql"]) + return nil, util.NewAgentError(fmt.Sprintf("unable to cast parameter 'sql' to string: %v", paramsMap["sql"]), nil) } // 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, sql, nil) + logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", resourceType, sqlStr)) + + resp, err := source.RunSQL(ctx, sqlStr, nil) + if err != nil { + return nil, util.ProcessGeneralError(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/firebird/firebirdsql/firebirdsql.go b/internal/tools/firebird/firebirdsql/firebirdsql.go index 73c455ccb6..fbadc1c2a1 100644 --- a/internal/tools/firebird/firebirdsql/firebirdsql.go +++ b/internal/tools/firebird/firebirdsql/firebirdsql.go @@ -18,12 +18,14 @@ import ( "context" "database/sql" "fmt" + "net/http" "strings" "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" ) @@ -98,21 +100,21 @@ 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() statement, 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) } namedArgs := make([]any, 0, len(newParams)) @@ -127,7 +129,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para namedArgs = append(namedArgs, value) } } - return source.RunSQL(ctx, statement, namedArgs) + + resp, err := source.RunSQL(ctx, statement, namedArgs) + if err != nil { + return nil, util.ProcessGeneralError(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/http/http.go b/internal/tools/http/http.go index 0b3b49e383..59bea51ea0 100644 --- a/internal/tools/http/http.go +++ b/internal/tools/http/http.go @@ -17,18 +17,18 @@ import ( "bytes" "context" "fmt" + "maps" "net/http" "net/url" "slices" "strings" - - "maps" "text/template" 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" ) @@ -226,10 +226,10 @@ func getHeaders(headerParams parameters.Parameters, defaultHeaders map[string]st return allHeaders, nil } -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() @@ -237,27 +237,35 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para // Calculate request body requestBody, err := getRequestBody(t.BodyParams, t.RequestBody, paramsMap) if err != nil { - return nil, fmt.Errorf("error populating request body: %s", err) + return nil, util.NewAgentError("error populating request body", err) } // Calculate URL urlString, err := getURL(source.HttpBaseURL(), t.Path, t.PathParams, t.QueryParams, source.HttpQueryParams(), paramsMap) if err != nil { - return nil, fmt.Errorf("error populating path parameters: %s", err) + return nil, util.NewAgentError("error populating path parameters", err) } - req, _ := http.NewRequest(string(t.Method), urlString, strings.NewReader(requestBody)) + req, err := http.NewRequestWithContext(ctx, string(t.Method), urlString, strings.NewReader(requestBody)) + if err != nil { + return nil, util.NewClientServerError("error creating http request", http.StatusInternalServerError, err) + } // Calculate request headers allHeaders, err := getHeaders(t.HeaderParams, t.Headers, paramsMap) if err != nil { - return nil, fmt.Errorf("error populating request headers: %s", err) + return nil, util.NewAgentError("error populating request headers", err) } // Set request headers for k, v := range allHeaders { req.Header.Set(k, v) } - return source.RunRequest(req) + + resp, err := source.RunRequest(req) + if err != nil { + return nil, util.ProcessGeneralError(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/looker/lookeradddashboardelement/lookeradddashboardelement.go b/internal/tools/looker/lookeradddashboardelement/lookeradddashboardelement.go index 5b1103c102..5c6a6e5880 100644 --- a/internal/tools/looker/lookeradddashboardelement/lookeradddashboardelement.go +++ b/internal/tools/looker/lookeradddashboardelement/lookeradddashboardelement.go @@ -16,6 +16,7 @@ package lookeradddashboardelement import ( "context" "fmt" + "net/http" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" @@ -134,58 +135,74 @@ var ( visType string = "vis" ) -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) } logger, err := util.LoggerFromContext(ctx) if err != nil { - return nil, fmt.Errorf("unable to get logger from ctx: %s", err) + return nil, util.NewClientServerError("unable to get logger from ctx", http.StatusInternalServerError, err) } logger.DebugContext(ctx, "params = ", params) wq, err := lookercommon.ProcessQueryArgs(ctx, params) if err != nil { - return nil, fmt.Errorf("error building query request: %w", err) + return nil, util.NewAgentError("error building query request", err) } paramsMap := params.AsMap() - dashboard_id := paramsMap["dashboard_id"].(string) - title := paramsMap["title"].(string) - visConfig := paramsMap["vis_config"].(map[string]any) + dashboard_id, ok := paramsMap["dashboard_id"].(string) + if !ok { + return nil, util.NewAgentError("dashboard_id parameter missing or invalid", nil) + } + + title, ok := paramsMap["title"].(string) + if !ok { + title = "" + } + + visConfig, ok := paramsMap["vis_config"].(map[string]any) + if !ok { + visConfig = make(map[string]any) + } wq.VisConfig = &visConfig sdk, err := source.GetLookerSDK(string(accessToken)) if err != nil { - return nil, fmt.Errorf("error getting sdk: %w", err) + return nil, util.NewClientServerError("error getting sdk", http.StatusInternalServerError, err) } qresp, err := sdk.CreateQuery(*wq, "id", source.LookerApiSettings()) if err != nil { - return nil, fmt.Errorf("error making create query request: %w", err) + return nil, util.ProcessGeneralError(err) } dashFilters := []any{} if v, ok := paramsMap["dashboard_filters"]; ok { if v != nil { - dashFilters = paramsMap["dashboard_filters"].([]any) + if df, ok := v.([]any); ok { + dashFilters = df + } } } var filterables []v4.ResultMakerFilterables for _, m := range dashFilters { - f := m.(map[string]any) + f, ok := m.(map[string]any) + if !ok { + return nil, util.NewAgentError("invalid dashboard filter structure", nil) + } name, ok := f["dashboard_filter_name"].(string) if !ok { - return nil, fmt.Errorf("error processing dashboard filter: %w", err) + return nil, util.NewAgentError("error processing dashboard filter: missing dashboard_filter_name", nil) } field, ok := f["field"].(string) if !ok { - return nil, fmt.Errorf("error processing dashboard filter: %w", err) + return nil, util.NewAgentError("error processing dashboard filter: missing field", nil) } listener := v4.ResultMakerFilterablesListen{ DashboardFilterName: &name, @@ -233,7 +250,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para resp, err := sdk.CreateDashboardElement(req, source.LookerApiSettings()) if err != nil { - return nil, fmt.Errorf("error making create dashboard element request: %w", err) + return nil, util.ProcessGeneralError(err) } logger.DebugContext(ctx, "resp = %v", resp) diff --git a/internal/tools/looker/lookeradddashboardfilter/lookeradddashboardfilter.go b/internal/tools/looker/lookeradddashboardfilter/lookeradddashboardfilter.go index e3da8838f8..71ca790850 100644 --- a/internal/tools/looker/lookeradddashboardfilter/lookeradddashboardfilter.go +++ b/internal/tools/looker/lookeradddashboardfilter/lookeradddashboardfilter.go @@ -16,6 +16,7 @@ package lookeradddashboardfilter import ( "context" "fmt" + "net/http" yaml "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" @@ -128,33 +129,54 @@ 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) } logger, err := util.LoggerFromContext(ctx) if err != nil { - return nil, fmt.Errorf("unable to get logger from ctx: %s", err) + return nil, util.NewClientServerError("unable to get logger from ctx", http.StatusInternalServerError, err) } logger.DebugContext(ctx, "params = ", params) paramsMap := params.AsMap() - dashboard_id := paramsMap["dashboard_id"].(string) - name := paramsMap["name"].(string) - title := paramsMap["title"].(string) - filterType := paramsMap["filter_type"].(string) + dashboard_id, ok := paramsMap["dashboard_id"].(string) + if !ok { + return nil, util.NewAgentError("dashboard_id parameter missing or invalid", nil) + } + name, ok := paramsMap["name"].(string) + if !ok { + return nil, util.NewAgentError("name parameter missing or invalid", nil) + } + title, ok := paramsMap["title"].(string) + if !ok { + return nil, util.NewAgentError("title parameter missing or invalid", nil) + } + filterType, ok := paramsMap["filter_type"].(string) + if !ok { + return nil, util.NewAgentError("filter_type parameter missing or invalid", nil) + } + switch filterType { case "date_filter": case "number_filter": case "string_filter": case "field_filter": default: - return nil, fmt.Errorf("invalid filter type: %s. Must be one of date_filter, number_filter, string_filter, field_filter", filterType) + return nil, util.NewAgentError(fmt.Sprintf("invalid filter type: %s. Must be one of date_filter, number_filter, string_filter, field_filter", filterType), nil) + } + + allowMultipleValues, ok := paramsMap["allow_multiple_values"].(bool) + if !ok { + // defaults should handle this, but safe fallback + allowMultipleValues = true + } + required, ok := paramsMap["required"].(bool) + if !ok { + required = false } - allowMultipleValues := paramsMap["allow_multiple_values"].(bool) - required := paramsMap["required"].(bool) req := v4.WriteCreateDashboardFilter{ DashboardId: dashboard_id, @@ -165,9 +187,8 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para Required: &required, } - if v, ok := paramsMap["default_value"]; ok { - if v != nil { - defaultValue := paramsMap["default_value"].(string) + if v, ok := paramsMap["default_value"]; ok && v != nil { + if defaultValue, ok := v.(string); ok { req.DefaultValue = &defaultValue } } @@ -175,15 +196,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if filterType == "field_filter" { model, ok := paramsMap["model"].(string) if !ok || model == "" { - return nil, fmt.Errorf("model must be specified for field_filter type") + return nil, util.NewAgentError("model must be specified for field_filter type", nil) } explore, ok := paramsMap["explore"].(string) if !ok || explore == "" { - return nil, fmt.Errorf("explore must be specified for field_filter type") + return nil, util.NewAgentError("explore must be specified for field_filter type", nil) } dimension, ok := paramsMap["dimension"].(string) if !ok || dimension == "" { - return nil, fmt.Errorf("dimension must be specified for field_filter type") + return nil, util.NewAgentError("dimension must be specified for field_filter type", nil) } req.Model = &model @@ -193,12 +214,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para sdk, err := source.GetLookerSDK(string(accessToken)) if err != nil { - return nil, fmt.Errorf("error getting sdk: %w", err) + return nil, util.NewClientServerError("error getting sdk", http.StatusInternalServerError, err) } resp, err := sdk.CreateDashboardFilter(req, "name", source.LookerApiSettings()) if err != nil { - return nil, fmt.Errorf("error making create dashboard filter request: %s", err) + return nil, util.ProcessGeneralError(err) } logger.DebugContext(ctx, "resp = %v", resp) diff --git a/internal/tools/looker/lookerconversationalanalytics/lookerconversationalanalytics.go b/internal/tools/looker/lookerconversationalanalytics/lookerconversationalanalytics.go index 3c548abc49..3eb28c4d5e 100644 --- a/internal/tools/looker/lookerconversationalanalytics/lookerconversationalanalytics.go +++ b/internal/tools/looker/lookerconversationalanalytics/lookerconversationalanalytics.go @@ -215,10 +215,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) } var tokenStr string @@ -226,11 +226,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para // Get credentials for the API call // Use cloud-platform token source for Gemini Data Analytics API if t.TokenSource == nil { - return nil, fmt.Errorf("cloud-platform token source is missing") + return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil) } token, err := t.TokenSource.Token() if err != nil { - return nil, fmt.Errorf("failed to get token from cloud-platform token source: %w", err) + return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err) } tokenStr = token.AccessToken @@ -286,7 +286,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para // Call the streaming API response, err := getStream(ctx, caURL, payload, headers) if err != nil { - return nil, fmt.Errorf("failed to get response from conversational analytics API: %w", err) + return nil, util.NewClientServerError("failed to get response from conversational analytics API", http.StatusInternalServerError, err) } return response, nil