diff --git a/cmd/root.go b/cmd/root.go index e183ed9af..826698afe 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -21,6 +21,7 @@ import ( "os" "strings" + "github.com/googleapis/genai-toolbox/internal/log" "github.com/googleapis/genai-toolbox/internal/server" "github.com/spf13/cobra" "gopkg.in/yaml.v3" @@ -61,6 +62,7 @@ type Command struct { *cobra.Command cfg server.ServerConfig + logger log.Logger tools_file string } @@ -68,8 +70,9 @@ type Command struct { func NewCommand() *Command { cmd := &Command{ Command: &cobra.Command{ - Use: "toolbox", - Version: versionString, + Use: "toolbox", + Version: versionString, + SilenceErrors: true, }, } @@ -77,7 +80,9 @@ func NewCommand() *Command { flags.StringVarP(&cmd.cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.") flags.IntVarP(&cmd.cfg.Port, "port", "p", 5000, "Port the server will listen on.") - flags.StringVar(&cmd.tools_file, "tools_file", "tools.yaml", "File path specifying the tool configuration") + flags.StringVar(&cmd.tools_file, "tools_file", "tools.yaml", "File path specifying the tool configuration.") + flags.Var(&cmd.cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.") + flags.Var(&cmd.cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.") // wrap RunE command so that we have access to original Command object cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) } @@ -104,24 +109,48 @@ func run(cmd *Command) error { ctx, cancel := context.WithCancel(cmd.Context()) defer cancel() + // Handle logger separately from config + switch strings.ToLower(cmd.cfg.LoggingFormat.String()) { + case "json": + logger, err := log.NewStructuredLogger(os.Stdout, os.Stderr, cmd.cfg.LogLevel.String()) + if err != nil { + return fmt.Errorf("unable to initialize logger: %w", err) + } + cmd.logger = logger + default: + logger, err := log.NewStdLogger(os.Stdout, os.Stderr, cmd.cfg.LogLevel.String()) + if err != nil { + return fmt.Errorf("unable to initialize logger: %w", err) + } + cmd.logger = logger + } + // Read tool file contents buf, err := os.ReadFile(cmd.tools_file) if err != nil { - return fmt.Errorf("unable to read tool file at %q: %w", cmd.tools_file, err) + errMsg := fmt.Errorf("unable to read tool file at %q: %w", cmd.tools_file, err) + cmd.logger.Error(errMsg.Error()) + return errMsg } cmd.cfg.SourceConfigs, cmd.cfg.ToolConfigs, cmd.cfg.ToolsetConfigs, err = parseToolsFile(buf) if err != nil { - return fmt.Errorf("unable to parse tool file at %q: %w", cmd.tools_file, err) + errMsg := fmt.Errorf("unable to parse tool file at %q: %w", cmd.tools_file, err) + cmd.logger.Error(errMsg.Error()) + return errMsg } // run server - s, err := server.NewServer(cmd.cfg) + s, err := server.NewServer(cmd.cfg, cmd.logger) if err != nil { - return fmt.Errorf("toolbox failed to start with the following error: %w", err) + errMsg := fmt.Errorf("toolbox failed to start with the following error: %w", err) + cmd.logger.Error(errMsg.Error()) + return errMsg } err = s.ListenAndServe(ctx) if err != nil { - return fmt.Errorf("toolbox crashed with the following error: %w", err) + errMsg := fmt.Errorf("toolbox crashed with the following error: %w", err) + cmd.logger.Error(errMsg.Error()) + return errMsg } return nil diff --git a/cmd/root_test.go b/cmd/root_test.go index 596178ac0..03b1029ff 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -30,6 +30,16 @@ import ( "github.com/spf13/cobra" ) +func withDefaults(c server.ServerConfig) server.ServerConfig { + if c.Address == "" { + c.Address = "127.0.0.1" + } + if c.Port == 0 { + c.Port = 5000 + } + return c +} + func invokeCommand(args []string) (*Command, string, error) { c := NewCommand() @@ -70,7 +80,7 @@ func TestVersion(t *testing.T) { } } -func TestAddrPort(t *testing.T) { +func TestServerConfigFlags(t *testing.T) { tcs := []struct { desc string args []string @@ -79,42 +89,49 @@ func TestAddrPort(t *testing.T) { { desc: "default values", args: []string{}, - want: server.ServerConfig{ - Address: "127.0.0.1", - Port: 5000, - }, + want: withDefaults(server.ServerConfig{}), }, { desc: "address short", args: []string{"-a", "127.0.1.1"}, - want: server.ServerConfig{ + want: withDefaults(server.ServerConfig{ Address: "127.0.1.1", - Port: 5000, - }, + }), }, { desc: "address long", args: []string{"--address", "0.0.0.0"}, - want: server.ServerConfig{ + want: withDefaults(server.ServerConfig{ Address: "0.0.0.0", - Port: 5000, - }, + }), }, { desc: "port short", args: []string{"-p", "5052"}, - want: server.ServerConfig{ - Address: "127.0.0.1", - Port: 5052, - }, + want: withDefaults(server.ServerConfig{ + Port: 5052, + }), }, { desc: "port long", args: []string{"--port", "5050"}, - want: server.ServerConfig{ - Address: "127.0.0.1", - Port: 5050, - }, + want: withDefaults(server.ServerConfig{ + Port: 5050, + }), + }, + { + desc: "logging format", + args: []string{"--logging-format", "JSON"}, + want: withDefaults(server.ServerConfig{ + LoggingFormat: "JSON", + }), + }, + { + desc: "debug logs", + args: []string{"--log-level", "WARN"}, + want: withDefaults(server.ServerConfig{ + LogLevel: "WARN", + }), }, } for _, tc := range tcs { @@ -131,6 +148,54 @@ func TestAddrPort(t *testing.T) { } } +func TestFailServerConfigFlags(t *testing.T) { + tcs := []struct { + desc string + args []string + }{ + { + desc: "logging format", + args: []string{"--logging-format", "fail"}, + }, + { + desc: "debug logs", + args: []string{"--log-level", "fail"}, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + _, _, err := invokeCommand(tc.args) + if err == nil { + t.Fatalf("expected an error, but got nil") + } + }) + } +} + +func TestDefaultLoggingFormat(t *testing.T) { + c, _, err := invokeCommand([]string{}) + if err != nil { + t.Fatalf("unexpected error invoking command: %s", err) + } + got := c.cfg.LoggingFormat.String() + want := "standard" + if got != want { + t.Fatalf("unexpected default logging format flag: got %v, want %v", got, want) + } +} + +func TestDefaultLogLevel(t *testing.T) { + c, _, err := invokeCommand([]string{}) + if err != nil { + t.Fatalf("unexpected error invoking command: %s", err) + } + got := c.cfg.LogLevel.String() + want := "info" + if got != want { + t.Fatalf("unexpected default log level flag: got %v, want %v", got, want) + } +} + func TestToolFileFlag(t *testing.T) { tcs := []struct { desc string diff --git a/go.mod b/go.mod index f34bf6596..f8bc31ebf 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( cloud.google.com/go/alloydbconn v1.13.0 cloud.google.com/go/cloudsqlconn v1.12.1 github.com/go-chi/chi/v5 v5.1.0 + github.com/go-chi/httplog/v2 v2.1.1 github.com/go-chi/render v1.0.3 github.com/google/go-cmp v0.6.0 github.com/jackc/pgx/v5 v5.7.1 diff --git a/go.sum b/go.sum index e6951137a..01d479a83 100644 --- a/go.sum +++ b/go.sum @@ -36,6 +36,8 @@ github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2 github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-chi/chi/v5 v5.1.0 h1:acVI1TYaD+hhedDJ3r54HyA6sExp3HfXq7QWEEY/xMw= github.com/go-chi/chi/v5 v5.1.0/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8= +github.com/go-chi/httplog/v2 v2.1.1 h1:ojojiu4PIaoeJ/qAO4GWUxJqvYUTobeo7zmuHQJAxRk= +github.com/go-chi/httplog/v2 v2.1.1/go.mod h1:/XXdxicJsp4BA5fapgIC3VuTD+z0Z/VzukoB3VDc1YE= github.com/go-chi/render v1.0.3 h1:AsXqd2a1/INaIfUSKq3G5uA8weYx20FOsM7uSoCyyt4= github.com/go-chi/render v1.0.3/go.mod h1:/gr3hVkmYR0YlEy3LxCuVRFzEu9Ruok+gFqbIofjao0= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= diff --git a/internal/log/log.go b/internal/log/log.go index a4eb815f5..443dbdd07 100644 --- a/internal/log/log.go +++ b/internal/log/log.go @@ -31,7 +31,7 @@ type StdLogger struct { func NewStdLogger(outW, errW io.Writer, logLevel string) (Logger, error) { //Set log level var programLevel = new(slog.LevelVar) - slogLevel, err := severityToLevel(logLevel) + slogLevel, err := SeverityToLevel(logLevel) if err != nil { return nil, err } @@ -73,7 +73,7 @@ const ( ) // Returns severity level based on string. -func severityToLevel(s string) (slog.Level, error) { +func SeverityToLevel(s string) (slog.Level, error) { switch strings.ToUpper(s) { case Debug: return slog.LevelDebug, nil @@ -113,7 +113,7 @@ type StructuredLogger struct { func NewStructuredLogger(outW, errW io.Writer, logLevel string) (Logger, error) { //Set log level var programLevel = new(slog.LevelVar) - slogLevel, err := severityToLevel(logLevel) + slogLevel, err := SeverityToLevel(logLevel) if err != nil { return nil, err } diff --git a/internal/log/log_test.go b/internal/log/log_test.go index e0d4ed703..1b274ea32 100644 --- a/internal/log/log_test.go +++ b/internal/log/log_test.go @@ -53,7 +53,7 @@ func TestSeverityToLevel(t *testing.T) { } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { - got, err := severityToLevel(tc.in) + got, err := SeverityToLevel(tc.in) if err != nil { t.Fatalf("unexpected error: %s", err) } @@ -66,7 +66,7 @@ func TestSeverityToLevel(t *testing.T) { } func TestSeverityToLevelError(t *testing.T) { - _, err := severityToLevel("fail") + _, err := SeverityToLevel("fail") if err == nil { t.Fatalf("expected error on incorrect level") } diff --git a/internal/server/api.go b/internal/server/api.go index 9fb52510a..b44075492 100644 --- a/internal/server/api.go +++ b/internal/server/api.go @@ -48,7 +48,8 @@ func toolsetHandler(s *Server, w http.ResponseWriter, r *http.Request) { toolsetName := chi.URLParam(r, "toolsetName") toolset, ok := s.toolsets[toolsetName] if !ok { - _ = render.Render(w, r, newErrResponse(fmt.Errorf("Toolset %q does not exist", toolsetName), http.StatusNotFound)) + err := fmt.Errorf("Toolset %q does not exist", toolsetName) + _ = render.Render(w, r, newErrResponse(err, http.StatusNotFound)) return } render.JSON(w, r, toolset.Manifest) diff --git a/internal/server/api_test.go b/internal/server/api_test.go index 35dbee665..0299412fe 100644 --- a/internal/server/api_test.go +++ b/internal/server/api_test.go @@ -20,8 +20,10 @@ import ( "io" "net/http" "net/http/httptest" + "os" "testing" + "github.com/googleapis/genai-toolbox/internal/log" "github.com/googleapis/genai-toolbox/internal/tools" ) @@ -78,7 +80,11 @@ func TestToolsetEndpoint(t *testing.T) { toolsets[name] = m } - server := Server{conf: ServerConfig{}, tools: toolsMap, toolsets: toolsets} + testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info") + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + server := Server{conf: ServerConfig{}, logger: testLogger, tools: toolsMap, toolsets: toolsets} r, err := apiRouter(&server) if err != nil { t.Fatalf("unable to initialize router: %s", err) @@ -189,7 +195,11 @@ func TestToolGetEndpoint(t *testing.T) { } toolsMap := map[string]tools.Tool{tool1.Name: tool1, tool2.Name: tool2} - server := Server{conf: ServerConfig{Version: "0.0.0"}, tools: toolsMap} + testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info") + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + server := Server{conf: ServerConfig{Version: "0.0.0"}, logger: testLogger, tools: toolsMap} r, err := apiRouter(&server) if err != nil { t.Fatalf("unable to initialize router: %s", err) diff --git a/internal/server/config.go b/internal/server/config.go index effb244ec..1730a47cd 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -15,6 +15,7 @@ package server import ( "fmt" + "strings" "github.com/googleapis/genai-toolbox/internal/sources" alloydbpgsrc "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" @@ -38,6 +39,62 @@ type ServerConfig struct { ToolConfigs ToolConfigs // ToolsetConfigs defines what tools are available. ToolsetConfigs ToolsetConfigs + // LoggingFormat defines whether structured loggings are used. + LoggingFormat logFormat + // LogLevel defines the levels to log + LogLevel StringLevel +} + +type logFormat string + +// String is used by both fmt.Print and by Cobra in help text +func (f *logFormat) String() string { + if string(*f) != "" { + return strings.ToLower(string(*f)) + } + return "standard" +} + +// validate logging format flag +func (f *logFormat) Set(v string) error { + switch strings.ToLower(v) { + case "standard", "json": + *f = logFormat(v) + return nil + default: + return fmt.Errorf(`log format must be one of "standard", or "json"`) + } +} + +// Type is used in Cobra help text +func (f *logFormat) Type() string { + return "logFormat" +} + +type StringLevel string + +// String is used by both fmt.Print and by Cobra in help text +func (s *StringLevel) String() string { + if string(*s) != "" { + return strings.ToLower(string(*s)) + } + return "info" +} + +// validate log level flag +func (s *StringLevel) Set(v string) error { + switch strings.ToLower(v) { + case "debug", "info", "warn", "error": + *s = StringLevel(v) + return nil + default: + return fmt.Errorf(`log level must be one of "debug", "info", "warn", or "error"`) + } +} + +// Type is used in Cobra help text +func (s *StringLevel) Type() string { + return "stringLevel" } // SourceConfigs is a type used to allow unmarshal of the data source config map diff --git a/internal/server/server.go b/internal/server/server.go index 8932b6e44..301cd7174 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -24,14 +24,17 @@ import ( "github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5/middleware" + "github.com/go-chi/httplog/v2" + logLib "github.com/googleapis/genai-toolbox/internal/log" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/googleapis/genai-toolbox/internal/tools" ) // Server contains info for running an instance of Toolbox. Should be instantiated with NewServer(). type Server struct { - conf ServerConfig - root chi.Router + conf ServerConfig + root chi.Router + logger logLib.Logger sources map[string]sources.Source tools map[string]tools.Tool @@ -39,9 +42,36 @@ type Server struct { } // NewServer returns a Server object based on provided Config. -func NewServer(cfg ServerConfig) (*Server, error) { +func NewServer(cfg ServerConfig, log logLib.Logger) (*Server, error) { + logLevel, err := logLib.SeverityToLevel(cfg.LogLevel.String()) + if err != nil { + return nil, fmt.Errorf("unable to initialize http log: %w", err) + } + var httpOpts httplog.Options + switch cfg.LoggingFormat.String() { + case "json": + httpOpts = httplog.Options{ + JSON: true, + LogLevel: logLevel, + Concise: true, + RequestHeaders: true, + MessageFieldName: "message", + SourceFieldName: "logging.googleapis.com/sourceLocation", + TimeFieldName: "timestamp", + LevelFieldName: "severity", + } + default: + httpOpts = httplog.Options{ + LogLevel: logLevel, + Concise: true, + RequestHeaders: true, + MessageFieldName: "message", + } + } + + logger := httplog.NewLogger("httplog", httpOpts) r := chi.NewRouter() - r.Use(middleware.Logger) + r.Use(httplog.RequestLogger(logger)) r.Use(middleware.Recoverer) r.Get("/", func(w http.ResponseWriter, r *http.Request) { @@ -57,7 +87,7 @@ func NewServer(cfg ServerConfig) (*Server, error) { } sourcesMap[name] = s } - fmt.Printf("Initalized %d sources.\n", len(sourcesMap)) + log.Info(fmt.Sprintf("Initalized %d sources.", len(sourcesMap))) // initalize and validate the tools toolsMap := make(map[string]tools.Tool) @@ -68,7 +98,7 @@ func NewServer(cfg ServerConfig) (*Server, error) { } toolsMap[name] = t } - fmt.Printf("Initalized %d tools.\n", len(toolsMap)) + log.Info(fmt.Sprintf("Initalized %d tools.", len(toolsMap))) // create a default toolset that contains all tools allToolNames := make([]string, 0, len(toolsMap)) @@ -88,11 +118,12 @@ func NewServer(cfg ServerConfig) (*Server, error) { } toolsetsMap[name] = t } - fmt.Printf("Initalized %d toolsets.\n", len(toolsetsMap)) + log.Info(fmt.Sprintf("Initalized %d toolsets.", len(toolsetsMap))) s := &Server{ conf: cfg, root: r, + logger: log, sources: sourcesMap, tools: toolsMap, toolsets: toolsetsMap, diff --git a/internal/server/server_test.go b/internal/server/server_test.go index f59158c05..2f12cdc79 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -17,10 +17,12 @@ package server_test import ( "context" "net" + "os" "strconv" "testing" "time" + "github.com/googleapis/genai-toolbox/internal/log" "github.com/googleapis/genai-toolbox/internal/server" ) @@ -47,9 +49,15 @@ func TestServe(t *testing.T) { Address: addr, Port: port, } - s, err := server.NewServer(cfg) + + testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info") + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + + s, err := server.NewServer(cfg, testLogger) if err != nil { - t.Fatalf("Unable to initialize server! %v", err) + t.Fatalf("unable to initialize server! %v", err) } // start server in background @@ -63,7 +71,7 @@ func TestServe(t *testing.T) { }() if !tryDial(net.JoinHostPort(addr, strconv.Itoa(port)), 10) { - t.Fatalf("Unable to dial server!") + t.Fatalf("unable to dial server!") } }