feat: opentelemetry tracing (getzep#271)

* otel implementation

* error if collector endpoint not configured

* refacvtor http calls; add telemetry; refactor tests

* reverse out openai timeouts

.cursorignore

@@ -0,0 +1 @@
+test_data

cmd/zep/run.go

@@ -22,12 +22,19 @@ import (
 	"github.com/getzep/zep/pkg/llms"
 	"github.com/getzep/zep/pkg/models"
 	"github.com/getzep/zep/pkg/server"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
+	"go.opentelemetry.io/otel/sdk/resource"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 )
 
 const (
-	ErrStoreTypeNotSet   = "store.type must be set"
-	ErrPostgresDSNNotSet = "store.postgres.dsn must be set"
-	StoreTypePostgres    = "postgres"
+	ErrStoreTypeNotSet             = "store.type must be set"
+	ErrPostgresDSNNotSet           = "store.postgres.dsn must be set"
+	ErrOtelEnabledButExporterEmpty = "OpenTelemtry is enabled but OTEL_EXPORTER_OTLP_ENDPOINT is not set"
+	StoreTypePostgres              = "postgres"
 )
 
 // run is the entrypoint for the zep server
@@ -44,6 +51,16 @@ func run() {
 	config.SetLogLevel(cfg)
 	appState := NewAppState(cfg)
 
+	if cfg.OpenTelemetry.Enabled {
+		cleanup := initTracer()
+		defer func() {
+			err := cleanup(context.Background())
+			if err != nil {
+				log.Errorf("Failed to cleanup tracer: %v", err)
+			}
+		}()
+	}
+
 	srv := server.Create(appState)
 
 	log.Infof("Listening on: %s", srv.Addr)
@@ -52,7 +69,7 @@ func run() {
 	}
 	err = srv.ListenAndServe()
 	if err != nil {
-		log.Fatal(err)
+		log.Panic(err)
 	}
 }
 
@@ -226,3 +243,36 @@ func dumpConfigToJSON(cfg *config.Config) string {
 
 	return string(b)
 }
+
+func initTracer() func(context.Context) error {
+	if os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT") == "" {
+		log.Fatal(ErrOtelEnabledButExporterEmpty)
+	}
+	exporter, err := otlptrace.New(
+		context.Background(),
+		otlptracehttp.NewClient(),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	resources, err := resource.New(context.Background(),
+		resource.WithFromEnv(),
+		resource.WithProcess(),
+		resource.WithOS(),
+		resource.WithContainer(),
+		resource.WithHost(),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	otel.SetTracerProvider(
+		sdktrace.NewTracerProvider(
+			sdktrace.WithSampler(sdktrace.AlwaysSample()),
+			sdktrace.WithBatcher(exporter),
+			sdktrace.WithResource(resources),
+		),
+	)
+	return exporter.Shutdown
+}

config.yaml

@@ -76,6 +76,10 @@ data:
   purge_every: 60
 log:
   level: "info"
+opentelemetry:
+  enabled: false
+  endpoint:
+  attributes: {}
 # Custom Prompts Configuration
 # Allows customization of extractor prompts.
 custom_prompts:

config/models.go

@@ -11,6 +11,7 @@ type Config struct {
 	Server        ServerConfig        `mapstructure:"server"`
 	Log           LogConfig           `mapstructure:"log"`
 	Auth          AuthConfig          `mapstructure:"auth"`
+	OpenTelemetry OpenTelemetryConfig `mapstructure:"opentelemetry"`
 	DataConfig    DataConfig          `mapstructure:"data"`
 	Development   bool                `mapstructure:"development"`
 	CustomPrompts CustomPromptsConfig `mapstructure:"custom_prompts"`
@@ -66,6 +67,10 @@ type LogConfig struct {
 	Level string `mapstructure:"level"`
 }
 
+type OpenTelemetryConfig struct {
+	Enabled bool `mapstructure:"enabled"`
+}
+
 type AuthConfig struct {
 	Secret   string `mapstructure:"secret"`
 	Required bool   `mapstructure:"required"`

go.mod

@@ -6,7 +6,6 @@ toolchain go1.21.3
 
 require (
 	dario.cat/mergo v1.0.0
-	github.com/avast/retry-go/v4 v4.5.0
 	github.com/brianvoe/gofakeit/v6 v6.23.2
 	github.com/chi-middleware/logrus-logger v0.2.0
 	github.com/go-chi/chi/v5 v5.0.10
@@ -40,11 +39,20 @@ require (
 	github.com/invopop/jsonschema v0.12.0
 	github.com/jackc/pgx/v5 v5.4.3
 	github.com/ma-hartma/watermill-logrus-adapter v0.0.0-20220319171828-0856b297f1c2
-	github.com/sony/gobreaker v0.5.0
+	github.com/riandyrn/otelchi v0.5.1
 	github.com/tmc/langchaingo v0.0.0-20230929160525-e16b77704b8d
 	github.com/uptrace/bun/dbfixture v1.1.16
 	github.com/uptrace/bun/extra/bundebug v1.1.16
+	github.com/uptrace/bun/extra/bunotel v1.1.16
 	github.com/viterin/vek v0.4.2
+	github.com/voi-oss/watermill-opentelemetry v0.1.3
+	go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.46.0
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.46.0
+	go.opentelemetry.io/otel v1.20.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0
+	go.opentelemetry.io/otel/sdk v1.20.0
+	go.opentelemetry.io/otel/trace v1.20.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -54,20 +62,26 @@ require (
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/cenkalti/backoff/v3 v3.2.2 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/chewxy/math32 v1.10.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.2.0 // indirect
 	github.com/dlclark/regexp2 v1.10.0 // indirect
 	github.com/fatih/color v1.15.0 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
+	github.com/go-logr/logr v1.3.0 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-openapi/jsonpointer v0.20.0 // indirect
 	github.com/go-openapi/jsonreference v0.20.2 // indirect
 	github.com/go-openapi/spec v0.20.9 // indirect
 	github.com/go-openapi/swag v0.22.4 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
@@ -99,23 +113,32 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
 	github.com/shopspring/decimal v1.3.1 // indirect
+	github.com/sony/gobreaker v0.5.0 // indirect
 	github.com/spf13/afero v1.10.0 // indirect
 	github.com/spf13/cast v1.5.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/sv-tools/openapi v0.2.2 // indirect
 	github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect
+	github.com/uptrace/opentelemetry-go-extra/otelsql v0.2.2 // indirect
 	github.com/viterin/partial v1.1.0 // indirect
 	github.com/vmihailenco/msgpack/v5 v5.3.5 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
+	go.opentelemetry.io/contrib v1.0.0 // indirect
+	go.opentelemetry.io/otel/metric v1.20.0 // indirect
+	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
 	golang.org/x/crypto v0.14.0 // indirect
 	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
 	golang.org/x/net v0.17.0 // indirect
-	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/sys v0.14.0 // indirect
 	golang.org/x/text v0.13.0 // indirect
 	golang.org/x/tools v0.14.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20230711160842-782d3b101e98 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20230711160842-782d3b101e98 // indirect
+	google.golang.org/grpc v1.58.2 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	mellium.im/sasl v0.3.1 // indirect