mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-02-10 15:15:21 -05:00
252 lines
7.7 KiB
Go
252 lines
7.7 KiB
Go
// Copyright 2026 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package cli
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"slices"
|
|
"strings"
|
|
|
|
"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/telemetry"
|
|
"github.com/googleapis/genai-toolbox/internal/util"
|
|
)
|
|
|
|
type IOStreams struct {
|
|
In io.Reader
|
|
Out io.Writer
|
|
ErrOut io.Writer
|
|
}
|
|
|
|
// ToolboxOptions holds dependencies shared by all commands.
|
|
type ToolboxOptions struct {
|
|
IOStreams IOStreams
|
|
Logger log.Logger
|
|
Cfg server.ServerConfig
|
|
ToolsFile string
|
|
ToolsFiles []string
|
|
ToolsFolder string
|
|
PrebuiltConfigs []string
|
|
}
|
|
|
|
// Option defines a function that modifies the ToolboxOptions struct.
|
|
type Option func(*ToolboxOptions)
|
|
|
|
// NewToolboxOptions creates a new instance with defaults, then applies any
|
|
// provided options.
|
|
func NewToolboxOptions(opts ...Option) *ToolboxOptions {
|
|
o := &ToolboxOptions{
|
|
IOStreams: IOStreams{
|
|
In: os.Stdin,
|
|
Out: os.Stdout,
|
|
ErrOut: os.Stderr,
|
|
},
|
|
}
|
|
|
|
for _, opt := range opts {
|
|
opt(o)
|
|
}
|
|
return o
|
|
}
|
|
|
|
// Apply allows you to update an EXISTING ToolboxOptions instance.
|
|
// This is useful for "late binding".
|
|
func (o *ToolboxOptions) Apply(opts ...Option) {
|
|
for _, opt := range opts {
|
|
opt(o)
|
|
}
|
|
}
|
|
|
|
// WithIOStreams updates the IO streams.
|
|
func WithIOStreams(out, err io.Writer) Option {
|
|
return func(o *ToolboxOptions) {
|
|
o.IOStreams.Out = out
|
|
o.IOStreams.ErrOut = err
|
|
}
|
|
}
|
|
|
|
// Setup create logger and telemetry instrumentations.
|
|
func (opts *ToolboxOptions) Setup(ctx context.Context) (context.Context, func(context.Context) error, error) {
|
|
// If stdio, set logger's out stream (usually DEBUG and INFO logs) to
|
|
// errStream
|
|
loggerOut := opts.IOStreams.Out
|
|
if opts.Cfg.Stdio {
|
|
loggerOut = opts.IOStreams.ErrOut
|
|
}
|
|
|
|
// Handle logger separately from config
|
|
logger, err := log.NewLogger(opts.Cfg.LoggingFormat.String(), opts.Cfg.LogLevel.String(), loggerOut, opts.IOStreams.ErrOut)
|
|
if err != nil {
|
|
return ctx, nil, fmt.Errorf("unable to initialize logger: %w", err)
|
|
}
|
|
|
|
ctx = util.WithLogger(ctx, logger)
|
|
opts.Logger = logger
|
|
|
|
// Set up OpenTelemetry
|
|
otelShutdown, err := telemetry.SetupOTel(ctx, opts.Cfg.Version, opts.Cfg.TelemetryOTLP, opts.Cfg.TelemetryGCP, opts.Cfg.TelemetryServiceName)
|
|
if err != nil {
|
|
errMsg := fmt.Errorf("error setting up OpenTelemetry: %w", err)
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return ctx, nil, errMsg
|
|
}
|
|
|
|
shutdownFunc := func(ctx context.Context) error {
|
|
err := otelShutdown(ctx)
|
|
if err != nil {
|
|
errMsg := fmt.Errorf("error shutting down OpenTelemetry: %w", err)
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
instrumentation, err := telemetry.CreateTelemetryInstrumentation(opts.Cfg.Version)
|
|
if err != nil {
|
|
errMsg := fmt.Errorf("unable to create telemetry instrumentation: %w", err)
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return ctx, shutdownFunc, errMsg
|
|
}
|
|
|
|
ctx = util.WithInstrumentation(ctx, instrumentation)
|
|
|
|
return ctx, shutdownFunc, nil
|
|
}
|
|
|
|
// LoadConfig checks and merge files that should be loaded into the server
|
|
func (opts *ToolboxOptions) LoadConfig(ctx context.Context) (bool, error) {
|
|
// Determine if Custom Files should be loaded
|
|
// Check for explicit custom flags
|
|
isCustomConfigured := opts.ToolsFile != "" || len(opts.ToolsFiles) > 0 || opts.ToolsFolder != ""
|
|
|
|
// Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags)
|
|
useDefaultToolsFile := len(opts.PrebuiltConfigs) == 0 && !isCustomConfigured
|
|
|
|
if useDefaultToolsFile {
|
|
opts.ToolsFile = "tools.yaml"
|
|
isCustomConfigured = true
|
|
}
|
|
|
|
logger, err := util.LoggerFromContext(ctx)
|
|
if err != nil {
|
|
return isCustomConfigured, err
|
|
}
|
|
|
|
var allToolsFiles []ToolsFile
|
|
|
|
// Load Prebuilt Configuration
|
|
|
|
if len(opts.PrebuiltConfigs) > 0 {
|
|
slices.Sort(opts.PrebuiltConfigs)
|
|
sourcesList := strings.Join(opts.PrebuiltConfigs, ", ")
|
|
logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList)
|
|
logger.InfoContext(ctx, logMsg)
|
|
|
|
for _, configName := range opts.PrebuiltConfigs {
|
|
buf, err := prebuiltconfigs.Get(configName)
|
|
if err != nil {
|
|
logger.ErrorContext(ctx, err.Error())
|
|
return isCustomConfigured, err
|
|
}
|
|
|
|
// Parse into ToolsFile struct
|
|
parsed, err := parseToolsFile(ctx, buf)
|
|
if err != nil {
|
|
errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err)
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return isCustomConfigured, errMsg
|
|
}
|
|
allToolsFiles = append(allToolsFiles, parsed)
|
|
}
|
|
}
|
|
|
|
// Load Custom Configurations
|
|
if isCustomConfigured {
|
|
// Enforce exclusivity among custom flags (tools-file vs tools-files vs tools-folder)
|
|
if (opts.ToolsFile != "" && len(opts.ToolsFiles) > 0) ||
|
|
(opts.ToolsFile != "" && opts.ToolsFolder != "") ||
|
|
(len(opts.ToolsFiles) > 0 && opts.ToolsFolder != "") {
|
|
errMsg := fmt.Errorf("--tools-file, --tools-files, and --tools-folder flags cannot be used simultaneously")
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return isCustomConfigured, errMsg
|
|
}
|
|
|
|
var customTools ToolsFile
|
|
var err error
|
|
|
|
if len(opts.ToolsFiles) > 0 {
|
|
// Use tools-files
|
|
logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(opts.ToolsFiles)))
|
|
customTools, err = LoadAndMergeToolsFiles(ctx, opts.ToolsFiles)
|
|
} else if opts.ToolsFolder != "" {
|
|
// Use tools-folder
|
|
logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", opts.ToolsFolder))
|
|
customTools, err = LoadAndMergeToolsFolder(ctx, opts.ToolsFolder)
|
|
} else {
|
|
// Use single file (tools-file or default `tools.yaml`)
|
|
buf, readFileErr := os.ReadFile(opts.ToolsFile)
|
|
if readFileErr != nil {
|
|
errMsg := fmt.Errorf("unable to read tool file at %q: %w", opts.ToolsFile, readFileErr)
|
|
logger.ErrorContext(ctx, errMsg.Error())
|
|
return isCustomConfigured, errMsg
|
|
}
|
|
customTools, err = parseToolsFile(ctx, buf)
|
|
if err != nil {
|
|
err = fmt.Errorf("unable to parse tool file at %q: %w", opts.ToolsFile, err)
|
|
}
|
|
}
|
|
|
|
if err != nil {
|
|
logger.ErrorContext(ctx, err.Error())
|
|
return isCustomConfigured, err
|
|
}
|
|
allToolsFiles = append(allToolsFiles, customTools)
|
|
}
|
|
|
|
// Modify version string based on loaded configurations
|
|
if len(opts.PrebuiltConfigs) > 0 {
|
|
tag := "prebuilt"
|
|
if isCustomConfigured {
|
|
tag = "custom"
|
|
}
|
|
// prebuiltConfigs is already sorted above
|
|
for _, configName := range opts.PrebuiltConfigs {
|
|
opts.Cfg.Version += fmt.Sprintf("+%s.%s", tag, configName)
|
|
}
|
|
}
|
|
|
|
// Merge Everything
|
|
// This will error if custom tools collide with prebuilt tools
|
|
finalToolsFile, err := mergeToolsFiles(allToolsFiles...)
|
|
if err != nil {
|
|
logger.ErrorContext(ctx, err.Error())
|
|
return isCustomConfigured, err
|
|
}
|
|
|
|
opts.Cfg.SourceConfigs = finalToolsFile.Sources
|
|
opts.Cfg.AuthServiceConfigs = finalToolsFile.AuthServices
|
|
opts.Cfg.EmbeddingModelConfigs = finalToolsFile.EmbeddingModels
|
|
opts.Cfg.ToolConfigs = finalToolsFile.Tools
|
|
opts.Cfg.ToolsetConfigs = finalToolsFile.Toolsets
|
|
opts.Cfg.PromptConfigs = finalToolsFile.Prompts
|
|
|
|
return isCustomConfigured, nil
|
|
}
|