feat: add --log-level and --logging-format flags (#97)

Logging support 4 different types of logging (debug, info, warn, error).
The default logging level is Info.

User will be able to set flag for log level (allowed values: "debug",
"info", "warn", "error"), example:
`go run . --log-level debug`

User will be able to set flag for logging format (allowed values:
"standard", "JSON"), example:
`go run . --logging-format json`

**sample http request log - std:**
server
```
2024-11-12T15:08:11.451377-08:00 INFO "Initalized 0 sources.\n"
```
httplog
```
2024-11-26T15:15:53.947287-08:00 INFO Response: 200 OK service: "httplog" httpRequest: {url: "http://127.0.0.1:5000/" method: "GET" path: "/" remoteIP: "127.0.0.1:64216" proto: "HTTP/1.1" requestID: "macbookpro.roam.interna/..." scheme: "http" header: {user-agent: "curl/8.7.1" accept: "*/*"}} httpResponse: {status: 200 bytes: 22 elapsed: 0.012417}
```

**sample http request log - structured:**
server
```
{
  "timestamp":"2024-11-04T16:45:11.987299-08:00",
  "severity":"ERROR",
  "logging.googleapis.com/sourceLocation":{
    "function":"github.com/googleapis/genai-toolbox/internal/log.(*StructuredLogger).Errorf",
    "file":"/Users/yuanteoh/github/genai-toolbox/internal/log/log.go","line":157
  },
  "message":"unable to parse tool file at \"tools.yaml\": \"cloud-sql-postgres1\" is not a valid kind of data source"
}
```
httplog
```
{
  "timestamp":"2024-11-26T15:12:49.290974-08:00",
  "severity":"INFO",
  "logging.googleapis.com/sourceLocation":{
      "function":"github.com/go-chi/httplog/v2.(*RequestLoggerEntry).Write",
      "file":"/Users/yuanteoh/go/pkg/mod/github.com/go-chi/httplog/v2@v2.1.1/httplog.go","line":173
  },
  "message":"Response: 200 OK",
  "service":"httplog",
  "httpRequest":{
      "url":"http://127.0.0.1:5000/",
      "method":"GET",
      "path":"/",
      "remoteIP":"127.0.0.1:64140",
      "proto":"HTTP/1.1",
      "requestID":"yuanteoh-macbookpro.roam.internal/NBrtYBu3q9-000001",
      "scheme":"http",
      "header":{"user-agent":"curl/8.7.1","accept":"*/*"}
  },
  "httpResponse":{"status":200,"bytes":22,"elapsed":0.0115}
}
```

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

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

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

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=

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
 	}

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")
 	}

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)

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)

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

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,

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("Unable to initialize server! %v", err)
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	s, err := server.NewServer(cfg, testLogger)
+	if err != nil {
+		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!")
 	}
 
 }