Skip to content

Commit

Permalink
feat: add --log-level and --logging-format flags (#97)
Browse files Browse the repository at this point in the history
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}
}
```
  • Loading branch information
Yuan325 authored Dec 3, 2024
1 parent bf614ed commit 9a0f618
Show file tree
Hide file tree
Showing 11 changed files with 249 additions and 45 deletions.
45 changes: 37 additions & 8 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -61,23 +62,27 @@ type Command struct {
*cobra.Command

cfg server.ServerConfig
logger log.Logger
tools_file string
}

// NewCommand returns a Command object representing an invocation of the CLI.
func NewCommand() *Command {
cmd := &Command{
Command: &cobra.Command{
Use: "toolbox",
Version: versionString,
Use: "toolbox",
Version: versionString,
SilenceErrors: true,
},
}

flags := cmd.Flags()
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) }
Expand All @@ -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
Expand Down
103 changes: 84 additions & 19 deletions cmd/root_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
6 changes: 3 additions & 3 deletions internal/log/log.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down
4 changes: 2 additions & 2 deletions internal/log/log_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand All @@ -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")
}
Expand Down
3 changes: 2 additions & 1 deletion internal/server/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
14 changes: 12 additions & 2 deletions internal/server/api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading

0 comments on commit 9a0f618

Please sign in to comment.