mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 08:28:11 -05:00
Compare commits
11 Commits
debug-mac
...
invoke-too
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
decff260fb | ||
|
|
dd61496420 | ||
|
|
9cd494b18a | ||
|
|
005a1e624c | ||
|
|
72a7282797 | ||
|
|
29fe3b93cd | ||
|
|
fb3f66acf4 | ||
|
|
1f95eb134b | ||
|
|
4c240ac3c9 | ||
|
|
c6ab74c5da | ||
|
|
04e2529ba9 |
1
.hugo/assets/scss/_styles_project.scss
Normal file
1
.hugo/assets/scss/_styles_project.scss
Normal file
@@ -0,0 +1 @@
|
||||
@import 'td/code-dark';
|
||||
@@ -48,7 +48,11 @@ enableRobotsTXT = true
|
||||
pre = "<i class='fa-brands fa-github'></i>"
|
||||
|
||||
[markup.goldmark.renderer]
|
||||
unsafe= true
|
||||
unsafe= true
|
||||
|
||||
[markup.highlight]
|
||||
noClasses = false
|
||||
style = "tango"
|
||||
|
||||
[outputFormats]
|
||||
[outputFormats.LLMS]
|
||||
|
||||
@@ -151,6 +151,8 @@ execute `toolbox` to start the server:
|
||||
```sh
|
||||
./toolbox --tools-file "tools.yaml"
|
||||
```
|
||||
> [!NOTE]
|
||||
> Toolbox enables dynamic reloading by default. To disable, use the `--disable-reload` flag.
|
||||
|
||||
You can use `toolbox help` for a full list of flags! To stop the server, send a
|
||||
terminate signal (`ctrl+c` on most platforms).
|
||||
@@ -165,7 +167,7 @@ Once your server is up and running, you can load the tools into your
|
||||
application. See below the list of Client SDKs for using various frameworks:
|
||||
|
||||
<details open>
|
||||
<summary>Python</summary>
|
||||
<summary>Python (<a href="https://github.com/googleapis/mcp-toolbox-sdk-python">Github</a>)</summary>
|
||||
<br>
|
||||
<blockquote>
|
||||
|
||||
@@ -256,7 +258,7 @@ For more detailed instructions on using the Toolbox Core SDK, see the
|
||||
</details>
|
||||
</blockquote>
|
||||
<details>
|
||||
<summary>Javascript/Typescript</summary>
|
||||
<summary>Javascript/Typescript (<a href="https://github.com/googleapis/mcp-toolbox-sdk-js">Github</a>)</summary>
|
||||
<br>
|
||||
<blockquote>
|
||||
|
||||
|
||||
237
cmd/root.go
237
cmd/root.go
@@ -19,20 +19,26 @@ import (
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"io"
|
||||
"maps"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/fsnotify/fsnotify"
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
|
||||
// Import tool packages for side effect of registration
|
||||
@@ -178,6 +184,8 @@ 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.")
|
||||
flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", "Use a prebuilt tool configuration by source type. Cannot be used with --tools-file. Allowed: 'alloydb-postgres', 'bigquery', 'cloud-sql-mysql', 'cloud-sql-postgres', 'cloud-sql-mssql', 'postgres', 'spanner', 'spanner-postgres'.")
|
||||
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.")
|
||||
|
||||
// wrap RunE command so that we have access to original Command object
|
||||
cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) }
|
||||
@@ -347,7 +355,7 @@ func loadAndMergeToolsFolder(ctx context.Context, folderPath string) (ToolsFile,
|
||||
|
||||
// Combine both file lists
|
||||
allFiles := append(yamlFiles, ymlFiles...)
|
||||
|
||||
|
||||
if len(allFiles) == 0 {
|
||||
return ToolsFile{}, fmt.Errorf("no YAML files found in directory %q", folderPath)
|
||||
}
|
||||
@@ -356,6 +364,177 @@ func loadAndMergeToolsFolder(ctx context.Context, folderPath string) (ToolsFile,
|
||||
return loadAndMergeToolsFiles(ctx, allFiles)
|
||||
}
|
||||
|
||||
func handleDynamicReload(ctx context.Context, toolsFile ToolsFile, s *server.Server) error {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, 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, toolsMap, toolsetsMap)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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]tools.Tool, map[string]tools.Toolset, error,
|
||||
) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
instrumentation, err := util.InstrumentationFromContext(ctx)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
logger.DebugContext(ctx, "Attempting to parse and validate reloaded tools file.")
|
||||
|
||||
ctx, span := instrumentation.Tracer.Start(ctx, "toolbox/server/reload")
|
||||
defer span.End()
|
||||
|
||||
reloadedConfig := server.ServerConfig{
|
||||
Version: versionString,
|
||||
SourceConfigs: toolsFile.Sources,
|
||||
AuthServiceConfigs: toolsFile.AuthServices,
|
||||
ToolConfigs: toolsFile.Tools,
|
||||
ToolsetConfigs: toolsFile.Toolsets,
|
||||
}
|
||||
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, 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, err
|
||||
}
|
||||
|
||||
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, nil
|
||||
}
|
||||
|
||||
// watchChanges checks for changes in the provided yaml tools file(s) or folder.
|
||||
func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
w, err := fsnotify.NewWatcher()
|
||||
if err != nil {
|
||||
logger.WarnContext(ctx, "error setting up new watcher %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
defer w.Close()
|
||||
|
||||
watchingFolder := false
|
||||
var folderToWatch string
|
||||
|
||||
// if watchedFiles is empty, indicates that user passed entire folder instead
|
||||
if len(watchedFiles) == 0 {
|
||||
watchingFolder = true
|
||||
|
||||
// validate that watchDirs only has single element
|
||||
if len(watchDirs) > 1 {
|
||||
logger.WarnContext(ctx, "error setting watcher, expected single tools folder if no file(s) are defined.")
|
||||
return
|
||||
}
|
||||
|
||||
for onlyKey := range watchDirs {
|
||||
folderToWatch = onlyKey
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
for dir := range watchDirs {
|
||||
err := w.Add(dir)
|
||||
if err != nil {
|
||||
logger.WarnContext(ctx, fmt.Sprintf("Error adding path %s to watcher: %s", dir, err))
|
||||
break
|
||||
}
|
||||
logger.DebugContext(ctx, fmt.Sprintf("Added directory %s to watcher.", dir))
|
||||
}
|
||||
|
||||
// debounce timer is used to prevent multiple writes triggering multiple reloads
|
||||
debounceDelay := 100 * time.Millisecond
|
||||
debounce := time.NewTimer(1 * time.Minute)
|
||||
debounce.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
logger.DebugContext(ctx, "file watcher context cancelled")
|
||||
return
|
||||
case err, ok := <-w.Errors:
|
||||
if !ok {
|
||||
logger.WarnContext(ctx, "file watcher was closed unexpectedly")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
logger.WarnContext(ctx, "file watcher error %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
case e, ok := <-w.Events:
|
||||
if !ok {
|
||||
logger.WarnContext(ctx, "file watcher already closed")
|
||||
return
|
||||
}
|
||||
|
||||
// only check for write events which indicate user saved a new tools file
|
||||
if !e.Has(fsnotify.Write) {
|
||||
continue
|
||||
}
|
||||
|
||||
cleanedFilename := filepath.Clean(e.Name)
|
||||
logger.DebugContext(ctx, fmt.Sprintf("WRITE event detected in %s", cleanedFilename))
|
||||
|
||||
folderChanged := watchingFolder &&
|
||||
(strings.HasSuffix(cleanedFilename, ".yaml") || strings.HasSuffix(cleanedFilename, ".yml"))
|
||||
|
||||
if folderChanged || watchedFiles[cleanedFilename] {
|
||||
// indicates the write event is on a relevant file
|
||||
debounce.Reset(debounceDelay)
|
||||
}
|
||||
|
||||
case <-debounce.C:
|
||||
debounce.Stop()
|
||||
var reloadedToolsFile ToolsFile
|
||||
|
||||
if watchingFolder {
|
||||
logger.DebugContext(ctx, "Reloading tools folder.")
|
||||
reloadedToolsFile, err = loadAndMergeToolsFolder(ctx, folderToWatch)
|
||||
if err != nil {
|
||||
logger.WarnContext(ctx, "error loading tools folder %s", err)
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
logger.DebugContext(ctx, "Reloading tools file(s).")
|
||||
reloadedToolsFile, err = loadAndMergeToolsFiles(ctx, slices.Collect(maps.Keys(watchedFiles)))
|
||||
if err != nil {
|
||||
logger.WarnContext(ctx, "error loading tools files %s", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
err = handleDynamicReload(ctx, reloadedToolsFile, s)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to parse reloaded tools file at %q: %w", reloadedToolsFile, err)
|
||||
logger.WarnContext(ctx, errMsg.Error())
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// updateLogLevel checks if Toolbox have to update the existing log level set by users.
|
||||
// stdio doesn't support "debug" and "info" logs.
|
||||
func updateLogLevel(stdio bool, logLevel string) bool {
|
||||
@@ -370,6 +549,33 @@ func updateLogLevel(stdio bool, logLevel string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func resolveWatcherInputs(toolsFile string, toolsFiles []string, toolsFolder string) (map[string]bool, map[string]bool) {
|
||||
var relevantFiles []string
|
||||
|
||||
// map for efficiently checking if a file is relevant
|
||||
watchedFiles := make(map[string]bool)
|
||||
|
||||
// dirs that will be added to watcher (fsnotify prefers watching directory then filtering for file)
|
||||
watchDirs := make(map[string]bool)
|
||||
|
||||
if len(toolsFiles) > 0 {
|
||||
relevantFiles = toolsFiles
|
||||
} else if toolsFolder != "" {
|
||||
watchDirs[filepath.Clean(toolsFolder)] = true
|
||||
} else {
|
||||
relevantFiles = []string{toolsFile}
|
||||
}
|
||||
|
||||
// extract parent dir for relevant files and dedup
|
||||
for _, f := range relevantFiles {
|
||||
cleanFile := filepath.Clean(f)
|
||||
watchedFiles[cleanFile] = true
|
||||
watchDirs[filepath.Dir(cleanFile)] = true
|
||||
}
|
||||
|
||||
return watchDirs, watchedFiles
|
||||
}
|
||||
|
||||
func run(cmd *Command) error {
|
||||
if updateLogLevel(cmd.cfg.Stdio, cmd.cfg.LogLevel.String()) {
|
||||
cmd.cfg.LogLevel = server.StringLevel(log.Warn)
|
||||
@@ -466,6 +672,7 @@ func run(cmd *Command) error {
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
}
|
||||
|
||||
// Use multiple tools files
|
||||
cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(cmd.tools_files)))
|
||||
var err error
|
||||
@@ -481,6 +688,7 @@ func run(cmd *Command) error {
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
}
|
||||
|
||||
// Use tools folder
|
||||
cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", cmd.tools_folder))
|
||||
var err error
|
||||
@@ -494,6 +702,7 @@ func run(cmd *Command) error {
|
||||
if cmd.tools_file == "" {
|
||||
cmd.tools_file = "tools.yaml"
|
||||
}
|
||||
|
||||
// Read single tool file contents
|
||||
buf, err := os.ReadFile(cmd.tools_file)
|
||||
if err != nil {
|
||||
@@ -516,9 +725,23 @@ func run(cmd *Command) error {
|
||||
cmd.logger.WarnContext(ctx, "`authSources` is deprecated, use `authServices` instead")
|
||||
cmd.cfg.AuthServiceConfigs = authSourceConfigs
|
||||
}
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to parse tool file at %q: %w", cmd.tools_file, err)
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
}
|
||||
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(versionString)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("unable to create telemetry instrumentation: %w", err)
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
return errMsg
|
||||
}
|
||||
|
||||
ctx = util.WithInstrumentation(ctx, instrumentation)
|
||||
|
||||
// start server
|
||||
s, err := server.NewServer(ctx, cmd.cfg, cmd.logger)
|
||||
s, err := server.NewServer(ctx, cmd.cfg)
|
||||
if err != nil {
|
||||
errMsg := fmt.Errorf("toolbox failed to initialize: %w", err)
|
||||
cmd.logger.ErrorContext(ctx, errMsg.Error())
|
||||
@@ -543,6 +766,9 @@ func run(cmd *Command) error {
|
||||
return errMsg
|
||||
}
|
||||
cmd.logger.InfoContext(ctx, "Server ready to serve!")
|
||||
if cmd.cfg.UI {
|
||||
cmd.logger.InfoContext(ctx, "Toolbox UI is up and running at: http://localhost:5000/ui")
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer close(srvErr)
|
||||
@@ -553,6 +779,13 @@ func run(cmd *Command) error {
|
||||
}()
|
||||
}
|
||||
|
||||
watchDirs, watchedFiles := resolveWatcherInputs(cmd.tools_file, cmd.tools_files, cmd.tools_folder)
|
||||
|
||||
if !cmd.cfg.DisableReload {
|
||||
// start watching the file(s) or folder for changes to trigger dynamic reloading
|
||||
go watchChanges(ctx, watchDirs, watchedFiles, s)
|
||||
}
|
||||
|
||||
// wait for either the server to error out or the command's context to be canceled
|
||||
select {
|
||||
case err := <-srvErr:
|
||||
|
||||
194
cmd/root_test.go
194
cmd/root_test.go
@@ -16,23 +16,33 @@ package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/auth/google"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
|
||||
httpsrc "github.com/googleapis/genai-toolbox/internal/sources/http"
|
||||
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/http"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
@@ -174,6 +184,13 @@ func TestServerConfigFlags(t *testing.T) {
|
||||
Stdio: true,
|
||||
}),
|
||||
},
|
||||
{
|
||||
desc: "disable reload",
|
||||
args: []string{"--disable-reload"},
|
||||
want: withDefaults(server.ServerConfig{
|
||||
DisableReload: true,
|
||||
}),
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
@@ -965,6 +982,183 @@ func TestEnvVarReplacement(t *testing.T) {
|
||||
|
||||
}
|
||||
|
||||
// normalizeFilepaths is a helper function to allow same filepath formats for Mac and Windows.
|
||||
// this prevents needing multiple "want" cases for TestResolveWatcherInputs
|
||||
func normalizeFilepaths(m map[string]bool) map[string]bool {
|
||||
newMap := make(map[string]bool)
|
||||
for k, v := range m {
|
||||
newMap[filepath.ToSlash(k)] = v
|
||||
}
|
||||
return newMap
|
||||
}
|
||||
|
||||
func TestResolveWatcherInputs(t *testing.T) {
|
||||
tcs := []struct {
|
||||
description string
|
||||
toolsFile string
|
||||
toolsFiles []string
|
||||
toolsFolder string
|
||||
wantWatchDirs map[string]bool
|
||||
wantWatchedFiles map[string]bool
|
||||
}{
|
||||
{
|
||||
description: "single tools file",
|
||||
toolsFile: "tools_folder/example_tools.yaml",
|
||||
toolsFiles: []string{},
|
||||
toolsFolder: "",
|
||||
wantWatchDirs: map[string]bool{"tools_folder": true},
|
||||
wantWatchedFiles: map[string]bool{"tools_folder/example_tools.yaml": true},
|
||||
},
|
||||
{
|
||||
description: "default tools file (root dir)",
|
||||
toolsFile: "tools.yaml",
|
||||
toolsFiles: []string{},
|
||||
toolsFolder: "",
|
||||
wantWatchDirs: map[string]bool{".": true},
|
||||
wantWatchedFiles: map[string]bool{"tools.yaml": true},
|
||||
},
|
||||
{
|
||||
description: "multiple files in different folders",
|
||||
toolsFile: "",
|
||||
toolsFiles: []string{"tools_folder/example_tools.yaml", "tools_folder2/example_tools.yaml"},
|
||||
toolsFolder: "",
|
||||
wantWatchDirs: map[string]bool{"tools_folder": true, "tools_folder2": true},
|
||||
wantWatchedFiles: map[string]bool{
|
||||
"tools_folder/example_tools.yaml": true,
|
||||
"tools_folder2/example_tools.yaml": true,
|
||||
},
|
||||
},
|
||||
{
|
||||
description: "multiple files in same folder",
|
||||
toolsFile: "",
|
||||
toolsFiles: []string{"tools_folder/example_tools.yaml", "tools_folder/example_tools2.yaml"},
|
||||
toolsFolder: "",
|
||||
wantWatchDirs: map[string]bool{"tools_folder": true},
|
||||
wantWatchedFiles: map[string]bool{
|
||||
"tools_folder/example_tools.yaml": true,
|
||||
"tools_folder/example_tools2.yaml": true,
|
||||
},
|
||||
},
|
||||
{
|
||||
description: "multiple files in different levels",
|
||||
toolsFile: "",
|
||||
toolsFiles: []string{
|
||||
"tools_folder/example_tools.yaml",
|
||||
"tools_folder/special_tools/example_tools2.yaml"},
|
||||
toolsFolder: "",
|
||||
wantWatchDirs: map[string]bool{"tools_folder": true, "tools_folder/special_tools": true},
|
||||
wantWatchedFiles: map[string]bool{
|
||||
"tools_folder/example_tools.yaml": true,
|
||||
"tools_folder/special_tools/example_tools2.yaml": true,
|
||||
},
|
||||
},
|
||||
{
|
||||
description: "tools folder",
|
||||
toolsFile: "",
|
||||
toolsFiles: []string{},
|
||||
toolsFolder: "tools_folder",
|
||||
wantWatchDirs: map[string]bool{"tools_folder": true},
|
||||
wantWatchedFiles: map[string]bool{},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.description, func(t *testing.T) {
|
||||
gotWatchDirs, gotWatchedFiles := resolveWatcherInputs(tc.toolsFile, tc.toolsFiles, tc.toolsFolder)
|
||||
|
||||
normalizedGotWatchDirs := normalizeFilepaths(gotWatchDirs)
|
||||
normalizedGotWatchedFiles := normalizeFilepaths(gotWatchedFiles)
|
||||
|
||||
if diff := cmp.Diff(tc.wantWatchDirs, normalizedGotWatchDirs); diff != "" {
|
||||
t.Errorf("incorrect watchDirs: diff %v", diff)
|
||||
}
|
||||
if diff := cmp.Diff(tc.wantWatchedFiles, normalizedGotWatchedFiles); diff != "" {
|
||||
t.Errorf("incorrect watchedFiles: diff %v", diff)
|
||||
}
|
||||
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// helper function for testing file detection in dynamic reloading
|
||||
func tmpFileWithCleanup(content []byte) (string, func(), error) {
|
||||
f, err := os.CreateTemp("", "*")
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
cleanup := func() { os.Remove(f.Name()) }
|
||||
|
||||
if _, err := f.Write(content); err != nil {
|
||||
cleanup()
|
||||
return "", nil, err
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
cleanup()
|
||||
return "", nil, err
|
||||
}
|
||||
return f.Name(), cleanup, err
|
||||
}
|
||||
|
||||
func TestSingleEdit(t *testing.T) {
|
||||
ctx, cancelCtx := context.WithTimeout(context.Background(), time.Minute)
|
||||
defer cancelCtx()
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
defer pw.Close()
|
||||
defer pr.Close()
|
||||
|
||||
fileToWatch, cleanup, err := tmpFileWithCleanup([]byte("initial content"))
|
||||
if err != nil {
|
||||
t.Fatalf("error editing tools file %s", err)
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
logger, err := log.NewStdLogger(pw, pw, "DEBUG")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to setup logger %s", err)
|
||||
}
|
||||
ctx = util.WithLogger(ctx, logger)
|
||||
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(versionString)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to setup instrumentation %s", err)
|
||||
}
|
||||
ctx = util.WithInstrumentation(ctx, instrumentation)
|
||||
|
||||
mockServer := &server.Server{}
|
||||
|
||||
cleanFileToWatch := filepath.Clean(fileToWatch)
|
||||
watchDir := filepath.Dir(cleanFileToWatch)
|
||||
|
||||
watchedFiles := map[string]bool{cleanFileToWatch: true}
|
||||
watchDirs := map[string]bool{watchDir: true}
|
||||
|
||||
go watchChanges(ctx, watchDirs, watchedFiles, mockServer)
|
||||
|
||||
// escape backslash so regex doesn't fail on windows filepaths
|
||||
regexEscapedPathFile := strings.ReplaceAll(cleanFileToWatch, `\`, `\\\\*\\`)
|
||||
regexEscapedPathFile = path.Clean(regexEscapedPathFile)
|
||||
|
||||
regexEscapedPathDir := strings.ReplaceAll(watchDir, `\`, `\\\\*\\`)
|
||||
regexEscapedPathDir = path.Clean(regexEscapedPathDir)
|
||||
|
||||
begunWatchingDir := regexp.MustCompile(fmt.Sprintf(`DEBUG "Added directory %s to watcher."`, regexEscapedPathDir))
|
||||
_, err = testutils.WaitForString(ctx, begunWatchingDir, pr)
|
||||
if err != nil {
|
||||
t.Fatalf("timeout or error waiting for watcher to start: %s", err)
|
||||
}
|
||||
|
||||
err = os.WriteFile(fileToWatch, []byte("modification"), 0777)
|
||||
if err != nil {
|
||||
t.Fatalf("error writing to file: %v", err)
|
||||
}
|
||||
|
||||
detectedFileChange := regexp.MustCompile(fmt.Sprintf(`DEBUG "WRITE event detected in %s"`, regexEscapedPathFile))
|
||||
_, err = testutils.WaitForString(ctx, detectedFileChange, pr)
|
||||
if err != nil {
|
||||
t.Fatalf("timeout or error waiting for file to detect write: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrebuiltTools(t *testing.T) {
|
||||
alloydb_config, _ := prebuiltconfigs.Get("alloydb-postgres")
|
||||
bigquery_config, _ := prebuiltconfigs.Get("bigquery")
|
||||
|
||||
@@ -123,6 +123,9 @@ execute `toolbox` to start the server:
|
||||
```sh
|
||||
./toolbox --tools-file "tools.yaml"
|
||||
```
|
||||
{{< notice note >}}
|
||||
Toolbox enables dynamic reloading by default. To disable, use the `--disable-reload` flag.
|
||||
{{< /notice >}}
|
||||
|
||||
You can use `toolbox help` for a full list of flags! To stop the server, send a
|
||||
terminate signal (`ctrl+c` on most platforms).
|
||||
|
||||
@@ -257,6 +257,9 @@ In this section, we will download Toolbox, configure our tools in a
|
||||
```bash
|
||||
./toolbox --tools-file "tools.yaml"
|
||||
```
|
||||
{{< notice note >}}
|
||||
Toolbox enables dynamic reloading by default. To disable, use the `--disable-reload` flag.
|
||||
{{< /notice >}}
|
||||
|
||||
## Step 3: Connect your agent to Toolbox
|
||||
|
||||
|
||||
@@ -63,6 +63,10 @@ When running with stdio, Toolbox will listen via stdio instead of acting as a
|
||||
remote HTTP server. Logs will be set to the `warn` level by default. `debug` and
|
||||
`info` logs are not supported with stdio.
|
||||
|
||||
{{< notice note >}}
|
||||
Toolbox enables dynamic reloading by default. To disable, use the `--disable-reload` flag.
|
||||
{{< /notice >}}
|
||||
|
||||
### Connecting via HTTP
|
||||
|
||||
Toolbox supports the HTTP transport protocol with and without SSE.
|
||||
|
||||
@@ -79,7 +79,7 @@ database are in the same VPC network.
|
||||
|
||||
Create a `tools.yaml` file that contains your configuration for Toolbox. For
|
||||
details, see the
|
||||
[configuration](https://github.com/googleapis/genai-toolbox/blob/main/README.md#configuration)
|
||||
[configuration](https://googleapis.github.io/genai-toolbox/resources/sources/)
|
||||
section.
|
||||
|
||||
## Deploy to Cloud Run
|
||||
|
||||
@@ -15,8 +15,10 @@ It's compatible with the following sources:
|
||||
|
||||
- [bigquery](../sources/bigquery.md)
|
||||
|
||||
bigquery-get-dataset-info takes a dataset parameter to specify the dataset
|
||||
on the given source.
|
||||
`bigquery-get-dataset-info` takes a `dataset` parameter to specify the dataset
|
||||
on the given source. It also optionally accepts a `project` parameter to
|
||||
define the Google Cloud project ID. If the `project` parameter is not provided,
|
||||
the tool defaults to using the project defined in the source configuration.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -15,8 +15,10 @@ It's compatible with the following sources:
|
||||
|
||||
- [bigquery](../sources/bigquery.md)
|
||||
|
||||
bigquery-get-table-info takes dataset and table parameters to specify
|
||||
the target table.
|
||||
`bigquery-get-table-info` takes `dataset` and `table` parameters to specify
|
||||
the target table. It also optionally accepts a `project` parameter to define
|
||||
the Google Cloud project ID. If the `project` parameter is not provided, the
|
||||
tool defaults to using the project defined in the source configuration.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -15,8 +15,9 @@ It's compatible with the following sources:
|
||||
|
||||
- [bigquery](../sources/bigquery.md)
|
||||
|
||||
bigquery-list-dataset-ids requires no input parameters beyond the configured
|
||||
source.
|
||||
`bigquery-list-dataset-ids` optionally accepts a `project` parameter to define
|
||||
the Google Cloud project ID. If the `project` parameter is not provided, the
|
||||
tool defaults to using the project defined in the source configuration.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -15,8 +15,10 @@ It's compatible with the following sources:
|
||||
|
||||
- [bigquery](../sources/bigquery.md)
|
||||
|
||||
bigquery-get-dataset-info takes a dataset parameter to specify the dataset
|
||||
from which to list table IDs.
|
||||
`bigquery-get-dataset-info` takes a required `dataset` parameter to specify the dataset
|
||||
from which to list table IDs. It also optionally accepts a `project` parameter to
|
||||
define the Google Cloud project ID. If the `project` parameter is not provided, the
|
||||
tool defaults to using the project defined in the source configuration.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -292,6 +292,9 @@ to use BigQuery, and then run the Toolbox server.
|
||||
```bash
|
||||
./toolbox --tools-file "tools.yaml"
|
||||
```
|
||||
{{< notice note >}}
|
||||
Toolbox enables dynamic reloading by default. To disable, use the `--disable-reload` flag.
|
||||
{{< /notice >}}
|
||||
|
||||
## Step 3: Connect your agent to Toolbox
|
||||
|
||||
|
||||
4
go.mod
4
go.mod
@@ -14,9 +14,11 @@ require (
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.29.0
|
||||
github.com/couchbase/gocb/v2 v2.10.0
|
||||
github.com/couchbase/tools-common/http v1.0.9
|
||||
github.com/fsnotify/fsnotify v1.9.0
|
||||
github.com/go-chi/chi/v5 v5.2.2
|
||||
github.com/go-chi/httplog/v2 v2.1.1
|
||||
github.com/go-chi/render v1.0.3
|
||||
github.com/go-goquery/goquery v1.0.1
|
||||
github.com/go-playground/validator/v10 v10.27.0
|
||||
github.com/go-sql-driver/mysql v1.9.3
|
||||
github.com/goccy/go-yaml v1.18.0
|
||||
@@ -59,7 +61,9 @@ require (
|
||||
github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp v1.5.3 // indirect
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 // indirect
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 // indirect
|
||||
github.com/PuerkitoBio/goquery v1.10.3 // indirect
|
||||
github.com/ajg/form v1.5.1 // indirect
|
||||
github.com/andybalholm/cascadia v1.3.3 // indirect
|
||||
github.com/apache/arrow/go/v15 v15.0.2 // indirect
|
||||
github.com/cenkalti/backoff/v5 v5.0.2 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
|
||||
42
go.sum
42
go.sum
@@ -661,6 +661,8 @@ github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapp
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0/go.mod h1:cSgYe11MCNYunTnRXrKiR/tHc0eoKjICUuWpNZoVCOo=
|
||||
github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk=
|
||||
github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
|
||||
github.com/PuerkitoBio/goquery v1.10.3 h1:pFYcNSqHxBD06Fpj/KsbStFRsgRATgnf3LeXiUkhzPo=
|
||||
github.com/PuerkitoBio/goquery v1.10.3/go.mod h1:tMUX0zDMHXYlAQk6p35XxQMqMweEKB7iK7iLNd4RH4Y=
|
||||
github.com/ajg/form v1.5.1 h1:t9c7v8JUKu/XxOGBU0yjNpaMloxGEJhUkqFRq0ibGeU=
|
||||
github.com/ajg/form v1.5.1/go.mod h1:uL1WgH+h2mgNtvBq0339dVnzXdBETtL2LeUXaIv25UY=
|
||||
github.com/ajstarks/deck v0.0.0-20200831202436-30c9fc6549a9/go.mod h1:JynElWSGnm/4RlzPXRlREEwqTHAN3T56Bv2ITsFT3gY=
|
||||
@@ -668,6 +670,8 @@ github.com/ajstarks/deck/generate v0.0.0-20210309230005-c3f852c02e19/go.mod h1:T
|
||||
github.com/ajstarks/svgo v0.0.0-20180226025133-644b8db467af/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw=
|
||||
github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b/go.mod h1:1KcenG0jGWcpt8ov532z81sp/kMMUG485J2InIOyADM=
|
||||
github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
|
||||
github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM=
|
||||
github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA=
|
||||
github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
|
||||
github.com/apache/arrow/go/v10 v10.0.1/go.mod h1:YvhnlEePVnBS4+0z3fhPfUy7W1Ikj0Ih0vcRo/gZ1M0=
|
||||
github.com/apache/arrow/go/v11 v11.0.0/go.mod h1:Eg5OsL5H+e299f7u5ssuXsuHQVEGC4xei5aX110hRiI=
|
||||
@@ -765,6 +769,8 @@ github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2
|
||||
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
|
||||
github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
|
||||
github.com/fogleman/gg v1.3.0/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
|
||||
github.com/gabriel-vasile/mimetype v1.4.8 h1:FfZ3gj38NjllZIeJAmMhr+qKL8Wu+nOoI3GqacKw1NM=
|
||||
github.com/gabriel-vasile/mimetype v1.4.8/go.mod h1:ByKUIKGjh1ODkGM1asKUbQZOLGrPjydw3hYPU2YU9t8=
|
||||
github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
|
||||
@@ -782,6 +788,8 @@ github.com/go-fonts/stix v0.1.0/go.mod h1:w/c1f0ldAUlJmLBvlbkvVXLAD+tAMqobIIQpmn
|
||||
github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU=
|
||||
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
|
||||
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
|
||||
github.com/go-goquery/goquery v1.0.1 h1:kpchVA1LdOFWdRpkDPESVdlb1JQI6ixsJ5MiNUITO7U=
|
||||
github.com/go-goquery/goquery v1.0.1/go.mod h1:W5s8OWbqWf6lG0LkXWBeh7U1Y/X5XTI0Br65MHF8uJk=
|
||||
github.com/go-jose/go-jose/v4 v4.0.5 h1:M6T8+mKZl/+fNNuFHvGIzDz7BTLQPIounk/b9dw3AaE=
|
||||
github.com/go-jose/go-jose/v4 v4.0.5/go.mod h1:s3P1lRrkT8igV8D9OjyL4WRyHvjB6a4JSllnOrmmBOA=
|
||||
github.com/go-kit/log v0.1.0/go.mod h1:zbhenjAZHb184qTLMA9ZjW7ThYL0H2mk7Q6pNt4vbaY=
|
||||
@@ -883,6 +891,7 @@ github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
|
||||
github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
|
||||
github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
@@ -1176,6 +1185,10 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
|
||||
golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.0.0-20211108221036-ceb1ce70b4fa/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
|
||||
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
|
||||
golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
|
||||
golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
|
||||
golang.org/x/exp v0.0.0-20180321215751-8460e604b9de/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
@@ -1237,6 +1250,9 @@ golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91
|
||||
golang.org/x/mod v0.7.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.9.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.25.0 h1:n7a+ZbQKQA/Ysbyb0/6IbB1H/X41mKgbhfv7AfG/44w=
|
||||
golang.org/x/mod v0.25.0/go.mod h1:IXM97Txy2VM4PJ3gI61r1YEk/gAj6zAHN3AdZt6S9Ww=
|
||||
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
@@ -1296,6 +1312,11 @@ golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
|
||||
golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
|
||||
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
|
||||
golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
|
||||
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
|
||||
golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
|
||||
golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
|
||||
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
|
||||
@@ -1345,6 +1366,10 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ
|
||||
golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
|
||||
golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8=
|
||||
golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
|
||||
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
@@ -1426,8 +1451,14 @@ golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
|
||||
golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.2.0/go.mod h1:TVmDHMZPmdnySmBfhjOoOdhjzdE1h4u1VwSiw2l1Nuc=
|
||||
@@ -1436,6 +1467,11 @@ golang.org/x/term v0.4.0/go.mod h1:9P2UbLfCdcvo3p/nzKvsmas4TnlujnuoV9hGgYzW1lQ=
|
||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||
golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
|
||||
golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=
|
||||
golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
|
||||
golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
|
||||
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
|
||||
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
|
||||
golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
@@ -1452,6 +1488,10 @@ golang.org/x/text v0.6.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
|
||||
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
|
||||
golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
|
||||
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
|
||||
golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
@@ -1526,6 +1566,8 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
|
||||
golang.org/x/tools v0.3.0/go.mod h1:/rWhSS2+zyEVwoJf8YAX6L2f0ntZ7Kn/mGgAWcipA5k=
|
||||
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
|
||||
golang.org/x/tools v0.7.0/go.mod h1:4pg6aUX35JBAogB10C9AtvVL+qowtN4pT3CGSQex14s=
|
||||
golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/tools v0.33.0 h1:4qz2S3zmRxbGIhDIAgjxvFutSvH5EfnsYrRBj0UI0bc=
|
||||
golang.org/x/tools v0.33.0/go.mod h1:CIJMaWEY88juyUfo7UbgPqbC8rU2OqfAV1h2Qp0oMYI=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
||||
@@ -75,7 +75,7 @@ func toolsetHandler(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
)
|
||||
}()
|
||||
|
||||
toolset, ok := s.toolsets[toolsetName]
|
||||
toolset, ok := s.ResourceMgr.GetToolset(toolsetName)
|
||||
if !ok {
|
||||
err = fmt.Errorf("toolset %q does not exist", toolsetName)
|
||||
s.logger.DebugContext(ctx, err.Error())
|
||||
@@ -111,7 +111,7 @@ func toolGetHandler(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
metric.WithAttributes(attribute.String("toolbox.operation.status", status)),
|
||||
)
|
||||
}()
|
||||
tool, ok := s.tools[toolName]
|
||||
tool, ok := s.ResourceMgr.GetTool(toolName)
|
||||
if !ok {
|
||||
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
|
||||
s.logger.DebugContext(ctx, err.Error())
|
||||
@@ -156,7 +156,7 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
)
|
||||
}()
|
||||
|
||||
tool, ok := s.tools[toolName]
|
||||
tool, ok := s.ResourceMgr.GetTool(toolName)
|
||||
if !ok {
|
||||
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
|
||||
s.logger.DebugContext(ctx, err.Error())
|
||||
@@ -167,7 +167,7 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
|
||||
// Tool authentication
|
||||
// claimsFromAuth maps the name of the authservice to the claims retrieved from it.
|
||||
claimsFromAuth := make(map[string]map[string]any)
|
||||
for _, aS := range s.authServices {
|
||||
for _, aS := range s.ResourceMgr.GetAuthServiceMap() {
|
||||
claims, err := aS.GetClaimsFromHeader(ctx, r.Header)
|
||||
if err != nil {
|
||||
s.logger.DebugContext(ctx, err.Error())
|
||||
|
||||
@@ -147,14 +147,23 @@ func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, tools
|
||||
t.Fatalf("unable to setup otel: %s", err)
|
||||
}
|
||||
|
||||
instrumentation, err := CreateTelemetryInstrumentation(fakeVersionString)
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(fakeVersionString)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create custom metrics: %s", err)
|
||||
}
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
server := Server{version: fakeVersionString, logger: testLogger, instrumentation: instrumentation, sseManager: sseManager, tools: tools, toolsets: toolsets}
|
||||
resourceManager := NewResourceManager(nil, nil, tools, toolsets)
|
||||
|
||||
server := Server{
|
||||
version: fakeVersionString,
|
||||
logger: testLogger,
|
||||
instrumentation: instrumentation,
|
||||
sseManager: sseManager,
|
||||
ResourceMgr: resourceManager,
|
||||
}
|
||||
|
||||
var r chi.Router
|
||||
switch router {
|
||||
case "api":
|
||||
|
||||
@@ -53,6 +53,10 @@ type ServerConfig struct {
|
||||
TelemetryServiceName string
|
||||
// Stdio indicates if Toolbox is listening via MCP stdio.
|
||||
Stdio bool
|
||||
// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.
|
||||
DisableReload bool
|
||||
// UI indicates if Toolbox UI endpoints (/ui) are available
|
||||
UI bool
|
||||
}
|
||||
|
||||
type logFormat string
|
||||
|
||||
@@ -479,12 +479,12 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
|
||||
}
|
||||
return v, res, err
|
||||
default:
|
||||
toolset, ok := s.toolsets[toolsetName]
|
||||
toolset, ok := s.ResourceMgr.GetToolset(toolsetName)
|
||||
if !ok {
|
||||
err = fmt.Errorf("toolset does not exist")
|
||||
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||
}
|
||||
res, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, s.tools, body)
|
||||
res, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, s.ResourceMgr.GetToolsMap(), body)
|
||||
return "", res, err
|
||||
}
|
||||
}
|
||||
|
||||
@@ -693,14 +693,22 @@ func TestStdioSession(t *testing.T) {
|
||||
}
|
||||
}()
|
||||
|
||||
instrumentation, err := CreateTelemetryInstrumentation(fakeVersionString)
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(fakeVersionString)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create custom metrics: %s", err)
|
||||
}
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
server := &Server{version: fakeVersionString, logger: testLogger, instrumentation: instrumentation, sseManager: sseManager, tools: toolsMap, toolsets: toolsets}
|
||||
resourceManager := NewResourceManager(nil, nil, toolsMap, toolsets)
|
||||
|
||||
server := &Server{
|
||||
version: fakeVersionString,
|
||||
logger: testLogger,
|
||||
instrumentation: instrumentation,
|
||||
sseManager: sseManager,
|
||||
ResourceMgr: resourceManager,
|
||||
}
|
||||
|
||||
in := bufio.NewReader(pr)
|
||||
stdioSession := NewStdioSession(server, in, pw)
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
@@ -29,6 +30,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
@@ -42,26 +44,225 @@ type Server struct {
|
||||
listener net.Listener
|
||||
root chi.Router
|
||||
logger log.Logger
|
||||
instrumentation *Instrumentation
|
||||
instrumentation *telemetry.Instrumentation
|
||||
sseManager *sseManager
|
||||
ResourceMgr *ResourceManager
|
||||
}
|
||||
|
||||
// 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
|
||||
tools map[string]tools.Tool
|
||||
toolsets map[string]tools.Toolset
|
||||
}
|
||||
|
||||
// NewServer returns a Server object based on provided Config.
|
||||
func NewServer(ctx context.Context, cfg ServerConfig, l log.Logger) (*Server, error) {
|
||||
instrumentation, err := CreateTelemetryInstrumentation(cfg.Version)
|
||||
func NewResourceManager(
|
||||
sourcesMap map[string]sources.Source,
|
||||
authServicesMap map[string]auth.AuthService,
|
||||
toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset,
|
||||
) *ResourceManager {
|
||||
resourceMgr := &ResourceManager{
|
||||
mu: sync.RWMutex{},
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
}
|
||||
|
||||
return resourceMgr
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetSource(sourceName string) (sources.Source, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
source, ok := r.sources[sourceName]
|
||||
return source, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetAuthService(authServiceName string) (auth.AuthService, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
authService, ok := r.authServices[authServiceName]
|
||||
return authService, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetTool(toolName string) (tools.Tool, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
tool, ok := r.tools[toolName]
|
||||
return tool, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetToolset(toolsetName string) (tools.Toolset, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
toolset, ok := r.toolsets[toolsetName]
|
||||
return toolset, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) SetResources(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.sources = sourcesMap
|
||||
r.authServices = authServicesMap
|
||||
r.tools = toolsMap
|
||||
r.toolsets = toolsetsMap
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetAuthServiceMap() map[string]auth.AuthService {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.authServices
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetToolsMap() map[string]tools.Tool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.tools
|
||||
}
|
||||
|
||||
func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
map[string]sources.Source,
|
||||
map[string]auth.AuthService,
|
||||
map[string]tools.Tool,
|
||||
map[string]tools.Toolset,
|
||||
error,
|
||||
) {
|
||||
ctx = util.WithUserAgent(ctx, cfg.Version)
|
||||
instrumentation, err := util.InstrumentationFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to create telemetry instrumentation: %w", err)
|
||||
panic(err)
|
||||
}
|
||||
|
||||
l, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// initialize and validate the sources from configs
|
||||
sourcesMap := make(map[string]sources.Source)
|
||||
for name, sc := range cfg.SourceConfigs {
|
||||
s, err := func() (sources.Source, error) {
|
||||
childCtx, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/source/init",
|
||||
trace.WithAttributes(attribute.String("source_kind", sc.SourceConfigKind())),
|
||||
trace.WithAttributes(attribute.String("source_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
s, err := sc.Initialize(childCtx, instrumentation.Tracer)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize source %q: %w", name, err)
|
||||
}
|
||||
return s, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
sourcesMap[name] = s
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d sources.", len(sourcesMap)))
|
||||
|
||||
// initialize and validate the auth services from configs
|
||||
authServicesMap := make(map[string]auth.AuthService)
|
||||
for name, sc := range cfg.AuthServiceConfigs {
|
||||
a, err := func() (auth.AuthService, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/auth/init",
|
||||
trace.WithAttributes(attribute.String("auth_kind", sc.AuthServiceConfigKind())),
|
||||
trace.WithAttributes(attribute.String("auth_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
a, err := sc.Initialize()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize auth service %q: %w", name, err)
|
||||
}
|
||||
return a, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
authServicesMap[name] = a
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d authServices.", len(authServicesMap)))
|
||||
|
||||
// initialize and validate the tools from configs
|
||||
toolsMap := make(map[string]tools.Tool)
|
||||
for name, tc := range cfg.ToolConfigs {
|
||||
t, err := func() (tools.Tool, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/tool/init",
|
||||
trace.WithAttributes(attribute.String("tool_kind", tc.ToolConfigKind())),
|
||||
trace.WithAttributes(attribute.String("tool_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
t, err := tc.Initialize(sourcesMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize tool %q: %w", name, err)
|
||||
}
|
||||
return t, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
toolsMap[name] = t
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d tools.", len(toolsMap)))
|
||||
|
||||
// create a default toolset that contains all tools
|
||||
allToolNames := make([]string, 0, len(toolsMap))
|
||||
for name := range toolsMap {
|
||||
allToolNames = append(allToolNames, name)
|
||||
}
|
||||
if cfg.ToolsetConfigs == nil {
|
||||
cfg.ToolsetConfigs = make(ToolsetConfigs)
|
||||
}
|
||||
cfg.ToolsetConfigs[""] = tools.ToolsetConfig{Name: "", ToolNames: allToolNames}
|
||||
|
||||
// initialize and validate the toolsets from configs
|
||||
toolsetsMap := make(map[string]tools.Toolset)
|
||||
for name, tc := range cfg.ToolsetConfigs {
|
||||
t, err := func() (tools.Toolset, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/toolset/init",
|
||||
trace.WithAttributes(attribute.String("toolset_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
t, err := tc.Initialize(cfg.Version, toolsMap)
|
||||
if err != nil {
|
||||
return tools.Toolset{}, fmt.Errorf("unable to initialize toolset %q: %w", name, err)
|
||||
}
|
||||
return t, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
toolsetsMap[name] = t
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d toolsets.", len(toolsetsMap)))
|
||||
|
||||
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, nil
|
||||
}
|
||||
|
||||
// NewServer returns a Server object based on provided Config.
|
||||
func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
instrumentation, err := util.InstrumentationFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, span := instrumentation.Tracer.Start(ctx, "toolbox/server/init")
|
||||
defer span.End()
|
||||
|
||||
ctx = util.WithUserAgent(ctx, cfg.Version)
|
||||
l, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// set up http serving
|
||||
r := chi.NewRouter()
|
||||
@@ -97,116 +298,18 @@ func NewServer(ctx context.Context, cfg ServerConfig, l log.Logger) (*Server, er
|
||||
httpLogger := httplog.NewLogger("httplog", httpOpts)
|
||||
r.Use(httplog.RequestLogger(httpLogger))
|
||||
|
||||
// initialize and validate the sources from configs
|
||||
sourcesMap := make(map[string]sources.Source)
|
||||
for name, sc := range cfg.SourceConfigs {
|
||||
s, err := func() (sources.Source, error) {
|
||||
childCtx, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/source/init",
|
||||
trace.WithAttributes(attribute.String("source_kind", sc.SourceConfigKind())),
|
||||
trace.WithAttributes(attribute.String("source_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
s, err := sc.Initialize(childCtx, instrumentation.Tracer)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize source %q: %w", name, err)
|
||||
}
|
||||
return s, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sourcesMap[name] = s
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, err := InitializeConfigs(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize configs: %w", err)
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d sources.", len(sourcesMap)))
|
||||
|
||||
// initialize and validate the auth services from configs
|
||||
authServicesMap := make(map[string]auth.AuthService)
|
||||
for name, sc := range cfg.AuthServiceConfigs {
|
||||
a, err := func() (auth.AuthService, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/auth/init",
|
||||
trace.WithAttributes(attribute.String("auth_kind", sc.AuthServiceConfigKind())),
|
||||
trace.WithAttributes(attribute.String("auth_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
a, err := sc.Initialize()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize auth service %q: %w", name, err)
|
||||
}
|
||||
return a, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
authServicesMap[name] = a
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d authServices.", len(authServicesMap)))
|
||||
|
||||
// initialize and validate the tools from configs
|
||||
toolsMap := make(map[string]tools.Tool)
|
||||
for name, tc := range cfg.ToolConfigs {
|
||||
t, err := func() (tools.Tool, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/tool/init",
|
||||
trace.WithAttributes(attribute.String("tool_kind", tc.ToolConfigKind())),
|
||||
trace.WithAttributes(attribute.String("tool_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
t, err := tc.Initialize(sourcesMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize tool %q: %w", name, err)
|
||||
}
|
||||
return t, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
toolsMap[name] = t
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d tools.", len(toolsMap)))
|
||||
|
||||
// create a default toolset that contains all tools
|
||||
allToolNames := make([]string, 0, len(toolsMap))
|
||||
for name := range toolsMap {
|
||||
allToolNames = append(allToolNames, name)
|
||||
}
|
||||
if cfg.ToolsetConfigs == nil {
|
||||
cfg.ToolsetConfigs = make(ToolsetConfigs)
|
||||
}
|
||||
cfg.ToolsetConfigs[""] = tools.ToolsetConfig{Name: "", ToolNames: allToolNames}
|
||||
|
||||
// initialize and validate the toolsets from configs
|
||||
toolsetsMap := make(map[string]tools.Toolset)
|
||||
for name, tc := range cfg.ToolsetConfigs {
|
||||
t, err := func() (tools.Toolset, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/toolset/init",
|
||||
trace.WithAttributes(attribute.String("toolset_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
t, err := tc.Initialize(cfg.Version, toolsMap)
|
||||
if err != nil {
|
||||
return tools.Toolset{}, fmt.Errorf("unable to initialize toolset %q: %w", name, err)
|
||||
}
|
||||
return t, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
toolsetsMap[name] = t
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d toolsets.", len(toolsetsMap)))
|
||||
|
||||
addr := net.JoinHostPort(cfg.Address, strconv.Itoa(cfg.Port))
|
||||
srv := &http.Server{Addr: addr, Handler: r}
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := NewResourceManager(sourcesMap, authServicesMap, toolsMap, toolsetsMap)
|
||||
|
||||
s := &Server{
|
||||
version: cfg.Version,
|
||||
srv: srv,
|
||||
@@ -214,11 +317,7 @@ func NewServer(ctx context.Context, cfg ServerConfig, l log.Logger) (*Server, er
|
||||
logger: l,
|
||||
instrumentation: instrumentation,
|
||||
sseManager: sseManager,
|
||||
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
ResourceMgr: resourceManager,
|
||||
}
|
||||
// control plane
|
||||
apiR, err := apiRouter(s)
|
||||
@@ -231,6 +330,13 @@ func NewServer(ctx context.Context, cfg ServerConfig, l log.Logger) (*Server, er
|
||||
return nil, err
|
||||
}
|
||||
r.Mount("/mcp", mcpR)
|
||||
if cfg.UI {
|
||||
webR, err := webRouter()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.Mount("/ui", webR)
|
||||
}
|
||||
// default endpoint for validating server is running
|
||||
r.Get("/", func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte("🧰 Hello, World! 🧰"))
|
||||
|
||||
@@ -23,9 +23,16 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
)
|
||||
|
||||
func TestServe(t *testing.T) {
|
||||
@@ -54,8 +61,16 @@ func TestServe(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
ctx = util.WithLogger(ctx, testLogger)
|
||||
|
||||
s, err := server.NewServer(ctx, cfg, testLogger)
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(cfg.Version)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
|
||||
ctx = util.WithInstrumentation(ctx, instrumentation)
|
||||
|
||||
s, err := server.NewServer(ctx, cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to initialize server: %v", err)
|
||||
}
|
||||
@@ -93,3 +108,67 @@ func TestServe(t *testing.T) {
|
||||
t.Fatalf("version missing from output: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateServer(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("error setting up logger: %s", err)
|
||||
}
|
||||
|
||||
addr, port := "127.0.0.1", 5000
|
||||
cfg := server.ServerConfig{
|
||||
Version: "0.0.0",
|
||||
Address: addr,
|
||||
Port: port,
|
||||
}
|
||||
|
||||
instrumentation, err := telemetry.CreateTelemetryInstrumentation(cfg.Version)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
|
||||
ctx = util.WithInstrumentation(ctx, instrumentation)
|
||||
|
||||
s, err := server.NewServer(ctx, cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("error setting up server: %s", err)
|
||||
}
|
||||
|
||||
newSources := map[string]sources.Source{
|
||||
"example-source": &alloydbpg.Source{
|
||||
Name: "example-alloydb-source",
|
||||
Kind: "alloydb-postgres",
|
||||
},
|
||||
}
|
||||
newAuth := map[string]auth.AuthService{"example-auth": nil}
|
||||
newTools := map[string]tools.Tool{"example-tool": nil}
|
||||
newToolsets := map[string]tools.Toolset{
|
||||
"example-toolset": {
|
||||
Name: "example-toolset", Tools: []*tools.Tool{},
|
||||
},
|
||||
}
|
||||
s.ResourceMgr.SetResources(newSources, newAuth, newTools, newToolsets)
|
||||
if err != nil {
|
||||
t.Errorf("error updating server: %s", err)
|
||||
}
|
||||
|
||||
gotSource, _ := s.ResourceMgr.GetSource("example-source")
|
||||
if diff := cmp.Diff(gotSource, newSources["example-source"]); diff != "" {
|
||||
t.Errorf("error updating server, sources (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
gotAuthService, _ := s.ResourceMgr.GetAuthService("example-auth")
|
||||
if diff := cmp.Diff(gotAuthService, newAuth["example-auth"]); diff != "" {
|
||||
t.Errorf("error updating server, authServices (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
gotTool, _ := s.ResourceMgr.GetTool("example-tool")
|
||||
if diff := cmp.Diff(gotTool, newTools["example-tool"]); diff != "" {
|
||||
t.Errorf("error updating server, tools (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
gotToolset, _ := s.ResourceMgr.GetToolset("example-toolset")
|
||||
if diff := cmp.Diff(gotToolset, newToolsets["example-toolset"]); diff != "" {
|
||||
t.Errorf("error updating server, toolset (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
BIN
internal/server/static/assets/mcptoolboxlogo.png
Normal file
BIN
internal/server/static/assets/mcptoolboxlogo.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 57 KiB |
231
internal/server/static/css/style.css
Normal file
231
internal/server/static/css/style.css
Normal file
@@ -0,0 +1,231 @@
|
||||
body {
|
||||
display: flex;
|
||||
height: 100vh;
|
||||
margin: 0;
|
||||
font-family: 'Trebuchet MS';
|
||||
background-color: #f8f9fa;
|
||||
}
|
||||
|
||||
.left-nav {
|
||||
flex: 0 0 200px;
|
||||
background-color: #fff;
|
||||
box-shadow: 4px 0px 12px rgba(0, 0, 0, 0.15);
|
||||
z-index: 10;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 15px;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.second-nav {
|
||||
flex: 0 0 250px;
|
||||
background-color: #fff;
|
||||
box-shadow: 4px 0px 12px rgba(0, 0, 0, 0.15);
|
||||
z-index: 10;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 15px;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.nav-logo {
|
||||
width: 100%;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
.nav-logo img {
|
||||
max-width: 100%;
|
||||
height: auto;
|
||||
display: block;
|
||||
}
|
||||
|
||||
.left-nav ul {
|
||||
font-family: 'Verdana';
|
||||
list-style: none;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
.left-nav ul li {
|
||||
margin-bottom: 5px;
|
||||
}
|
||||
|
||||
.left-nav ul li a {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
padding: 12px;
|
||||
text-decoration: none;
|
||||
color: #333;
|
||||
border-radius: 0;
|
||||
}
|
||||
|
||||
.left-nav ul li a:hover {
|
||||
background-color: #e9e9e9;
|
||||
border-radius: 35px;
|
||||
}
|
||||
|
||||
.left-nav ul li a.active {
|
||||
background-color: #d0d0d0;
|
||||
font-weight: bold;
|
||||
border-radius: 35px;
|
||||
}
|
||||
|
||||
.main-content-area {
|
||||
flex: 1;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
|
||||
.top-bar {
|
||||
background-color: #fff;
|
||||
padding: 30px 30px;
|
||||
display: flex;
|
||||
justify-content: flex-end;
|
||||
align-items: center;
|
||||
border-bottom: 1px solid #eee;
|
||||
}
|
||||
|
||||
.top-bar span {
|
||||
font-weight: bold;
|
||||
font-size: 1.2em;
|
||||
font-family: 'Trebuchet MS';
|
||||
color: #333;
|
||||
}
|
||||
|
||||
.content {
|
||||
padding: 20px;
|
||||
flex-grow: 1;
|
||||
overflow-y: auto;
|
||||
}
|
||||
|
||||
.tool-button {
|
||||
/* --- Mimicking .left-nav ul li a styles --- */
|
||||
display: flex;
|
||||
align-items: center;
|
||||
padding: 12px;
|
||||
text-decoration: none;
|
||||
color: #333;
|
||||
background-color: transparent; /* Reset default button background */
|
||||
border: none; /* Reset default button border */
|
||||
border-radius: 0; /* Start with sharp corners */
|
||||
width: 100%;
|
||||
text-align: left;
|
||||
cursor: pointer;
|
||||
font-family: inherit; /* Inherit 'Verdana' from ul */
|
||||
font-size: inherit; /* Inherit font size */
|
||||
|
||||
/* Transition to match a tags */
|
||||
transition: background-color 0.1s ease-in-out, border-radius 0.1s ease-in-out;
|
||||
}
|
||||
|
||||
.tool-button:hover {
|
||||
/* Mimic .left-nav ul li a:hover */
|
||||
background-color: #e9e9e9;
|
||||
border-radius: 35px;
|
||||
}
|
||||
|
||||
.tool-button:focus {
|
||||
outline: none;
|
||||
/* Optional: You might want a subtle focus indicator, e.g.
|
||||
box-shadow: 0 0 0 2px rgba(208, 208, 208, 0.5); */
|
||||
}
|
||||
|
||||
.tool-button.active {
|
||||
/* Mimic .left-nav ul li a.active */
|
||||
background-color: #d0d0d0;
|
||||
font-weight: bold;
|
||||
border-radius: 35px;
|
||||
}
|
||||
|
||||
.tool-button.active:hover {
|
||||
background-color: #d0d0d0; /* Keep active color when hovered */
|
||||
}
|
||||
#secondary-panel-content ul {
|
||||
list-style: none; /* This is the main property to remove bullet points */
|
||||
padding: 0; /* Reset default browser padding on the ul */
|
||||
margin: 0; /* Reset default browser margin on the ul */
|
||||
width: 100%; /* Ensure it takes full width if needed */
|
||||
}
|
||||
|
||||
.tool-details-grid {
|
||||
display: grid;
|
||||
grid-template-columns: 2fr 1fr; /* 2/3 for info, 1/3 for params */
|
||||
gap: 20px;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
.tool-info {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 15px;
|
||||
}
|
||||
|
||||
.tool-params {
|
||||
background-color: #f9f9f9;
|
||||
padding: 15px;
|
||||
border-radius: 4px;
|
||||
border: 1px solid #ddd;
|
||||
}
|
||||
|
||||
.tool-box {
|
||||
background-color: #fff;
|
||||
padding: 15px;
|
||||
border-radius: 4px;
|
||||
border: 1px solid #eee;
|
||||
box-shadow: 0 1px 3px rgba(0,0,0,0.05);
|
||||
}
|
||||
|
||||
.tool-box h5 {
|
||||
margin-top: 0;
|
||||
margin-bottom: 8px;
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
.param-item {
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.param-item label {
|
||||
display: block;
|
||||
margin-bottom: 4px;
|
||||
font-weight: 500;
|
||||
}
|
||||
|
||||
.param-item input[type="text"],
|
||||
.param-item input[type="number"],
|
||||
.param-item select {
|
||||
width: calc(100% - 12px);
|
||||
padding: 6px;
|
||||
border: 1px solid #ccc;
|
||||
border-radius: 4px;
|
||||
}
|
||||
|
||||
.run-tool-btn {
|
||||
padding: 10px 15px;
|
||||
background-color: #4285f4; /* Google Blue */
|
||||
color: white;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
cursor: pointer;
|
||||
font-size: 1em;
|
||||
margin-top: 10px;
|
||||
}
|
||||
|
||||
.run-tool-btn:hover {
|
||||
background-color: #357ae8;
|
||||
}
|
||||
|
||||
.tool-response {
|
||||
width: 100%;
|
||||
margin-top: 20px;
|
||||
}
|
||||
|
||||
.tool-response textarea {
|
||||
width: calc(100% - 24px);
|
||||
min-height: 150px;
|
||||
padding: 12px;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 4px;
|
||||
font-family: monospace;
|
||||
}
|
||||
30
internal/server/static/index.html
Normal file
30
internal/server/static/index.html
Normal file
@@ -0,0 +1,30 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Toolbox UI</title>
|
||||
<link rel="stylesheet" href="/ui/css/style.css">
|
||||
</head>
|
||||
<body>
|
||||
<nav class="left-nav">
|
||||
<div class="nav-logo">
|
||||
<img src="/ui/assets/mcptoolboxlogo.png" alt="App Logo">
|
||||
</div>
|
||||
<ul>
|
||||
<li><a href="/ui/sources">Sources</a></li>
|
||||
<li><a href="/ui/authservices">Auth Services</a></li>
|
||||
<li><a href="/ui/tools">Tools</a></li>
|
||||
<li><a href="/ui/toolsets">Toolsets</a></li>
|
||||
</ul>
|
||||
</nav>
|
||||
|
||||
<div class="main-content-area", align="center">
|
||||
<div class="top-bar">
|
||||
</div>
|
||||
<main class="content">
|
||||
<h1>Homepage</h1>
|
||||
</main>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
255
internal/server/static/js/toolDisplay.js
Normal file
255
internal/server/static/js/toolDisplay.js
Normal file
@@ -0,0 +1,255 @@
|
||||
// helper function to create form inputs for parameters
|
||||
function createParamInput(param, toolId) {
|
||||
const paramItem = document.createElement('div');
|
||||
paramItem.className = 'param-item';
|
||||
|
||||
const label = document.createElement('label');
|
||||
const inputId = `param-${toolId}-${param.name}`;
|
||||
label.setAttribute('for', inputId);
|
||||
label.textContent = param.label;
|
||||
paramItem.appendChild(label);
|
||||
|
||||
let inputElement;
|
||||
if (param.type === 'select') {
|
||||
inputElement = document.createElement('select');
|
||||
param.options.forEach(optionValue => {
|
||||
const option = document.createElement('option');
|
||||
option.value = optionValue;
|
||||
option.textContent = optionValue;
|
||||
if (optionValue === param.defaultValue) {
|
||||
option.selected = true;
|
||||
}
|
||||
inputElement.appendChild(option);
|
||||
});
|
||||
} else if (param.type === 'textarea') { // Handle textarea for arrays
|
||||
inputElement = document.createElement('textarea');
|
||||
inputElement.rows = 3;
|
||||
inputElement.value = param.defaultValue || '';
|
||||
if (param.valueType && param.valueType.startsWith('array')) {
|
||||
inputElement.placeholder = 'E.g., ["item1", "item2"] or [1, 2, 3]';
|
||||
}
|
||||
} else { // text, number, etc.
|
||||
inputElement = document.createElement('input');
|
||||
inputElement.type = param.type;
|
||||
inputElement.value = param.defaultValue || '';
|
||||
}
|
||||
// Common properties
|
||||
inputElement.id = inputId;
|
||||
inputElement.name = param.name;
|
||||
paramItem.appendChild(inputElement);
|
||||
return paramItem;
|
||||
}
|
||||
|
||||
function displayResults(results, responseArea, prettify) {
|
||||
if (results === null || results === undefined) {
|
||||
// responseArea.value = ''; // Keep placeholder or old error message
|
||||
return;
|
||||
}
|
||||
try {
|
||||
if (prettify) {
|
||||
responseArea.value = JSON.stringify(JSON.parse(results.result), null, 2);
|
||||
} else {
|
||||
responseArea.value = JSON.stringify(JSON.parse(results.result));
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error stringifying results:", error);
|
||||
responseArea.value = "Error displaying results.";
|
||||
}
|
||||
}
|
||||
|
||||
// function to run the tool (calls API version of endpoint)
|
||||
async function handleRunTool(toolId, form, responseArea, parameters, prettifyCheckbox, updateLastResults) {
|
||||
responseArea.value = 'Running tool...';
|
||||
updateLastResults(null); // Clear last results before new run
|
||||
const formData = new FormData(form);
|
||||
const typedParams = {};
|
||||
|
||||
for (const param of parameters) {
|
||||
const rawValue = formData.get(param.name);
|
||||
|
||||
if (rawValue === null || rawValue === undefined || rawValue === '') {
|
||||
if (param.required) {
|
||||
console.warn(`Required parameter ${param.name} is missing.`);
|
||||
}
|
||||
typedParams[param.name] = null;
|
||||
continue;
|
||||
}
|
||||
|
||||
const valueType = param.valueType;
|
||||
|
||||
try {
|
||||
if (valueType && valueType.startsWith('array<')) {
|
||||
const elementType = valueType.substring(6, valueType.length - 1);
|
||||
let parsedArray;
|
||||
try {
|
||||
parsedArray = JSON.parse(rawValue);
|
||||
} catch (e) {
|
||||
throw new Error(`Invalid JSON format for ${param.name}. ${e.message}`);
|
||||
}
|
||||
|
||||
if (!Array.isArray(parsedArray)) {
|
||||
throw new Error(`Input for ${param.name} must be a JSON array (e.g., ["a", "b"]).`);
|
||||
}
|
||||
|
||||
if (elementType === 'number') {
|
||||
typedParams[param.name] = parsedArray.map((item, index) => {
|
||||
const num = Number(item);
|
||||
if (isNaN(num)) {
|
||||
throw new Error(`Invalid number "${item}" found in array for ${param.name} at index ${index}.`);
|
||||
}
|
||||
return num;
|
||||
});
|
||||
} else if (elementType === 'boolean') {
|
||||
typedParams[param.name] = parsedArray.map(item => item === true || String(item).toLowerCase() === 'true');
|
||||
} else { // string or other types
|
||||
typedParams[param.name] = parsedArray;
|
||||
}
|
||||
} else {
|
||||
switch (valueType) {
|
||||
case 'number':
|
||||
const num = Number(rawValue);
|
||||
if (isNaN(num)) {
|
||||
throw new Error(`Invalid number input for ${param.name}: ${rawValue}`);
|
||||
}
|
||||
typedParams[param.name] = num;
|
||||
break;
|
||||
case 'boolean':
|
||||
typedParams[param.name] = rawValue === 'true';
|
||||
break;
|
||||
case 'string':
|
||||
default:
|
||||
typedParams[param.name] = rawValue;
|
||||
break;
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error processing parameter:', param.name, error);
|
||||
responseArea.value = `Error for ${param.name}: ${error.message}`;
|
||||
return; // Stop processing
|
||||
}
|
||||
}
|
||||
|
||||
console.log('Running tool:', toolId, 'with typed params:', typedParams);
|
||||
try {
|
||||
const response = await fetch(`/api/tool/${toolId}/invoke`, {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json'},
|
||||
body: JSON.stringify(typedParams)
|
||||
});
|
||||
if (!response.ok) {
|
||||
const errorBody = await response.text();
|
||||
throw new Error(`HTTP error ${response.status}: ${errorBody}`);
|
||||
}
|
||||
const results = await response.json();
|
||||
updateLastResults(results); // Update the stored results
|
||||
displayResults(results, responseArea, prettifyCheckbox.checked); // Display formatted results
|
||||
} catch (error) {
|
||||
console.error('Error running tool:', error);
|
||||
responseArea.value = `Error: ${error.message}`;
|
||||
updateLastResults(null); // Clear results on error
|
||||
}
|
||||
}
|
||||
|
||||
// renders the tool display area
|
||||
export function renderToolInterface(tool, containerElement) {
|
||||
containerElement.innerHTML = '';
|
||||
const toolId = tool.id;
|
||||
|
||||
let lastResults = null; // Store the most recent successful result object
|
||||
|
||||
// Function to update lastResults, closure to keep it private to this scope
|
||||
const updateLastResults = (newResults) => {
|
||||
lastResults = newResults;
|
||||
};
|
||||
|
||||
const gridContainer = document.createElement('div');
|
||||
gridContainer.className = 'tool-details-grid';
|
||||
|
||||
const toolInfoContainer = document.createElement('div');
|
||||
toolInfoContainer.className = 'tool-info';
|
||||
|
||||
const nameBox = document.createElement('div');
|
||||
nameBox.className = 'tool-box tool-name';
|
||||
nameBox.innerHTML = `<h5>Name:</h5><p>${tool.name}</p>`;
|
||||
toolInfoContainer.appendChild(nameBox);
|
||||
|
||||
const descBox = document.createElement('div');
|
||||
descBox.className = 'tool-box tool-description';
|
||||
descBox.innerHTML = `<h5>Description:</h5><p>${tool.description}</p>`;
|
||||
toolInfoContainer.appendChild(descBox);
|
||||
|
||||
gridContainer.appendChild(toolInfoContainer);
|
||||
|
||||
const paramsContainer = document.createElement('div');
|
||||
paramsContainer.className = 'tool-params tool-box';
|
||||
paramsContainer.innerHTML = '<h5>Parameters:</h5>';
|
||||
const form = document.createElement('form');
|
||||
form.id = `tool-params-form-${toolId}`;
|
||||
|
||||
tool.parameters.forEach(param => {
|
||||
form.appendChild(createParamInput(param, toolId));
|
||||
});
|
||||
paramsContainer.appendChild(form);
|
||||
|
||||
const runButton = document.createElement('button');
|
||||
runButton.className = 'run-tool-btn';
|
||||
runButton.textContent = 'Run Tool';
|
||||
paramsContainer.appendChild(runButton);
|
||||
|
||||
gridContainer.appendChild(paramsContainer);
|
||||
containerElement.appendChild(gridContainer);
|
||||
|
||||
// Response Area
|
||||
const responseContainer = document.createElement('div');
|
||||
responseContainer.className = 'tool-response tool-box';
|
||||
|
||||
const responseHeader = document.createElement('h5');
|
||||
responseHeader.textContent = 'Response:';
|
||||
responseContainer.appendChild(responseHeader);
|
||||
|
||||
// Prettify Checkbox
|
||||
const prettifyId = `prettify-${toolId}`;
|
||||
const prettifyLabel = document.createElement('label');
|
||||
prettifyLabel.setAttribute('for', prettifyId);
|
||||
prettifyLabel.textContent = 'Prettify JSON';
|
||||
prettifyLabel.style.display = 'inline-block';
|
||||
prettifyLabel.style.marginLeft = '10px';
|
||||
prettifyLabel.style.verticalAlign = 'middle';
|
||||
prettifyLabel.style.cursor = 'pointer';
|
||||
|
||||
const prettifyCheckbox = document.createElement('input');
|
||||
prettifyCheckbox.type = 'checkbox';
|
||||
prettifyCheckbox.id = prettifyId;
|
||||
prettifyCheckbox.checked = true; // Default to pretty
|
||||
prettifyCheckbox.style.verticalAlign = 'middle';
|
||||
prettifyCheckbox.style.cursor = 'pointer';
|
||||
|
||||
const prettifyDiv = document.createElement('div');
|
||||
prettifyDiv.style.marginBottom = '5px';
|
||||
prettifyDiv.appendChild(prettifyCheckbox);
|
||||
prettifyDiv.appendChild(prettifyLabel);
|
||||
responseContainer.appendChild(prettifyDiv);
|
||||
|
||||
const responseAreaId = `tool-response-area-${toolId}`;
|
||||
const responseArea = document.createElement('textarea');
|
||||
responseArea.id = responseAreaId;
|
||||
responseArea.readOnly = true;
|
||||
responseArea.placeholder = 'Results will appear here...';
|
||||
responseArea.style.width = 'calc(100% - 12px)';
|
||||
responseArea.rows = 10;
|
||||
responseContainer.appendChild(responseArea);
|
||||
|
||||
containerElement.appendChild(responseContainer);
|
||||
|
||||
// Event Listeners
|
||||
prettifyCheckbox.addEventListener('change', () => {
|
||||
if (lastResults) {
|
||||
displayResults(lastResults, responseArea, prettifyCheckbox.checked);
|
||||
}
|
||||
});
|
||||
|
||||
runButton.addEventListener('click', (event) => {
|
||||
event.preventDefault();
|
||||
handleRunTool(toolId, form, responseArea, tool.parameters, prettifyCheckbox, updateLastResults);
|
||||
});
|
||||
}
|
||||
156
internal/server/static/js/tools.js
Normal file
156
internal/server/static/js/tools.js
Normal file
@@ -0,0 +1,156 @@
|
||||
import { renderToolInterface } from "./toolDisplay.js";
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
const toolDisplayArea = document.getElementById('tool-display-area');
|
||||
const secondaryPanelContent = document.getElementById('secondary-panel-content');
|
||||
|
||||
if (!secondaryPanelContent || !toolDisplayArea) {
|
||||
console.error('Required DOM elements not found.');
|
||||
return;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetches the list of tools from the API and renders them in the secondary panel.
|
||||
*/
|
||||
async function loadTools() {
|
||||
secondaryPanelContent.innerHTML = '<p>Fetching tools...</p>';
|
||||
try {
|
||||
// This endpoint should list tools, the structure you provided seems to be for a single tool
|
||||
const response = await fetch('/api/toolset');
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
const apiResponse = await response.json();
|
||||
renderToolList(apiResponse);
|
||||
} catch (error) {
|
||||
console.error('Failed to load tools:', error);
|
||||
secondaryPanelContent.innerHTML = '<p class="error">Failed to load tools. Please try again later.</p>';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Renders the list of tools in the secondary navigation panel.
|
||||
* @param {object} apiResponse - The parsed JSON response from the API.
|
||||
*/
|
||||
function renderToolList(apiResponse) {
|
||||
secondaryPanelContent.innerHTML = '';
|
||||
|
||||
if (!apiResponse || typeof apiResponse.tools !== 'object' || apiResponse.tools === null) {
|
||||
console.error('Error: Expected an object with a "tools" property, but received:', apiResponse);
|
||||
secondaryPanelContent.textContent = 'Error: Invalid response format from toolset API.';
|
||||
return;
|
||||
}
|
||||
|
||||
const toolsObject = apiResponse.tools;
|
||||
const toolNames = Object.keys(toolsObject);
|
||||
|
||||
if (toolNames.length === 0) {
|
||||
secondaryPanelContent.textContent = 'No tools found.';
|
||||
return;
|
||||
}
|
||||
|
||||
const ul = document.createElement('ul');
|
||||
toolNames.forEach(toolName => {
|
||||
const li = document.createElement('li');
|
||||
const button = document.createElement('button');
|
||||
button.textContent = toolName;
|
||||
button.dataset.toolname = toolName;
|
||||
button.classList.add('tool-button');
|
||||
button.addEventListener('click', handleToolClick);
|
||||
li.appendChild(button);
|
||||
ul.appendChild(li);
|
||||
});
|
||||
secondaryPanelContent.appendChild(ul);
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles the click event on a tool button in the secondary panel.
|
||||
* @param {MouseEvent} event - The click event.
|
||||
*/
|
||||
function handleToolClick(event) {
|
||||
const toolName = event.target.dataset.toolname;
|
||||
if (toolName) {
|
||||
const currentActive = secondaryPanelContent.querySelector('.tool-button.active');
|
||||
if (currentActive) {
|
||||
currentActive.classList.remove('active');
|
||||
}
|
||||
event.target.classList.add('active');
|
||||
fetchToolDetails(toolName);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetches details for a specific tool from the API and renders the UI.
|
||||
* @param {string} toolName - The name of the tool.
|
||||
*/
|
||||
async function fetchToolDetails(toolName) {
|
||||
toolDisplayArea.innerHTML = '<p>Loading tool details...</p>';
|
||||
|
||||
try {
|
||||
const response = await fetch(`/api/tool/${encodeURIComponent(toolName)}`);
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
const apiResponse = await response.json();
|
||||
|
||||
if (!apiResponse.tools || !apiResponse.tools[toolName]) {
|
||||
throw new Error(`Tool "${toolName}" data not found in API response.`);
|
||||
}
|
||||
const toolObject = apiResponse.tools[toolName];
|
||||
|
||||
const toolInterfaceData = {
|
||||
id: toolName,
|
||||
name: toolName,
|
||||
description: toolObject.description || "No description provided.",
|
||||
parameters: (toolObject.parameters || []).map(param => {
|
||||
let inputType = 'text'; // Default HTML input type
|
||||
let options;
|
||||
const apiType = param.type ? param.type.toLowerCase() : 'string';
|
||||
let valueType = 'string'; // Data type for API payload
|
||||
let label = param.description || param.name;
|
||||
|
||||
if (apiType === 'integer' || apiType === 'number') {
|
||||
inputType = 'number';
|
||||
valueType = 'number';
|
||||
} else if (apiType === 'boolean') {
|
||||
inputType = 'select';
|
||||
options = ['true', 'false'];
|
||||
valueType = 'boolean';
|
||||
} else if (apiType === 'array') {
|
||||
inputType = 'textarea'; // Use textarea for array inputs
|
||||
const itemType = param.items && param.items.type ? param.items.type.toLowerCase() : 'string';
|
||||
valueType = `array<${itemType}>`;
|
||||
label += ' (JSON Array string)'; // Hint to the user
|
||||
} else if (param.enum && Array.isArray(param.enum)) {
|
||||
inputType = 'select';
|
||||
options = param.enum;
|
||||
valueType = 'string';
|
||||
}
|
||||
console.log(param.name, inputType, label, apiType, valueType)
|
||||
|
||||
return {
|
||||
name: param.name,
|
||||
type: inputType, // For HTML input element type
|
||||
valueType: valueType, // For API request payload type
|
||||
label: label,
|
||||
required: param.required || false,
|
||||
options: options,
|
||||
defaultValue: param.default,
|
||||
};
|
||||
})
|
||||
};
|
||||
|
||||
console.log("Transformed toolInterfaceData:", toolInterfaceData);
|
||||
|
||||
renderToolInterface(toolInterfaceData, toolDisplayArea);
|
||||
|
||||
} catch (error) {
|
||||
console.error(`Failed to load details for tool "${toolName}":`, error);
|
||||
toolDisplayArea.innerHTML = `<p class="error">Failed to load details for ${toolName}. ${error.message}</p>`;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Initial load of tools list
|
||||
loadTools();
|
||||
});
|
||||
40
internal/server/static/tools.html
Normal file
40
internal/server/static/tools.html
Normal file
@@ -0,0 +1,40 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Page Three</title>
|
||||
<link rel="stylesheet" href="css/style.css">
|
||||
</head>
|
||||
<body>
|
||||
<nav class="left-nav">
|
||||
<div class="nav-logo">
|
||||
<img src="/ui/assets/mcptoolboxlogo.png" alt="App Logo">
|
||||
</div>
|
||||
<ul>
|
||||
<li><a href="/ui/sources">Sources</a></li>
|
||||
<li><a href="/ui/authservices">Auth Services</a></li>
|
||||
<li><a href="/ui/tools" class="active">Tools</a></li>
|
||||
<li><a href="/ui/toolsets">Toolsets</a></li>
|
||||
</ul>
|
||||
</nav>
|
||||
|
||||
<aside class="second-nav">
|
||||
<h4>My Tools</h4>
|
||||
<div id="secondary-panel-content">
|
||||
<p>Fetching tools...</p>
|
||||
</div>
|
||||
</aside>
|
||||
|
||||
<div class="main-content-area">
|
||||
<div class="top-bar">
|
||||
<span>Toolbox UI Inspector</span>
|
||||
</div>
|
||||
<main class="content" id="tool-display-area">
|
||||
<h1>Welcome to the Main Page</h1>
|
||||
<p>This is the main content area. Click a tab on the left to navigate.</p>
|
||||
</main>
|
||||
</div>
|
||||
<script type="module" src="js/tools.js"></script>
|
||||
</body>
|
||||
</html>
|
||||
53
internal/server/web.go
Normal file
53
internal/server/web.go
Normal file
@@ -0,0 +1,53 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"embed"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/go-chi/chi/v5/middleware"
|
||||
)
|
||||
|
||||
//go:embed all:static
|
||||
var staticContent embed.FS
|
||||
|
||||
// webRouter creates a router that represents the routes under /ui
|
||||
func webRouter() (chi.Router, error) {
|
||||
r := chi.NewRouter()
|
||||
r.Use(middleware.StripSlashes)
|
||||
|
||||
// direct routes for html pages to provide clean URLs
|
||||
r.Get("/", func(w http.ResponseWriter, r *http.Request) { serveHTML(w, r, "static/index.html") })
|
||||
r.Get("/tools", func(w http.ResponseWriter, r *http.Request) { serveHTML(w, r, "static/tools.html") })
|
||||
|
||||
// handler for all other static files/assets
|
||||
staticFS, _ := fs.Sub(staticContent, "static")
|
||||
r.Handle("/*", http.StripPrefix("/ui", http.FileServer(http.FS(staticFS))))
|
||||
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func serveHTML(w http.ResponseWriter, r *http.Request, filepath string) {
|
||||
file, err := staticContent.Open(filepath)
|
||||
if err != nil {
|
||||
http.Error(w, "File not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
fileBytes, err := io.ReadAll(file)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Error reading file: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
fileInfo, err := file.Stat()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
http.ServeContent(w, r, fileInfo.Name(), fileInfo.ModTime(), bytes.NewReader(fileBytes))
|
||||
}
|
||||
77
internal/server/web_test.go
Normal file
77
internal/server/web_test.go
Normal file
@@ -0,0 +1,77 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/go-goquery/goquery"
|
||||
)
|
||||
|
||||
func TestWebEndpoint(t *testing.T) {
|
||||
router, err := webRouter()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create webRouter: %v", err)
|
||||
}
|
||||
|
||||
ts := httptest.NewServer(router)
|
||||
defer ts.Close()
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
wantStatus int
|
||||
wantContentType string
|
||||
wantPageTitle string
|
||||
}{
|
||||
{
|
||||
name: "web index page GET",
|
||||
method: http.MethodGet,
|
||||
path: "/",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContentType: "text/html; charset=utf-8",
|
||||
wantPageTitle: "Toolbox UI",
|
||||
},
|
||||
}
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
url := ts.URL + tc.path
|
||||
req, err := http.NewRequest(tc.method, url, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create request: %v", err)
|
||||
}
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to send request: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to read response body: %v", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != tc.wantStatus {
|
||||
t.Errorf("Unexpected status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatus, string(body))
|
||||
}
|
||||
|
||||
if contentType := resp.Header.Get("Content-Type"); contentType != tc.wantContentType {
|
||||
t.Errorf("Unexpected Content-Type header: got %s, want %s", contentType, tc.wantContentType)
|
||||
}
|
||||
|
||||
doc, err := goquery.NewDocumentFromReader(strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to parse HTML: %v", err)
|
||||
}
|
||||
gotPageTitle := doc.Find("title").Text()
|
||||
|
||||
if gotPageTitle != tc.wantPageTitle {
|
||||
t.Errorf("Unexpected page title: got %q, want %q", gotPageTitle, tc.wantPageTitle)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package server
|
||||
package telemetry
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
@@ -15,9 +15,12 @@
|
||||
package testutils
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
@@ -42,3 +45,63 @@ func ContextWithNewLogger() (context.Context, error) {
|
||||
}
|
||||
return util.WithLogger(ctx, logger), nil
|
||||
}
|
||||
|
||||
// WaitForString waits until the server logs a single line that matches the provided regex.
|
||||
// returns the output of whatever the server sent so far.
|
||||
func WaitForString(ctx context.Context, re *regexp.Regexp, pr io.ReadCloser) (string, error) {
|
||||
in := bufio.NewReader(pr)
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// read lines in background, sending result of each read over a channel
|
||||
// this allows us to use in.ReadString without blocking
|
||||
type result struct {
|
||||
s string
|
||||
err error
|
||||
}
|
||||
output := make(chan result)
|
||||
go func() {
|
||||
defer close(output)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// if the context is canceled, the orig thread will send back the error
|
||||
// so we can just exit the goroutine here
|
||||
return
|
||||
default:
|
||||
// otherwise read a line from the output
|
||||
s, err := in.ReadString('\n')
|
||||
if err != nil {
|
||||
output <- result{err: err}
|
||||
return
|
||||
}
|
||||
output <- result{s: s}
|
||||
// if that last string matched, exit the goroutine
|
||||
if re.MatchString(s) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// collect the output until the ctx is canceled, an error was hit,
|
||||
// or match was found (which is indicated the channel is closed)
|
||||
var sb strings.Builder
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// if ctx is done, return that error
|
||||
return sb.String(), ctx.Err()
|
||||
case o, ok := <-output:
|
||||
if !ok {
|
||||
// match was found!
|
||||
return sb.String(), nil
|
||||
}
|
||||
if o.err != nil {
|
||||
// error was found!
|
||||
return sb.String(), o.err
|
||||
}
|
||||
sb.WriteString(o.s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,8 @@ import (
|
||||
)
|
||||
|
||||
const kind string = "bigquery-get-dataset-info"
|
||||
const projectKey string = "project"
|
||||
const datasetKey string = "dataset"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
@@ -78,8 +80,9 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
datasetParameter := tools.NewStringParameter("dataset", "The dataset to get metadata information.")
|
||||
parameters := tools.Parameters{datasetParameter}
|
||||
projectParameter := tools.NewStringParameterWithDefault(projectKey, s.BigQueryClient().Project(), "The Google Cloud project ID containing the dataset.")
|
||||
datasetParameter := tools.NewStringParameter(datasetKey, "The dataset to get metadata information.")
|
||||
parameters := tools.Parameters{projectParameter, datasetParameter}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
@@ -116,14 +119,18 @@ type Tool struct {
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
|
||||
|
||||
sliceParams := params.AsSlice()
|
||||
datasetId, ok := sliceParams[0].(string)
|
||||
mapParams := params.AsMap()
|
||||
projectId, ok := mapParams[projectKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.Dataset(datasetId)
|
||||
datasetId, ok := mapParams[datasetKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.DatasetInProject(projectId, datasetId)
|
||||
|
||||
metadata, err := dsHandle.Metadata(ctx)
|
||||
if err != nil {
|
||||
|
||||
@@ -26,6 +26,9 @@ import (
|
||||
)
|
||||
|
||||
const kind string = "bigquery-get-table-info"
|
||||
const projectKey string = "project"
|
||||
const datasetKey string = "dataset"
|
||||
const tableKey string = "table"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
@@ -78,9 +81,10 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
datasetParameter := tools.NewStringParameter("dataset", "The table's parent dataset.")
|
||||
tableParameter := tools.NewStringParameter("table", "The table to get metadata information.")
|
||||
parameters := tools.Parameters{datasetParameter, tableParameter}
|
||||
projectParameter := tools.NewStringParameterWithDefault(projectKey, s.BigQueryClient().Project(), "The Google Cloud project ID containing the dataset and table.")
|
||||
datasetParameter := tools.NewStringParameter(datasetKey, "The table's parent dataset.")
|
||||
tableParameter := tools.NewStringParameter(tableKey, "The table to get metadata information.")
|
||||
parameters := tools.Parameters{projectParameter, datasetParameter, tableParameter}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
@@ -117,22 +121,28 @@ type Tool struct {
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
|
||||
sliceParams := params.AsSlice()
|
||||
datasetId, ok := sliceParams[0].(string)
|
||||
mapParams := params.AsMap()
|
||||
projectId, ok := mapParams[projectKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
|
||||
}
|
||||
tableId, ok := sliceParams[1].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unable to get cast %s", sliceParams[1])
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.Dataset(datasetId)
|
||||
datasetId, ok := mapParams[datasetKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
|
||||
}
|
||||
|
||||
tableId, ok := mapParams[tableKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", tableKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.DatasetInProject(projectId, datasetId)
|
||||
tableHandle := dsHandle.Table(tableId)
|
||||
|
||||
metadata, err := tableHandle.Metadata(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get metadata for table %s.%s.%s: %w", t.Client.Project(), datasetId, tableId, err)
|
||||
return nil, fmt.Errorf("failed to get metadata for table %s.%s.%s: %w", projectId, datasetId, tableId, err)
|
||||
}
|
||||
|
||||
return []any{metadata}, nil
|
||||
|
||||
@@ -27,6 +27,7 @@ import (
|
||||
)
|
||||
|
||||
const kind string = "bigquery-list-dataset-ids"
|
||||
const projectKey string = "project"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
@@ -79,7 +80,9 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
parameters := tools.Parameters{}
|
||||
projectParameter := tools.NewStringParameterWithDefault(projectKey, s.BigQueryClient().Project(), "The Google Cloud project to list dataset ids.")
|
||||
|
||||
parameters := tools.Parameters{projectParameter}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
@@ -116,7 +119,13 @@ type Tool struct {
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
|
||||
mapParams := params.AsMap()
|
||||
projectId, ok := mapParams[projectKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
|
||||
}
|
||||
datasetIterator := t.Client.Datasets(ctx)
|
||||
datasetIterator.ProjectID = projectId
|
||||
|
||||
var datasetIds []any
|
||||
for {
|
||||
|
||||
@@ -27,6 +27,8 @@ import (
|
||||
)
|
||||
|
||||
const kind string = "bigquery-list-table-ids"
|
||||
const projectKey string = "project"
|
||||
const datasetKey string = "dataset"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
@@ -79,8 +81,9 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
datasetParameter := tools.NewStringParameter("dataset", "The dataset to list table ids.")
|
||||
parameters := tools.Parameters{datasetParameter}
|
||||
projectParameter := tools.NewStringParameterWithDefault(projectKey, s.BigQueryClient().Project(), "The Google Cloud project ID containing the dataset.")
|
||||
datasetParameter := tools.NewStringParameter(datasetKey, "The dataset to list table ids.")
|
||||
parameters := tools.Parameters{projectParameter, datasetParameter}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
@@ -117,14 +120,18 @@ type Tool struct {
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
|
||||
|
||||
sliceParams := params.AsSlice()
|
||||
datasetId, ok := sliceParams[0].(string)
|
||||
mapParams := params.AsMap()
|
||||
projectId, ok := mapParams[projectKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.Dataset(datasetId)
|
||||
datasetId, ok := mapParams[datasetKey].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
|
||||
}
|
||||
|
||||
dsHandle := t.Client.DatasetInProject(projectId, datasetId)
|
||||
|
||||
var tableIds []any
|
||||
tableIterator := dsHandle.Tables(ctx)
|
||||
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
"github.com/go-playground/validator/v10"
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
||||
)
|
||||
|
||||
// DecodeJSON decodes a given reader into an interface using the json decoder.
|
||||
@@ -98,10 +99,25 @@ func WithLogger(ctx context.Context, logger log.Logger) context.Context {
|
||||
return context.WithValue(ctx, loggerKey, logger)
|
||||
}
|
||||
|
||||
// LoggerFromContext retreives the logger or return an error
|
||||
// LoggerFromContext retrieves the logger or return an error
|
||||
func LoggerFromContext(ctx context.Context) (log.Logger, error) {
|
||||
if logger, ok := ctx.Value(loggerKey).(log.Logger); ok {
|
||||
return logger, nil
|
||||
}
|
||||
return nil, fmt.Errorf("unable to retrieve logger")
|
||||
}
|
||||
|
||||
const instrumentationKey contextKey = "instrumentation"
|
||||
|
||||
// WithInstrumentation adds an instrumentation into the context as a value
|
||||
func WithInstrumentation(ctx context.Context, instrumentation *telemetry.Instrumentation) context.Context {
|
||||
return context.WithValue(ctx, instrumentationKey, instrumentation)
|
||||
}
|
||||
|
||||
// InstrumentationFromContext retrieves the instrumentation or return an error
|
||||
func InstrumentationFromContext(ctx context.Context) (*telemetry.Instrumentation, error) {
|
||||
if instrumentation, ok := ctx.Value(instrumentationKey).(*telemetry.Instrumentation); ok {
|
||||
return instrumentation, nil
|
||||
}
|
||||
return nil, fmt.Errorf("unable to retrieve instrumentation")
|
||||
}
|
||||
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -90,7 +91,7 @@ func TestAlloyDBAINLToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
|
||||
"cloud.google.com/go/alloydbconn"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
@@ -134,17 +135,17 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, AlloyDBPostgresToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, AlloyDBPostgresToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddPgExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, AlloyDBPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -157,7 +158,7 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -166,8 +167,8 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetPostgresWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
|
||||
bigqueryapi "cloud.google.com/go/bigquery"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"golang.org/x/oauth2/google"
|
||||
"google.golang.org/api/googleapi"
|
||||
@@ -101,17 +102,17 @@ func TestBigQueryToolEndpoints(t *testing.T) {
|
||||
)
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := getBigQueryParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := getBigQueryParamToolInfo(tableNameParam)
|
||||
teardownTable1 := setupBigQueryTable(t, ctx, client, createStatement1, insertStatement1, datasetName, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := getBigQueryAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := getBigQueryAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := setupBigQueryTable(t, ctx, client, createStatement2, insertStatement2, datasetName, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, BigqueryToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, BigqueryToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = addBigQueryPrebuiltToolsConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := getBigQueryTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, BigqueryToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -124,7 +125,7 @@ func TestBigQueryToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -137,8 +138,8 @@ func TestBigQueryToolEndpoints(t *testing.T) {
|
||||
failInvocationWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: googleapi: Error 400: Syntax error: Unexpected identifier \"SELEC\" at [1:1]`
|
||||
datasetInfoWant := "\"Location\":\"US\",\"DefaultTableExpiration\":0,\"Labels\":null,\"Access\":"
|
||||
tableInfoWant := "[{\"Name\":\"\",\"Location\":\"US\",\"Description\":\"\",\"Schema\":[{\"Name\":\"id\""
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
templateParamTestConfig := tests.NewTemplateParameterTestConfig(
|
||||
tests.WithCreateColArray(`["id INT64", "name STRING", "age INT64"]`),
|
||||
@@ -153,18 +154,20 @@ func TestBigQueryToolEndpoints(t *testing.T) {
|
||||
}
|
||||
|
||||
// getBigQueryParamToolInfo returns statements and param for my-param-tool for bigquery kind
|
||||
func getBigQueryParamToolInfo(tableName string) (string, string, string, []bigqueryapi.QueryParameter) {
|
||||
func getBigQueryParamToolInfo(tableName string) (string, string, string, string, []bigqueryapi.QueryParameter) {
|
||||
createStatement := fmt.Sprintf(`
|
||||
CREATE TABLE IF NOT EXISTS %s (id INT64, name STRING);`, tableName)
|
||||
insertStatement := fmt.Sprintf(`
|
||||
INSERT INTO %s (id, name) VALUES (?, ?), (?, ?), (?, ?);`, tableName)
|
||||
INSERT INTO %s (id, name) VALUES (?, ?), (?, ?), (?, ?), (?, NULL);`, tableName)
|
||||
toolStatement := fmt.Sprintf(`SELECT * FROM %s WHERE id = ? OR name = ? ORDER BY id;`, tableName)
|
||||
toolStatement2 := fmt.Sprintf(`SELECT * FROM %s WHERE id = ? ORDER BY id;`, tableName)
|
||||
params := []bigqueryapi.QueryParameter{
|
||||
{Value: int64(1)}, {Value: "Alice"},
|
||||
{Value: int64(2)}, {Value: "Jane"},
|
||||
{Value: int64(3)}, {Value: "Sid"},
|
||||
{Value: int64(4)},
|
||||
}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// getBigQueryAuthToolInfo returns statements and param of my-auth-tool for bigquery kind
|
||||
@@ -499,6 +502,21 @@ func runBigQueryListDatasetToolInvokeTest(t *testing.T, datasetWant string) {
|
||||
isErr: false,
|
||||
want: datasetWant,
|
||||
},
|
||||
{
|
||||
name: "invoke my-list-dataset-ids-tool with project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-dataset-ids-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s\"}", BigqueryProject))),
|
||||
isErr: false,
|
||||
want: datasetWant,
|
||||
},
|
||||
{
|
||||
name: "invoke my-list-dataset-ids-tool with non-existent project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-dataset-ids-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s-%s\"}", BigqueryProject, uuid.NewString()))),
|
||||
isErr: true,
|
||||
},
|
||||
{
|
||||
name: "invoke my-auth-list-dataset-ids-tool",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-dataset-ids-tool/invoke",
|
||||
@@ -583,6 +601,21 @@ func runBigQueryGetDatasetInfoToolInvokeTest(t *testing.T, datasetName, datasetI
|
||||
want: datasetInfoWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-get-dataset-info-tool with correct project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-get-dataset-info-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s\", \"dataset\":\"%s\"}", BigqueryProject, datasetName))),
|
||||
want: datasetInfoWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-get-dataset-info-tool with non-existent project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-get-dataset-info-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s-%s\", \"dataset\":\"%s\"}", BigqueryProject, uuid.NewString(), datasetName))),
|
||||
isErr: true,
|
||||
},
|
||||
{
|
||||
name: "invoke my-auth-get-dataset-info-tool without body",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-get-dataset-info-tool/invoke",
|
||||
@@ -703,6 +736,21 @@ func runBigQueryListTableIdsToolInvokeTest(t *testing.T, datasetName, tablename_
|
||||
want: tablename_want,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-list-table-ids-tool with correct project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-table-ids-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s\", \"dataset\":\"%s\"}", BigqueryProject, datasetName))),
|
||||
want: tablename_want,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-list-table-ids-tool with non-existent project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-table-ids-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s-%s\", \"dataset\":\"%s\"}", BigqueryProject, uuid.NewString(), datasetName))),
|
||||
isErr: true,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-list-table-ids-tool with invalid auth token",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-list-table-ids-tool/invoke",
|
||||
@@ -808,6 +856,21 @@ func runBigQueryGetTableInfoToolInvokeTest(t *testing.T, datasetName, tableName,
|
||||
want: tableInfoWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-get-table-info-tool with correct project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-get-table-info-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s\", \"dataset\":\"%s\", \"table\":\"%s\"}", BigqueryProject, datasetName, tableName))),
|
||||
want: tableInfoWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-get-table-info-tool with non-existent project",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-get-table-info-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"project\":\"%s-%s\", \"dataset\":\"%s\", \"table\":\"%s\"}", BigqueryProject, uuid.NewString(), datasetName, tableName))),
|
||||
isErr: true,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-get-table-info-tool with invalid auth token",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-get-table-info-tool/invoke",
|
||||
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
|
||||
"cloud.google.com/go/bigtable"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
@@ -78,6 +79,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
|
||||
// Do not change the shape of statement without checking tests/common_test.go.
|
||||
// The structure and value of seed data has to match https://github.com/googleapis/genai-toolbox/blob/4dba0df12dc438eca3cb476ef52aa17cdf232c12/tests/common_test.go#L200-L251
|
||||
paramTestStatement := fmt.Sprintf("SELECT TO_INT64(cf['id']) as id, CAST(cf['name'] AS string) as name, FROM %s WHERE TO_INT64(cf['id']) = @id OR CAST(cf['name'] AS string) = @name;", tableName)
|
||||
paramTestStatement2 := fmt.Sprintf("SELECT TO_INT64(cf['id']) as id, CAST(cf['name'] AS string) as name, FROM %s WHERE TO_INT64(cf['id']) = @id;", tableName)
|
||||
teardownTable1 := setupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableName, columnFamilyName, muts, rowKeys)
|
||||
defer teardownTable1(t)
|
||||
|
||||
@@ -92,7 +94,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
|
||||
defer teardownTableTmpl(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, BigtableToolKind, paramTestStatement, authToolStatement)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, BigtableToolKind, paramTestStatement, paramTestStatement2, authToolStatement)
|
||||
toolsFile = addTemplateParamConfig(t, toolsFile)
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
@@ -103,7 +105,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -114,8 +116,9 @@ func TestBigtableToolEndpoints(t *testing.T) {
|
||||
// Actual test parameters are set in https://github.com/googleapis/genai-toolbox/blob/52b09a67cb40ac0c5f461598b4673136699a3089/tests/tool_test.go#L250
|
||||
select1Want := "[{\"$col1\":1}]"
|
||||
failInvocationWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to prepare statement: rpc error: code = InvalidArgument desc = Syntax error: Unexpected identifier \"SELEC\" [at 1:1]"}],"isError":true}}`
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, _, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
invokeParamWantNull := `[{"id":4,"name":""}]`
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
|
||||
templateParamTestConfig := tests.NewTemplateParameterTestConfig(
|
||||
@@ -139,7 +142,7 @@ func getTestData(columnFamilyName string) ([]*bigtable.Mutation, []string) {
|
||||
muts := []*bigtable.Mutation{}
|
||||
rowKeys := []string{}
|
||||
|
||||
var ids [3][]byte
|
||||
var ids [4][]byte
|
||||
for i := range ids {
|
||||
ids[i] = convertToBytes(i + 1)
|
||||
}
|
||||
@@ -163,6 +166,10 @@ func getTestData(columnFamilyName string) ([]*bigtable.Mutation, []string) {
|
||||
"name": []byte("Sid"),
|
||||
"id": ids[2],
|
||||
},
|
||||
"row-04": {
|
||||
"name": nil,
|
||||
"id": ids[3],
|
||||
},
|
||||
} {
|
||||
mut := bigtable.NewMutation()
|
||||
for col, v := range mutData {
|
||||
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
"cloud.google.com/go/cloudsqlconn"
|
||||
"cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -128,17 +129,17 @@ func TestCloudSQLMSSQLToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetMSSQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetMSSQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupMsSQLTable(t, ctx, db, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetMSSQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetMSSQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupMsSQLTable(t, ctx, db, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLMSSQLToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLMSSQLToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddMSSQLExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMSSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLMSSQLToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -151,7 +152,7 @@ func TestCloudSQLMSSQLToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -160,8 +161,8 @@ func TestCloudSQLMSSQLToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetMSSQLWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"cloud.google.com/go/cloudsqlconn"
|
||||
"cloud.google.com/go/cloudsqlconn/mysql/mysql"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -115,17 +116,17 @@ func TestCloudSQLMySQLToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetMySQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetMySQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupMySQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetMySQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetMySQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupMySQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLMySQLToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLMySQLToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddMySqlExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMySQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLMySQLToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -138,7 +139,7 @@ func TestCloudSQLMySQLToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -147,8 +148,8 @@ func TestCloudSQLMySQLToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetMySQLWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
|
||||
"cloud.google.com/go/cloudsqlconn"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
@@ -119,17 +120,17 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddPgExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -142,7 +143,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -151,8 +152,8 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetPostgresWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -28,7 +28,7 @@ import (
|
||||
)
|
||||
|
||||
// GetToolsConfig returns a mock tools config file
|
||||
func GetToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, authToolStatement string) map[string]any {
|
||||
func GetToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, paramToolStatement2, authToolStatement string) map[string]any {
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := map[string]any{
|
||||
"sources": map[string]any{
|
||||
@@ -65,6 +65,19 @@ func GetToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, a
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-param-tool2": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test invocation with params.",
|
||||
"statement": paramToolStatement2,
|
||||
"parameters": []any{
|
||||
map[string]any{
|
||||
"name": "id",
|
||||
"type": "integer",
|
||||
"description": "user ID",
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-auth-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
@@ -261,12 +274,13 @@ func AddMSSQLExecuteSqlConfig(t *testing.T, config map[string]any) map[string]an
|
||||
}
|
||||
|
||||
// GetPostgresSQLParamToolInfo returns statements and param for my-param-tool postgres-sql kind
|
||||
func GetPostgresSQLParamToolInfo(tableName string) (string, string, string, []any) {
|
||||
func GetPostgresSQLParamToolInfo(tableName string) (string, string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id SERIAL PRIMARY KEY, name TEXT);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES ($1), ($2), ($3);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES ($1), ($2), ($3), ($4);", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = $1 OR name = $2;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid"}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
toolStatement2 := fmt.Sprintf("SELECT * FROM %s WHERE id = $1;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid", nil}
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// GetPostgresSQLAuthToolInfo returns statements and param of my-auth-tool for postgres-sql kind
|
||||
@@ -286,12 +300,13 @@ func GetPostgresSQLTmplToolStatement() (string, string) {
|
||||
}
|
||||
|
||||
// GetMSSQLParamToolInfo returns statements and param for my-param-tool mssql-sql kind
|
||||
func GetMSSQLParamToolInfo(tableName string) (string, string, string, []any) {
|
||||
func GetMSSQLParamToolInfo(tableName string) (string, string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT IDENTITY(1,1) PRIMARY KEY, name VARCHAR(255));", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (@alice), (@jane), (@sid);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (@alice), (@jane), (@sid), (@nil);", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = @id OR name = @p2;", tableName)
|
||||
params := []any{sql.Named("alice", "Alice"), sql.Named("jane", "Jane"), sql.Named("sid", "Sid")}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
toolStatement2 := fmt.Sprintf("SELECT * FROM %s WHERE id = @id;", tableName)
|
||||
params := []any{sql.Named("alice", "Alice"), sql.Named("jane", "Jane"), sql.Named("sid", "Sid"), sql.Named("nil", nil)}
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// GetMSSQLAuthToolInfo returns statements and param of my-auth-tool for mssql-sql kind
|
||||
@@ -311,12 +326,13 @@ func GetMSSQLTmplToolStatement() (string, string) {
|
||||
}
|
||||
|
||||
// GetMySQLParamToolInfo returns statements and param for my-param-tool mysql-sql kind
|
||||
func GetMySQLParamToolInfo(tableName string) (string, string, string, []any) {
|
||||
func GetMySQLParamToolInfo(tableName string) (string, string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT NOT NULL AUTO_INCREMENT PRIMARY KEY, name VARCHAR(255));", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (?), (?), (?);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (?), (?), (?), (?);", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = ? OR name = ?;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid"}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
toolStatement2 := fmt.Sprintf("SELECT * FROM %s WHERE id = ?;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid", nil}
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// GetMySQLAuthToolInfo returns statements and param of my-auth-tool for mysql-sql kind
|
||||
@@ -335,10 +351,11 @@ func GetMySQLTmplToolStatement() (string, string) {
|
||||
return tmplSelectCombined, tmplSelectFilterCombined
|
||||
}
|
||||
|
||||
func GetNonSpannerInvokeParamWant() (string, string) {
|
||||
func GetNonSpannerInvokeParamWant() (string, string, string) {
|
||||
invokeParamWant := "[{\"id\":1,\"name\":\"Alice\"},{\"id\":3,\"name\":\"Sid\"}]"
|
||||
invokeParamWantNull := "[{\"id\":4,\"name\":null}]"
|
||||
mcpInvokeParamWant := `{"jsonrpc":"2.0","id":"my-param-tool","result":{"content":[{"type":"text","text":"{\"id\":1,\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":3,\"name\":\"Sid\"}"}]}}`
|
||||
return invokeParamWant, mcpInvokeParamWant
|
||||
return invokeParamWant, invokeParamWantNull, mcpInvokeParamWant
|
||||
}
|
||||
|
||||
// GetPostgresWants return the expected wants for postgres
|
||||
@@ -453,12 +470,13 @@ func SetupMySQLTable(t *testing.T, ctx context.Context, pool *sql.DB, createStat
|
||||
}
|
||||
|
||||
// GetRedisWants return the expected wants for redis
|
||||
func GetRedisValkeyWants() (string, string, string, string) {
|
||||
func GetRedisValkeyWants() (string, string, string, string, string) {
|
||||
select1Want := "[\"PONG\"]"
|
||||
failInvocationWant := `unknown command 'SELEC 1;', with args beginning with: \""}]}}`
|
||||
invokeParamWant := "[{\"id\":\"1\",\"name\":\"Alice\"},{\"id\":\"3\",\"name\":\"Sid\"}]"
|
||||
invokeParamWantNull := `[{"id":"4","name":""}]`
|
||||
mcpInvokeParamWant := `{"jsonrpc":"2.0","id":"my-param-tool","result":{"content":[{"type":"text","text":"{\"id\":\"1\",\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":\"3\",\"name\":\"Sid\"}"}]}}`
|
||||
return select1Want, failInvocationWant, invokeParamWant, mcpInvokeParamWant
|
||||
return select1Want, failInvocationWant, invokeParamWant, invokeParamWantNull, mcpInvokeParamWant
|
||||
}
|
||||
|
||||
func GetRedisValkeyToolsConfig(sourceConfig map[string]any, toolKind string) map[string]any {
|
||||
@@ -497,6 +515,19 @@ func GetRedisValkeyToolsConfig(sourceConfig map[string]any, toolKind string) map
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-param-tool2": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test invocation with params.",
|
||||
"commands": [][]string{{"HGETALL", "row4"}},
|
||||
"parameters": []any{
|
||||
map[string]any{
|
||||
"name": "id",
|
||||
"type": "integer",
|
||||
"description": "user ID",
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-auth-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
@@ -534,7 +565,5 @@ func GetRedisValkeyToolsConfig(sourceConfig map[string]any, toolKind string) map
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
return toolsFile
|
||||
|
||||
}
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
|
||||
"github.com/couchbase/gocb/v2"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -102,7 +103,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) {
|
||||
collectionNameTemplateParam := "template_param_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// Set up data for param tool
|
||||
paramToolStatement, params1 := getCouchbaseParamToolInfo(collectionNameParam)
|
||||
paramToolStatement1, paramToolStatement2, params1 := getCouchbaseParamToolInfo(collectionNameParam)
|
||||
teardownCollection1 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameParam, params1)
|
||||
defer teardownCollection1(t)
|
||||
|
||||
@@ -117,7 +118,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) {
|
||||
defer teardownCollection3(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolKind, paramToolStatement, authToolStatement)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, couchbaseToolKind, tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll)
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
@@ -128,7 +129,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -139,8 +140,8 @@ func TestCouchbaseToolEndpoints(t *testing.T) {
|
||||
select1Want := "[{\"$1\":1}]"
|
||||
failMcpInvocationWant := "{\"jsonrpc\":\"2.0\",\"id\":\"invoke-fail-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"unable to execute query: parsing failure | {\\\"statement\\\":\\\"SELEC 1;\\\""
|
||||
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failMcpInvocationWant)
|
||||
|
||||
templateParamTestConfig := tests.NewTemplateParameterTestConfig(
|
||||
@@ -230,18 +231,22 @@ func setupCouchbaseCollection(t *testing.T, ctx context.Context, cluster *gocb.C
|
||||
}
|
||||
|
||||
// getCouchbaseParamToolInfo returns statements and params for my-param-tool couchbase-sql kind
|
||||
func getCouchbaseParamToolInfo(collectionName string) (string, []map[string]any) {
|
||||
func getCouchbaseParamToolInfo(collectionName string) (string, string, []map[string]any) {
|
||||
// N1QL uses positional or named parameters with $ prefix
|
||||
toolStatement := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
|
||||
"%s.* FROM %s WHERE meta().id = TOSTRING($id) OR name = $name order by meta().id",
|
||||
collectionName, collectionName)
|
||||
toolStatement2 := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
|
||||
"%s.* FROM %s WHERE meta().id = TOSTRING($id) order by meta().id",
|
||||
collectionName, collectionName)
|
||||
|
||||
params := []map[string]any{
|
||||
{"name": "Alice"},
|
||||
{"name": "Jane"},
|
||||
{"name": "Sid"},
|
||||
{"name": nil},
|
||||
}
|
||||
return toolStatement, params
|
||||
return toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// getCouchbaseAuthToolInfo returns statements and param of my-auth-tool for couchbase-sql kind
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -78,7 +79,7 @@ func TestDgraphToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
@@ -59,6 +60,8 @@ func multiTool(w http.ResponseWriter, r *http.Request) {
|
||||
handleTool0(w, r)
|
||||
case "tool1":
|
||||
handleTool1(w, r)
|
||||
case "tool1a":
|
||||
handleTool1a(w, r)
|
||||
case "tool2":
|
||||
handleTool2(w, r)
|
||||
case "tool3":
|
||||
@@ -130,6 +133,27 @@ func handleTool1(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
|
||||
// handler function for the test server
|
||||
func handleTool1a(w http.ResponseWriter, r *http.Request) {
|
||||
// expect GET method
|
||||
if r.Method != http.MethodGet {
|
||||
errorMessage := fmt.Sprintf("expected GET method but got: %s", string(r.Method))
|
||||
http.Error(w, errorMessage, http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
id := r.URL.Query().Get("id")
|
||||
if id == "4" {
|
||||
response := `[{"id":4,"name":null}]`
|
||||
_, err := w.Write([]byte(response))
|
||||
if err != nil {
|
||||
http.Error(w, "Failed to write response", http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
|
||||
// handler function for the test server
|
||||
func handleTool2(w http.ResponseWriter, r *http.Request) {
|
||||
// expect GET method
|
||||
@@ -254,16 +278,16 @@ func TestHttpToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
}
|
||||
|
||||
select1Want := `["Hello","World"]`
|
||||
invokeParamWant, _ := tests.GetNonSpannerInvokeParamWant()
|
||||
invokeParamWant, invokeParamWantNull, _ := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolGetTest(t)
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
runAdvancedHTTPInvokeTest(t)
|
||||
}
|
||||
|
||||
@@ -383,6 +407,16 @@ func getHTTPToolsConfig(sourceConfig map[string]any, toolKind string) map[string
|
||||
"bodyParams": []tools.Parameter{tools.NewStringParameter("name", "user name")},
|
||||
"headers": map[string]string{"Content-Type": "application/json"},
|
||||
},
|
||||
"my-param-tool2": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"method": "GET",
|
||||
"path": "/tool1a",
|
||||
"description": "some description",
|
||||
"queryParams": []tools.Parameter{
|
||||
tools.NewIntParameter("id", "user ID")},
|
||||
"headers": map[string]string{"Content-Type": "application/json"},
|
||||
},
|
||||
"my-auth-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -101,17 +102,17 @@ func TestMSSQLToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetMSSQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetMSSQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupMsSQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetMSSQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetMSSQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupMsSQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, MSSQLToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, MSSQLToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddMSSQLExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMSSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, MSSQLToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -124,7 +125,7 @@ func TestMSSQLToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -133,8 +134,8 @@ func TestMSSQLToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetMSSQLWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -92,17 +93,17 @@ func TestMySQLToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetMySQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetMySQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupMySQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetMySQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetMySQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupMySQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, MySQLToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, MySQLToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddMySqlExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMySQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, MySQLToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -115,7 +116,7 @@ func TestMySQLToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -124,8 +125,8 @@ func TestMySQLToolEndpoints(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetMySQLWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -87,7 +88,7 @@ func TestNeo4jToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
@@ -98,17 +99,17 @@ func TestPostgres(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := tests.GetPostgresSQLParamToolInfo(tableNameParam)
|
||||
teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, PostgresToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, PostgresToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = tests.AddPgExecuteSqlConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, PostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
@@ -121,7 +122,7 @@ func TestPostgres(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -130,8 +131,8 @@ func TestPostgres(t *testing.T) {
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, createTableStatement := tests.GetPostgresWants()
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
@@ -90,7 +91,7 @@ func TestRedisToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -98,17 +99,18 @@ func TestRedisToolEndpoints(t *testing.T) {
|
||||
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, invokeParamWant, mcpInvokeParamWant := tests.GetRedisValkeyWants()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
select1Want, failInvocationWant, invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetRedisValkeyWants()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
}
|
||||
|
||||
func setupRedisDB(t *testing.T, ctx context.Context, client *redis.Client) func(*testing.T) {
|
||||
keys := []string{"row1", "row2", "row3"}
|
||||
keys := []string{"row1", "row2", "row3", "row4"}
|
||||
commands := [][]any{
|
||||
{"HSET", keys[0], "id", 1, "name", "Alice"},
|
||||
{"HSET", keys[1], "id", 2, "name", "Jane"},
|
||||
{"HSET", keys[2], "id", 3, "name", "Sid"},
|
||||
{"HSET", keys[3], "id", 4, "name", nil},
|
||||
{"HSET", tests.ServiceAccountEmail, "name", "Alice"},
|
||||
}
|
||||
for _, c := range commands {
|
||||
|
||||
@@ -15,13 +15,10 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
|
||||
@@ -133,63 +130,3 @@ func (c *CmdExec) Close() {
|
||||
c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// WaitForString waits until the server logs a single line that matches the provided regex.
|
||||
// returns the output of whatever the server sent so far.
|
||||
func (c *CmdExec) WaitForString(ctx context.Context, re *regexp.Regexp) (string, error) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
in := bufio.NewReader(c.Out)
|
||||
|
||||
// read lines in background, sending result of each read over a channel
|
||||
// this allows us to use in.ReadString without blocking
|
||||
type result struct {
|
||||
s string
|
||||
err error
|
||||
}
|
||||
output := make(chan result)
|
||||
go func() {
|
||||
defer close(output)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// if the context is canceled, the orig thread will send back the error
|
||||
// so we can just exit the goroutine here
|
||||
return
|
||||
default:
|
||||
// otherwise read a line from the output
|
||||
s, err := in.ReadString('\n')
|
||||
if err != nil {
|
||||
output <- result{err: err}
|
||||
return
|
||||
}
|
||||
output <- result{s: s}
|
||||
// if that last string matched, exit the goroutine
|
||||
if re.MatchString(s) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// collect the output until the ctx is canceled, an error was hit,
|
||||
// or match was found (which is indicated the channel is closed)
|
||||
var sb strings.Builder
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// if ctx is done, return that error
|
||||
return sb.String(), ctx.Err()
|
||||
case o, ok := <-output:
|
||||
if !ok {
|
||||
// match was found!
|
||||
return sb.String(), nil
|
||||
}
|
||||
if o.err != nil {
|
||||
// error was found!
|
||||
return sb.String(), o.err
|
||||
}
|
||||
sb.WriteString(o.s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cloud.google.com/go/cloudsqlconn"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
// RunSourceConnection test for source connection
|
||||
@@ -57,7 +58,7 @@ func RunSourceConnectionTest(t *testing.T, sourceConfig map[string]any, toolKind
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
return fmt.Errorf("toolbox didn't start successfully: %s", err)
|
||||
|
||||
@@ -31,6 +31,7 @@ import (
|
||||
database "cloud.google.com/go/spanner/admin/database/apiv1"
|
||||
"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
@@ -90,7 +91,7 @@ func initSpannerClients(ctx context.Context, project, instance, dbname string) (
|
||||
|
||||
func TestSpannerToolEndpoints(t *testing.T) {
|
||||
sourceConfig := getSpannerVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
var args []string
|
||||
@@ -107,7 +108,7 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := getSpannerParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := getSpannerParamToolInfo(tableNameParam)
|
||||
dbString := fmt.Sprintf(
|
||||
"projects/%s/instances/%s/databases/%s",
|
||||
SpannerProject,
|
||||
@@ -118,7 +119,7 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := getSpannerAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := getSpannerAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := setupSpannerTable(t, ctx, adminClient, dataClient, createStatement2, insertStatement2, tableNameAuth, dbString, params2)
|
||||
defer teardownTable2(t)
|
||||
|
||||
@@ -128,7 +129,7 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
defer teardownTableTmpl(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, SpannerToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, SpannerToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
toolsFile = addSpannerExecuteSqlConfig(t, toolsFile)
|
||||
toolsFile = addSpannerReadOnlyConfig(t, toolsFile)
|
||||
toolsFile = addTemplateParamConfig(t, toolsFile)
|
||||
@@ -141,7 +142,7 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -152,10 +153,11 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
select1Want := "[{\"\":\"1\"}]"
|
||||
accessSchemaWant := "[{\"schema_name\":\"INFORMATION_SCHEMA\"}]"
|
||||
invokeParamWant := "[{\"id\":\"1\",\"name\":\"Alice\"},{\"id\":\"3\",\"name\":\"Sid\"}]"
|
||||
invokeParamWantNull := `[{"id":"4","name":null}]`
|
||||
mcpInvokeParamWant := `{"jsonrpc":"2.0","id":"my-param-tool","result":{"content":[{"type":"text","text":"{\"id\":\"1\",\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":\"3\",\"name\":\"Sid\"}"}]}}`
|
||||
failInvocationWant := `"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute client: unable to parse row: spanner: code = \"InvalidArgument\", desc = \"Syntax error: Unexpected identifier \\\\\\\"SELEC\\\\\\\" [at 1:1]\\\\nSELEC 1;\\\\n^\"`
|
||||
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
runSpannerSchemaToolInvokeTest(t, accessSchemaWant)
|
||||
runSpannerExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam, tableNameAuth)
|
||||
@@ -169,12 +171,13 @@ func TestSpannerToolEndpoints(t *testing.T) {
|
||||
}
|
||||
|
||||
// getSpannerToolInfo returns statements and param for my-param-tool for spanner-sql kind
|
||||
func getSpannerParamToolInfo(tableName string) (string, string, string, map[string]any) {
|
||||
func getSpannerParamToolInfo(tableName string) (string, string, string, string, map[string]any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT64, name STRING(MAX)) PRIMARY KEY (id)", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (id, name) VALUES (1, @name1), (2, @name2), (3, @name3)", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (id, name) VALUES (1, @name1), (2, @name2), (3, @name3), (4, @name4)", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = @id OR name = @name", tableName)
|
||||
params := map[string]any{"name1": "Alice", "name2": "Jane", "name3": "Sid"}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
toolStatement2 := fmt.Sprintf("SELECT * FROM %s WHERE id = @id", tableName)
|
||||
params := map[string]any{"name1": "Alice", "name2": "Jane", "name3": "Sid", "name4": nil}
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
// getSpannerAuthToolInfo returns statements and param of my-auth-tool for spanner-sql kind
|
||||
@@ -230,13 +233,13 @@ func setupSpannerTable(t *testing.T, ctx context.Context, adminClient *database.
|
||||
Statements: []string{fmt.Sprintf("DROP TABLE %s", tableName)},
|
||||
})
|
||||
if err != nil {
|
||||
t.Errorf("unable to start drop table operation: %s", err)
|
||||
t.Errorf("unable to start drop %s operation: %s", tableName, err)
|
||||
return
|
||||
}
|
||||
|
||||
err = op.Wait(ctx)
|
||||
if err != nil {
|
||||
t.Errorf("Teardown failed: %s", err)
|
||||
opErr := op.Wait(ctx)
|
||||
if opErr != nil {
|
||||
t.Errorf("Teardown failed: %s", opErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -438,7 +441,7 @@ func runSpannerExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWa
|
||||
name: "invoke my-exec-sql-tool insert entry",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"INSERT INTO %s (id, name) VALUES (4, 'test_name')\"}", tableNameParam))),
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"INSERT INTO %s (id, name) VALUES (5, 'test_name')\"}", tableNameParam))),
|
||||
want: "null",
|
||||
isErr: false,
|
||||
},
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
@@ -80,12 +81,13 @@ func setupSQLiteTestDB(t *testing.T, ctx context.Context, db *sql.DB, createStat
|
||||
}
|
||||
}
|
||||
|
||||
func getSQLiteParamToolInfo(tableName string) (string, string, string, []any) {
|
||||
func getSQLiteParamToolInfo(tableName string) (string, string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (id INTEGER PRIMARY KEY, name TEXT);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (?), (?), (?);", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name) VALUES (?), (?), (?), (?);", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = ? OR name = ?;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid"}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
toolStatement2 := fmt.Sprintf("SELECT * FROM %s WHERE id = ?;", tableName)
|
||||
params := []any{"Alice", "Jane", "Sid", nil}
|
||||
return createStatement, insertStatement, toolStatement, toolStatement2, params
|
||||
}
|
||||
|
||||
func getSQLiteAuthToolInfo(tableName string) (string, string, string, []any) {
|
||||
@@ -123,15 +125,15 @@ func TestSQLiteToolEndpoint(t *testing.T) {
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createStatement1, insertStatement1, toolStatement1, params1 := getSQLiteParamToolInfo(tableNameParam)
|
||||
createStatement1, insertStatement1, paramToolStatement1, paramToolStatement2, params1 := getSQLiteParamToolInfo(tableNameParam)
|
||||
setupSQLiteTestDB(t, ctx, db, createStatement1, insertStatement1, tableNameParam, params1)
|
||||
|
||||
// set up data for auth tool
|
||||
createStatement2, insertStatement2, toolStatement2, params2 := getSQLiteAuthToolInfo(tableNameAuth)
|
||||
createStatement2, insertStatement2, authToolStatement, params2 := getSQLiteAuthToolInfo(tableNameAuth)
|
||||
setupSQLiteTestDB(t, ctx, db, createStatement2, insertStatement2, tableNameAuth, params2)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, SQLiteToolKind, toolStatement1, toolStatement2)
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, SQLiteToolKind, paramToolStatement1, paramToolStatement2, authToolStatement)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := getSQLiteTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, SQLiteToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
|
||||
@@ -143,7 +145,7 @@ func TestSQLiteToolEndpoint(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -153,8 +155,8 @@ func TestSQLiteToolEndpoint(t *testing.T) {
|
||||
|
||||
select1Want := "[{\"1\":1}]"
|
||||
failInvocationWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: SQL logic error: near \"SELEC\": syntax error (1)"}],"isError":true}}`
|
||||
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.NewTemplateParameterTestConfig())
|
||||
}
|
||||
|
||||
@@ -76,7 +76,7 @@ func RunToolGetTest(t *testing.T) {
|
||||
}
|
||||
|
||||
// RunToolInvoke runs the tool invoke endpoint
|
||||
func RunToolInvokeTest(t *testing.T, select1Want, invokeParamWant string) {
|
||||
func RunToolInvokeTest(t *testing.T, select1Want, invokeParamWant, invokeParamWantNull string) {
|
||||
// Get ID token
|
||||
idToken, err := GetGoogleIdToken(ClientId)
|
||||
if err != nil {
|
||||
@@ -108,6 +108,14 @@ func RunToolInvokeTest(t *testing.T, select1Want, invokeParamWant string) {
|
||||
want: invokeParamWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "invoke my-param-tool2 with nil response",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-param-tool2/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(`{"id": 4}`)),
|
||||
want: invokeParamWantNull,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-param-tool without parameters",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-param-tool/invoke",
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/valkey-io/valkey-go"
|
||||
)
|
||||
@@ -93,7 +94,7 @@ func TestValkeyToolEndpoints(t *testing.T) {
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
@@ -101,17 +102,18 @@ func TestValkeyToolEndpoints(t *testing.T) {
|
||||
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, invokeParamWant, mcpInvokeParamWant := tests.GetRedisValkeyWants()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
select1Want, failInvocationWant, invokeParamWant, invokeParamWantNull, mcpInvokeParamWant := tests.GetRedisValkeyWants()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
}
|
||||
|
||||
func setupValkeyDB(t *testing.T, ctx context.Context, client valkey.Client) func(*testing.T) {
|
||||
keys := []string{"row1", "row2", "row3"}
|
||||
keys := []string{"row1", "row2", "row3", "row4"}
|
||||
commands := [][]string{
|
||||
{"HSET", keys[0], "name", "Alice", "id", "1"},
|
||||
{"HSET", keys[1], "name", "Jane", "id", "2"},
|
||||
{"HSET", keys[2], "name", "Sid", "id", "3"},
|
||||
{"HSET", keys[3], "name", "", "id", "4"},
|
||||
{"HSET", tests.ServiceAccountEmail, "name", "Alice"},
|
||||
}
|
||||
builtCmds := make(valkey.Commands, len(commands))
|
||||
|
||||
Reference in New Issue
Block a user