mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-19 04:18:18 -05:00
Compare commits
1 Commits
js-sdk-doc
...
looker-val
| Author | SHA1 | Date | |
|---|---|---|---|
|
b85fcdffc3 |
@@ -875,8 +875,8 @@ steps:
|
||||
total_coverage=$(go tool cover -func=oracle_coverage.out | grep "total:" | awk '{print $3}')
|
||||
echo "Oracle total coverage: $total_coverage"
|
||||
coverage_numeric=$(echo "$total_coverage" | sed 's/%//')
|
||||
if awk -v cov="$coverage_numeric" 'BEGIN {exit !(cov < 20)}'; then
|
||||
echo "Coverage failure: $total_coverage is below 20%."
|
||||
if awk -v cov="$coverage_numeric" 'BEGIN {exit !(cov < 30)}'; then
|
||||
echo "Coverage failure: $total_coverage is below 30%."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
|
||||
61
cmd/root.go
61
cmd/root.go
@@ -33,7 +33,6 @@ import (
|
||||
"github.com/fsnotify/fsnotify"
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
@@ -155,6 +154,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerrundashboard"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerrunlook"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerupdateprojectfile"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookervalidateproject"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/mindsdb/mindsdbexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/mindsdb/mindsdbsql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbaggregate"
|
||||
@@ -198,7 +198,6 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistroles"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistsequences"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresliststoredprocedure"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttablespaces"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttablestats"
|
||||
@@ -387,13 +386,12 @@ func NewCommand(opts ...Option) *Command {
|
||||
}
|
||||
|
||||
type ToolsFile struct {
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
AuthSources server.AuthServiceConfigs `yaml:"authSources"` // Deprecated: Kept for compatibility.
|
||||
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
|
||||
EmbeddingModels server.EmbeddingModelConfigs `yaml:"embeddingModels"`
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
|
||||
Prompts server.PromptConfigs `yaml:"prompts"`
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
AuthSources server.AuthServiceConfigs `yaml:"authSources"` // Deprecated: Kept for compatibility.
|
||||
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
|
||||
Prompts server.PromptConfigs `yaml:"prompts"`
|
||||
}
|
||||
|
||||
// parseEnv replaces environment variables ${ENV_NAME} with their values.
|
||||
@@ -442,12 +440,11 @@ func parseToolsFile(ctx context.Context, raw []byte) (ToolsFile, error) {
|
||||
// All resource names (sources, authServices, tools, toolsets) must be unique across all files.
|
||||
func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
|
||||
merged := ToolsFile{
|
||||
Sources: make(server.SourceConfigs),
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
EmbeddingModels: make(server.EmbeddingModelConfigs),
|
||||
Tools: make(server.ToolConfigs),
|
||||
Toolsets: make(server.ToolsetConfigs),
|
||||
Prompts: make(server.PromptConfigs),
|
||||
Sources: make(server.SourceConfigs),
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
Tools: make(server.ToolConfigs),
|
||||
Toolsets: make(server.ToolsetConfigs),
|
||||
Prompts: make(server.PromptConfigs),
|
||||
}
|
||||
|
||||
var conflicts []string
|
||||
@@ -483,15 +480,6 @@ func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// Check for conflicts and merge embeddingModels
|
||||
for name, em := range file.EmbeddingModels {
|
||||
if _, exists := merged.EmbeddingModels[name]; exists {
|
||||
conflicts = append(conflicts, fmt.Sprintf("embedding model '%s' (file #%d)", name, fileIndex+1))
|
||||
} else {
|
||||
merged.EmbeddingModels[name] = em
|
||||
}
|
||||
}
|
||||
|
||||
// Check for conflicts and merge tools
|
||||
for name, tool := range file.Tools {
|
||||
if _, exists := merged.Tools[name]; exists {
|
||||
@@ -596,14 +584,14 @@ func handleDynamicReload(ctx context.Context, toolsFile ToolsFile, s *server.Ser
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to validate reloaded edits: %w", err)
|
||||
logger.WarnContext(ctx, errMsg.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
s.ResourceMgr.SetResources(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
s.ResourceMgr.SetResources(sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -611,7 +599,7 @@ func handleDynamicReload(ctx context.Context, toolsFile ToolsFile, s *server.Ser
|
||||
// validateReloadEdits checks that the reloaded tools file configs can initialized without failing
|
||||
func validateReloadEdits(
|
||||
ctx context.Context, toolsFile ToolsFile,
|
||||
) (map[string]sources.Source, map[string]auth.AuthService, map[string]embeddingmodels.EmbeddingModel, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
|
||||
) (map[string]sources.Source, map[string]auth.AuthService, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
|
||||
) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
@@ -629,23 +617,22 @@ func validateReloadEdits(
|
||||
defer span.End()
|
||||
|
||||
reloadedConfig := server.ServerConfig{
|
||||
Version: versionString,
|
||||
SourceConfigs: toolsFile.Sources,
|
||||
AuthServiceConfigs: toolsFile.AuthServices,
|
||||
EmbeddingModelConfigs: toolsFile.EmbeddingModels,
|
||||
ToolConfigs: toolsFile.Tools,
|
||||
ToolsetConfigs: toolsFile.Toolsets,
|
||||
PromptConfigs: toolsFile.Prompts,
|
||||
Version: versionString,
|
||||
SourceConfigs: toolsFile.Sources,
|
||||
AuthServiceConfigs: toolsFile.AuthServices,
|
||||
ToolConfigs: toolsFile.Tools,
|
||||
ToolsetConfigs: toolsFile.Toolsets,
|
||||
PromptConfigs: toolsFile.Prompts,
|
||||
}
|
||||
|
||||
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to initialize reloaded configs: %w", err)
|
||||
logger.WarnContext(ctx, errMsg.Error())
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
|
||||
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
}
|
||||
|
||||
// watchChanges checks for changes in the provided yaml tools file(s) or folder.
|
||||
|
||||
@@ -32,7 +32,6 @@ import (
|
||||
"github.com/google/go-cmp/cmp"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/auth/google"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
@@ -1504,7 +1503,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"alloydb_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "alloydb_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats", "list_stored_procedure"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1534,7 +1533,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"cloud_sql_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "cloud_sql_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats", "list_stored_procedure"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1614,7 +1613,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"looker_tools": tools.ToolsetConfig{
|
||||
Name: "looker_tools",
|
||||
ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "run_dashboard", "make_dashboard", "add_dashboard_element", "add_dashboard_filter", "generate_embed_url", "health_pulse", "health_analyze", "health_vacuum", "dev_mode", "get_projects", "get_project_files", "get_project_file", "create_project_file", "update_project_file", "delete_project_file", "get_connections", "get_connection_schemas", "get_connection_databases", "get_connection_tables", "get_connection_table_columns"},
|
||||
ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "run_dashboard", "make_dashboard", "add_dashboard_element", "add_dashboard_filter", "generate_embed_url", "health_pulse", "health_analyze", "health_vacuum", "dev_mode", "get_projects", "get_project_files", "get_project_file", "create_project_file", "update_project_file", "delete_project_file", "validate_project","get_connections", "get_connection_schemas", "get_connection_databases", "get_connection_tables", "get_connection_table_columns"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1634,7 +1633,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats", "list_stored_procedure"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1831,10 +1830,9 @@ func TestFileLoadingErrors(t *testing.T) {
|
||||
|
||||
func TestMergeToolsFiles(t *testing.T) {
|
||||
file1 := ToolsFile{
|
||||
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
|
||||
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}},
|
||||
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}},
|
||||
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
|
||||
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
|
||||
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}},
|
||||
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}},
|
||||
}
|
||||
file2 := ToolsFile{
|
||||
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
|
||||
@@ -1856,12 +1854,11 @@ func TestMergeToolsFiles(t *testing.T) {
|
||||
name: "merge two distinct files",
|
||||
files: []ToolsFile{file1, file2},
|
||||
want: ToolsFile{
|
||||
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
|
||||
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
|
||||
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}, "tool2": http.Config{Name: "tool2"}},
|
||||
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}, "set2": tools.ToolsetConfig{Name: "set2"}},
|
||||
Prompts: server.PromptConfigs{},
|
||||
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
|
||||
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
|
||||
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
|
||||
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}, "tool2": http.Config{Name: "tool2"}},
|
||||
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}, "set2": tools.ToolsetConfig{Name: "set2"}},
|
||||
Prompts: server.PromptConfigs{},
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
@@ -1874,24 +1871,22 @@ func TestMergeToolsFiles(t *testing.T) {
|
||||
name: "merge single file",
|
||||
files: []ToolsFile{file1},
|
||||
want: ToolsFile{
|
||||
Sources: file1.Sources,
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
|
||||
Tools: file1.Tools,
|
||||
Toolsets: file1.Toolsets,
|
||||
Prompts: server.PromptConfigs{},
|
||||
Sources: file1.Sources,
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
Tools: file1.Tools,
|
||||
Toolsets: file1.Toolsets,
|
||||
Prompts: server.PromptConfigs{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "merge empty list",
|
||||
files: []ToolsFile{},
|
||||
want: ToolsFile{
|
||||
Sources: make(server.SourceConfigs),
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
EmbeddingModels: make(server.EmbeddingModelConfigs),
|
||||
Tools: make(server.ToolConfigs),
|
||||
Toolsets: make(server.ToolsetConfigs),
|
||||
Prompts: server.PromptConfigs{},
|
||||
Sources: make(server.SourceConfigs),
|
||||
AuthServices: make(server.AuthServiceConfigs),
|
||||
Tools: make(server.ToolConfigs),
|
||||
Toolsets: make(server.ToolsetConfigs),
|
||||
Prompts: server.PromptConfigs{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -68,7 +68,6 @@ The BigQuery MCP server is configured using environment variables.
|
||||
export BIGQUERY_PROJECT="<your-gcp-project-id>"
|
||||
export BIGQUERY_LOCATION="<your-dataset-location>" # Optional
|
||||
export BIGQUERY_USE_CLIENT_OAUTH="true" # Optional
|
||||
export BIGQUERY_SCOPES="<comma-separated-scopes>" # Optional
|
||||
```
|
||||
|
||||
Add the following configuration to your MCP client (e.g., `settings.json` for Gemini CLI, `mcp_config.json` for Antigravity):
|
||||
|
||||
@@ -28,11 +28,11 @@ require (
|
||||
go.opentelemetry.io/otel/metric v1.38.0 // indirect
|
||||
go.opentelemetry.io/otel/sdk v1.38.0 // indirect
|
||||
go.opentelemetry.io/otel/trace v1.38.0 // indirect
|
||||
golang.org/x/crypto v0.45.0 // indirect
|
||||
golang.org/x/net v0.47.0 // indirect
|
||||
golang.org/x/crypto v0.43.0 // indirect
|
||||
golang.org/x/net v0.46.0 // indirect
|
||||
golang.org/x/oauth2 v0.32.0 // indirect
|
||||
golang.org/x/sys v0.38.0 // indirect
|
||||
golang.org/x/text v0.31.0 // indirect
|
||||
golang.org/x/sys v0.37.0 // indirect
|
||||
golang.org/x/text v0.30.0 // indirect
|
||||
google.golang.org/api v0.255.0 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda // indirect
|
||||
google.golang.org/grpc v1.76.0 // indirect
|
||||
|
||||
@@ -88,18 +88,18 @@ go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJr
|
||||
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
|
||||
golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q=
|
||||
golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4=
|
||||
golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
|
||||
golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
|
||||
golang.org/x/crypto v0.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04=
|
||||
golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0=
|
||||
golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4=
|
||||
golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
|
||||
golang.org/x/oauth2 v0.32.0 h1:jsCblLleRMDrxMN29H3z/k1KliIvpLgCkE6R8FXXNgY=
|
||||
golang.org/x/oauth2 v0.32.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
|
||||
golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
|
||||
golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
|
||||
golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
|
||||
golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
|
||||
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
|
||||
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
|
||||
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
|
||||
golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
|
||||
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
|
||||
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
|
||||
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
|
||||
|
||||
@@ -1813,9 +1813,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/qs": {
|
||||
"version": "6.14.1",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz",
|
||||
"integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==",
|
||||
"version": "6.14.0",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.14.0.tgz",
|
||||
"integrity": "sha512-YWWTjgABSKcvs/nWBi9PycY/JiPJqOD4JA6o9Sej2AtvSGarXxKC3OQSk4pAarbdQlKAh5D4FCQkJNkW+GAn3w==",
|
||||
"license": "BSD-3-Clause",
|
||||
"dependencies": {
|
||||
"side-channel": "^1.1.0"
|
||||
|
||||
@@ -3376,23 +3376,22 @@
|
||||
}
|
||||
},
|
||||
"node_modules/body-parser": {
|
||||
"version": "1.20.4",
|
||||
"resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.4.tgz",
|
||||
"integrity": "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA==",
|
||||
"license": "MIT",
|
||||
"version": "1.20.3",
|
||||
"resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.3.tgz",
|
||||
"integrity": "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==",
|
||||
"dependencies": {
|
||||
"bytes": "~3.1.2",
|
||||
"bytes": "3.1.2",
|
||||
"content-type": "~1.0.5",
|
||||
"debug": "2.6.9",
|
||||
"depd": "2.0.0",
|
||||
"destroy": "~1.2.0",
|
||||
"http-errors": "~2.0.1",
|
||||
"iconv-lite": "~0.4.24",
|
||||
"on-finished": "~2.4.1",
|
||||
"qs": "~6.14.0",
|
||||
"raw-body": "~2.5.3",
|
||||
"destroy": "1.2.0",
|
||||
"http-errors": "2.0.0",
|
||||
"iconv-lite": "0.4.24",
|
||||
"on-finished": "2.4.1",
|
||||
"qs": "6.13.0",
|
||||
"raw-body": "2.5.2",
|
||||
"type-is": "~1.6.18",
|
||||
"unpipe": "~1.0.0"
|
||||
"unpipe": "1.0.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 0.8",
|
||||
@@ -3407,40 +3406,11 @@
|
||||
"ms": "2.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/body-parser/node_modules/http-errors": {
|
||||
"version": "2.0.1",
|
||||
"resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
|
||||
"integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"depd": "~2.0.0",
|
||||
"inherits": "~2.0.4",
|
||||
"setprototypeof": "~1.2.0",
|
||||
"statuses": "~2.0.2",
|
||||
"toidentifier": "~1.0.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
},
|
||||
"funding": {
|
||||
"type": "opencollective",
|
||||
"url": "https://opencollective.com/express"
|
||||
}
|
||||
},
|
||||
"node_modules/body-parser/node_modules/ms": {
|
||||
"version": "2.0.0",
|
||||
"resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz",
|
||||
"integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A=="
|
||||
},
|
||||
"node_modules/body-parser/node_modules/statuses": {
|
||||
"version": "2.0.2",
|
||||
"resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
|
||||
"integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
}
|
||||
},
|
||||
"node_modules/buffer-equal-constant-time": {
|
||||
"version": "1.0.1",
|
||||
"resolved": "https://registry.npmjs.org/buffer-equal-constant-time/-/buffer-equal-constant-time-1.0.1.tgz",
|
||||
@@ -3464,7 +3434,6 @@
|
||||
"version": "3.1.2",
|
||||
"resolved": "https://registry.npmjs.org/bytes/-/bytes-3.1.2.tgz",
|
||||
"integrity": "sha512-/Nf7TyzTx6S3yRJObOAV7956r8cr2+Oj8AC5dt8wSP3BQAoeX58NoHyCU8P8zGkNXStjTSi6fzO6F0pBdcYbEg==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
}
|
||||
@@ -3861,39 +3830,38 @@
|
||||
}
|
||||
},
|
||||
"node_modules/express": {
|
||||
"version": "4.22.1",
|
||||
"resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz",
|
||||
"integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==",
|
||||
"license": "MIT",
|
||||
"version": "4.21.2",
|
||||
"resolved": "https://registry.npmjs.org/express/-/express-4.21.2.tgz",
|
||||
"integrity": "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA==",
|
||||
"dependencies": {
|
||||
"accepts": "~1.3.8",
|
||||
"array-flatten": "1.1.1",
|
||||
"body-parser": "~1.20.3",
|
||||
"content-disposition": "~0.5.4",
|
||||
"body-parser": "1.20.3",
|
||||
"content-disposition": "0.5.4",
|
||||
"content-type": "~1.0.4",
|
||||
"cookie": "~0.7.1",
|
||||
"cookie-signature": "~1.0.6",
|
||||
"cookie": "0.7.1",
|
||||
"cookie-signature": "1.0.6",
|
||||
"debug": "2.6.9",
|
||||
"depd": "2.0.0",
|
||||
"encodeurl": "~2.0.0",
|
||||
"escape-html": "~1.0.3",
|
||||
"etag": "~1.8.1",
|
||||
"finalhandler": "~1.3.1",
|
||||
"fresh": "~0.5.2",
|
||||
"http-errors": "~2.0.0",
|
||||
"finalhandler": "1.3.1",
|
||||
"fresh": "0.5.2",
|
||||
"http-errors": "2.0.0",
|
||||
"merge-descriptors": "1.0.3",
|
||||
"methods": "~1.1.2",
|
||||
"on-finished": "~2.4.1",
|
||||
"on-finished": "2.4.1",
|
||||
"parseurl": "~1.3.3",
|
||||
"path-to-regexp": "~0.1.12",
|
||||
"path-to-regexp": "0.1.12",
|
||||
"proxy-addr": "~2.0.7",
|
||||
"qs": "~6.14.0",
|
||||
"qs": "6.13.0",
|
||||
"range-parser": "~1.2.1",
|
||||
"safe-buffer": "5.2.1",
|
||||
"send": "~0.19.0",
|
||||
"serve-static": "~1.16.2",
|
||||
"send": "0.19.0",
|
||||
"serve-static": "1.16.2",
|
||||
"setprototypeof": "1.2.0",
|
||||
"statuses": "~2.0.1",
|
||||
"statuses": "2.0.1",
|
||||
"type-is": "~1.6.18",
|
||||
"utils-merge": "1.0.1",
|
||||
"vary": "~1.1.2"
|
||||
@@ -4936,7 +4904,6 @@
|
||||
"version": "0.4.24",
|
||||
"resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz",
|
||||
"integrity": "sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"safer-buffer": ">= 2.1.2 < 3"
|
||||
},
|
||||
@@ -5694,12 +5661,11 @@
|
||||
}
|
||||
},
|
||||
"node_modules/qs": {
|
||||
"version": "6.14.1",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz",
|
||||
"integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==",
|
||||
"license": "BSD-3-Clause",
|
||||
"version": "6.13.0",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.13.0.tgz",
|
||||
"integrity": "sha512-+38qI9SOr8tfZ4QmJNplMUxqjbe7LKvvZgWdExBOmd+egZTtjLB67Gu0HRX3u/XOq7UU2Nx6nsjvS16Z9uwfpg==",
|
||||
"dependencies": {
|
||||
"side-channel": "^1.1.0"
|
||||
"side-channel": "^1.0.6"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=0.6"
|
||||
@@ -5717,49 +5683,19 @@
|
||||
}
|
||||
},
|
||||
"node_modules/raw-body": {
|
||||
"version": "2.5.3",
|
||||
"resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz",
|
||||
"integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==",
|
||||
"license": "MIT",
|
||||
"version": "2.5.2",
|
||||
"resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.2.tgz",
|
||||
"integrity": "sha512-8zGqypfENjCIqGhgXToC8aB2r7YrBX+AQAfIPs/Mlk+BtPTztOvTS01NRW/3Eh60J+a48lt8qsCzirQ6loCVfA==",
|
||||
"dependencies": {
|
||||
"bytes": "~3.1.2",
|
||||
"http-errors": "~2.0.1",
|
||||
"iconv-lite": "~0.4.24",
|
||||
"unpipe": "~1.0.0"
|
||||
"bytes": "3.1.2",
|
||||
"http-errors": "2.0.0",
|
||||
"iconv-lite": "0.4.24",
|
||||
"unpipe": "1.0.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
}
|
||||
},
|
||||
"node_modules/raw-body/node_modules/http-errors": {
|
||||
"version": "2.0.1",
|
||||
"resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
|
||||
"integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"depd": "~2.0.0",
|
||||
"inherits": "~2.0.4",
|
||||
"setprototypeof": "~1.2.0",
|
||||
"statuses": "~2.0.2",
|
||||
"toidentifier": "~1.0.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
},
|
||||
"funding": {
|
||||
"type": "opencollective",
|
||||
"url": "https://opencollective.com/express"
|
||||
}
|
||||
},
|
||||
"node_modules/raw-body/node_modules/statuses": {
|
||||
"version": "2.0.2",
|
||||
"resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
|
||||
"integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">= 0.8"
|
||||
}
|
||||
},
|
||||
"node_modules/readable-stream": {
|
||||
"version": "3.6.2",
|
||||
"resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-3.6.2.tgz",
|
||||
@@ -5877,8 +5813,7 @@
|
||||
"node_modules/safer-buffer": {
|
||||
"version": "2.1.2",
|
||||
"resolved": "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz",
|
||||
"integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==",
|
||||
"license": "MIT"
|
||||
"integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg=="
|
||||
},
|
||||
"node_modules/semver": {
|
||||
"version": "7.7.2",
|
||||
|
||||
@@ -66,6 +66,40 @@
|
||||
"node": ">=20"
|
||||
}
|
||||
},
|
||||
"node_modules/@langchain/core/node_modules/langsmith": {
|
||||
"version": "0.4.2",
|
||||
"resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.4.2.tgz",
|
||||
"integrity": "sha512-BvBeFgSmR9esl8x5wsiDlALiHKKPybw2wE2Hh6x1tgSZki46H9c9KI9/06LARbPhyyDu/TZU7exfg6fnhdj1Qg==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@types/uuid": "^10.0.0",
|
||||
"chalk": "^4.1.2",
|
||||
"console-table-printer": "^2.12.1",
|
||||
"p-queue": "^6.6.2",
|
||||
"semver": "^7.6.3",
|
||||
"uuid": "^10.0.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": "*",
|
||||
"@opentelemetry/exporter-trace-otlp-proto": "*",
|
||||
"@opentelemetry/sdk-trace-base": "*",
|
||||
"openai": "*"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"@opentelemetry/api": {
|
||||
"optional": true
|
||||
},
|
||||
"@opentelemetry/exporter-trace-otlp-proto": {
|
||||
"optional": true
|
||||
},
|
||||
"@opentelemetry/sdk-trace-base": {
|
||||
"optional": true
|
||||
},
|
||||
"openai": {
|
||||
"optional": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/@langchain/google-genai": {
|
||||
"version": "2.1.3",
|
||||
"resolved": "https://registry.npmjs.org/@langchain/google-genai/-/google-genai-2.1.3.tgz",
|
||||
@@ -854,14 +888,13 @@
|
||||
}
|
||||
},
|
||||
"node_modules/langchain": {
|
||||
"version": "1.2.3",
|
||||
"resolved": "https://registry.npmjs.org/langchain/-/langchain-1.2.3.tgz",
|
||||
"integrity": "sha512-3k986xJuqg4az53JxV5LnGlOzIXF1d9Kq6Y9s7XjitvzhpsbFuTDV5/kiF4cx3pkNGyw0mUXC4tLz9RxucO0hw==",
|
||||
"license": "MIT",
|
||||
"version": "1.0.2",
|
||||
"resolved": "https://registry.npmjs.org/langchain/-/langchain-1.0.2.tgz",
|
||||
"integrity": "sha512-He/xvjVl8DHESvdaW6Dpyba72OaLCAfS2CyOm1aWrlJ4C38dKXyTIxphtld8hiii6MWX7qMSmu2EaUwWBx2STg==",
|
||||
"dependencies": {
|
||||
"@langchain/langgraph": "^1.0.0",
|
||||
"@langchain/langgraph-checkpoint": "^1.0.0",
|
||||
"langsmith": ">=0.4.0 <1.0.0",
|
||||
"langsmith": "~0.3.74",
|
||||
"uuid": "^10.0.0",
|
||||
"zod": "^3.25.76 || ^4"
|
||||
},
|
||||
@@ -869,19 +902,19 @@
|
||||
"node": ">=20"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@langchain/core": "1.1.8"
|
||||
"@langchain/core": "^1.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/langsmith": {
|
||||
"version": "0.4.3",
|
||||
"resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.4.3.tgz",
|
||||
"integrity": "sha512-vuBAagBZulXj0rpZhUTxmHhrYIBk53z8e2Q8ty4OHVkahN4ul7Im3OZxD9jsXZB0EuncK1xRYtY8J3BW4vj1zw==",
|
||||
"license": "MIT",
|
||||
"version": "0.3.77",
|
||||
"resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.3.77.tgz",
|
||||
"integrity": "sha512-wbS/9IX/hOAsOEOtPj8kCS8H0tFHaelwQ97gTONRtIfoPPLd9MMUmhk0KQB5DdsGAI5abg966+f0dZ/B+YRRzg==",
|
||||
"dependencies": {
|
||||
"@types/uuid": "^10.0.0",
|
||||
"chalk": "^4.1.2",
|
||||
"console-table-printer": "^2.12.1",
|
||||
"p-queue": "^6.6.2",
|
||||
"p-retry": "4",
|
||||
"semver": "^7.6.3",
|
||||
"uuid": "^10.0.0"
|
||||
},
|
||||
|
||||
@@ -105,8 +105,6 @@ See [Usage Examples](../reference/cli.md#examples).
|
||||
* `BIGQUERY_LOCATION`: (Optional) The dataset location.
|
||||
* `BIGQUERY_USE_CLIENT_OAUTH`: (Optional) If `true`, forwards the client's
|
||||
OAuth access token for authentication. Defaults to `false`.
|
||||
* `BIGQUERY_SCOPES`: (Optional) A comma-separated list of OAuth scopes to
|
||||
use for authentication.
|
||||
* **Permissions:**
|
||||
* **BigQuery User** (`roles/bigquery.user`) to execute queries and view
|
||||
metadata.
|
||||
@@ -436,6 +434,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
||||
* `create_project_file`: Create a new LookML file.
|
||||
* `update_project_file`: Update an existing LookML file.
|
||||
* `delete_project_file`: Delete a LookML file.
|
||||
* `validate_project`: Check the syntax of a LookML project.
|
||||
* `get_connections`: Get the available connections in a Looker instance.
|
||||
* `get_connection_schemas`: Get the available schemas in a connection.
|
||||
* `get_connection_databases`: Get the available databases in a connection.
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
---
|
||||
title: "EmbeddingModels"
|
||||
type: docs
|
||||
weight: 2
|
||||
description: >
|
||||
EmbeddingModels represent services that transform text into vector embeddings for semantic search.
|
||||
---
|
||||
|
||||
EmbeddingModels represent services that generate vector representations of text
|
||||
data. In the MCP Toolbox, these models enable **Semantic Queries**,
|
||||
allowing [Tools](../tools/) to automatically convert human-readable text into
|
||||
numerical vectors before using them in a query.
|
||||
|
||||
This is primarily used in two scenarios:
|
||||
|
||||
- **Vector Ingestion**: Converting a text parameter into a vector string during
|
||||
an `INSERT` operation.
|
||||
|
||||
- **Semantic Search**: Converting a natural language query into a vector to
|
||||
perform similarity searches.
|
||||
|
||||
## Example
|
||||
|
||||
The following configuration defines an embedding model and applies it to
|
||||
specific tool parameters.
|
||||
|
||||
{{< notice tip >}}
|
||||
Use environment variable replacement with the format ${ENV_NAME}
|
||||
instead of hardcoding your API keys into the configuration file.
|
||||
{{< /notice >}}
|
||||
|
||||
### Step 1 - Define an Embedding Model
|
||||
|
||||
Define an embedding model in the `embeddingModels` section:
|
||||
|
||||
```yaml
|
||||
embeddingModels:
|
||||
gemini-model: # Name of the embedding model
|
||||
kind: gemini
|
||||
model: gemini-embedding-001
|
||||
apiKey: ${GOOGLE_API_KEY}
|
||||
dimension: 768
|
||||
|
||||
```
|
||||
|
||||
### Step 2 - Embed Tool Parameters
|
||||
|
||||
Use the defined embedding model, embed your query parameters using the
|
||||
`embeddedBy` field. Only string-typed
|
||||
parameters can be embedded:
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
# Vector ingestion tool
|
||||
insert_embedding:
|
||||
kind: postgres-sql
|
||||
source: my-pg-instance
|
||||
statement: |
|
||||
INSERT INTO documents (content, embedding)
|
||||
VALUES ($1, $2);
|
||||
parameters:
|
||||
- name: content
|
||||
type: string
|
||||
- name: vector_string
|
||||
type: string
|
||||
description: The text to be vectorized and stored.
|
||||
embeddedBy: gemini-model # refers to the name of a defined embedding model
|
||||
|
||||
# Semantic search tool
|
||||
search_embedding:
|
||||
kind: postgres-sql
|
||||
source: my-pg-instance
|
||||
statement: |
|
||||
SELECT id, content, embedding <-> $1 AS distance
|
||||
FROM documents
|
||||
ORDER BY distance LIMIT 1
|
||||
parameters:
|
||||
- name: semantic_search_string
|
||||
type: string
|
||||
description: The search query that will be converted to a vector.
|
||||
embeddedBy: gemini-model # refers to the name of a defined embedding model
|
||||
```
|
||||
|
||||
## Kinds of Embedding Models
|
||||
@@ -1,73 +0,0 @@
|
||||
---
|
||||
title: "Gemini Embedding"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
Use Google's Gemini models to generate high-performance text embeddings for vector databases.
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
Google Gemini provides state-of-the-art embedding models that convert text into
|
||||
high-dimensional vectors.
|
||||
|
||||
### Authentication
|
||||
|
||||
Toolbox uses your [Application Default Credentials
|
||||
(ADC)][adc] to authorize with the
|
||||
Gemini API client.
|
||||
|
||||
Optionally, you can use an [API key][api-key] obtain an API
|
||||
Key from the [Google AI Studio][ai-studio].
|
||||
|
||||
We recommend using an API key for testing and using application default
|
||||
credentials for production.
|
||||
|
||||
[adc]: https://cloud.google.com/docs/authentication#adc
|
||||
[api-key]: https://ai.google.dev/gemini-api/docs/api-key#api-keys
|
||||
[ai-studio]: https://aistudio.google.com/app/apikey
|
||||
|
||||
## Behavior
|
||||
|
||||
### Automatic Vectorization
|
||||
|
||||
When a tool parameter is configured with `embeddedBy: <your-gemini-model-name>`,
|
||||
the Toolbox intercepts the raw text input from the client and sends it to the
|
||||
Gemini API. The resulting numerical array is then formatted before being passed
|
||||
to your database source.
|
||||
|
||||
### Dimension Matching
|
||||
|
||||
The `dimension` field must match the expected size of your database column
|
||||
(e.g., a `vector(768)` column in PostgreSQL). This setting is supported by newer
|
||||
models since 2024 only. You cannot set this value if using the earlier model
|
||||
(`models/embedding-001`). Check out [available Gemini models][modellist] for more
|
||||
information.
|
||||
|
||||
[modellist]:
|
||||
https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
embeddingModels:
|
||||
gemini-model:
|
||||
kind: gemini
|
||||
model: gemini-embedding-001
|
||||
apiKey: ${GOOGLE_API_KEY}
|
||||
dimension: 768
|
||||
```
|
||||
|
||||
{{< notice tip >}}
|
||||
Use environment variable replacement with the format ${ENV_NAME}
|
||||
instead of hardcoding your secrets into the configuration file.
|
||||
{{< /notice >}}
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-----------|:--------:|:------------:|--------------------------------------------------------------|
|
||||
| kind | string | true | Must be `gemini`. |
|
||||
| model | string | true | The Gemini model ID to use (e.g., `gemini-embedding-001`). |
|
||||
| apiKey | string | false | Your API Key from Google AI Studio. |
|
||||
| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). |
|
||||
@@ -94,10 +94,7 @@ cluster][alloydb-free-trial].
|
||||
instance.
|
||||
|
||||
- [`postgres-list-roles`](../tools/postgres/postgres-list-roles.md)
|
||||
Lists all the user-created roles in PostgreSQL database.
|
||||
|
||||
- [`postgres-list-stored-procedure`](../tools/postgres/postgres-list-stored-procedure.md)
|
||||
Lists all the stored procedure in PostgreSQL database.
|
||||
Lists all the user-created roles in PostgreSQL database..
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
|
||||
@@ -94,13 +94,6 @@ intend to run. Common roles include `roles/bigquery.user` (which includes
|
||||
permissions to run jobs and read data) or `roles/bigbigquery.dataViewer`.
|
||||
Follow this [guide][set-adc] to set up your ADC.
|
||||
|
||||
If you are running on Google Compute Engine (GCE) or Google Kubernetes Engine
|
||||
(GKE), you might need to explicitly set the access scopes for the service
|
||||
account. While you can configure scopes when creating the VM or node pool, you
|
||||
can also specify them in the source configuration using the `scopes` field.
|
||||
Common scopes include `https://www.googleapis.com/auth/bigquery` or
|
||||
`https://www.googleapis.com/auth/cloud-platform`.
|
||||
|
||||
### Authentication via User's OAuth Access Token
|
||||
|
||||
If the `useClientOAuth` parameter is set to `true`, Toolbox will instead use the
|
||||
@@ -131,9 +124,6 @@ sources:
|
||||
# - "my_dataset_1"
|
||||
# - "other_project.my_dataset_2"
|
||||
# impersonateServiceAccount: "service-account@project-id.iam.gserviceaccount.com" # Optional: Service account to impersonate
|
||||
# scopes: # Optional: List of OAuth scopes to request.
|
||||
# - "https://www.googleapis.com/auth/bigquery"
|
||||
# - "https://www.googleapis.com/auth/drive.readonly"
|
||||
```
|
||||
|
||||
Initialize a BigQuery source that uses the client's access token:
|
||||
@@ -150,9 +140,6 @@ sources:
|
||||
# - "my_dataset_1"
|
||||
# - "other_project.my_dataset_2"
|
||||
# impersonateServiceAccount: "service-account@project-id.iam.gserviceaccount.com" # Optional: Service account to impersonate
|
||||
# scopes: # Optional: List of OAuth scopes to request.
|
||||
# - "https://www.googleapis.com/auth/bigquery"
|
||||
# - "https://www.googleapis.com/auth/drive.readonly"
|
||||
```
|
||||
|
||||
## Reference
|
||||
@@ -165,5 +152,4 @@ sources:
|
||||
| writeMode | string | false | Controls the write behavior for tools. `allowed` (default): All queries are permitted. `blocked`: Only `SELECT` statements are allowed for the `bigquery-execute-sql` tool. `protected`: Enables session-based execution where all tools associated with this source instance share the same [BigQuery session](https://cloud.google.com/bigquery/docs/sessions-intro). This allows for stateful operations using temporary tables (e.g., `CREATE TEMP TABLE`). For `bigquery-execute-sql`, `SELECT` statements can be used on all tables, but write operations are restricted to the session's temporary dataset. For tools like `bigquery-sql`, `bigquery-forecast`, and `bigquery-analyze-contribution`, the `writeMode` restrictions do not apply, but they will operate within the shared session. **Note:** The `protected` mode cannot be used with `useClientOAuth: true`. It is also not recommended for multi-user server environments, as all users would share the same session. A session is terminated automatically after 24 hours of inactivity or after 7 days, whichever comes first. A new session is created on the next request, and any temporary data from the previous session will be lost. |
|
||||
| allowedDatasets | []string | false | An optional list of dataset IDs that tools using this source are allowed to access. If provided, any tool operation attempting to access a dataset not in this list will be rejected. To enforce this, two types of operations are also disallowed: 1) Dataset-level operations (e.g., `CREATE SCHEMA`), and 2) operations where table access cannot be statically analyzed (e.g., `EXECUTE IMMEDIATE`, `CREATE PROCEDURE`). If a single dataset is provided, it will be treated as the default for prebuilt tools. |
|
||||
| useClientOAuth | bool | false | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. **Note:** This cannot be used with `writeMode: protected`. |
|
||||
| scopes | []string | false | A list of OAuth 2.0 scopes to use for the credentials. If not provided, default scopes are used. |
|
||||
| impersonateServiceAccount | string | false | Service account email to impersonate when making BigQuery and Dataplex API calls. The authenticated principal must have the `roles/iam.serviceAccountTokenCreator` role on the target service account. [Learn More](https://cloud.google.com/iam/docs/service-account-impersonation) |
|
||||
|
||||
@@ -91,10 +91,7 @@ to a database by following these instructions][csql-pg-quickstart].
|
||||
instance.
|
||||
|
||||
- [`postgres-list-roles`](../tools/postgres/postgres-list-roles.md)
|
||||
Lists all the user-created roles in PostgreSQL database.
|
||||
|
||||
- [`postgres-list-stored-procedure`](../tools/postgres/postgres-list-stored-procedure.md)
|
||||
Lists all the stored procedure in PostgreSQL database.
|
||||
Lists all the user-created roles in PostgreSQL database..
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
|
||||
@@ -229,38 +229,22 @@ Finds resources that were created within, before, or after a given date or time.
|
||||
### Aspect Search
|
||||
To search for entries based on their attached aspects, use the following query syntax.
|
||||
|
||||
`has:x`
|
||||
Matches `x` as a substring of the full path to the aspect type of an aspect that is attached to the entry, in the format `projectid.location.ASPECT_TYPE_ID`
|
||||
aspect:x Matches x as a substring of the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
|
||||
aspect=x Matches x as the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
|
||||
aspect:xOPERATORvalue
|
||||
Searches for aspect field values. Matches x as a substring of the full path to the aspect type and field name of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID.FIELD_NAME
|
||||
|
||||
`has=x`
|
||||
Matches `x` as the full path to the aspect type of an aspect that is attached to the entry, in the format `projectid.location.ASPECT_TYPE_ID`
|
||||
The list of supported {OPERATOR}s depends on the type of field in the aspect, as follows:
|
||||
- String: = (exact match) and : (substring)
|
||||
- All number types: =, :, <, >, <=, >=, =>, =<
|
||||
- Enum: =
|
||||
- Datetime: same as for numbers, but the values to compare are treated as datetimes instead of numbers
|
||||
- Boolean: =
|
||||
|
||||
`xOPERATORvalue`
|
||||
Searches for aspect field values. Matches x as a substring of the full path to the aspect type and field name of an aspect that is attached to the entry, in the format `projectid.location.ASPECT_TYPE_ID.FIELD_NAME`
|
||||
|
||||
The list of supported operators depends on the type of field in the aspect, as follows:
|
||||
* **String**: `=` (exact match)
|
||||
* **All number types**: `=`, `:`, `<`, `>`, `<=`, `>=`, `=>`, `=<`
|
||||
* **Enum**: `=` (exact match only)
|
||||
* **Datetime**: same as for numbers, but the values to compare are treated as datetimes instead of numbers
|
||||
* **Boolean**: `=`
|
||||
|
||||
Only top-level fields of the aspect are searchable.
|
||||
|
||||
* Syntax for system aspect types:
|
||||
* `ASPECT_TYPE_ID.FIELD_NAME`
|
||||
* `dataplex-types.ASPECT_TYPE_ID.FIELD_NAME`
|
||||
* `dataplex-types.LOCATION.ASPECT_TYPE_ID.FIELD_NAME`
|
||||
For example, the following queries match entries where the value of the `type` field in the `bigquery-dataset` aspect is `default`:
|
||||
* `bigquery-dataset.type=default`
|
||||
* `dataplex-types.bigquery-dataset.type=default`
|
||||
* `dataplex-types.global.bigquery-dataset.type=default`
|
||||
* Syntax for custom aspect types:
|
||||
* If the aspect is created in the global region: `PROJECT_ID.ASPECT_TYPE_ID.FIELD_NAME`
|
||||
* If the aspect is created in a specific region: `PROJECT_ID.REGION.ASPECT_TYPE_ID.FIELD_NAME`
|
||||
For example, the following queries match entries where the value of the `is-enrolled` field in the `employee-info` aspect is `true`.
|
||||
* `example-project.us-central1.employee-info.is-enrolled=true`
|
||||
* `example-project.employee-info.is-enrolled=true`
|
||||
Only top-level fields of the aspect are searchable. For example, all of the following queries match entries where the value of the is-enrolled field in the employee-info aspect type is true. Other entries that match on the substring are also returned.
|
||||
- aspect:example-project.us-central1.employee-info.is-enrolled=true
|
||||
- aspect:example-project.us-central1.employee=true
|
||||
- aspect:employee=true
|
||||
|
||||
Example:-
|
||||
You can use following filters
|
||||
@@ -274,25 +258,6 @@ Logical AND and logical OR are supported. For example, foo OR bar.
|
||||
You can negate a predicate with a - (hyphen) or NOT prefix. For example, -name:foo returns resources with names that don't match the predicate foo.
|
||||
Logical operators are case-sensitive. `OR` and `AND` are acceptable whereas `or` and `and` are not.
|
||||
|
||||
### Abbreviated syntax
|
||||
|
||||
An abbreviated search syntax is also available, using `|` (vertical bar) for `OR` operators and `,` (comma) for `AND` operators.
|
||||
|
||||
For example, to search for entries inside one of many projects using the `OR` operator, you can use the following abbreviated syntax:
|
||||
|
||||
`projectid:(id1|id2|id3|id4)`
|
||||
|
||||
The same search without using abbreviated syntax looks like the following:
|
||||
|
||||
`projectid:id1 OR projectid:id2 OR projectid:id3 OR projectid:id4`
|
||||
|
||||
To search for entries with matching column names, use the following:
|
||||
|
||||
* **AND**: `column:(name1,name2,name3)`
|
||||
* **OR**: `column:(name1|name2|name3)`
|
||||
|
||||
This abbreviated syntax works for the qualified predicates except for `label` in keyword search.
|
||||
|
||||
### Request
|
||||
1. Always try to rewrite the prompt using search syntax.
|
||||
|
||||
|
||||
@@ -85,10 +85,7 @@ reputation for reliability, feature robustness, and performance.
|
||||
server.
|
||||
|
||||
- [`postgres-list-roles`](../tools/postgres/postgres-list-roles.md)
|
||||
Lists all the user-created roles in PostgreSQL database.
|
||||
|
||||
- [`postgres-list-stored-procedure`](../tools/postgres/postgres-list-stored-procedure.md)
|
||||
Lists all the stored procedure in PostgreSQL database.
|
||||
Lists all the user-created roles in PostgreSQL database..
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
|
||||
@@ -18,11 +18,9 @@ It's compatible with the following sources:
|
||||
|
||||
- [looker](../../sources/looker.md)
|
||||
|
||||
`looker-make-dashboard` takes three parameters:
|
||||
`looker-make-dashboard` takes one parameter:
|
||||
|
||||
1. the `title`
|
||||
2. the `description`
|
||||
3. an optional `folder` id. If not provided, the user's default folder will be used.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ It's compatible with the following sources:
|
||||
|
||||
- [looker](../../sources/looker.md)
|
||||
|
||||
`looker-make-look` takes twelve parameters:
|
||||
`looker-make-look` takes eleven parameters:
|
||||
|
||||
1. the `model`
|
||||
2. the `explore`
|
||||
@@ -31,7 +31,6 @@ It's compatible with the following sources:
|
||||
9. an optional `vis_config`
|
||||
10. the `title`
|
||||
11. an optional `description`
|
||||
12. an optional `folder` id. If not provided, the user's default folder will be used.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
47
docs/en/resources/tools/looker/looker-validate-project.md
Normal file
47
docs/en/resources/tools/looker/looker-validate-project.md
Normal file
@@ -0,0 +1,47 @@
|
||||
---
|
||||
title: "looker-validate-project"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
A "looker-validate-project" tool checks the syntax of a LookML project and reports any errors
|
||||
aliases:
|
||||
- /resources/tools/looker-validate-project
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
A "looker-validate-project" tool checks the syntax of a LookML project and reports any errors
|
||||
|
||||
It's compatible with the following sources:
|
||||
|
||||
- [looker](../../sources/looker.md)
|
||||
|
||||
`looker-validate-project` accepts a project_id parameter.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
validate_project:
|
||||
kind: looker-validate-project
|
||||
source: looker-source
|
||||
description: |
|
||||
This tool checks a LookML project for syntax errors.
|
||||
|
||||
Prerequisite: The Looker session must be in Development Mode. Use `dev_mode: true` first.
|
||||
|
||||
Parameters:
|
||||
- project_id (required): The unique ID of the LookML project.
|
||||
|
||||
Output:
|
||||
A list of error details including the file path and line number, and also a list of models
|
||||
that are not currently valid due to LookML errors.
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:--------:|:------------:|----------------------------------------------------|
|
||||
| kind | string | true | Must be "looker-validate-project". |
|
||||
| source | string | true | Name of the source Looker instance. |
|
||||
| description | string | true | Description of the tool that is passed to the LLM. |
|
||||
@@ -1,141 +0,0 @@
|
||||
---
|
||||
title: "postgres-list-stored-procedure"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
The "postgres-list-stored-procedure" tool retrieves metadata for stored procedures in PostgreSQL, including procedure definitions, owners, languages, and descriptions.
|
||||
aliases:
|
||||
- /resources/tools/postgres-list-stored-procedure
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `postgres-list-stored-procedure` tool queries PostgreSQL system catalogs (`pg_proc`, `pg_namespace`, `pg_roles`, and `pg_language`) to retrieve comprehensive metadata about stored procedures in the database. It filters for procedures (kind = 'p') and provides the full procedure definition along with ownership and language information.
|
||||
|
||||
Compatible sources:
|
||||
|
||||
- [alloydb-postgres](../../sources/alloydb-pg.md)
|
||||
- [cloud-sql-postgres](../../sources/cloud-sql-pg.md)
|
||||
- [postgres](../../sources/postgres.md)
|
||||
|
||||
The tool returns a JSON array where each element represents a stored procedure with its schema, name, owner, language, complete definition, and optional description. Results are sorted by schema name and procedure name, with a default limit of 20 procedures.
|
||||
|
||||
## Parameters
|
||||
|
||||
| parameter | type | required | default | description |
|
||||
|--------------|---------|----------|---------|-------------|
|
||||
| role_name | string | false | null | Optional: The owner name to filter stored procedures by (supports partial matching) |
|
||||
| schema_name | string | false | null | Optional: The schema name to filter stored procedures by (supports partial matching) |
|
||||
| limit | integer | false | 20 | Optional: The maximum number of stored procedures to return |
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
list_stored_procedure:
|
||||
kind: postgres-list-stored-procedure
|
||||
source: postgres-source
|
||||
description: "Retrieves stored procedure metadata including definitions and owners."
|
||||
```
|
||||
|
||||
### Example Requests
|
||||
|
||||
**List all stored procedures (default limit 20):**
|
||||
```json
|
||||
{}
|
||||
```
|
||||
|
||||
**Filter by specific owner (role):**
|
||||
```json
|
||||
{
|
||||
"role_name": "app_user"
|
||||
}
|
||||
```
|
||||
|
||||
**Filter by schema:**
|
||||
```json
|
||||
{
|
||||
"schema_name": "public"
|
||||
}
|
||||
```
|
||||
|
||||
**Filter by owner and schema with custom limit:**
|
||||
```json
|
||||
{
|
||||
"role_name": "postgres",
|
||||
"schema_name": "public",
|
||||
"limit": 50
|
||||
}
|
||||
```
|
||||
|
||||
**Filter by partial schema name:**
|
||||
```json
|
||||
{
|
||||
"schema_name": "audit"
|
||||
}
|
||||
```
|
||||
|
||||
### Example Response
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"schema_name": "public",
|
||||
"name": "process_payment",
|
||||
"owner": "postgres",
|
||||
"language": "plpgsql",
|
||||
"definition": "CREATE OR REPLACE PROCEDURE public.process_payment(p_order_id integer, p_amount numeric)\n LANGUAGE plpgsql\nAS $procedure$\nBEGIN\n UPDATE orders SET status = 'paid', amount = p_amount WHERE id = p_order_id;\n INSERT INTO payment_log (order_id, amount, timestamp) VALUES (p_order_id, p_amount, now());\n COMMIT;\nEND\n$procedure$",
|
||||
"description": "Processes payment for an order and logs the transaction"
|
||||
},
|
||||
{
|
||||
"schema_name": "public",
|
||||
"name": "cleanup_old_records",
|
||||
"owner": "postgres",
|
||||
"language": "plpgsql",
|
||||
"definition": "CREATE OR REPLACE PROCEDURE public.cleanup_old_records(p_days_old integer)\n LANGUAGE plpgsql\nAS $procedure$\nDECLARE\n v_deleted integer;\nBEGIN\n DELETE FROM audit_logs WHERE created_at < now() - (p_days_old || ' days')::interval;\n GET DIAGNOSTICS v_deleted = ROW_COUNT;\n RAISE NOTICE 'Deleted % records', v_deleted;\nEND\n$procedure$",
|
||||
"description": "Removes audit log records older than specified days"
|
||||
},
|
||||
{
|
||||
"schema_name": "audit",
|
||||
"name": "audit_table_changes",
|
||||
"owner": "app_user",
|
||||
"language": "plpgsql",
|
||||
"definition": "CREATE OR REPLACE PROCEDURE audit.audit_table_changes()\n LANGUAGE plpgsql\nAS $procedure$\nBEGIN\n INSERT INTO audit.change_log (table_name, operation, changed_at) VALUES (TG_TABLE_NAME, TG_OP, now());\nEND\n$procedure$",
|
||||
"description": null
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
## Output Fields Reference
|
||||
|
||||
| field | type | description |
|
||||
|-------------|---------|-------------|
|
||||
| schema_name | string | Name of the schema containing the stored procedure. |
|
||||
| name | string | Name of the stored procedure. |
|
||||
| owner | string | PostgreSQL role/user who owns the stored procedure. |
|
||||
| language | string | Programming language in which the procedure is written (e.g., plpgsql, sql, c). |
|
||||
| definition | string | Complete SQL definition of the stored procedure, including the CREATE PROCEDURE statement. |
|
||||
| description | string | Optional description or comment for the procedure (may be null if no comment is set). |
|
||||
|
||||
## Use Cases
|
||||
|
||||
- **Code review and auditing**: Export procedure definitions for version control or compliance audits.
|
||||
- **Documentation generation**: Automatically extract procedure metadata and descriptions for documentation.
|
||||
- **Permission auditing**: Identify procedures owned by specific users or in specific schemas.
|
||||
- **Migration planning**: Retrieve all procedure definitions when planning database migrations.
|
||||
- **Dependency analysis**: Review procedure definitions to understand dependencies and call chains.
|
||||
- **Security assessment**: Audit which roles own and can modify stored procedures.
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
- The tool filters at the database level using LIKE pattern matching, so partial matches are supported.
|
||||
- Procedure definitions can be large; consider using the `limit` parameter for large databases with many procedures.
|
||||
- Results are ordered by schema name and procedure name for consistent output.
|
||||
- The default limit of 20 procedures is suitable for most use cases; increase as needed.
|
||||
|
||||
## Notes
|
||||
|
||||
- Only stored **procedures** are returned; functions and other callable objects are excluded via the `prokind = 'p'` filter.
|
||||
- Filtering uses `LIKE` pattern matching, so filter values support partial matches (e.g., `role_name: "app"` will match "app_user", "app_admin", etc.).
|
||||
- The `definition` field contains the complete, runnable CREATE PROCEDURE statement.
|
||||
- The `description` field is populated from comments set via PostgreSQL's COMMENT command and may be null.
|
||||
@@ -1,68 +0,0 @@
|
||||
---
|
||||
title: "Javascript"
|
||||
type: docs
|
||||
weight: 7
|
||||
description: >
|
||||
Javascript SDKs to connect to the MCP Toolbox server.
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
The MCP Toolbox service provides a centralized way to manage and expose tools
|
||||
(like API connectors, database query tools, etc.) for use by GenAI applications.
|
||||
|
||||
These JS SDKs act as clients for that service. They handle the communication needed to:
|
||||
|
||||
* Fetch tool definitions from your running Toolbox instance.
|
||||
* Provide convenient JS objects or functions representing those tools.
|
||||
* Invoke the tools (calling the underlying APIs/services configured in Toolbox).
|
||||
* Handle authentication and parameter binding as needed.
|
||||
|
||||
By using these SDKs, you can easily leverage your Toolbox-managed tools directly
|
||||
within your JS applications or AI orchestration frameworks.
|
||||
|
||||
## Which Package Should I Use?
|
||||
|
||||
Choosing the right package depends on how you are building your application:
|
||||
|
||||
- [`@toolbox-sdk/core`](https://github.com/googleapis/mcp-toolbox-sdk-js/tree/main/packages/toolbox-core):
|
||||
This is a framework agnostic way to connect the tools to popular frameworks
|
||||
like Langchain, LlamaIndex and Genkit.
|
||||
- [`@toolbox-sdk/adk`](https://github.com/googleapis/mcp-toolbox-sdk-js/tree/main/packages/toolbox-adk):
|
||||
This package provides a seamless way to connect to [Google ADK TS](https://github.com/google/adk-js).
|
||||
|
||||
## Available Packages
|
||||
|
||||
This repository hosts the following TS packages. See the package-specific
|
||||
README for detailed installation and usage instructions:
|
||||
|
||||
| Package | Target Use Case | Integration | Path | Details (README) | Npm Version |
|
||||
| :------ | :---------- | :---------- | :---------------------- | :---------- | :---------
|
||||
| `toolbox-core` | Framework-agnostic / Custom applications | Use directly / Custom | `packages/toolbox-core/` | 📄 [View README](https://github.com/googleapis/mcp-toolbox-sdk-js/blob/main/packages/toolbox-core/README.md) |  |
|
||||
| `toolbox-adk` | ADK applications | ADK | `packages/toolbox-adk/` | 📄 [View README](https://github.com/googleapis/mcp-toolbox-sdk-js/blob/main/packages/toolbox-adk/README.md) |  |
|
||||
|
||||
|
||||
## Getting Started
|
||||
|
||||
To get started using Toolbox tools with an application, follow these general steps:
|
||||
|
||||
1. **Set up and Run the Toolbox Service:**
|
||||
|
||||
Before using the SDKs, you need the main MCP Toolbox service running. Follow
|
||||
the instructions here: [**Toolbox Getting Started
|
||||
Guide**](https://github.com/googleapis/genai-toolbox?tab=readme-ov-file#getting-started)
|
||||
|
||||
2. **Install the Appropriate SDK:**
|
||||
|
||||
Choose the package based on your needs (see "[Which Package Should I Use?](#which-package-should-i-use)" above) and install it:
|
||||
|
||||
```bash
|
||||
# For the core, framework-agnostic SDK
|
||||
npm install @toolbox-sdk/core
|
||||
```
|
||||
|
||||
|
||||
|
||||
{{< notice note >}}
|
||||
Source code for [js-sdk](https://github.com/googleapis/mcp-toolbox-sdk-js)
|
||||
{{< /notice >}}
|
||||
15
docs/en/sdks/go-sdk.md
Normal file
15
docs/en/sdks/go-sdk.md
Normal file
@@ -0,0 +1,15 @@
|
||||
---
|
||||
title: "Go SDK"
|
||||
weight: 2
|
||||
description: Go lang client SDK
|
||||
icon: fa-brands fa-golang
|
||||
manualLink: "https://github.com/googleapis/mcp-toolbox-sdk-go"
|
||||
manualLinkTarget: _blank
|
||||
---
|
||||
|
||||
<html>
|
||||
<head>
|
||||
<link rel="canonical" href="https://github.com/googleapis/mcp-toolbox-sdk-go"/>
|
||||
<meta http-equiv="refresh" content="0;url=https://github.com/googleapis/mcp-toolbox-sdk-go"/>
|
||||
</head>
|
||||
</html>
|
||||
@@ -1,25 +0,0 @@
|
||||
---
|
||||
title: "Go SDK"
|
||||
type: docs
|
||||
weight: 7
|
||||
description: >
|
||||
Go SDKs to connect to the MCP Toolbox server.
|
||||
---
|
||||
|
||||
|
||||
## Overview
|
||||
|
||||
The MCP Toolbox service provides a centralized way to manage and expose tools
|
||||
(like API connectors, database query tools, etc.) for use by GenAI applications.
|
||||
|
||||
The Go SDK act as clients for that service. They handle the communication needed to:
|
||||
|
||||
* Fetch tool definitions from your running Toolbox instance.
|
||||
* Provide convenient Go structs representing those tools.
|
||||
* Invoke the tools (calling the underlying APIs/services configured in Toolbox).
|
||||
* Handle authentication and parameter binding as needed.
|
||||
|
||||
By using the SDK, you can easily leverage your Toolbox-managed tools directly
|
||||
within your Go applications or AI orchestration frameworks.
|
||||
|
||||
[Github](https://github.com/googleapis/mcp-toolbox-sdk-go)
|
||||
15
docs/en/sdks/js-sdk.md
Normal file
15
docs/en/sdks/js-sdk.md
Normal file
@@ -0,0 +1,15 @@
|
||||
---
|
||||
title: "JS SDK"
|
||||
weight: 2
|
||||
description: Javascript client SDK
|
||||
icon: fa-brands fa-node-js
|
||||
manualLink: "https://github.com/googleapis/mcp-toolbox-sdk-js"
|
||||
manualLinkTarget: _blank
|
||||
---
|
||||
|
||||
<html>
|
||||
<head>
|
||||
<link rel="canonical" href="https://github.com/googleapis/mcp-toolbox-sdk-js"/>
|
||||
<meta http-equiv="refresh" content="0;url=https://github.com/googleapis/mcp-toolbox-sdk-js"/>
|
||||
</head>
|
||||
</html>
|
||||
15
docs/en/sdks/python-sdk.md
Normal file
15
docs/en/sdks/python-sdk.md
Normal file
@@ -0,0 +1,15 @@
|
||||
---
|
||||
title: "Python SDK"
|
||||
weight: 2
|
||||
description: Python client SDK
|
||||
icon: fa-brands fa-python
|
||||
manualLink: "https://github.com/googleapis/mcp-toolbox-sdk-python"
|
||||
manualLinkTarget: _blank
|
||||
---
|
||||
|
||||
<html>
|
||||
<head>
|
||||
<link rel="canonical" href="https://github.com/googleapis/mcp-toolbox-sdk-python"/>
|
||||
<meta http-equiv="refresh" content="0;url=https://github.com/googleapis/mcp-toolbox-sdk-python"/>
|
||||
</head>
|
||||
</html>
|
||||
@@ -1,25 +0,0 @@
|
||||
---
|
||||
title: "Python SDK"
|
||||
type: docs
|
||||
weight: 7
|
||||
description: >
|
||||
Python SDKs to connect to the MCP Toolbox server.
|
||||
---
|
||||
|
||||
|
||||
## Overview
|
||||
|
||||
The MCP Toolbox service provides a centralized way to manage and expose tools
|
||||
(like API connectors, database query tools, etc.) for use by GenAI applications.
|
||||
|
||||
These Python SDKs act as clients for that service. They handle the communication needed to:
|
||||
|
||||
* Fetch tool definitions from your running Toolbox instance.
|
||||
* Provide convenient Python objects or functions representing those tools.
|
||||
* Invoke the tools (calling the underlying APIs/services configured in Toolbox).
|
||||
* Handle authentication and parameter binding as needed.
|
||||
|
||||
By using these SDKs, you can easily leverage your Toolbox-managed tools directly
|
||||
within your Python applications or AI orchestration frameworks.
|
||||
|
||||
[Github](https://github.com/googleapis/mcp-toolbox-sdk-python)
|
||||
3
go.mod
3
go.mod
@@ -37,6 +37,7 @@ require (
|
||||
github.com/google/go-cmp v0.7.0
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jackc/pgx/v5 v5.7.6
|
||||
github.com/json-iterator/go v1.1.12
|
||||
github.com/looker-open-source/sdk-codegen/go v0.25.21
|
||||
github.com/microsoft/go-mssqldb v1.9.3
|
||||
github.com/nakagami/firebirdsql v0.9.15
|
||||
@@ -59,7 +60,6 @@ require (
|
||||
go.opentelemetry.io/otel/trace v1.38.0
|
||||
golang.org/x/oauth2 v0.33.0
|
||||
google.golang.org/api v0.256.0
|
||||
google.golang.org/genai v1.37.0
|
||||
google.golang.org/genproto v0.0.0-20251022142026-3a174f9686a8
|
||||
google.golang.org/protobuf v1.36.10
|
||||
modernc.org/sqlite v1.40.0
|
||||
@@ -138,7 +138,6 @@ require (
|
||||
github.com/jcmturner/goidentity/v6 v6.0.1 // indirect
|
||||
github.com/jcmturner/gokrb5/v8 v8.4.4 // indirect
|
||||
github.com/jcmturner/rpc/v2 v2.0.3 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0 // indirect
|
||||
github.com/klauspost/compress v1.18.0 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.2.11 // indirect
|
||||
|
||||
2
go.sum
2
go.sum
@@ -1869,8 +1869,6 @@ google.golang.org/appengine v1.6.1/go.mod h1:i06prIuMbXzDqacNJfV5OdTW448YApPu5ww
|
||||
google.golang.org/appengine v1.6.5/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
|
||||
google.golang.org/appengine v1.6.6/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
|
||||
google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
|
||||
google.golang.org/genai v1.37.0 h1:dgp71k1wQ+/+APdZrN3LFgAGnVnr5IdTF1Oj0Dg+BQc=
|
||||
google.golang.org/genai v1.37.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
|
||||
google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
|
||||
google.golang.org/genproto v0.0.0-20190307195333-5fe7a883aa19/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=
|
||||
google.golang.org/genproto v0.0.0-20190418145605-e7d98fc518a7/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=
|
||||
|
||||
@@ -1,59 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// 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
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package embeddingmodels
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// EmbeddingModelConfig is the interface for configuring embedding models.
|
||||
type EmbeddingModelConfig interface {
|
||||
EmbeddingModelConfigKind() string
|
||||
Initialize(context.Context) (EmbeddingModel, error)
|
||||
}
|
||||
|
||||
type EmbeddingModel interface {
|
||||
EmbeddingModelKind() string
|
||||
ToConfig() EmbeddingModelConfig
|
||||
EmbedParameters(context.Context, []string) ([][]float32, error)
|
||||
}
|
||||
|
||||
type VectorFormatter func(vectorFloats []float32) any
|
||||
|
||||
// FormatVectorForPgvector converts a slice of floats into a PostgreSQL vector literal string: '[x, y, z]'
|
||||
func FormatVectorForPgvector(vectorFloats []float32) any {
|
||||
if len(vectorFloats) == 0 {
|
||||
return "[]"
|
||||
}
|
||||
|
||||
// Pre-allocate the builder.
|
||||
var b strings.Builder
|
||||
b.Grow(len(vectorFloats) * 10)
|
||||
|
||||
b.WriteByte('[')
|
||||
for i, f := range vectorFloats {
|
||||
if i > 0 {
|
||||
b.WriteString(", ")
|
||||
}
|
||||
b.Write(strconv.AppendFloat(nil, float64(f), 'g', -1, 32))
|
||||
}
|
||||
b.WriteByte(']')
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
var _ VectorFormatter = FormatVectorForPgvector
|
||||
@@ -1,122 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// 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
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"google.golang.org/genai"
|
||||
)
|
||||
|
||||
const EmbeddingModelKind string = "gemini"
|
||||
|
||||
// validate interface
|
||||
var _ embeddingmodels.EmbeddingModelConfig = Config{}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Model string `yaml:"model" validate:"required"`
|
||||
ApiKey string `yaml:"apiKey"`
|
||||
Dimension int32 `yaml:"dimension"`
|
||||
}
|
||||
|
||||
// Returns the embedding model kind
|
||||
func (cfg Config) EmbeddingModelConfigKind() string {
|
||||
return EmbeddingModelKind
|
||||
}
|
||||
|
||||
// Initialize a Gemini embedding model
|
||||
func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingModel, error) {
|
||||
// Get client configs
|
||||
configs := &genai.ClientConfig{}
|
||||
if cfg.ApiKey != "" {
|
||||
configs.APIKey = cfg.ApiKey
|
||||
}
|
||||
|
||||
// Create new Gemini API client
|
||||
client, err := genai.NewClient(ctx, configs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to create Gemini API client")
|
||||
}
|
||||
|
||||
m := &EmbeddingModel{
|
||||
Config: cfg,
|
||||
Client: client,
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
var _ embeddingmodels.EmbeddingModel = EmbeddingModel{}
|
||||
|
||||
type EmbeddingModel struct {
|
||||
Client *genai.Client
|
||||
Config
|
||||
}
|
||||
|
||||
// Returns the embedding model kind
|
||||
func (m EmbeddingModel) EmbeddingModelKind() string {
|
||||
return EmbeddingModelKind
|
||||
}
|
||||
|
||||
func (m EmbeddingModel) ToConfig() embeddingmodels.EmbeddingModelConfig {
|
||||
return m.Config
|
||||
}
|
||||
|
||||
func (m EmbeddingModel) EmbedParameters(ctx context.Context, parameters []string) ([][]float32, error) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
|
||||
}
|
||||
|
||||
contents := convertStringsToContents(parameters)
|
||||
|
||||
embedConfig := &genai.EmbedContentConfig{
|
||||
TaskType: "SEMANTIC_SIMILARITY",
|
||||
}
|
||||
|
||||
if m.Dimension > 0 {
|
||||
embedConfig.OutputDimensionality = genai.Ptr(m.Dimension)
|
||||
}
|
||||
|
||||
result, err := m.Client.Models.EmbedContent(ctx, m.Model, contents, embedConfig)
|
||||
if err != nil {
|
||||
logger.ErrorContext(ctx, "Error calling EmbedContent for model %s: %v", m.Model, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
embeddings := make([][]float32, 0, len(result.Embeddings))
|
||||
for _, embedding := range result.Embeddings {
|
||||
embeddings = append(embeddings, embedding.Values)
|
||||
}
|
||||
|
||||
logger.InfoContext(ctx, "Successfully embedded %d text parameters using model %s", len(parameters), m.Model)
|
||||
|
||||
return embeddings, nil
|
||||
}
|
||||
|
||||
// convertStringsToContents takes a slice of strings and converts it into a slice of *genai.Content objects.
|
||||
func convertStringsToContents(texts []string) []*genai.Content {
|
||||
contents := make([]*genai.Content, 0, len(texts))
|
||||
|
||||
for _, text := range texts {
|
||||
content := genai.NewContentFromText(text, "")
|
||||
contents = append(contents, content)
|
||||
}
|
||||
return contents
|
||||
}
|
||||
@@ -1,130 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// 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
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gemini_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYamlGemini(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.EmbeddingModelConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
embeddingModels:
|
||||
my-gemini-model:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
`,
|
||||
want: map[string]embeddingmodels.EmbeddingModelConfig{
|
||||
"my-gemini-model": gemini.Config{
|
||||
Name: "my-gemini-model",
|
||||
Kind: gemini.EmbeddingModelKind,
|
||||
Model: "text-embedding-004",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "full example with optional fields",
|
||||
in: `
|
||||
embeddingModels:
|
||||
complex-gemini:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
apiKey: "test-api-key"
|
||||
dimension: 768
|
||||
`,
|
||||
want: map[string]embeddingmodels.EmbeddingModelConfig{
|
||||
"complex-gemini": gemini.Config{
|
||||
Name: "complex-gemini",
|
||||
Kind: gemini.EmbeddingModelKind,
|
||||
Model: "text-embedding-004",
|
||||
ApiKey: "test-api-key",
|
||||
Dimension: 768,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Models server.EmbeddingModelConfigs `yaml:"embeddingModels"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if !cmp.Equal(tc.want, got.Models) {
|
||||
t.Fatalf("incorrect parse: %v", cmp.Diff(tc.want, got.Models))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
func TestFailParseFromYamlGemini(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "missing required model field",
|
||||
in: `
|
||||
embeddingModels:
|
||||
bad-model:
|
||||
kind: gemini
|
||||
`,
|
||||
// Removed the specific model name from the prefix to match your output
|
||||
err: "unable to parse as \"gemini\": Key: 'Config.Model' Error:Field validation for 'Model' failed on the 'required' tag",
|
||||
},
|
||||
{
|
||||
desc: "unknown field",
|
||||
in: `
|
||||
embeddingModels:
|
||||
bad-field:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
invalid_param: true
|
||||
`,
|
||||
// Updated to match the specific line-starting format of your error output
|
||||
err: "unable to parse as \"gemini\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | kind: gemini\n 3 | model: text-embedding-004",
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Models server.EmbeddingModelConfigs `yaml:"embeddingModels"`
|
||||
}{}
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err == nil {
|
||||
t.Fatalf("expect parsing to fail")
|
||||
}
|
||||
if err.Error() != tc.err {
|
||||
t.Fatalf("unexpected error:\ngot: %q\nwant: %q", err.Error(), tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -224,10 +224,6 @@ tools:
|
||||
kind: postgres-list-roles
|
||||
source: alloydb-pg-source
|
||||
|
||||
list_stored_procedure:
|
||||
kind: postgres-list-stored-procedure
|
||||
source: alloydb-pg-source
|
||||
|
||||
toolsets:
|
||||
alloydb_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -258,4 +254,3 @@ toolsets:
|
||||
- list_database_stats
|
||||
- list_roles
|
||||
- list_table_stats
|
||||
- list_stored_procedure
|
||||
|
||||
@@ -18,7 +18,6 @@ sources:
|
||||
project: ${BIGQUERY_PROJECT}
|
||||
location: ${BIGQUERY_LOCATION:}
|
||||
useClientOAuth: ${BIGQUERY_USE_CLIENT_OAUTH:false}
|
||||
scopes: ${BIGQUERY_SCOPES:}
|
||||
|
||||
tools:
|
||||
analyze_contribution:
|
||||
|
||||
@@ -19,8 +19,8 @@ sources:
|
||||
region: ${CLOUD_SQL_MYSQL_REGION}
|
||||
instance: ${CLOUD_SQL_MYSQL_INSTANCE}
|
||||
database: ${CLOUD_SQL_MYSQL_DATABASE}
|
||||
user: ${CLOUD_SQL_MYSQL_USER:}
|
||||
password: ${CLOUD_SQL_MYSQL_PASSWORD:}
|
||||
user: ${CLOUD_SQL_MYSQL_USER}
|
||||
password: ${CLOUD_SQL_MYSQL_PASSWORD}
|
||||
ipType: ${CLOUD_SQL_MYSQL_IP_TYPE:PUBLIC}
|
||||
tools:
|
||||
execute_sql:
|
||||
|
||||
@@ -226,10 +226,6 @@ tools:
|
||||
kind: postgres-list-roles
|
||||
source: cloudsql-pg-source
|
||||
|
||||
list_stored_procedure:
|
||||
kind: postgres-list-stored-procedure
|
||||
source: cloudsql-pg-source
|
||||
|
||||
toolsets:
|
||||
cloud_sql_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -260,4 +256,3 @@ toolsets:
|
||||
- list_database_stats
|
||||
- list_roles
|
||||
- list_table_stats
|
||||
- list_stored_procedure
|
||||
|
||||
@@ -959,6 +959,21 @@ tools:
|
||||
Output:
|
||||
A confirmation message upon successful file deletion.
|
||||
|
||||
validate_project:
|
||||
kind: looker-validate-project
|
||||
source: looker-source
|
||||
description: |
|
||||
This tool checks a LookML project for syntax errors.
|
||||
|
||||
Prerequisite: The Looker session must be in Development Mode. Use `dev_mode: true` first.
|
||||
|
||||
Parameters:
|
||||
- project_id (required): The unique ID of the LookML project.
|
||||
|
||||
Output:
|
||||
A list of error details including the file path and line number, and also a list of models
|
||||
that are not currently valid due to LookML errors.
|
||||
|
||||
get_connections:
|
||||
kind: looker-get-connections
|
||||
source: looker-source
|
||||
@@ -1072,6 +1087,7 @@ toolsets:
|
||||
- create_project_file
|
||||
- update_project_file
|
||||
- delete_project_file
|
||||
- validate_project
|
||||
- get_connections
|
||||
- get_connection_schemas
|
||||
- get_connection_databases
|
||||
|
||||
@@ -225,10 +225,6 @@ tools:
|
||||
kind: postgres-list-roles
|
||||
source: postgresql-source
|
||||
|
||||
list_stored_procedure:
|
||||
kind: postgres-list-stored-procedure
|
||||
source: postgresql-source
|
||||
|
||||
toolsets:
|
||||
postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -259,4 +255,3 @@ toolsets:
|
||||
- list_database_stats
|
||||
- list_roles
|
||||
- list_table_stats
|
||||
- list_stored_procedure
|
||||
|
||||
@@ -246,14 +246,6 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
s.logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
|
||||
|
||||
params, err = tool.EmbedParams(ctx, params, s.ResourceMgr.GetEmbeddingModelMap())
|
||||
if err != nil {
|
||||
err = fmt.Errorf("error embedding parameters: %w", err)
|
||||
s.logger.DebugContext(ctx, err.Error())
|
||||
_ = render.Render(w, r, newErrResponse(err, http.StatusBadRequest))
|
||||
return
|
||||
}
|
||||
|
||||
res, err := tool.Invoke(ctx, s.ResourceMgr, params, accessToken)
|
||||
|
||||
// Determine what error to return to the users.
|
||||
|
||||
@@ -24,7 +24,6 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||
@@ -65,10 +64,6 @@ func (t MockTool) ParseParams(data map[string]any, claimsMap map[string]map[stri
|
||||
return parameters.ParseParams(t.Params, data, claimsMap)
|
||||
}
|
||||
|
||||
func (t MockTool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Params, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t MockTool) Manifest() tools.Manifest {
|
||||
pMs := make([]parameters.ParameterManifest, 0, len(t.Params))
|
||||
for _, p := range t.Params {
|
||||
@@ -281,7 +276,7 @@ func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, tools
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(nil, nil, nil, tools, toolsets, prompts, promptsets)
|
||||
resourceManager := resources.NewResourceManager(nil, nil, tools, toolsets, prompts, promptsets)
|
||||
|
||||
server := Server{
|
||||
version: fakeVersionString,
|
||||
|
||||
@@ -21,8 +21,6 @@ import (
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth/google"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -40,8 +38,6 @@ type ServerConfig struct {
|
||||
SourceConfigs SourceConfigs
|
||||
// AuthServiceConfigs defines what sources of authentication are available for tools.
|
||||
AuthServiceConfigs AuthServiceConfigs
|
||||
// EmbeddingModelConfigs defines a models used to embed parameters.
|
||||
EmbeddingModelConfigs EmbeddingModelConfigs
|
||||
// ToolConfigs defines what tools are available.
|
||||
ToolConfigs ToolConfigs
|
||||
// ToolsetConfigs defines what tools are available.
|
||||
@@ -209,50 +205,6 @@ func (c *AuthServiceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(i
|
||||
return nil
|
||||
}
|
||||
|
||||
// EmbeddingModelConfigs is a type used to allow unmarshal of the embedding model config map
|
||||
type EmbeddingModelConfigs map[string]embeddingmodels.EmbeddingModelConfig
|
||||
|
||||
// validate interface
|
||||
var _ yaml.InterfaceUnmarshalerContext = &EmbeddingModelConfigs{}
|
||||
|
||||
func (c *EmbeddingModelConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
|
||||
*c = make(EmbeddingModelConfigs)
|
||||
// Parse the 'kind' fields for each embedding model
|
||||
var raw map[string]util.DelayedUnmarshaler
|
||||
if err := unmarshal(&raw); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for name, u := range raw {
|
||||
// Unmarshal to a general type that ensure it capture all fields
|
||||
var v map[string]any
|
||||
if err := u.Unmarshal(&v); err != nil {
|
||||
return fmt.Errorf("unable to unmarshal embedding model %q: %w", name, err)
|
||||
}
|
||||
|
||||
kind, ok := v["kind"]
|
||||
if !ok {
|
||||
return fmt.Errorf("missing 'kind' field for embedding model %q", name)
|
||||
}
|
||||
|
||||
dec, err := util.NewStrictDecoder(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error creating decoder: %w", err)
|
||||
}
|
||||
switch kind {
|
||||
case gemini.EmbeddingModelKind:
|
||||
actual := gemini.Config{Name: name}
|
||||
if err := dec.DecodeContext(ctx, &actual); err != nil {
|
||||
return fmt.Errorf("unable to parse as %q: %w", kind, err)
|
||||
}
|
||||
(*c)[name] = actual
|
||||
default:
|
||||
return fmt.Errorf("%q is not a valid kind of auth source", kind)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ToolConfigs is a type used to allow unmarshal of the tool configs
|
||||
type ToolConfigs map[string]tools.ToolConfig
|
||||
|
||||
|
||||
@@ -1107,7 +1107,7 @@ func TestStdioSession(t *testing.T) {
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
|
||||
resourceManager := resources.NewResourceManager(nil, nil, toolsMap, toolsets, promptsMap, promptsets)
|
||||
|
||||
server := &Server{
|
||||
version: fakeVersionString,
|
||||
|
||||
@@ -18,7 +18,6 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -26,33 +25,30 @@ import (
|
||||
|
||||
// ResourceManager contains available resources for the server. Should be initialized with NewResourceManager().
|
||||
type ResourceManager struct {
|
||||
mu sync.RWMutex
|
||||
sources map[string]sources.Source
|
||||
authServices map[string]auth.AuthService
|
||||
embeddingModels map[string]embeddingmodels.EmbeddingModel
|
||||
tools map[string]tools.Tool
|
||||
toolsets map[string]tools.Toolset
|
||||
prompts map[string]prompts.Prompt
|
||||
promptsets map[string]prompts.Promptset
|
||||
mu sync.RWMutex
|
||||
sources map[string]sources.Source
|
||||
authServices map[string]auth.AuthService
|
||||
tools map[string]tools.Tool
|
||||
toolsets map[string]tools.Toolset
|
||||
prompts map[string]prompts.Prompt
|
||||
promptsets map[string]prompts.Promptset
|
||||
}
|
||||
|
||||
func NewResourceManager(
|
||||
sourcesMap map[string]sources.Source,
|
||||
authServicesMap map[string]auth.AuthService,
|
||||
embeddingModelsMap map[string]embeddingmodels.EmbeddingModel,
|
||||
toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset,
|
||||
promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset,
|
||||
|
||||
) *ResourceManager {
|
||||
resourceMgr := &ResourceManager{
|
||||
mu: sync.RWMutex{},
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
embeddingModels: embeddingModelsMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
prompts: promptsMap,
|
||||
promptsets: promptsetsMap,
|
||||
mu: sync.RWMutex{},
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
prompts: promptsMap,
|
||||
promptsets: promptsetsMap,
|
||||
}
|
||||
|
||||
return resourceMgr
|
||||
@@ -72,13 +68,6 @@ func (r *ResourceManager) GetAuthService(authServiceName string) (auth.AuthServi
|
||||
return authService, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetEmbeddingModel(embeddingModelName string) (embeddingmodels.EmbeddingModel, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
model, ok := r.embeddingModels[embeddingModelName]
|
||||
return model, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetTool(toolName string) (tools.Tool, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
@@ -107,12 +96,11 @@ func (r *ResourceManager) GetPromptset(promptsetName string) (prompts.Promptset,
|
||||
return promptset, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) SetResources(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset, promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset) {
|
||||
func (r *ResourceManager) SetResources(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset, promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.sources = sourcesMap
|
||||
r.authServices = authServicesMap
|
||||
r.embeddingModels = embeddingModelsMap
|
||||
r.tools = toolsMap
|
||||
r.toolsets = toolsetsMap
|
||||
r.prompts = promptsMap
|
||||
@@ -129,16 +117,6 @@ func (r *ResourceManager) GetAuthServiceMap() map[string]auth.AuthService {
|
||||
return copiedMap
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetEmbeddingModelMap() map[string]embeddingmodels.EmbeddingModel {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
copiedMap := make(map[string]embeddingmodels.EmbeddingModel, len(r.embeddingModels))
|
||||
for k, v := range r.embeddingModels {
|
||||
copiedMap[k] = v
|
||||
}
|
||||
return copiedMap
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetToolsMap() map[string]tools.Tool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
@@ -37,7 +36,6 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
newAuth := map[string]auth.AuthService{"example-auth": nil}
|
||||
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
|
||||
newTools := map[string]tools.Tool{"example-tool": nil}
|
||||
newToolsets := map[string]tools.Toolset{
|
||||
"example-toolset": {
|
||||
@@ -56,7 +54,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
Prompts: []*prompts.Prompt{},
|
||||
},
|
||||
}
|
||||
resMgr := resources.NewResourceManager(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
resMgr := resources.NewResourceManager(newSources, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
|
||||
gotSource, _ := resMgr.GetSource("example-source")
|
||||
if diff := cmp.Diff(gotSource, newSources["example-source"]); diff != "" {
|
||||
@@ -97,7 +95,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resMgr.SetResources(updateSource, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
resMgr.SetResources(updateSource, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
gotSource, _ = resMgr.GetSource("example-source2")
|
||||
if diff := cmp.Diff(gotSource, updateSource["example-source2"]); diff != "" {
|
||||
t.Errorf("error updating server, sources (-want +got):\n%s", diff)
|
||||
|
||||
@@ -30,7 +30,6 @@ import (
|
||||
"github.com/go-chi/cors"
|
||||
"github.com/go-chi/httplog/v2"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||
@@ -57,7 +56,6 @@ type Server struct {
|
||||
func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
map[string]sources.Source,
|
||||
map[string]auth.AuthService,
|
||||
map[string]embeddingmodels.EmbeddingModel,
|
||||
map[string]tools.Tool,
|
||||
map[string]tools.Toolset,
|
||||
map[string]prompts.Prompt,
|
||||
@@ -93,7 +91,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return s, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
sourcesMap[name] = s
|
||||
}
|
||||
@@ -121,7 +119,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return a, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
authServicesMap[name] = a
|
||||
}
|
||||
@@ -131,34 +129,6 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d authServices: %s", len(authServicesMap), strings.Join(authServiceNames, ", ")))
|
||||
|
||||
// Initialize and validate embedding models from configs.
|
||||
embeddingModelsMap := make(map[string]embeddingmodels.EmbeddingModel)
|
||||
for name, ec := range cfg.EmbeddingModelConfigs {
|
||||
em, err := func() (embeddingmodels.EmbeddingModel, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/embeddingmodel/init",
|
||||
trace.WithAttributes(attribute.String("model_kind", ec.EmbeddingModelConfigKind())),
|
||||
trace.WithAttributes(attribute.String("model_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
em, err := ec.Initialize(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize embedding model %q: %w", name, err)
|
||||
}
|
||||
return em, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
embeddingModelsMap[name] = em
|
||||
}
|
||||
embeddingModelNames := make([]string, 0, len(embeddingModelsMap))
|
||||
for name := range embeddingModelsMap {
|
||||
embeddingModelNames = append(embeddingModelNames, name)
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d embeddingModels: %s", len(embeddingModelsMap), strings.Join(embeddingModelNames, ", ")))
|
||||
|
||||
// initialize and validate the tools from configs
|
||||
toolsMap := make(map[string]tools.Tool)
|
||||
for name, tc := range cfg.ToolConfigs {
|
||||
@@ -177,7 +147,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return t, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
toolsMap[name] = t
|
||||
}
|
||||
@@ -214,7 +184,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return t, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
toolsetsMap[name] = t
|
||||
}
|
||||
@@ -246,7 +216,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return p, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
promptsMap[name] = p
|
||||
}
|
||||
@@ -283,7 +253,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return p, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
promptsetsMap[name] = p
|
||||
}
|
||||
@@ -297,7 +267,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d promptsets: %s", len(promptsetsMap), strings.Join(promptsetNames, ", ")))
|
||||
|
||||
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
}
|
||||
|
||||
// NewServer returns a Server object based on provided Config.
|
||||
@@ -350,7 +320,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
httpLogger := httplog.NewLogger("httplog", httpOpts)
|
||||
r.Use(httplog.RequestLogger(httpLogger))
|
||||
|
||||
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize configs: %w", err)
|
||||
}
|
||||
@@ -360,7 +330,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
|
||||
s := &Server{
|
||||
version: cfg.Version,
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
@@ -145,7 +144,6 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
newAuth := map[string]auth.AuthService{"example-auth": nil}
|
||||
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
|
||||
newTools := map[string]tools.Tool{"example-tool": nil}
|
||||
newToolsets := map[string]tools.Toolset{
|
||||
"example-toolset": {
|
||||
@@ -164,7 +162,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
Prompts: []*prompts.Prompt{},
|
||||
},
|
||||
}
|
||||
s.ResourceMgr.SetResources(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
s.ResourceMgr.SetResources(newSources, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
if err != nil {
|
||||
t.Errorf("error updating server: %s", err)
|
||||
}
|
||||
|
||||
@@ -24,7 +24,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
@@ -105,22 +104,22 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.Pool.Query(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
return nil, fmt.Errorf("unable to execute query: %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, statement, params)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
v, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
vMap := make(map[string]any)
|
||||
for i, f := range fields {
|
||||
row.Add(f.Name, v[i])
|
||||
vMap[f.Name] = v[i]
|
||||
}
|
||||
out = append(out, row)
|
||||
out = append(out, vMap)
|
||||
}
|
||||
// this will catch actual query execution errors
|
||||
if err := results.Err(); err != nil {
|
||||
|
||||
@@ -17,9 +17,7 @@ package bigquery
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -28,24 +26,18 @@ import (
|
||||
dataplexapi "cloud.google.com/go/dataplex/apiv1"
|
||||
"github.com/goccy/go-yaml"
|
||||
"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/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/google"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/googleapi"
|
||||
"google.golang.org/api/impersonate"
|
||||
"google.golang.org/api/iterator"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
|
||||
const SourceKind string = "bigquery"
|
||||
|
||||
// CloudPlatformScope is a broad scope for Google Cloud Platform services.
|
||||
const CloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform"
|
||||
|
||||
const (
|
||||
// No write operations are allowed.
|
||||
WriteModeBlocked string = "blocked"
|
||||
@@ -80,42 +72,14 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources
|
||||
|
||||
type Config struct {
|
||||
// BigQuery configs
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Project string `yaml:"project" validate:"required"`
|
||||
Location string `yaml:"location"`
|
||||
WriteMode string `yaml:"writeMode"`
|
||||
AllowedDatasets StringOrStringSlice `yaml:"allowedDatasets"`
|
||||
UseClientOAuth bool `yaml:"useClientOAuth"`
|
||||
ImpersonateServiceAccount string `yaml:"impersonateServiceAccount"`
|
||||
Scopes StringOrStringSlice `yaml:"scopes"`
|
||||
}
|
||||
|
||||
// StringOrStringSlice is a custom type that can unmarshal both a single string
|
||||
// (which it splits by comma) and a sequence of strings into a string slice.
|
||||
type StringOrStringSlice []string
|
||||
|
||||
// UnmarshalYAML implements the yaml.Unmarshaler interface.
|
||||
func (s *StringOrStringSlice) UnmarshalYAML(unmarshal func(any) error) error {
|
||||
var v any
|
||||
if err := unmarshal(&v); err != nil {
|
||||
return err
|
||||
}
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
*s = strings.Split(val, ",")
|
||||
return nil
|
||||
case []any:
|
||||
for _, item := range val {
|
||||
if str, ok := item.(string); ok {
|
||||
*s = append(*s, str)
|
||||
} else {
|
||||
return fmt.Errorf("element in sequence is not a string: %v", item)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("cannot unmarshal %T into StringOrStringSlice", v)
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Project string `yaml:"project" validate:"required"`
|
||||
Location string `yaml:"location"`
|
||||
WriteMode string `yaml:"writeMode"`
|
||||
AllowedDatasets []string `yaml:"allowedDatasets"`
|
||||
UseClientOAuth bool `yaml:"useClientOAuth"`
|
||||
ImpersonateServiceAccount string `yaml:"impersonateServiceAccount"`
|
||||
}
|
||||
|
||||
func (r Config) SourceConfigKind() string {
|
||||
@@ -164,7 +128,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
|
||||
|
||||
} else {
|
||||
// Initializes a BigQuery Google SQL source
|
||||
client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount, r.Scopes)
|
||||
client, restService, tokenSource, err = initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location, r.ImpersonateServiceAccount)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from ADC: %w", err)
|
||||
}
|
||||
@@ -427,26 +391,19 @@ func (s *Source) BigQueryTokenSource() oauth2.TokenSource {
|
||||
return s.TokenSource
|
||||
}
|
||||
|
||||
func (s *Source) BigQueryTokenSourceWithScope(ctx context.Context, scopes []string) (oauth2.TokenSource, error) {
|
||||
if len(scopes) == 0 {
|
||||
scopes = s.Scopes
|
||||
if len(scopes) == 0 {
|
||||
scopes = []string{CloudPlatformScope}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Source) BigQueryTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error) {
|
||||
if s.ImpersonateServiceAccount != "" {
|
||||
// Create impersonated credentials token source with the requested scopes
|
||||
// Create impersonated credentials token source with the requested scope
|
||||
ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
|
||||
TargetPrincipal: s.ImpersonateServiceAccount,
|
||||
Scopes: scopes,
|
||||
Scopes: []string{scope},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create impersonated credentials for %q with scopes %v: %w", s.ImpersonateServiceAccount, scopes, err)
|
||||
return nil, fmt.Errorf("failed to create impersonated credentials for %q with scope %q: %w", s.ImpersonateServiceAccount, scope, err)
|
||||
}
|
||||
return ts, nil
|
||||
}
|
||||
return google.DefaultTokenSource(ctx, scopes...)
|
||||
return google.DefaultTokenSource(ctx, scope)
|
||||
}
|
||||
|
||||
func (s *Source) GetMaxQueryResultRows() int {
|
||||
@@ -492,7 +449,7 @@ func (s *Source) lazyInitDataplexClient(ctx context.Context, tracer trace.Tracer
|
||||
|
||||
return func() (*dataplexapi.CatalogClient, DataplexClientCreator, error) {
|
||||
once.Do(func() {
|
||||
c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientOAuth, s.ImpersonateServiceAccount, s.Scopes)
|
||||
c, cc, e := initDataplexConnection(ctx, tracer, s.Name, s.Project, s.UseClientOAuth, s.ImpersonateServiceAccount)
|
||||
if e != nil {
|
||||
err = fmt.Errorf("failed to initialize dataplex client: %w", e)
|
||||
return
|
||||
@@ -526,131 +483,6 @@ func (s *Source) lazyInitDataplexClient(ctx context.Context, tracer trace.Tracer
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Source) RetrieveClientAndService(accessToken tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error) {
|
||||
bqClient := s.BigQueryClient()
|
||||
restService := s.BigQueryRestService()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if s.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, restService, err = s.BigQueryClientCreator()(tokenStr, true)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
return bqClient, restService, nil
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, bqClient *bigqueryapi.Client, statement, statementType string, params []bigqueryapi.QueryParameter, connProps []*bigqueryapi.ConnectionProperty) (any, error) {
|
||||
query := bqClient.Query(statement)
|
||||
query.Location = bqClient.Location
|
||||
if params != nil {
|
||||
query.Parameters = params
|
||||
}
|
||||
if connProps != nil {
|
||||
query.ConnectionProperties = connProps
|
||||
}
|
||||
|
||||
// This block handles SELECT statements, which return a row set.
|
||||
// We iterate through the results, convert each row into a map of
|
||||
// column names to values, and return the collection of rows.
|
||||
job, err := query.Run(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
it, err := job.Read(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to read query results: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for {
|
||||
var val []bigqueryapi.Value
|
||||
err = it.Next(&val)
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to iterate through query results: %w", err)
|
||||
}
|
||||
schema := it.Schema
|
||||
row := orderedmap.Row{}
|
||||
for i, field := range schema {
|
||||
row.Add(field.Name, NormalizeValue(val[i]))
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
// If the query returned any rows, return them directly.
|
||||
if len(out) > 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This handles the standard case for a SELECT query that successfully
|
||||
// executes but returns zero rows.
|
||||
if statementType == "SELECT" {
|
||||
return "The query returned 0 rows.", nil
|
||||
}
|
||||
// This is the fallback for a successful query that doesn't return content.
|
||||
// In most cases, this will be for DML/DDL statements like INSERT, UPDATE, CREATE, etc.
|
||||
// However, it is also possible that this was a query that was expected to return rows
|
||||
// but returned none, a case that we cannot distinguish here.
|
||||
return "Query executed successfully and returned no content.", nil
|
||||
}
|
||||
|
||||
// NormalizeValue converts BigQuery specific types to standard JSON-compatible types.
|
||||
// Specifically, it handles *big.Rat (used for NUMERIC/BIGNUMERIC) by converting
|
||||
// them to decimal strings with up to 38 digits of precision, trimming trailing zeros.
|
||||
// It recursively handles slices (arrays) and maps (structs) using reflection.
|
||||
func NormalizeValue(v any) any {
|
||||
if v == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Handle *big.Rat specifically.
|
||||
if rat, ok := v.(*big.Rat); ok {
|
||||
// Convert big.Rat to a decimal string.
|
||||
// Use a precision of 38 digits (enough for BIGNUMERIC and NUMERIC)
|
||||
// and trim trailing zeros to match BigQuery's behavior.
|
||||
s := rat.FloatString(38)
|
||||
if strings.Contains(s, ".") {
|
||||
s = strings.TrimRight(s, "0")
|
||||
s = strings.TrimRight(s, ".")
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Use reflection for slices and maps to handle various underlying types.
|
||||
rv := reflect.ValueOf(v)
|
||||
switch rv.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
// Preserve []byte as is, so json.Marshal encodes it as Base64 string (BigQuery BYTES behavior).
|
||||
if rv.Type().Elem().Kind() == reflect.Uint8 {
|
||||
return v
|
||||
}
|
||||
newSlice := make([]any, rv.Len())
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
newSlice[i] = NormalizeValue(rv.Index(i).Interface())
|
||||
}
|
||||
return newSlice
|
||||
case reflect.Map:
|
||||
// Ensure keys are strings to produce a JSON-compatible map.
|
||||
if rv.Type().Key().Kind() != reflect.String {
|
||||
return v
|
||||
}
|
||||
newMap := make(map[string]any, rv.Len())
|
||||
iter := rv.MapRange()
|
||||
for iter.Next() {
|
||||
newMap[iter.Key().String()] = NormalizeValue(iter.Value().Interface())
|
||||
}
|
||||
return newMap
|
||||
}
|
||||
|
||||
return v
|
||||
}
|
||||
|
||||
func initBigQueryConnection(
|
||||
ctx context.Context,
|
||||
tracer trace.Tracer,
|
||||
@@ -658,7 +490,6 @@ func initBigQueryConnection(
|
||||
project string,
|
||||
location string,
|
||||
impersonateServiceAccount string,
|
||||
scopes []string,
|
||||
) (*bigqueryapi.Client, *bigqueryrestapi.Service, oauth2.TokenSource, error) {
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
defer span.End()
|
||||
@@ -671,21 +502,12 @@ func initBigQueryConnection(
|
||||
var tokenSource oauth2.TokenSource
|
||||
var opts []option.ClientOption
|
||||
|
||||
var credScopes []string
|
||||
if len(scopes) > 0 {
|
||||
credScopes = scopes
|
||||
} else if impersonateServiceAccount != "" {
|
||||
credScopes = []string{CloudPlatformScope}
|
||||
} else {
|
||||
credScopes = []string{bigqueryapi.Scope}
|
||||
}
|
||||
|
||||
if impersonateServiceAccount != "" {
|
||||
// Create impersonated credentials token source
|
||||
// Create impersonated credentials token source with cloud-platform scope
|
||||
// This broader scope is needed for tools like conversational analytics
|
||||
cloudPlatformTokenSource, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
|
||||
TargetPrincipal: impersonateServiceAccount,
|
||||
Scopes: credScopes,
|
||||
Scopes: []string{"https://www.googleapis.com/auth/cloud-platform"},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, nil, nil, fmt.Errorf("failed to create impersonated credentials for %q: %w", impersonateServiceAccount, err)
|
||||
@@ -697,9 +519,9 @@ func initBigQueryConnection(
|
||||
}
|
||||
} else {
|
||||
// Use default credentials
|
||||
cred, err := google.FindDefaultCredentials(ctx, credScopes...)
|
||||
cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope)
|
||||
if err != nil {
|
||||
return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scopes %v: %w", credScopes, err)
|
||||
return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err)
|
||||
}
|
||||
tokenSource = cred.TokenSource
|
||||
opts = []option.ClientOption{
|
||||
@@ -790,7 +612,6 @@ func initDataplexConnection(
|
||||
project string,
|
||||
useClientOAuth bool,
|
||||
impersonateServiceAccount string,
|
||||
scopes []string,
|
||||
) (*dataplexapi.CatalogClient, DataplexClientCreator, error) {
|
||||
var client *dataplexapi.CatalogClient
|
||||
var clientCreator DataplexClientCreator
|
||||
@@ -809,16 +630,11 @@ func initDataplexConnection(
|
||||
} else {
|
||||
var opts []option.ClientOption
|
||||
|
||||
credScopes := scopes
|
||||
if len(credScopes) == 0 {
|
||||
credScopes = []string{CloudPlatformScope}
|
||||
}
|
||||
|
||||
if impersonateServiceAccount != "" {
|
||||
// Create impersonated credentials token source
|
||||
ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
|
||||
TargetPrincipal: impersonateServiceAccount,
|
||||
Scopes: credScopes,
|
||||
Scopes: []string{"https://www.googleapis.com/auth/cloud-platform"},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to create impersonated credentials for %q: %w", impersonateServiceAccount, err)
|
||||
@@ -829,7 +645,7 @@ func initDataplexConnection(
|
||||
}
|
||||
} else {
|
||||
// Use default credentials
|
||||
cred, err := google.FindDefaultCredentials(ctx, credScopes...)
|
||||
cred, err := google.FindDefaultCredentials(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err)
|
||||
}
|
||||
|
||||
@@ -15,8 +15,6 @@
|
||||
package bigquery_test
|
||||
|
||||
import (
|
||||
"math/big"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
@@ -132,28 +130,6 @@ func TestParseFromYamlBigQuery(t *testing.T) {
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "with custom scopes example",
|
||||
in: `
|
||||
sources:
|
||||
my-instance:
|
||||
kind: bigquery
|
||||
project: my-project
|
||||
location: us
|
||||
scopes:
|
||||
- https://www.googleapis.com/auth/bigquery
|
||||
- https://www.googleapis.com/auth/cloud-platform
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-instance": bigquery.Config{
|
||||
Name: "my-instance",
|
||||
Kind: bigquery.SourceKind,
|
||||
Project: "my-project",
|
||||
Location: "us",
|
||||
Scopes: []string{"https://www.googleapis.com/auth/bigquery", "https://www.googleapis.com/auth/cloud-platform"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
@@ -219,105 +195,3 @@ func TestFailParseFromYaml(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input any
|
||||
expected any
|
||||
}{
|
||||
{
|
||||
name: "big.Rat 1/3 (NUMERIC scale 9)",
|
||||
input: new(big.Rat).SetFrac64(1, 3), // 0.33333333333...
|
||||
expected: "0.33333333333333333333333333333333333333", // FloatString(38)
|
||||
},
|
||||
{
|
||||
name: "big.Rat 19/2 (9.5)",
|
||||
input: new(big.Rat).SetFrac64(19, 2),
|
||||
expected: "9.5",
|
||||
},
|
||||
{
|
||||
name: "big.Rat 12341/10 (1234.1)",
|
||||
input: new(big.Rat).SetFrac64(12341, 10),
|
||||
expected: "1234.1",
|
||||
},
|
||||
{
|
||||
name: "big.Rat 10/1 (10)",
|
||||
input: new(big.Rat).SetFrac64(10, 1),
|
||||
expected: "10",
|
||||
},
|
||||
{
|
||||
name: "string",
|
||||
input: "hello",
|
||||
expected: "hello",
|
||||
},
|
||||
{
|
||||
name: "int",
|
||||
input: 123,
|
||||
expected: 123,
|
||||
},
|
||||
{
|
||||
name: "nested slice of big.Rat",
|
||||
input: []any{
|
||||
new(big.Rat).SetFrac64(19, 2),
|
||||
new(big.Rat).SetFrac64(1, 4),
|
||||
},
|
||||
expected: []any{"9.5", "0.25"},
|
||||
},
|
||||
{
|
||||
name: "nested map of big.Rat",
|
||||
input: map[string]any{
|
||||
"val1": new(big.Rat).SetFrac64(19, 2),
|
||||
"val2": new(big.Rat).SetFrac64(1, 2),
|
||||
},
|
||||
expected: map[string]any{
|
||||
"val1": "9.5",
|
||||
"val2": "0.5",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "complex nested structure",
|
||||
input: map[string]any{
|
||||
"list": []any{
|
||||
map[string]any{
|
||||
"rat": new(big.Rat).SetFrac64(3, 2),
|
||||
},
|
||||
},
|
||||
},
|
||||
expected: map[string]any{
|
||||
"list": []any{
|
||||
map[string]any{
|
||||
"rat": "1.5",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "slice of *big.Rat",
|
||||
input: []*big.Rat{
|
||||
new(big.Rat).SetFrac64(19, 2),
|
||||
new(big.Rat).SetFrac64(1, 4),
|
||||
},
|
||||
expected: []any{"9.5", "0.25"},
|
||||
},
|
||||
{
|
||||
name: "slice of strings",
|
||||
input: []string{"a", "b"},
|
||||
expected: []any{"a", "b"},
|
||||
},
|
||||
{
|
||||
name: "byte slice (BYTES)",
|
||||
input: []byte("hello"),
|
||||
expected: []byte("hello"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := bigquery.NormalizeValue(tt.input)
|
||||
if !reflect.DeepEqual(got, tt.expected) {
|
||||
t.Errorf("NormalizeValue() = %v, want %v", got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
@@ -89,94 +88,6 @@ func (s *Source) BigtableClient() *bigtable.Client {
|
||||
return s.Client
|
||||
}
|
||||
|
||||
func getBigtableType(paramType string) (bigtable.SQLType, error) {
|
||||
switch paramType {
|
||||
case "boolean":
|
||||
return bigtable.BoolSQLType{}, nil
|
||||
case "string":
|
||||
return bigtable.StringSQLType{}, nil
|
||||
case "integer":
|
||||
return bigtable.Int64SQLType{}, nil
|
||||
case "float":
|
||||
return bigtable.Float64SQLType{}, nil
|
||||
case "array":
|
||||
return bigtable.ArraySQLType{}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unknow param type %s", paramType)
|
||||
}
|
||||
}
|
||||
|
||||
func getMapParamsType(tparams parameters.Parameters) (map[string]bigtable.SQLType, error) {
|
||||
btParamTypes := make(map[string]bigtable.SQLType)
|
||||
for _, p := range tparams {
|
||||
if p.GetType() == "array" {
|
||||
itemType, err := getBigtableType(p.Manifest().Items.Type)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
btParamTypes[p.GetName()] = bigtable.ArraySQLType{
|
||||
ElemType: itemType,
|
||||
}
|
||||
continue
|
||||
}
|
||||
paramType, err := getBigtableType(p.GetType())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
btParamTypes[p.GetName()] = paramType
|
||||
}
|
||||
return btParamTypes, nil
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, configParam parameters.Parameters, params parameters.ParamValues) (any, error) {
|
||||
mapParamsType, err := getMapParamsType(configParam)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fail to get map params: %w", err)
|
||||
}
|
||||
|
||||
ps, err := s.BigtableClient().PrepareStatement(
|
||||
ctx,
|
||||
statement,
|
||||
mapParamsType,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to prepare statement: %w", err)
|
||||
}
|
||||
|
||||
bs, err := ps.Bind(params.AsMap())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to bind: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
var rowErr error
|
||||
err = bs.Execute(ctx, func(resultRow bigtable.ResultRow) bool {
|
||||
vMap := make(map[string]any)
|
||||
cols := resultRow.Metadata.Columns
|
||||
|
||||
for _, c := range cols {
|
||||
var columValue any
|
||||
if err = resultRow.GetByName(c.Name, &columValue); err != nil {
|
||||
rowErr = err
|
||||
return false
|
||||
}
|
||||
vMap[c.Name] = columValue
|
||||
}
|
||||
|
||||
out = append(out, vMap)
|
||||
|
||||
return true
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute client: %w", err)
|
||||
}
|
||||
if rowErr != nil {
|
||||
return nil, fmt.Errorf("error processing row: %w", rowErr)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initBigtableClient(ctx context.Context, tracer trace.Tracer, name, project, instance string) (*bigtable.Client, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -21,7 +21,6 @@ import (
|
||||
gocql "github.com/apache/cassandra-gocql-driver/v2"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -90,32 +89,10 @@ func (s *Source) ToConfig() sources.SourceConfig {
|
||||
}
|
||||
|
||||
// SourceKind implements sources.Source.
|
||||
func (s *Source) SourceKind() string {
|
||||
func (s Source) SourceKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params parameters.ParamValues) (any, error) {
|
||||
sliceParams := params.AsSlice()
|
||||
iter := s.CassandraSession().Query(statement, sliceParams...).IterContext(ctx)
|
||||
|
||||
// Create a slice to store the out
|
||||
var out []map[string]interface{}
|
||||
|
||||
// Scan results into a map and append to the slice
|
||||
for {
|
||||
row := make(map[string]interface{}) // Create a new map for each row
|
||||
if !iter.MapScan(row) {
|
||||
break // No more rows
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := iter.Close(); err != nil {
|
||||
return nil, fmt.Errorf("unable to parse rows: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
var _ sources.Source = &Source{}
|
||||
|
||||
func initCassandraSession(ctx context.Context, tracer trace.Tracer, c Config) (*gocql.Session, error) {
|
||||
|
||||
@@ -24,7 +24,6 @@ import (
|
||||
_ "github.com/ClickHouse/clickhouse-go/v2"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -100,69 +99,6 @@ func (s *Source) ClickHousePool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params parameters.ParamValues) (any, error) {
|
||||
var sliceParams []any
|
||||
if params != nil {
|
||||
sliceParams = params.AsSlice()
|
||||
}
|
||||
results, err := s.ClickHousePool().QueryContext(ctx, statement, sliceParams...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
// ClickHouse driver may return specific types that need handling
|
||||
switch colTypes[i].DatabaseTypeName() {
|
||||
case "String", "FixedString":
|
||||
if rawValues[i] != nil {
|
||||
// Handle potential []byte to string conversion if needed
|
||||
if b, ok := rawValues[i].([]byte); ok {
|
||||
vMap[name] = string(b)
|
||||
} else {
|
||||
vMap[name] = rawValues[i]
|
||||
}
|
||||
} else {
|
||||
vMap[name] = nil
|
||||
}
|
||||
default:
|
||||
vMap[name] = rawValues[i]
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered by results.Scan: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func validateConfig(protocol string) error {
|
||||
validProtocols := map[string]bool{"http": true, "https": true}
|
||||
|
||||
|
||||
@@ -14,11 +14,8 @@
|
||||
package cloudgda
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
@@ -134,43 +131,3 @@ func (s *Source) GetClient(ctx context.Context, accessToken string) (*http.Clien
|
||||
func (s *Source) UseClientAuthorization() bool {
|
||||
return s.UseClientOAuth
|
||||
}
|
||||
|
||||
func (s *Source) RunQuery(ctx context.Context, tokenStr string, bodyBytes []byte) (any, error) {
|
||||
// The API endpoint itself always uses the "global" location.
|
||||
apiLocation := "global"
|
||||
apiParent := fmt.Sprintf("projects/%s/locations/%s", s.GetProjectID(), apiLocation)
|
||||
apiURL := fmt.Sprintf("%s/v1beta/%s:queryData", s.GetBaseURL(), apiParent)
|
||||
|
||||
client, err := s.GetClient(ctx, tokenStr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get HTTP client: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewBuffer(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -15,9 +15,7 @@ package cloudmonitoring
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
@@ -133,44 +131,3 @@ func (s *Source) GetClient(ctx context.Context, accessToken string) (*http.Clien
|
||||
func (s *Source) UseClientAuthorization() bool {
|
||||
return s.UseClientOAuth
|
||||
}
|
||||
|
||||
func (s *Source) RunQuery(projectID, query string) (any, error) {
|
||||
url := fmt.Sprintf("%s/v1/projects/%s/location/global/prometheus/api/v1/query", s.BaseURL(), projectID)
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Add("query", query)
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
req.Header.Set("User-Agent", s.UserAgent())
|
||||
|
||||
resp, err := s.Client().Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("request failed: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
if len(body) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal json: %w, body: %s", err, string(body))
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -15,16 +15,10 @@ package cloudsqladmin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"text/template"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
@@ -36,8 +30,6 @@ import (
|
||||
|
||||
const SourceKind string = "cloud-sql-admin"
|
||||
|
||||
var targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
|
||||
|
||||
// validate interface
|
||||
var _ sources.SourceConfig = Config{}
|
||||
|
||||
@@ -138,304 +130,3 @@ func (s *Source) GetService(ctx context.Context, accessToken string) (*sqladmin.
|
||||
func (s *Source) UseClientAuthorization() bool {
|
||||
return s.UseClientOAuth
|
||||
}
|
||||
|
||||
func (s *Source) CloneInstance(ctx context.Context, project, sourceInstanceName, destinationInstanceName, pointInTime, preferredZone, preferredSecondaryZone, accessToken string) (any, error) {
|
||||
cloneContext := &sqladmin.CloneContext{
|
||||
DestinationInstanceName: destinationInstanceName,
|
||||
}
|
||||
|
||||
if pointInTime != "" {
|
||||
cloneContext.PointInTime = pointInTime
|
||||
}
|
||||
if preferredZone != "" {
|
||||
cloneContext.PreferredZone = preferredZone
|
||||
}
|
||||
if preferredSecondaryZone != "" {
|
||||
cloneContext.PreferredSecondaryZone = preferredSecondaryZone
|
||||
}
|
||||
|
||||
rb := &sqladmin.InstancesCloneRequest{
|
||||
CloneContext: cloneContext,
|
||||
}
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := service.Instances.Clone(project, sourceInstanceName, rb).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error cloning instance: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Source) CreateDatabase(ctx context.Context, name, project, instance, accessToken string) (any, error) {
|
||||
database := sqladmin.Database{
|
||||
Name: name,
|
||||
Project: project,
|
||||
Instance: instance,
|
||||
}
|
||||
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Databases.Insert(project, instance, &database).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating database: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Source) CreateUsers(ctx context.Context, project, instance, name, password string, iamUser bool, accessToken string) (any, error) {
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
user := sqladmin.User{
|
||||
Name: name,
|
||||
}
|
||||
|
||||
if iamUser {
|
||||
user.Type = "CLOUD_IAM_USER"
|
||||
} else {
|
||||
user.Type = "BUILT_IN"
|
||||
if password == "" {
|
||||
return nil, fmt.Errorf("missing 'password' parameter for non-IAM user")
|
||||
}
|
||||
user.Password = password
|
||||
}
|
||||
|
||||
resp, err := service.Users.Insert(project, instance, &user).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating user: %w", err)
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Source) GetInstance(ctx context.Context, projectId, instanceId, accessToken string) (any, error) {
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Instances.Get(projectId, instanceId).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting instance: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Source) ListDatabase(ctx context.Context, project, instance, accessToken string) (any, error) {
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Databases.List(project, instance).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error listing databases: %w", err)
|
||||
}
|
||||
|
||||
if resp.Items == nil {
|
||||
return []any{}, nil
|
||||
}
|
||||
|
||||
type databaseInfo struct {
|
||||
Name string `json:"name"`
|
||||
Charset string `json:"charset"`
|
||||
Collation string `json:"collation"`
|
||||
}
|
||||
|
||||
var databases []databaseInfo
|
||||
for _, item := range resp.Items {
|
||||
databases = append(databases, databaseInfo{
|
||||
Name: item.Name,
|
||||
Charset: item.Charset,
|
||||
Collation: item.Collation,
|
||||
})
|
||||
}
|
||||
return databases, nil
|
||||
}
|
||||
|
||||
func (s *Source) ListInstance(ctx context.Context, project, accessToken string) (any, error) {
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Instances.List(project).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error listing instances: %w", err)
|
||||
}
|
||||
|
||||
if resp.Items == nil {
|
||||
return []any{}, nil
|
||||
}
|
||||
|
||||
type instanceInfo struct {
|
||||
Name string `json:"name"`
|
||||
InstanceType string `json:"instanceType"`
|
||||
}
|
||||
|
||||
var instances []instanceInfo
|
||||
for _, item := range resp.Items {
|
||||
instances = append(instances, instanceInfo{
|
||||
Name: item.Name,
|
||||
InstanceType: item.InstanceType,
|
||||
})
|
||||
}
|
||||
return instances, nil
|
||||
}
|
||||
|
||||
func (s *Source) CreateInstance(ctx context.Context, project, name, dbVersion, rootPassword string, settings sqladmin.Settings, accessToken string) (any, error) {
|
||||
instance := sqladmin.DatabaseInstance{
|
||||
Name: name,
|
||||
DatabaseVersion: dbVersion,
|
||||
RootPassword: rootPassword,
|
||||
Settings: &settings,
|
||||
Project: project,
|
||||
}
|
||||
|
||||
service, err := s.GetService(ctx, accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Instances.Insert(project, &instance).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating instance: %w", err)
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Source) GetWaitForOperations(ctx context.Context, service *sqladmin.Service, project, operation, connectionMessageTemplate string, delay time.Duration) (any, error) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
op, err := service.Operations.Get(project, operation).Do()
|
||||
if err != nil {
|
||||
logger.DebugContext(ctx, fmt.Sprintf("error getting operation: %s, retrying in %v", err, delay))
|
||||
} else {
|
||||
if op.Status == "DONE" {
|
||||
if op.Error != nil {
|
||||
var errorBytes []byte
|
||||
errorBytes, err = json.Marshal(op.Error)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("operation finished with error but could not marshal error object: %w", err)
|
||||
}
|
||||
return nil, fmt.Errorf("operation finished with error: %s", string(errorBytes))
|
||||
}
|
||||
|
||||
var opBytes []byte
|
||||
opBytes, err = op.MarshalJSON()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not marshal operation: %w", err)
|
||||
}
|
||||
|
||||
var data map[string]any
|
||||
if err := json.Unmarshal(opBytes, &data); err != nil {
|
||||
return nil, fmt.Errorf("could not unmarshal operation: %w", err)
|
||||
}
|
||||
|
||||
if msg, ok := generateCloudSQLConnectionMessage(ctx, s, logger, data, connectionMessageTemplate); ok {
|
||||
return msg, nil
|
||||
}
|
||||
return string(opBytes), nil
|
||||
}
|
||||
logger.DebugContext(ctx, fmt.Sprintf("operation not complete, retrying in %v", delay))
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func generateCloudSQLConnectionMessage(ctx context.Context, source *Source, logger log.Logger, opResponse map[string]any, connectionMessageTemplate string) (string, bool) {
|
||||
operationType, ok := opResponse["operationType"].(string)
|
||||
if !ok || operationType != "CREATE_DATABASE" {
|
||||
return "", false
|
||||
}
|
||||
|
||||
targetLink, ok := opResponse["targetLink"].(string)
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
|
||||
matches := targetLinkRegex.FindStringSubmatch(targetLink)
|
||||
if len(matches) < 4 {
|
||||
return "", false
|
||||
}
|
||||
project := matches[1]
|
||||
instance := matches[2]
|
||||
database := matches[3]
|
||||
|
||||
dbInstance, err := fetchInstanceData(ctx, source, project, instance)
|
||||
if err != nil {
|
||||
logger.DebugContext(ctx, fmt.Sprintf("error fetching instance data: %v", err))
|
||||
return "", false
|
||||
}
|
||||
|
||||
region := dbInstance.Region
|
||||
if region == "" {
|
||||
return "", false
|
||||
}
|
||||
|
||||
databaseVersion := dbInstance.DatabaseVersion
|
||||
if databaseVersion == "" {
|
||||
return "", false
|
||||
}
|
||||
|
||||
var dbType string
|
||||
if strings.Contains(databaseVersion, "POSTGRES") {
|
||||
dbType = "postgres"
|
||||
} else if strings.Contains(databaseVersion, "MYSQL") {
|
||||
dbType = "mysql"
|
||||
} else if strings.Contains(databaseVersion, "SQLSERVER") {
|
||||
dbType = "mssql"
|
||||
} else {
|
||||
return "", false
|
||||
}
|
||||
|
||||
tmpl, err := template.New("cloud-sql-connection").Parse(connectionMessageTemplate)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("template parsing error: %v", err), false
|
||||
}
|
||||
|
||||
data := struct {
|
||||
Project string
|
||||
Region string
|
||||
Instance string
|
||||
DBType string
|
||||
DBTypeUpper string
|
||||
Database string
|
||||
}{
|
||||
Project: project,
|
||||
Region: region,
|
||||
Instance: instance,
|
||||
DBType: dbType,
|
||||
DBTypeUpper: strings.ToUpper(dbType),
|
||||
Database: database,
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
if err := tmpl.Execute(&b, data); err != nil {
|
||||
return fmt.Sprintf("template execution error: %v", err), false
|
||||
}
|
||||
|
||||
return b.String(), true
|
||||
}
|
||||
|
||||
func fetchInstanceData(ctx context.Context, source *Source, project, instance string) (*sqladmin.DatabaseInstance, error) {
|
||||
service, err := source.GetService(ctx, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := service.Instances.Get(project, instance).Do()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting instance: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -108,48 +107,6 @@ func (s *Source) MSSQLDB() *sql.DB {
|
||||
return s.Db
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.MSSQLDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
// If Columns() errors, it might be a DDL/DML without an OUTPUT clause.
|
||||
// We proceed, and results.Err() will catch actual query execution errors.
|
||||
// 'out' will remain nil if cols is empty or err is not nil here.
|
||||
var out []any
|
||||
if err == nil && len(cols) > 0 {
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
for results.Next() {
|
||||
scanErr := results.Scan(values...)
|
||||
if scanErr != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", scanErr)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
row.Add(name, rawValues[i])
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
}
|
||||
|
||||
// Check for errors from iterating over rows or from the query execution itself.
|
||||
// results.Close() is handled by defer.
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during query execution or row processing: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initCloudSQLMssqlConnection(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipType, user, pass, dbname string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -24,9 +24,7 @@ import (
|
||||
"cloud.google.com/go/cloudsqlconn/mysql/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -102,60 +100,6 @@ func (s *Source) MySQLPool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.MySQLPool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
row.Add(name, nil)
|
||||
continue
|
||||
}
|
||||
|
||||
convertedValue, err := mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
row.Add(name, convertedValue)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func getConnectionConfig(ctx context.Context, user, pass string) (string, string, bool, error) {
|
||||
useIAM := true
|
||||
|
||||
|
||||
@@ -23,7 +23,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
@@ -100,33 +99,6 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.PostgresPool().Query(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
var out []any
|
||||
for results.Next() {
|
||||
values, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, f := range fields {
|
||||
row.Add(f.Name, values[i])
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
// this will catch actual query execution errors
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func getConnectionConfig(ctx context.Context, user, pass, dbname string) (string, bool, error) {
|
||||
userAgent, err := util.UserAgentFromContext(ctx)
|
||||
if err != nil {
|
||||
|
||||
@@ -17,7 +17,6 @@ package couchbase
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
@@ -25,7 +24,6 @@ import (
|
||||
tlsutil "github.com/couchbase/tools-common/http/tls"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -112,27 +110,6 @@ func (s *Source) CouchbaseQueryScanConsistency() uint {
|
||||
return s.QueryScanConsistency
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(statement string, params parameters.ParamValues) (any, error) {
|
||||
results, err := s.CouchbaseScope().Query(statement, &gocb.QueryOptions{
|
||||
ScanConsistency: gocb.QueryScanConsistency(s.CouchbaseQueryScanConsistency()),
|
||||
NamedParameters: params.AsMap(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
var result json.RawMessage
|
||||
err := results.Row(&result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error processing row: %w", err)
|
||||
}
|
||||
out = append(out, result)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r Config) createCouchbaseOptions() (gocb.ClusterOptions, error) {
|
||||
cbOpts := gocb.ClusterOptions{}
|
||||
|
||||
|
||||
@@ -26,7 +26,6 @@ import (
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -115,28 +114,6 @@ func (s *Source) DgraphClient() *DgraphClient {
|
||||
return s.Client
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(statement string, params parameters.ParamValues, isQuery bool, timeout string) (any, error) {
|
||||
paramsMap := params.AsMapWithDollarPrefix()
|
||||
resp, err := s.DgraphClient().ExecuteQuery(statement, paramsMap, isQuery, timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := checkError(resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Data map[string]interface{} `json:"data"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(resp, &result); err != nil {
|
||||
return nil, fmt.Errorf("error parsing JSON: %v", err)
|
||||
}
|
||||
|
||||
return result.Data, nil
|
||||
}
|
||||
|
||||
func initDgraphHttpClient(ctx context.Context, tracer trace.Tracer, r Config) (*DgraphClient, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, r.Name)
|
||||
@@ -308,7 +285,7 @@ func (hc *DgraphClient) doLogin(creds map[string]interface{}) error {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := checkError(resp); err != nil {
|
||||
if err := CheckError(resp); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -393,7 +370,7 @@ func getUrl(baseUrl, resource string, params url.Values) (string, error) {
|
||||
return u.String(), nil
|
||||
}
|
||||
|
||||
func checkError(resp []byte) error {
|
||||
func CheckError(resp []byte) error {
|
||||
var errResp struct {
|
||||
Errors []struct {
|
||||
Message string `json:"message"`
|
||||
|
||||
@@ -15,9 +15,7 @@
|
||||
package elasticsearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
@@ -151,80 +149,3 @@ func (s *Source) ToConfig() sources.SourceConfig {
|
||||
func (s *Source) ElasticsearchClient() EsClient {
|
||||
return s.Client
|
||||
}
|
||||
|
||||
type EsqlColumn struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
type EsqlResult struct {
|
||||
Columns []EsqlColumn `json:"columns"`
|
||||
Values [][]any `json:"values"`
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, format, query string, params []map[string]any) (any, error) {
|
||||
bodyStruct := struct {
|
||||
Query string `json:"query"`
|
||||
Params []map[string]any `json:"params,omitempty"`
|
||||
}{
|
||||
Query: query,
|
||||
Params: params,
|
||||
}
|
||||
body, err := json.Marshal(bodyStruct)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal query body: %w", err)
|
||||
}
|
||||
|
||||
res, err := esapi.EsqlQueryRequest{
|
||||
Body: bytes.NewReader(body),
|
||||
Format: format,
|
||||
FilterPath: []string{"columns", "values"},
|
||||
Instrument: s.ElasticsearchClient().InstrumentationEnabled(),
|
||||
}.Do(ctx, s.ElasticsearchClient())
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
|
||||
if res.IsError() {
|
||||
// Try to extract error message from response
|
||||
var esErr json.RawMessage
|
||||
err = util.DecodeJSON(res.Body, &esErr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("elasticsearch error: status %s", res.Status())
|
||||
}
|
||||
return esErr, nil
|
||||
}
|
||||
|
||||
var result EsqlResult
|
||||
err = util.DecodeJSON(res.Body, &result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode response body: %w", err)
|
||||
}
|
||||
|
||||
output := EsqlToMap(result)
|
||||
|
||||
return output, nil
|
||||
}
|
||||
|
||||
// EsqlToMap converts the esqlResult to a slice of maps.
|
||||
func EsqlToMap(result EsqlResult) []map[string]any {
|
||||
output := make([]map[string]any, 0, len(result.Values))
|
||||
for _, value := range result.Values {
|
||||
row := make(map[string]any)
|
||||
if value == nil {
|
||||
output = append(output, row)
|
||||
continue
|
||||
}
|
||||
for i, col := range result.Columns {
|
||||
if i < len(value) {
|
||||
row[col.Name] = value[i]
|
||||
} else {
|
||||
row[col.Name] = nil
|
||||
}
|
||||
}
|
||||
output = append(output, row)
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
package elasticsearch_test
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
@@ -65,155 +64,3 @@ func TestParseFromYamlElasticsearch(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTool_esqlToMap(t1 *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
result elasticsearch.EsqlResult
|
||||
want []map[string]any
|
||||
}{
|
||||
{
|
||||
name: "simple case with two rows",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "first_name", Type: "text"},
|
||||
{Name: "last_name", Type: "text"},
|
||||
},
|
||||
Values: [][]any{
|
||||
{"John", "Doe"},
|
||||
{"Jane", "Smith"},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{"first_name": "John", "last_name": "Doe"},
|
||||
{"first_name": "Jane", "last_name": "Smith"},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "different data types",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
{Name: "active", Type: "boolean"},
|
||||
{Name: "score", Type: "float"},
|
||||
},
|
||||
Values: [][]any{
|
||||
{1, true, 95.5},
|
||||
{2, false, 88.0},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{"id": 1, "active": true, "score": 95.5},
|
||||
{"id": 2, "active": false, "score": 88.0},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "no rows",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
{Name: "name", Type: "text"},
|
||||
},
|
||||
Values: [][]any{},
|
||||
},
|
||||
want: []map[string]any{},
|
||||
},
|
||||
{
|
||||
name: "null values",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
{Name: "name", Type: "text"},
|
||||
},
|
||||
Values: [][]any{
|
||||
{1, nil},
|
||||
{2, "Alice"},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{"id": 1, "name": nil},
|
||||
{"id": 2, "name": "Alice"},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing values in a row",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
{Name: "name", Type: "text"},
|
||||
{Name: "age", Type: "integer"},
|
||||
},
|
||||
Values: [][]any{
|
||||
{1, "Bob"},
|
||||
{2, "Charlie", 30},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{"id": 1, "name": "Bob", "age": nil},
|
||||
{"id": 2, "name": "Charlie", "age": 30},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "all null row",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
{Name: "name", Type: "text"},
|
||||
},
|
||||
Values: [][]any{
|
||||
nil,
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "empty columns",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{},
|
||||
Values: [][]any{
|
||||
{},
|
||||
{},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{},
|
||||
{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "more values than columns",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{
|
||||
{Name: "id", Type: "integer"},
|
||||
},
|
||||
Values: [][]any{
|
||||
{1, "extra"},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{"id": 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "no columns but with values",
|
||||
result: elasticsearch.EsqlResult{
|
||||
Columns: []elasticsearch.EsqlColumn{},
|
||||
Values: [][]any{
|
||||
{1, "data"},
|
||||
},
|
||||
},
|
||||
want: []map[string]any{
|
||||
{},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t1.Run(tt.name, func(t1 *testing.T) {
|
||||
if got := elasticsearch.EsqlToMap(tt.result); !reflect.DeepEqual(got, tt.want) {
|
||||
t1.Errorf("esqlToMap() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -96,53 +96,6 @@ func (s *Source) FirebirdDB() *sql.DB {
|
||||
return s.Db
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
rows, err := s.FirebirdDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get columns: %w", err)
|
||||
}
|
||||
|
||||
values := make([]any, len(cols))
|
||||
scanArgs := make([]any, len(values))
|
||||
for i := range values {
|
||||
scanArgs[i] = &values[i]
|
||||
}
|
||||
|
||||
var out []any
|
||||
for rows.Next() {
|
||||
|
||||
err = rows.Scan(scanArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
|
||||
vMap := make(map[string]any)
|
||||
for i, col := range cols {
|
||||
if b, ok := values[i].([]byte); ok {
|
||||
vMap[col] = string(b)
|
||||
} else {
|
||||
vMap[col] = values[i]
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("error iterating rows: %w", err)
|
||||
}
|
||||
|
||||
// In most cases, DML/DDL statements like INSERT, UPDATE, CREATE, etc. might return no rows
|
||||
// However, it is also possible that this was a query that was expected to return rows
|
||||
// but returned none, a case that we cannot distinguish here.
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initFirebirdConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*sql.DB, error) {
|
||||
_, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
defer span.End()
|
||||
|
||||
@@ -23,7 +23,6 @@ import (
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -102,61 +101,6 @@ func (s *Source) MySQLPool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
// MindsDB now supports MySQL prepared statements natively
|
||||
results, err := s.MindsDBPool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
// MindsDB uses mysql driver
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initMindsDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -23,7 +23,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
_ "github.com/microsoft/go-mssqldb"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
@@ -105,48 +104,6 @@ func (s *Source) MSSQLDB() *sql.DB {
|
||||
return s.Db
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.MSSQLDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
// If Columns() errors, it might be a DDL/DML without an OUTPUT clause.
|
||||
// We proceed, and results.Err() will catch actual query execution errors.
|
||||
// 'out' will remain nil if cols is empty or err is not nil here.
|
||||
var out []any
|
||||
if err == nil && len(cols) > 0 {
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
for results.Next() {
|
||||
scanErr := results.Scan(values...)
|
||||
if scanErr != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", scanErr)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
row.Add(name, rawValues[i])
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
}
|
||||
|
||||
// Check for errors from iterating over rows or from the query execution itself.
|
||||
// results.Close() is handled by defer.
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during query execution or row processing: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initMssqlConnection(
|
||||
ctx context.Context,
|
||||
tracer trace.Tracer,
|
||||
|
||||
@@ -24,9 +24,7 @@ import (
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -102,60 +100,6 @@ func (s *Source) MySQLPool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.MySQLPool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
row.Add(name, nil)
|
||||
continue
|
||||
}
|
||||
|
||||
convertedValue, err := mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
row.Add(name, convertedValue)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initMySQLConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string, queryParams map[string]string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -20,19 +20,14 @@ import (
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jexecutecypher/classifier"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jschema/helpers"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/neo4j/neo4j-go-driver/v5/neo4j"
|
||||
neo4jconf "github.com/neo4j/neo4j-go-driver/v5/neo4j/config"
|
||||
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
const SourceKind string = "neo4j"
|
||||
|
||||
var sourceClassifier *classifier.QueryClassifier = classifier.NewQueryClassifier()
|
||||
|
||||
// validate interface
|
||||
var _ sources.SourceConfig = Config{}
|
||||
|
||||
@@ -107,79 +102,6 @@ func (s *Source) Neo4jDatabase() string {
|
||||
return s.Database
|
||||
}
|
||||
|
||||
func (s *Source) RunQuery(ctx context.Context, cypherStr string, params map[string]any, readOnly, dryRun bool) (any, error) {
|
||||
// validate the cypher query before executing
|
||||
cf := sourceClassifier.Classify(cypherStr)
|
||||
if cf.Error != nil {
|
||||
return nil, cf.Error
|
||||
}
|
||||
|
||||
if cf.Type == classifier.WriteQuery && readOnly {
|
||||
return nil, fmt.Errorf("this tool is read-only and cannot execute write queries")
|
||||
}
|
||||
|
||||
if dryRun {
|
||||
// Add EXPLAIN to the beginning of the query to validate it without executing
|
||||
cypherStr = "EXPLAIN " + cypherStr
|
||||
}
|
||||
|
||||
config := neo4j.ExecuteQueryWithDatabase(s.Neo4jDatabase())
|
||||
results, err := neo4j.ExecuteQuery[*neo4j.EagerResult](ctx, s.Neo4jDriver(), cypherStr, params,
|
||||
neo4j.EagerResultTransformer, config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
// If dry run, return the summary information only
|
||||
if dryRun {
|
||||
summary := results.Summary
|
||||
plan := summary.Plan()
|
||||
execPlan := map[string]any{
|
||||
"queryType": cf.Type.String(),
|
||||
"statementType": summary.StatementType(),
|
||||
"operator": plan.Operator(),
|
||||
"arguments": plan.Arguments(),
|
||||
"identifiers": plan.Identifiers(),
|
||||
"childrenCount": len(plan.Children()),
|
||||
}
|
||||
if len(plan.Children()) > 0 {
|
||||
execPlan["children"] = addPlanChildren(plan)
|
||||
}
|
||||
return []map[string]any{execPlan}, nil
|
||||
}
|
||||
|
||||
var out []map[string]any
|
||||
keys := results.Keys
|
||||
records := results.Records
|
||||
for _, record := range records {
|
||||
vMap := make(map[string]any)
|
||||
for col, value := range record.Values {
|
||||
vMap[keys[col]] = helpers.ConvertValue(value)
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Recursive function to add plan children
|
||||
func addPlanChildren(p neo4j.Plan) []map[string]any {
|
||||
var children []map[string]any
|
||||
for _, child := range p.Children() {
|
||||
childMap := map[string]any{
|
||||
"operator": child.Operator(),
|
||||
"arguments": child.Arguments(),
|
||||
"identifiers": child.Identifiers(),
|
||||
"children_count": len(child.Children()),
|
||||
}
|
||||
if len(child.Children()) > 0 {
|
||||
childMap["children"] = addPlanChildren(child)
|
||||
}
|
||||
children = append(children, childMap)
|
||||
}
|
||||
return children
|
||||
}
|
||||
|
||||
func initNeo4jDriver(ctx context.Context, tracer trace.Tracer, uri, user, password, name string) (neo4j.DriverWithContext, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -23,7 +23,6 @@ import (
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -98,60 +97,6 @@ func (s *Source) OceanBasePool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.OceanBasePool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
// oceanbase uses mysql driver
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initOceanBaseConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
|
||||
_, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
defer span.End()
|
||||
|
||||
@@ -4,7 +4,6 @@ package oracle
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
@@ -136,107 +135,6 @@ func (s *Source) OracleDB() *sql.DB {
|
||||
return s.DB
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
rows, err := s.OracleDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
// If Columns() errors, it might be a DDL/DML without an OUTPUT clause.
|
||||
// We proceed, and results.Err() will catch actual query execution errors.
|
||||
// 'out' will remain nil if cols is empty or err is not nil here.
|
||||
cols, _ := rows.Columns()
|
||||
|
||||
// Get Column types
|
||||
colTypes, err := rows.ColumnTypes()
|
||||
if err != nil {
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("query execution error: %w", err)
|
||||
}
|
||||
return []any{}, nil
|
||||
}
|
||||
|
||||
var out []any
|
||||
for rows.Next() {
|
||||
values := make([]any, len(cols))
|
||||
for i, colType := range colTypes {
|
||||
switch strings.ToUpper(colType.DatabaseTypeName()) {
|
||||
case "NUMBER", "FLOAT", "BINARY_FLOAT", "BINARY_DOUBLE":
|
||||
if _, scale, ok := colType.DecimalSize(); ok && scale == 0 {
|
||||
// Scale is 0, treat it as an integer.
|
||||
values[i] = new(sql.NullInt64)
|
||||
} else {
|
||||
// Scale is non-zero or unknown, treat
|
||||
// it as a float.
|
||||
values[i] = new(sql.NullFloat64)
|
||||
}
|
||||
case "DATE", "TIMESTAMP", "TIMESTAMP WITH TIME ZONE", "TIMESTAMP WITH LOCAL TIME ZONE":
|
||||
values[i] = new(sql.NullTime)
|
||||
case "JSON":
|
||||
values[i] = new(sql.RawBytes)
|
||||
default:
|
||||
values[i] = new(sql.NullString)
|
||||
}
|
||||
}
|
||||
|
||||
if err := rows.Scan(values...); err != nil {
|
||||
return nil, fmt.Errorf("unable to scan row: %w", err)
|
||||
}
|
||||
|
||||
vMap := make(map[string]any)
|
||||
for i, col := range cols {
|
||||
receiver := values[i]
|
||||
|
||||
switch v := receiver.(type) {
|
||||
case *sql.NullInt64:
|
||||
if v.Valid {
|
||||
vMap[col] = v.Int64
|
||||
} else {
|
||||
vMap[col] = nil
|
||||
}
|
||||
case *sql.NullFloat64:
|
||||
if v.Valid {
|
||||
vMap[col] = v.Float64
|
||||
} else {
|
||||
vMap[col] = nil
|
||||
}
|
||||
case *sql.NullString:
|
||||
if v.Valid {
|
||||
vMap[col] = v.String
|
||||
} else {
|
||||
vMap[col] = nil
|
||||
}
|
||||
case *sql.NullTime:
|
||||
if v.Valid {
|
||||
vMap[col] = v.Time
|
||||
} else {
|
||||
vMap[col] = nil
|
||||
}
|
||||
case *sql.RawBytes:
|
||||
if *v != nil {
|
||||
var unmarshaledData any
|
||||
if err := json.Unmarshal(*v, &unmarshaledData); err != nil {
|
||||
return nil, fmt.Errorf("unable to unmarshal json data for column %s", col)
|
||||
}
|
||||
vMap[col] = unmarshaledData
|
||||
} else {
|
||||
vMap[col] = nil
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected receiver type: %T", v)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during query execution or row processing: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initOracleConnection(ctx context.Context, tracer trace.Tracer, config Config) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, config.Name)
|
||||
|
||||
@@ -23,7 +23,6 @@ import (
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
@@ -99,33 +98,6 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.PostgresPool().Query(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
var out []any
|
||||
for results.Next() {
|
||||
values, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
row := orderedmap.Row{}
|
||||
for i, f := range fields {
|
||||
row.Add(f.Name, values[i])
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
// this will catch actual query execution errors
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string) (*pgxpool.Pool, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -152,50 +152,3 @@ func (s *Source) ToConfig() sources.SourceConfig {
|
||||
func (s *Source) RedisClient() RedisClient {
|
||||
return s.Client
|
||||
}
|
||||
|
||||
func (s *Source) RunCommand(ctx context.Context, cmds [][]any) (any, error) {
|
||||
// Execute commands
|
||||
responses := make([]*redis.Cmd, len(cmds))
|
||||
for i, cmd := range cmds {
|
||||
responses[i] = s.RedisClient().Do(ctx, cmd...)
|
||||
}
|
||||
// Parse responses
|
||||
out := make([]any, len(cmds))
|
||||
for i, resp := range responses {
|
||||
if err := resp.Err(); err != nil {
|
||||
// Add error from each command to `errSum`
|
||||
errString := fmt.Sprintf("error from executing command at index %d: %s", i, err)
|
||||
out[i] = errString
|
||||
continue
|
||||
}
|
||||
val, err := resp.Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting result: %s", err)
|
||||
}
|
||||
out[i] = convertRedisResult(val)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// convertRedisResult recursively converts redis results (map[any]any) to be
|
||||
// JSON-marshallable (map[string]any).
|
||||
// It converts map[any]any to map[string]any and handles nested structures.
|
||||
func convertRedisResult(v any) any {
|
||||
switch val := v.(type) {
|
||||
case map[any]any:
|
||||
m := make(map[string]any)
|
||||
for k, v := range val {
|
||||
m[fmt.Sprint(k)] = convertRedisResult(v)
|
||||
}
|
||||
return m
|
||||
case []any:
|
||||
s := make([]any, len(val))
|
||||
for i, v := range val {
|
||||
s[i] = convertRedisResult(v)
|
||||
}
|
||||
return s
|
||||
default:
|
||||
return v
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
@@ -107,59 +106,6 @@ func (s *Source) SingleStorePool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.SingleStorePool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initSingleStoreConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -16,16 +16,13 @@ package spanner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"cloud.google.com/go/spanner"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
const SourceKind string = "spanner"
|
||||
@@ -96,79 +93,6 @@ func (s *Source) DatabaseDialect() string {
|
||||
return s.Dialect.String()
|
||||
}
|
||||
|
||||
// processRows iterates over the spanner.RowIterator and converts each row to a map[string]any.
|
||||
func processRows(iter *spanner.RowIterator) ([]any, error) {
|
||||
var out []any
|
||||
defer iter.Stop()
|
||||
|
||||
for {
|
||||
row, err := iter.Next()
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
|
||||
rowMap := orderedmap.Row{}
|
||||
cols := row.ColumnNames()
|
||||
for i, c := range cols {
|
||||
if c == "object_details" { // for list graphs or list tables
|
||||
val := row.ColumnValue(i)
|
||||
if val == nil { // ColumnValue returns the Cloud Spanner Value of column i, or nil for invalid column.
|
||||
rowMap.Add(c, nil)
|
||||
} else {
|
||||
jsonString, ok := val.AsInterface().(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("column 'object_details' is not a string, but %T", val.AsInterface())
|
||||
}
|
||||
var details map[string]any
|
||||
if err := json.Unmarshal([]byte(jsonString), &details); err != nil {
|
||||
return nil, fmt.Errorf("unable to unmarshal JSON: %w", err)
|
||||
}
|
||||
rowMap.Add(c, details)
|
||||
}
|
||||
} else {
|
||||
rowMap.Add(c, row.ColumnValue(i))
|
||||
}
|
||||
}
|
||||
out = append(out, rowMap)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, readOnly bool, statement string, params map[string]any) (any, error) {
|
||||
var results []any
|
||||
var err error
|
||||
var opErr error
|
||||
stmt := spanner.Statement{
|
||||
SQL: statement,
|
||||
}
|
||||
if params != nil {
|
||||
stmt.Params = params
|
||||
}
|
||||
|
||||
if readOnly {
|
||||
iter := s.SpannerClient().Single().Query(ctx, stmt)
|
||||
results, opErr = processRows(iter)
|
||||
} else {
|
||||
_, opErr = s.SpannerClient().ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
|
||||
iter := txn.Query(ctx, stmt)
|
||||
results, err = processRows(iter)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
if opErr != nil {
|
||||
return nil, fmt.Errorf("unable to execute client: %w", opErr)
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func initSpannerClient(ctx context.Context, tracer trace.Tracer, name, project, instance, dbname string) (*spanner.Client, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -17,12 +17,10 @@ package sqlite
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
_ "modernc.org/sqlite" // Pure Go SQLite driver
|
||||
)
|
||||
@@ -93,66 +91,6 @@ func (s *Source) SQLiteDB() *sql.DB {
|
||||
return s.Db
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
// Execute the SQL query with parameters
|
||||
rows, err := s.SQLiteDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
// Get column names
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column names: %w", err)
|
||||
}
|
||||
|
||||
// The sqlite driver does not support ColumnTypes, so we can't get the
|
||||
// underlying database type of the columns. We'll have to rely on the
|
||||
// generic `any` type and then handle the JSON data separately.
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
// Prepare the result slice
|
||||
var out []any
|
||||
for rows.Next() {
|
||||
if err := rows.Scan(values...); err != nil {
|
||||
return nil, fmt.Errorf("unable to scan row: %w", err)
|
||||
}
|
||||
|
||||
// Create a map for this row
|
||||
row := orderedmap.Row{}
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
// Handle nil values
|
||||
if val == nil {
|
||||
row.Add(name, nil)
|
||||
continue
|
||||
}
|
||||
// Handle JSON data
|
||||
if jsonString, ok := val.(string); ok {
|
||||
var unmarshaledData any
|
||||
if json.Unmarshal([]byte(jsonString), &unmarshaledData) == nil {
|
||||
row.Add(name, unmarshaledData)
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Store the value in the map
|
||||
row.Add(name, val)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("error iterating rows: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initSQLiteConnection(ctx context.Context, tracer trace.Tracer, name, dbPath string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -17,7 +17,6 @@ package tidb
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
|
||||
@@ -105,79 +104,6 @@ func (s *Source) TiDBPool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.TiDBPool().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
// mysql driver return []uint8 type for "TEXT", "VARCHAR", and "NVARCHAR"
|
||||
// we'll need to cast it back to string
|
||||
switch colTypes[i].DatabaseTypeName() {
|
||||
case "JSON":
|
||||
// unmarshal JSON data before storing to prevent double
|
||||
// marshaling
|
||||
byteVal, ok := val.([]byte)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected []byte for JSON column, but got %T", val)
|
||||
}
|
||||
var unmarshaledData any
|
||||
if err := json.Unmarshal(byteVal, &unmarshaledData); err != nil {
|
||||
return nil, fmt.Errorf("unable to unmarshal json data %s", val)
|
||||
}
|
||||
vMap[name] = unmarshaledData
|
||||
case "TEXT", "VARCHAR", "NVARCHAR":
|
||||
byteVal, ok := val.([]byte)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected []byte for text-like column, but got %T", val)
|
||||
}
|
||||
vMap[name] = string(byteVal)
|
||||
default:
|
||||
vMap[name] = val
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func IsTiDBCloudHost(host string) bool {
|
||||
pattern := `gateway\d{2}\.(.+)\.(prod|dev|staging)\.(.+)\.tidbcloud\.com`
|
||||
match, err := regexp.MatchString(pattern, host)
|
||||
|
||||
@@ -102,56 +102,6 @@ func (s *Source) TrinoDB() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.TrinoDB().QueryContext(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve column names: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert byte arrays to strings for text fields
|
||||
if b, ok := val.([]byte); ok {
|
||||
vMap[name] = string(b)
|
||||
} else {
|
||||
vMap[name] = val
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initTrinoConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -125,37 +125,3 @@ func (s *Source) ToConfig() sources.SourceConfig {
|
||||
func (s *Source) ValkeyClient() valkey.Client {
|
||||
return s.Client
|
||||
}
|
||||
|
||||
func (s *Source) RunCommand(ctx context.Context, cmds [][]string) (any, error) {
|
||||
// Build commands
|
||||
builtCmds := make(valkey.Commands, len(cmds))
|
||||
|
||||
for i, cmd := range cmds {
|
||||
builtCmds[i] = s.ValkeyClient().B().Arbitrary(cmd...).Build()
|
||||
}
|
||||
|
||||
if len(builtCmds) == 0 {
|
||||
return nil, fmt.Errorf("no valid commands were built to execute")
|
||||
}
|
||||
|
||||
// Execute commands
|
||||
responses := s.ValkeyClient().DoMulti(ctx, builtCmds...)
|
||||
|
||||
// Parse responses
|
||||
out := make([]any, len(cmds))
|
||||
for i, resp := range responses {
|
||||
if err := resp.Error(); err != nil {
|
||||
// Store error message in the output for this command
|
||||
out[i] = fmt.Sprintf("error from executing command at index %d: %s", i, err)
|
||||
continue
|
||||
}
|
||||
val, err := resp.ToAny()
|
||||
if err != nil {
|
||||
out[i] = fmt.Sprintf("error parsing response: %s", err)
|
||||
continue
|
||||
}
|
||||
out[i] = val
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
@@ -99,35 +99,6 @@ func (s *Source) YugabyteDBPool() *pgxpool.Pool {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
|
||||
results, err := s.YugabyteDBPool().Query(ctx, statement, params...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
v, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, f := range fields {
|
||||
vMap[f.Name] = v[i]
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
// this will catch actual query execution errors
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func initYugabyteDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, loadBalance, topologyKeys, refreshInterval, explicitFallback, failedHostTTL string) (*pgxpool.Pool, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -167,10 +166,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -173,10 +172,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -178,10 +177,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -149,10 +148,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -153,10 +152,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -153,10 +152,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -143,10 +142,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -148,10 +147,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -19,7 +19,6 @@ import (
|
||||
"fmt"
|
||||
|
||||
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/parameters"
|
||||
@@ -148,10 +147,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -21,7 +21,6 @@ import (
|
||||
"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/parameters"
|
||||
@@ -272,10 +271,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
|
||||
@@ -20,7 +20,6 @@ import (
|
||||
"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/parameters"
|
||||
@@ -146,21 +145,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
allParamValues[i+2] = fmt.Sprintf("%s", param)
|
||||
}
|
||||
|
||||
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 resp, nil
|
||||
return source.RunSQL(ctx, t.Statement, allParamValues)
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -22,13 +22,13 @@ import (
|
||||
bigqueryapi "cloud.google.com/go/bigquery"
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
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/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
const kind string = "bigquery-analyze-contribution"
|
||||
@@ -49,12 +49,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryRestService() *bigqueryrestapi.Service
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
BigQuerySession() bigqueryds.BigQuerySessionProvider
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
RunSQL(context.Context, *bigqueryapi.Client, string, string, []bigqueryapi.QueryParameter, []*bigqueryapi.ConnectionProperty) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -166,9 +166,19 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return nil, fmt.Errorf("unable to cast input_data parameter %s", paramsMap["input_data"])
|
||||
}
|
||||
|
||||
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
restService := source.BigQueryRestService()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, restService, err = source.BigQueryClientCreator()(tokenStr, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
modelID := fmt.Sprintf("contribution_analysis_model_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
@@ -304,18 +314,49 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
}
|
||||
|
||||
getInsightsSQL := fmt.Sprintf("SELECT * FROM ML.GET_INSIGHTS(MODEL %s)", modelID)
|
||||
connProps := []*bigqueryapi.ConnectionProperty{{Key: "session_id", Value: sessionID}}
|
||||
return source.RunSQL(ctx, bqClient, getInsightsSQL, "SELECT", nil, connProps)
|
||||
|
||||
getInsightsQuery := bqClient.Query(getInsightsSQL)
|
||||
getInsightsQuery.ConnectionProperties = []*bigqueryapi.ConnectionProperty{{Key: "session_id", Value: sessionID}}
|
||||
|
||||
job, err := getInsightsQuery.Run(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute get insights query: %w", err)
|
||||
}
|
||||
it, err := job.Read(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to read query results: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for {
|
||||
var row map[string]bigqueryapi.Value
|
||||
err := it.Next(&row)
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to iterate through query results: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for key, value := range row {
|
||||
vMap[key] = value
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if len(out) > 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This handles the standard case for a SELECT query that successfully
|
||||
// executes but returns zero rows.
|
||||
return "The query returned 0 rows.", nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
123
internal/tools/bigquery/bigquerycommon/conversion_test.go
Normal file
123
internal/tools/bigquery/bigquerycommon/conversion_test.go
Normal file
@@ -0,0 +1,123 @@
|
||||
// Copyright 2025 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// 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
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package bigquerycommon
|
||||
|
||||
import (
|
||||
"math/big"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizeValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input any
|
||||
expected any
|
||||
}{
|
||||
{
|
||||
name: "big.Rat 1/3 (NUMERIC scale 9)",
|
||||
input: new(big.Rat).SetFrac64(1, 3), // 0.33333333333...
|
||||
expected: "0.33333333333333333333333333333333333333", // FloatString(38)
|
||||
},
|
||||
{
|
||||
name: "big.Rat 19/2 (9.5)",
|
||||
input: new(big.Rat).SetFrac64(19, 2),
|
||||
expected: "9.5",
|
||||
},
|
||||
{
|
||||
name: "big.Rat 12341/10 (1234.1)",
|
||||
input: new(big.Rat).SetFrac64(12341, 10),
|
||||
expected: "1234.1",
|
||||
},
|
||||
{
|
||||
name: "big.Rat 10/1 (10)",
|
||||
input: new(big.Rat).SetFrac64(10, 1),
|
||||
expected: "10",
|
||||
},
|
||||
{
|
||||
name: "string",
|
||||
input: "hello",
|
||||
expected: "hello",
|
||||
},
|
||||
{
|
||||
name: "int",
|
||||
input: 123,
|
||||
expected: 123,
|
||||
},
|
||||
{
|
||||
name: "nested slice of big.Rat",
|
||||
input: []any{
|
||||
new(big.Rat).SetFrac64(19, 2),
|
||||
new(big.Rat).SetFrac64(1, 4),
|
||||
},
|
||||
expected: []any{"9.5", "0.25"},
|
||||
},
|
||||
{
|
||||
name: "nested map of big.Rat",
|
||||
input: map[string]any{
|
||||
"val1": new(big.Rat).SetFrac64(19, 2),
|
||||
"val2": new(big.Rat).SetFrac64(1, 2),
|
||||
},
|
||||
expected: map[string]any{
|
||||
"val1": "9.5",
|
||||
"val2": "0.5",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "complex nested structure",
|
||||
input: map[string]any{
|
||||
"list": []any{
|
||||
map[string]any{
|
||||
"rat": new(big.Rat).SetFrac64(3, 2),
|
||||
},
|
||||
},
|
||||
},
|
||||
expected: map[string]any{
|
||||
"list": []any{
|
||||
map[string]any{
|
||||
"rat": "1.5",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "slice of *big.Rat",
|
||||
input: []*big.Rat{
|
||||
new(big.Rat).SetFrac64(19, 2),
|
||||
new(big.Rat).SetFrac64(1, 4),
|
||||
},
|
||||
expected: []any{"9.5", "0.25"},
|
||||
},
|
||||
{
|
||||
name: "slice of strings",
|
||||
input: []string{"a", "b"},
|
||||
expected: []any{"a", "b"},
|
||||
},
|
||||
{
|
||||
name: "byte slice (BYTES)",
|
||||
input: []byte("hello"),
|
||||
expected: []byte("hello"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := NormalizeValue(tt.input)
|
||||
if !reflect.DeepEqual(got, tt.expected) {
|
||||
t.Errorf("NormalizeValue() = %v, want %v", got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,8 @@ package bigquerycommon
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"reflect"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
@@ -118,3 +120,54 @@ func InitializeDatasetParameters(
|
||||
|
||||
return projectParam, datasetParam
|
||||
}
|
||||
|
||||
// NormalizeValue converts BigQuery specific types to standard JSON-compatible types.
|
||||
// Specifically, it handles *big.Rat (used for NUMERIC/BIGNUMERIC) by converting
|
||||
// them to decimal strings with up to 38 digits of precision, trimming trailing zeros.
|
||||
// It recursively handles slices (arrays) and maps (structs) using reflection.
|
||||
func NormalizeValue(v any) any {
|
||||
if v == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Handle *big.Rat specifically.
|
||||
if rat, ok := v.(*big.Rat); ok {
|
||||
// Convert big.Rat to a decimal string.
|
||||
// Use a precision of 38 digits (enough for BIGNUMERIC and NUMERIC)
|
||||
// and trim trailing zeros to match BigQuery's behavior.
|
||||
s := rat.FloatString(38)
|
||||
if strings.Contains(s, ".") {
|
||||
s = strings.TrimRight(s, "0")
|
||||
s = strings.TrimRight(s, ".")
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Use reflection for slices and maps to handle various underlying types.
|
||||
rv := reflect.ValueOf(v)
|
||||
switch rv.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
// Preserve []byte as is, so json.Marshal encodes it as Base64 string (BigQuery BYTES behavior).
|
||||
if rv.Type().Elem().Kind() == reflect.Uint8 {
|
||||
return v
|
||||
}
|
||||
newSlice := make([]any, rv.Len())
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
newSlice[i] = NormalizeValue(rv.Index(i).Interface())
|
||||
}
|
||||
return newSlice
|
||||
case reflect.Map:
|
||||
// Ensure keys are strings to produce a JSON-compatible map.
|
||||
if rv.Type().Key().Kind() != reflect.String {
|
||||
return v
|
||||
}
|
||||
newMap := make(map[string]any, rv.Len())
|
||||
iter := rv.MapRange()
|
||||
for iter.Next() {
|
||||
newMap[iter.Key().String()] = NormalizeValue(iter.Value().Interface())
|
||||
}
|
||||
return newMap
|
||||
}
|
||||
|
||||
return v
|
||||
}
|
||||
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
|
||||
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"
|
||||
@@ -56,7 +55,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryTokenSourceWithScope(ctx context.Context, scopes []string) (oauth2.TokenSource, error)
|
||||
BigQueryTokenSourceWithScope(ctx context.Context, scope string) (oauth2.TokenSource, error)
|
||||
BigQueryProject() string
|
||||
BigQueryLocation() string
|
||||
GetMaxQueryResultRows() int
|
||||
@@ -191,10 +190,10 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
} else {
|
||||
// Get a token source for the Gemini Data Analytics API.
|
||||
tokenSource, err := source.BigQueryTokenSourceWithScope(ctx, nil)
|
||||
// Get cloud-platform token source for Gemini Data Analytics API during initialization
|
||||
tokenSource, err := source.BigQueryTokenSourceWithScope(ctx, "https://www.googleapis.com/auth/cloud-platform")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get token source: %w", err)
|
||||
return nil, fmt.Errorf("failed to get cloud-platform token source: %w", err)
|
||||
}
|
||||
|
||||
// Use cloud-platform token source for Gemini Data Analytics API
|
||||
@@ -268,10 +267,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -22,14 +22,15 @@ import (
|
||||
|
||||
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"
|
||||
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/orderedmap"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
const kind string = "bigquery-execute-sql"
|
||||
@@ -52,11 +53,11 @@ type compatibleSource interface {
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQuerySession() bigqueryds.BigQuerySessionProvider
|
||||
BigQueryWriteMode() string
|
||||
BigQueryRestService() *bigqueryrestapi.Service
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
RunSQL(context.Context, *bigqueryapi.Client, string, string, []bigqueryapi.QueryParameter, []*bigqueryapi.ConnectionProperty) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -168,9 +169,19 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return nil, fmt.Errorf("unable to cast dry_run parameter %s", paramsMap["dry_run"])
|
||||
}
|
||||
|
||||
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
restService := source.BigQueryRestService()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, restService, err = source.BigQueryClientCreator()(tokenStr, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var connProps []*bigqueryapi.ConnectionProperty
|
||||
@@ -272,23 +283,67 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return "Dry run was requested, but no job information was returned.", nil
|
||||
}
|
||||
|
||||
query := bqClient.Query(sql)
|
||||
query.Location = bqClient.Location
|
||||
|
||||
query.ConnectionProperties = connProps
|
||||
|
||||
// Log the query executed for debugging.
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting logger: %s", err)
|
||||
}
|
||||
logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", kind, sql))
|
||||
return source.RunSQL(ctx, bqClient, sql, statementType, nil, connProps)
|
||||
|
||||
// This block handles SELECT statements, which return a row set.
|
||||
// We iterate through the results, convert each row into a map of
|
||||
// column names to values, and return the collection of rows.
|
||||
var out []any
|
||||
job, err := query.Run(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
it, err := job.Read(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to read query results: %w", err)
|
||||
}
|
||||
for {
|
||||
var val []bigqueryapi.Value
|
||||
err = it.Next(&val)
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to iterate through query results: %w", err)
|
||||
}
|
||||
schema := it.Schema
|
||||
row := orderedmap.Row{}
|
||||
for i, field := range schema {
|
||||
row.Add(field.Name, bqutil.NormalizeValue(val[i]))
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
// If the query returned any rows, return them directly.
|
||||
if len(out) > 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This handles the standard case for a SELECT query that successfully
|
||||
// executes but returns zero rows.
|
||||
if statementType == "SELECT" {
|
||||
return "The query returned 0 rows.", nil
|
||||
}
|
||||
// This is the fallback for a successful query that doesn't return content.
|
||||
// In most cases, this will be for DML/DDL statements like INSERT, UPDATE, CREATE, etc.
|
||||
// However, it is also possible that this was a query that was expected to return rows
|
||||
// but returned none, a case that we cannot distinguish here.
|
||||
return "Query executed successfully and returned no content.", nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -21,7 +21,6 @@ import (
|
||||
|
||||
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"
|
||||
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -29,6 +28,7 @@ import (
|
||||
"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"
|
||||
)
|
||||
|
||||
const kind string = "bigquery-forecast"
|
||||
@@ -49,12 +49,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryRestService() *bigqueryrestapi.Service
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
BigQuerySession() bigqueryds.BigQuerySessionProvider
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
RunSQL(context.Context, *bigqueryapi.Client, string, string, []bigqueryapi.QueryParameter, []*bigqueryapi.ConnectionProperty) (any, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -173,9 +173,19 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
}
|
||||
}
|
||||
|
||||
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
restService := source.BigQueryRestService()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, restService, err = source.BigQueryClientCreator()(tokenStr, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var historyDataSource string
|
||||
@@ -241,6 +251,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
idColsFormatted := fmt.Sprintf("['%s']", strings.Join(idCols, "', '"))
|
||||
idColsArg = fmt.Sprintf(", id_cols => %s", idColsFormatted)
|
||||
}
|
||||
|
||||
sql := fmt.Sprintf(`SELECT *
|
||||
FROM AI.FORECAST(
|
||||
%s,
|
||||
@@ -249,14 +260,16 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
horizon => %d%s)`,
|
||||
historyDataSource, dataCol, timestampCol, horizon, idColsArg)
|
||||
|
||||
// JobStatistics.QueryStatistics.StatementType
|
||||
query := bqClient.Query(sql)
|
||||
query.Location = bqClient.Location
|
||||
session, err := source.BigQuerySession()(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
|
||||
}
|
||||
var connProps []*bigqueryapi.ConnectionProperty
|
||||
if session != nil {
|
||||
// Add session ID to the connection properties for subsequent calls.
|
||||
connProps = []*bigqueryapi.ConnectionProperty{
|
||||
query.ConnectionProperties = []*bigqueryapi.ConnectionProperty{
|
||||
{Key: "session_id", Value: session.ID},
|
||||
}
|
||||
}
|
||||
@@ -268,17 +281,46 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
}
|
||||
logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", kind, sql))
|
||||
|
||||
return source.RunSQL(ctx, bqClient, sql, "SELECT", nil, connProps)
|
||||
// This block handles SELECT statements, which return a row set.
|
||||
// We iterate through the results, convert each row into a map of
|
||||
// column names to values, and return the collection of rows.
|
||||
var out []any
|
||||
job, err := query.Run(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
it, err := job.Read(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to read query results: %w", err)
|
||||
}
|
||||
for {
|
||||
var row map[string]bigqueryapi.Value
|
||||
err = it.Next(&row)
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to iterate through query results: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for key, value := range row {
|
||||
vMap[key] = value
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
// If the query returned any rows, return them directly.
|
||||
if len(out) > 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// This handles the standard case for a SELECT query that successfully
|
||||
return "The query returned 0 rows.", nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -20,12 +20,11 @@ import (
|
||||
|
||||
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"
|
||||
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/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
)
|
||||
|
||||
const kind string = "bigquery-get-dataset-info"
|
||||
@@ -48,10 +47,11 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryProject() string
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -138,9 +138,18 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
|
||||
}
|
||||
|
||||
bqClient, _, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, _, err = source.BigQueryClientCreator()(tokenStr, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if !source.IsDatasetAllowed(projectId, datasetId) {
|
||||
@@ -161,10 +170,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -20,12 +20,11 @@ import (
|
||||
|
||||
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"
|
||||
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/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
)
|
||||
|
||||
const kind string = "bigquery-get-table-info"
|
||||
@@ -49,10 +48,11 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryProject() string
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -151,9 +151,18 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
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)
|
||||
}
|
||||
|
||||
bqClient, _, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, _, err = source.BigQueryClientCreator()(tokenStr, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
dsHandle := bqClient.DatasetInProject(projectId, datasetId)
|
||||
@@ -171,10 +180,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -20,11 +20,10 @@ import (
|
||||
|
||||
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"
|
||||
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
@@ -47,9 +46,10 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryProject() string
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
UseClientAuthorization() bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -135,9 +135,17 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
|
||||
}
|
||||
|
||||
bqClient, _, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, _, err = source.BigQueryClientCreator()(tokenStr, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
datasetIterator := bqClient.Datasets(ctx)
|
||||
datasetIterator.ProjectID = projectId
|
||||
@@ -167,10 +175,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -20,12 +20,11 @@ import (
|
||||
|
||||
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"
|
||||
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/parameters"
|
||||
bigqueryrestapi "google.golang.org/api/bigquery/v2"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
@@ -48,11 +47,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
BigQueryClient() *bigqueryapi.Client
|
||||
BigQueryClientCreator() bigqueryds.BigqueryClientCreator
|
||||
BigQueryProject() string
|
||||
UseClientAuthorization() bool
|
||||
IsDatasetAllowed(projectID, datasetID string) bool
|
||||
BigQueryAllowedDatasets() []string
|
||||
RetrieveClientAndService(tools.AccessToken) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
@@ -145,9 +145,17 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
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)
|
||||
}
|
||||
|
||||
bqClient, _, err := source.RetrieveClientAndService(accessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
bqClient := source.BigQueryClient()
|
||||
// Initialize new client if using user OAuth token
|
||||
if source.UseClientAuthorization() {
|
||||
tokenStr, err := accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error parsing access token: %w", err)
|
||||
}
|
||||
bqClient, _, err = source.BigQueryClientCreator()(tokenStr, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
dsHandle := bqClient.DatasetInProject(projectId, datasetId)
|
||||
@@ -178,10 +186,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
dataplexapi "cloud.google.com/go/dataplex/apiv1"
|
||||
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"
|
||||
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -263,10 +262,6 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
|
||||
return parameters.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
// Returns the tool manifest
|
||||
return t.manifest
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user