mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-12 17:09:48 -05:00
Compare commits
1 Commits
refactor/c
...
multi_preb
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d833d1992d |
64
cmd/root.go
64
cmd/root.go
@@ -313,15 +313,15 @@ func Execute() {
|
||||
type Command struct {
|
||||
*cobra.Command
|
||||
|
||||
cfg server.ServerConfig
|
||||
logger log.Logger
|
||||
tools_file string
|
||||
tools_files []string
|
||||
tools_folder string
|
||||
prebuiltConfig string
|
||||
inStream io.Reader
|
||||
outStream io.Writer
|
||||
errStream io.Writer
|
||||
cfg server.ServerConfig
|
||||
logger log.Logger
|
||||
tools_file string
|
||||
tools_files []string
|
||||
tools_folder string
|
||||
prebuiltConfigs []string
|
||||
inStream io.Reader
|
||||
outStream io.Writer
|
||||
errStream io.Writer
|
||||
}
|
||||
|
||||
// NewCommand returns a Command object representing an invocation of the CLI.
|
||||
@@ -374,10 +374,10 @@ func NewCommand(opts ...Option) *Command {
|
||||
flags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
|
||||
// Fetch prebuilt tools sources to customize the help description
|
||||
prebuiltHelp := fmt.Sprintf(
|
||||
"Use a prebuilt tool configuration by source type. Allowed: '%s'.",
|
||||
"Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.",
|
||||
strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
|
||||
)
|
||||
flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", prebuiltHelp)
|
||||
flags.StringSliceVar(&cmd.prebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
|
||||
flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
|
||||
flags.BoolVar(&cmd.cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.")
|
||||
flags.BoolVar(&cmd.cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
|
||||
@@ -864,24 +864,32 @@ func run(cmd *Command) error {
|
||||
var allToolsFiles []ToolsFile
|
||||
|
||||
// Load Prebuilt Configuration
|
||||
if cmd.prebuiltConfig != "" {
|
||||
buf, err := prebuiltconfigs.Get(cmd.prebuiltConfig)
|
||||
if err != nil {
|
||||
cmd.logger.ErrorContext(ctx, err.Error())
|
||||
return err
|
||||
}
|
||||
logMsg := fmt.Sprint("Using prebuilt tool configuration for ", cmd.prebuiltConfig)
|
||||
cmd.logger.InfoContext(ctx, logMsg)
|
||||
// Append prebuilt.source to Version string for the User Agent
|
||||
cmd.cfg.Version += "+prebuilt." + cmd.prebuiltConfig
|
||||
|
||||
parsed, err := parseToolsFile(ctx, buf)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to parse prebuilt tool configuration: %w", err)
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
if len(cmd.prebuiltConfigs) > 0 {
|
||||
slices.Sort(cmd.prebuiltConfigs)
|
||||
sourcesList := strings.Join(cmd.prebuiltConfigs, ", ")
|
||||
logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList)
|
||||
cmd.logger.InfoContext(ctx, logMsg)
|
||||
|
||||
for _, configName := range cmd.prebuiltConfigs {
|
||||
buf, err := prebuiltconfigs.Get(configName)
|
||||
if err != nil {
|
||||
cmd.logger.ErrorContext(ctx, err.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
// Update version string
|
||||
cmd.cfg.Version += "+prebuilt." + configName
|
||||
|
||||
// Parse into ToolsFile struct
|
||||
parsed, err := parseToolsFile(ctx, buf)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err)
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
}
|
||||
allToolsFiles = append(allToolsFiles, parsed)
|
||||
}
|
||||
allToolsFiles = append(allToolsFiles, parsed)
|
||||
}
|
||||
|
||||
// Determine if Custom Files should be loaded
|
||||
@@ -889,7 +897,7 @@ func run(cmd *Command) error {
|
||||
isCustomConfigured := cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != ""
|
||||
|
||||
// Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags)
|
||||
useDefaultToolsFile := cmd.prebuiltConfig == "" && !isCustomConfigured
|
||||
useDefaultToolsFile := len(cmd.prebuiltConfigs) == 0 && !isCustomConfigured
|
||||
|
||||
if useDefaultToolsFile {
|
||||
cmd.tools_file = "tools.yaml"
|
||||
|
||||
@@ -420,17 +420,27 @@ func TestPrebuiltFlag(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
args []string
|
||||
want string
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
desc: "default value",
|
||||
args: []string{},
|
||||
want: "",
|
||||
want: []string{},
|
||||
},
|
||||
{
|
||||
desc: "custom pre built flag",
|
||||
args: []string{"--tools-file", "alloydb"},
|
||||
want: "alloydb",
|
||||
desc: "single prebuilt flag",
|
||||
args: []string{"--prebuilt", "alloydb"},
|
||||
want: []string{"alloydb"},
|
||||
},
|
||||
{
|
||||
desc: "multiple prebuilt flags",
|
||||
args: []string{"--prebuilt", "alloydb", "--prebuilt", "bigquery"},
|
||||
want: []string{"alloydb", "bigquery"},
|
||||
},
|
||||
{
|
||||
desc: "comma separated prebuilt flags",
|
||||
args: []string{"--prebuilt", "alloydb,bigquery"},
|
||||
want: []string{"alloydb", "bigquery"},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
@@ -439,8 +449,8 @@ func TestPrebuiltFlag(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error invoking command: %s", err)
|
||||
}
|
||||
if c.tools_file != tc.want {
|
||||
t.Fatalf("got %v, want %v", c.cfg, tc.want)
|
||||
if diff := cmp.Diff(c.prebuiltConfigs, tc.want); diff != "" {
|
||||
t.Fatalf("got %v, want %v, diff %s", c.prebuiltConfigs, tc.want, diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -2063,6 +2073,12 @@ authSources:
|
||||
return nil
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "sqlite called twice error",
|
||||
args: []string{"--prebuilt", "sqlite", "--prebuilt", "sqlite"},
|
||||
wantErr: true,
|
||||
errString: "resource conflicts detected",
|
||||
},
|
||||
{
|
||||
desc: "tool conflict error",
|
||||
args: []string{"--prebuilt", "sqlite", "--tools-file", toolConflictFile},
|
||||
|
||||
@@ -41,13 +41,13 @@ tools:
|
||||
|
||||
### Usage Flow
|
||||
|
||||
When using this tool, a `query` parameter containing a natural language query is provided to the tool (typically by an agent). The tool then interacts with the Gemini Data Analytics API using the context defined in your configuration.
|
||||
When using this tool, a `prompt` parameter containing a natural language query is provided to the tool (typically by an agent). The tool then interacts with the Gemini Data Analytics API using the context defined in your configuration.
|
||||
|
||||
The structure of the response depends on the `generationOptions` configured in your tool definition (e.g., enabling `generateQueryResult` will include the SQL query results).
|
||||
|
||||
See [Data Analytics API REST documentation](https://clouddocs.devsite.corp.google.com/gemini/docs/conversational-analytics-api/reference/rest/v1alpha/projects.locations/queryData?rep_location=global) for details.
|
||||
|
||||
**Example Input Query:**
|
||||
**Example Input Prompt:**
|
||||
|
||||
```text
|
||||
How many accounts who have region in Prague are eligible for loans? A3 contains the data of region.
|
||||
|
||||
@@ -28,22 +28,6 @@ import (
|
||||
|
||||
const kind string = "cloud-gemini-data-analytics-query"
|
||||
|
||||
// Guidance is the tool guidance string.
|
||||
const Guidance = `Tool guidance:
|
||||
Inputs:
|
||||
1. query: A natural language formulation of a database query.
|
||||
Outputs: (all optional)
|
||||
1. disambiguation_question: Clarification questions or comments where the tool needs the users' input.
|
||||
2. generated_query: The generated query for the user query.
|
||||
3. intent_explanation: An explanation for why the tool produced ` + "`generated_query`" + `.
|
||||
4. query_result: The result of executing ` + "`generated_query`" + `.
|
||||
5. natural_language_answer: The natural language answer that summarizes the ` + "`query`" + ` and ` + "`query_result`" + `.
|
||||
|
||||
Usage guidance:
|
||||
1. If ` + "`disambiguation_question`" + ` is produced, then solicit the needed inputs from the user and try the tool with a new ` + "`query`" + ` that has the needed clarification.
|
||||
2. If ` + "`natural_language_answer`" + ` is produced, use ` + "`intent_explanation`" + ` and ` + "`generated_query`" + ` to see if you need to clarify any assumptions for the user.
|
||||
3. If the tool output indicates failure or empty results, explain that clearly using the provided reasoning.`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
@@ -84,18 +68,11 @@ func (cfg Config) ToolConfigKind() string {
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// Define the parameters for the Gemini Data Analytics Query API
|
||||
// The query is the only input parameter.
|
||||
// The prompt is the only input parameter.
|
||||
allParameters := parameters.Parameters{
|
||||
parameters.NewStringParameterWithRequired("query", "A natural language formulation of a database query.", true),
|
||||
parameters.NewStringParameterWithRequired("prompt", "The natural language question to ask.", true),
|
||||
}
|
||||
// The input and outputs are for tool guidance, usage guidance is for multi-turn interaction.
|
||||
guidance := Guidance
|
||||
|
||||
if cfg.Description != "" {
|
||||
cfg.Description += "\n\n" + guidance
|
||||
} else {
|
||||
cfg.Description = guidance
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
|
||||
|
||||
return Tool{
|
||||
@@ -128,9 +105,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
}
|
||||
|
||||
paramsMap := params.AsMap()
|
||||
query, ok := paramsMap["query"].(string)
|
||||
prompt, ok := paramsMap["prompt"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("query parameter not found or not a string")
|
||||
return nil, fmt.Errorf("prompt parameter not found or not a string")
|
||||
}
|
||||
|
||||
// Parse the access token if provided
|
||||
@@ -148,7 +125,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
|
||||
payload := &QueryDataRequest{
|
||||
Parent: payloadParent,
|
||||
Prompt: query,
|
||||
Prompt: prompt,
|
||||
Context: t.Context,
|
||||
GenerationOptions: t.GenerationOptions,
|
||||
}
|
||||
|
||||
@@ -328,9 +328,9 @@ func TestInvoke(t *testing.T) {
|
||||
t.Fatalf("failed to initialize tool: %v", err)
|
||||
}
|
||||
|
||||
// Prepare parameters for invocation - ONLY query
|
||||
// Prepare parameters for invocation - ONLY prompt
|
||||
params := parameters.ParamValues{
|
||||
{Name: "query", Value: "How many accounts who have region in Prague are eligible for loans?"},
|
||||
{Name: "prompt", Value: "How many accounts who have region in Prague are eligible for loans?"},
|
||||
}
|
||||
|
||||
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
@@ -139,12 +139,12 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
||||
// 1. RunToolGetTestByName
|
||||
expectedManifest := map[string]any{
|
||||
toolName: map[string]any{
|
||||
"description": "Test GDA Tool\n\n" + cloudgda.Guidance,
|
||||
"description": "Test GDA Tool",
|
||||
"parameters": []any{
|
||||
map[string]any{
|
||||
"name": "query",
|
||||
"name": "prompt",
|
||||
"type": "string",
|
||||
"description": "A natural language formulation of a database query.",
|
||||
"description": "The natural language question to ask.",
|
||||
"required": true,
|
||||
"authSources": []any{},
|
||||
},
|
||||
@@ -155,7 +155,7 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTestByName(t, toolName, expectedManifest)
|
||||
|
||||
// 2. RunToolInvokeParametersTest
|
||||
params := []byte(`{"query": "test question"}`)
|
||||
params := []byte(`{"prompt": "test question"}`)
|
||||
tests.RunToolInvokeParametersTest(t, toolName, params, "\"queryResult\":\"SELECT * FROM table;\"")
|
||||
|
||||
// 3. Manual MCP Tool Call Test
|
||||
@@ -172,7 +172,7 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
||||
Params: map[string]any{
|
||||
"name": toolName,
|
||||
"arguments": map[string]any{
|
||||
"query": "test question",
|
||||
"prompt": "test question",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user