From be3172c2cf0670db16e9a8724886f1902dbd3866 Mon Sep 17 00:00:00 2001 From: pmady Date: Tue, 28 Apr 2026 16:25:09 -0500 Subject: [PATCH 1/6] go: initial traceai core module Scaffold the Go SDK with OTel TracerProvider setup, GenAI semantic convention keys, and the Instrumentor interface. Signed-off-by: pmady --- go/traceai/config.go | 205 +++++++++++++++++++++++++++++++++++++ go/traceai/go.mod | 30 ++++++ go/traceai/go.sum | 59 +++++++++++ go/traceai/instrumentor.go | 17 +++ go/traceai/semconv.go | 39 +++++++ 5 files changed, 350 insertions(+) create mode 100644 go/traceai/config.go create mode 100644 go/traceai/go.mod create mode 100644 go/traceai/go.sum create mode 100644 go/traceai/instrumentor.go create mode 100644 go/traceai/semconv.go diff --git a/go/traceai/config.go b/go/traceai/config.go new file mode 100644 index 0000000..036f2e4 --- /dev/null +++ b/go/traceai/config.go @@ -0,0 +1,205 @@ +// Package traceai sets up OTel tracing with Future AGI defaults. +package traceai + +import ( + "context" + "fmt" + "os" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.26.0" + "go.opentelemetry.io/otel/trace" +) + +const ( + envAPIKey = "FI_API_KEY" + envSecretKey = "FI_SECRET_KEY" + envBaseURL = "FI_BASE_URL" + envGRPCURL = "FI_GRPC_URL" + envProjectName = "FI_PROJECT_NAME" + + defaultBaseURL = "https://api.futureagi.com" + defaultGRPCURL = "https://grpc.futureagi.com" + defaultProjectName = "DEFAULT_PROJECT_NAME" + + headerAPIKey = "X-Api-Key" + headerSecretKey = "X-Secret-Key" +) + +// Transport specifies the OTLP export protocol. +type Transport int + +const ( + TransportHTTP Transport = iota + TransportGRPC +) + +// Config holds the options for initializing a traceAI TracerProvider. +type Config struct { + ProjectName string + Transport Transport + BaseURL string + GRPCURL string + APIKey string + SecretKey string + BatchExport bool + SetGlobal bool + ShutdownTimeout time.Duration +} + +// DefaultConfig returns a Config populated from environment variables. +func DefaultConfig() Config { + return Config{ + ProjectName: envOrDefault(envProjectName, defaultProjectName), + Transport: TransportHTTP, + BaseURL: envOrDefault(envBaseURL, defaultBaseURL), + GRPCURL: envOrDefault(envGRPCURL, defaultGRPCURL), + APIKey: os.Getenv(envAPIKey), + SecretKey: os.Getenv(envSecretKey), + BatchExport: true, + SetGlobal: true, + ShutdownTimeout: 10 * time.Second, + } +} + +type Option func(*Config) + +func WithProjectName(name string) Option { + return func(c *Config) { c.ProjectName = name } +} + +func WithTransport(t Transport) Option { + return func(c *Config) { c.Transport = t } +} + +func WithBaseURL(url string) Option { + return func(c *Config) { c.BaseURL = url } +} + +func WithCredentials(apiKey, secretKey string) Option { + return func(c *Config) { + c.APIKey = apiKey + c.SecretKey = secretKey + } +} + +func WithBatchExport(batch bool) Option { + return func(c *Config) { c.BatchExport = batch } +} + +func WithSetGlobal(global bool) Option { + return func(c *Config) { c.SetGlobal = global } +} + +func WithShutdownTimeout(d time.Duration) Option { + return func(c *Config) { c.ShutdownTimeout = d } +} + +// Provider wraps a TracerProvider with shutdown handling. +type Provider struct { + tp *sdktrace.TracerProvider + shutdownTimeout time.Duration +} + +// Register creates and configures a new traceAI TracerProvider. +// It returns a Provider whose Shutdown method should be called on process exit. +// +// provider, err := traceai.Register( +// traceai.WithProjectName("my-llm-service"), +// ) +// if err != nil { +// log.Fatal(err) +// } +// defer provider.Shutdown(context.Background()) +func Register(opts ...Option) (*Provider, error) { + cfg := DefaultConfig() + for _, o := range opts { + o(&cfg) + } + + ctx := context.Background() + + exporter, err := newExporter(ctx, cfg) + if err != nil { + return nil, fmt.Errorf("traceai: create exporter: %w", err) + } + + res, err := resource.Merge( + resource.Default(), + resource.NewWithAttributes( + semconv.SchemaURL, + semconv.ServiceName(cfg.ProjectName), + ), + ) + if err != nil { + return nil, fmt.Errorf("traceai: create resource: %w", err) + } + + var spanProcessor sdktrace.SpanProcessor + if cfg.BatchExport { + spanProcessor = sdktrace.NewBatchSpanProcessor(exporter) + } else { + spanProcessor = sdktrace.NewSimpleSpanProcessor(exporter) + } + + tp := sdktrace.NewTracerProvider( + sdktrace.WithResource(res), + sdktrace.WithSpanProcessor(spanProcessor), + ) + + if cfg.SetGlobal { + otel.SetTracerProvider(tp) + } + + return &Provider{ + tp: tp, + shutdownTimeout: cfg.ShutdownTimeout, + }, nil +} + +func (p *Provider) TracerProvider() trace.TracerProvider { + return p.tp +} + +// Shutdown flushes pending spans and shuts down the exporter. +func (p *Provider) Shutdown(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, p.shutdownTimeout) + defer cancel() + return p.tp.Shutdown(ctx) +} + +func newExporter(ctx context.Context, cfg Config) (sdktrace.SpanExporter, error) { + headers := make(map[string]string) + if cfg.APIKey != "" { + headers[headerAPIKey] = cfg.APIKey + } + if cfg.SecretKey != "" { + headers[headerSecretKey] = cfg.SecretKey + } + + switch cfg.Transport { + case TransportGRPC: + return otlptracegrpc.New(ctx, + otlptracegrpc.WithEndpoint(cfg.GRPCURL), + otlptracegrpc.WithHeaders(headers), + ) + default: + endpoint := cfg.BaseURL + "/tracer/v1/traces" + return otlptracehttp.New(ctx, + otlptracehttp.WithEndpointURL(endpoint), + otlptracehttp.WithHeaders(headers), + ) + } +} + +func envOrDefault(key, fallback string) string { + if v := os.Getenv(key); v != "" { + return v + } + return fallback +} diff --git a/go/traceai/go.mod b/go/traceai/go.mod new file mode 100644 index 0000000..3226a1a --- /dev/null +++ b/go/traceai/go.mod @@ -0,0 +1,30 @@ +module github.com/future-agi/traceAI/go/traceai + +go 1.22.0 + +require ( + go.opentelemetry.io/otel v1.35.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 + go.opentelemetry.io/otel/sdk v1.35.0 + go.opentelemetry.io/otel/trace v1.35.0 +) + +require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect + go.opentelemetry.io/auto/sdk v1.1.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 // indirect + go.opentelemetry.io/otel/metric v1.35.0 // indirect + go.opentelemetry.io/proto/otlp v1.5.0 // indirect + golang.org/x/net v0.35.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.22.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a // indirect + google.golang.org/grpc v1.71.0 // indirect + google.golang.org/protobuf v1.36.5 // indirect +) diff --git a/go/traceai/go.sum b/go/traceai/go.sum new file mode 100644 index 0000000..d7f8143 --- /dev/null +++ b/go/traceai/go.sum @@ -0,0 +1,59 @@ +github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8= +github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= +github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 h1:e9Rjr40Z98/clHv5Yg79Is0NtosR5LXRvdr7o/6NwbA= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1/go.mod h1:tIxuGz/9mpox++sgp9fJjHO0+q1X9/UOWd798aAm22M= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ= +go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 h1:1fTNlAIJZGWLP5FVu0fikVry1IsiUnXjf7QFvoNN3Xw= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0/go.mod h1:zjPK58DtkqQFn+YUMbx0M2XV3QgKU0gS9LeGohREyK4= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0 h1:m639+BofXTvcY1q8CGs4ItwQarYtJPOWmVobfM1HpVI= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0/go.mod h1:LjReUci/F4BUyv+y4dwnq3h/26iNOeC3wAIqgvTIZVo= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 h1:xJ2qHD0C1BeYVTLLR9sX12+Qb95kfeD/byKj6Ky1pXg= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0/go.mod h1:u5BF1xyjstDowA1R5QAO9JHzqK+ublenEW/dyqTjBVk= +go.opentelemetry.io/otel/metric v1.35.0 h1:0znxYu2SNyuMSQT4Y9WDWej0VpcsxkuklLa4/siN90M= +go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJGaE8CpgeAU3oE= +go.opentelemetry.io/otel/sdk v1.35.0 h1:iPctf8iprVySXSKJffSS79eOjl9pvxV9ZqOWT0QejKY= +go.opentelemetry.io/otel/sdk v1.35.0/go.mod h1:+ga1bZliga3DxJ3CQGg3updiaAJoNECOgJREo9KHGQg= +go.opentelemetry.io/otel/sdk/metric v1.34.0 h1:5CeK9ujjbFVL5c1PhLuStg1wxA7vQv7ce1EK0Gyvahk= +go.opentelemetry.io/otel/sdk/metric v1.34.0/go.mod h1:jQ/r8Ze28zRKoNRdkjCZxfs6YvBTG1+YIqyFVFYec5w= +go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs= +go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc= +go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU/3i4= +go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8= +golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= +golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= +google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a h1:nwKuGPlUAt+aR+pcrkfFRrTU1BVrSmYyYMxYbUIVHr0= +google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a/go.mod h1:3kWAYMk1I75K4vykHtKt2ycnOgpA6974V7bREqbsenU= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a h1:51aaUVRocpvUOSQKM6Q7VuoaktNIaMCLuhZB6DKksq4= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a/go.mod h1:uRxBH1mhmO8PGhU89cMcHaXKZqO+OfakD8QQO0oYwlQ= +google.golang.org/grpc v1.71.0 h1:kF77BGdPTQ4/JZWMlb9VpJ5pa25aqvVqogsxNHHdeBg= +google.golang.org/grpc v1.71.0/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec= +google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM= +google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/go/traceai/instrumentor.go b/go/traceai/instrumentor.go new file mode 100644 index 0000000..cd0d3b5 --- /dev/null +++ b/go/traceai/instrumentor.go @@ -0,0 +1,17 @@ +package traceai + +import "go.opentelemetry.io/otel/trace" + +// Instrumentor wraps an AI framework client with OTel tracing. +type Instrumentor interface { + Instrument(tp trace.TracerProvider) error + Uninstrument() error +} + +func Tracer(tp trace.TracerProvider, name string, opts ...trace.TracerOption) trace.Tracer { + if tp == nil { + // fall back to the global provider set by Register() + return trace.NewNoopTracerProvider().Tracer(name) + } + return tp.Tracer(name, opts...) +} diff --git a/go/traceai/semconv.go b/go/traceai/semconv.go new file mode 100644 index 0000000..62a77bc --- /dev/null +++ b/go/traceai/semconv.go @@ -0,0 +1,39 @@ +package traceai + +import "go.opentelemetry.io/otel/attribute" + +// GenAI semantic convention keys. +// See https://opentelemetry.io/docs/specs/semconv/gen-ai/ +const ( + AttrGenAISystem = attribute.Key("gen_ai.system") + AttrGenAIRequestModel = attribute.Key("gen_ai.request.model") + AttrGenAIResponseModel = attribute.Key("gen_ai.response.model") + AttrGenAIOperationName = attribute.Key("gen_ai.operation.name") + + AttrGenAIRequestMaxTokens = attribute.Key("gen_ai.request.max_tokens") + AttrGenAIRequestTemperature = attribute.Key("gen_ai.request.temperature") + AttrGenAIRequestTopP = attribute.Key("gen_ai.request.top_p") + AttrGenAIRequestStopSequences = attribute.Key("gen_ai.request.stop_sequences") + + AttrGenAIUsageInputTokens = attribute.Key("gen_ai.usage.input_tokens") + AttrGenAIUsageOutputTokens = attribute.Key("gen_ai.usage.output_tokens") + + AttrGenAIResponseFinishReasons = attribute.Key("gen_ai.response.finish_reasons") + AttrGenAIResponseID = attribute.Key("gen_ai.response.id") + + AttrGenAIPrompt = attribute.Key("gen_ai.prompt") + AttrGenAICompletion = attribute.Key("gen_ai.completion") +) + +const ( + GenAISystemOpenAI = "openai" + GenAISystemAnthropic = "anthropic" + GenAISystemCohere = "cohere" + GenAISystemGoogle = "google" +) + +const ( + OpChat = "chat" + OpCompletion = "completion" + OpEmbedding = "embedding" +) From 72c200688f200b995ad38091a532b0dd421c4b82 Mon Sep 17 00:00:00 2001 From: pmady Date: Fri, 1 May 2026 11:10:15 -0500 Subject: [PATCH 2/6] go: openai instrumentor middleware HTTP middleware that intercepts chat/completion/embedding calls and records them as OTel spans with GenAI semconv attributes. Signed-off-by: pmady --- go/traceai_openai/openai.go | 203 +++++++++++++++++++++++++++++++++++ go/traceai_openai/options.go | 26 +++++ 2 files changed, 229 insertions(+) create mode 100644 go/traceai_openai/openai.go create mode 100644 go/traceai_openai/options.go diff --git a/go/traceai_openai/openai.go b/go/traceai_openai/openai.go new file mode 100644 index 0000000..2aa407c --- /dev/null +++ b/go/traceai_openai/openai.go @@ -0,0 +1,203 @@ +// Package traceai_openai instruments the OpenAI Go client with OTel tracing. +// +// Usage: +// +// provider, _ := traceai.Register() +// defer provider.Shutdown(context.Background()) +// +// client := openai.NewClient( +// option.WithMiddleware(traceai_openai.Middleware()), +// ) +package traceai_openai + +import ( + "encoding/json" + "io" + "net/http" + "strings" + "time" + + "github.com/future-agi/traceAI/go/traceai" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" +) + +const instrumentationName = "traceai_openai" +const instrumentationVersion = "0.1.0" + +// Middleware returns an openai-go compatible middleware that records +// LLM calls as OTel spans. +func Middleware(opts ...Option) func(req *http.Request, next func(req *http.Request) (*http.Response, error)) (*http.Response, error) { + cfg := defaultOptions() + for _, o := range opts { + o(&cfg) + } + + tp := cfg.tracerProvider + if tp == nil { + tp = otel.GetTracerProvider() + } + + tracer := tp.Tracer( + instrumentationName, + trace.WithInstrumentationVersion(instrumentationVersion), + ) + + return func(req *http.Request, next func(req *http.Request) (*http.Response, error)) (*http.Response, error) { + opName := operationFromPath(req.URL.Path) + if opName == "" { + return next(req) + } + + spanName := "openai." + opName + ctx, span := tracer.Start(req.Context(), spanName, + trace.WithSpanKind(trace.SpanKindClient), + ) + defer span.End() + + span.SetAttributes( + traceai.AttrGenAISystem.String(traceai.GenAISystemOpenAI), + traceai.AttrGenAIOperationName.String(opName), + ) + + if req.Body != nil && cfg.captureContent { + bodyBytes, err := io.ReadAll(req.Body) + if err == nil { + req.Body = io.NopCloser(strings.NewReader(string(bodyBytes))) + extractRequestAttributes(span, bodyBytes) + } + } + + req = req.WithContext(ctx) + start := time.Now() + + resp, err := next(req) + + span.SetAttributes(attribute.Float64("gen_ai.request.duration_ms", + float64(time.Since(start).Milliseconds()))) + + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return resp, err + } + + if resp != nil && resp.StatusCode >= 400 { + span.SetStatus(codes.Error, resp.Status) + span.SetAttributes(attribute.Int("http.response.status_code", resp.StatusCode)) + return resp, err + } + + // read response for token counts + if resp != nil && resp.Body != nil && cfg.captureContent { + bodyBytes, readErr := io.ReadAll(resp.Body) + if readErr == nil { + resp.Body = io.NopCloser(strings.NewReader(string(bodyBytes))) + extractResponseAttributes(span, bodyBytes) + } + } + + span.SetStatus(codes.Ok, "") + return resp, err + } +} + +func operationFromPath(path string) string { + switch { + case strings.Contains(path, "/chat/completions"): + return traceai.OpChat + case strings.Contains(path, "/completions"): + return traceai.OpCompletion + case strings.Contains(path, "/embeddings"): + return traceai.OpEmbedding + default: + return "" + } +} + +type chatRequest struct { + Model string `json:"model"` + Messages json.RawMessage `json:"messages"` + MaxTokens *int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` +} + +func extractRequestAttributes(span trace.Span, body []byte) { + var req chatRequest + if err := json.Unmarshal(body, &req); err != nil { + return + } + + if req.Model != "" { + span.SetAttributes(traceai.AttrGenAIRequestModel.String(req.Model)) + } + if req.MaxTokens != nil { + span.SetAttributes(traceai.AttrGenAIRequestMaxTokens.Int(*req.MaxTokens)) + } + if req.Temperature != nil { + span.SetAttributes(traceai.AttrGenAIRequestTemperature.Float64(*req.Temperature)) + } + if req.TopP != nil { + span.SetAttributes(traceai.AttrGenAIRequestTopP.Float64(*req.TopP)) + } + if len(req.Messages) > 0 { + span.SetAttributes(traceai.AttrGenAIPrompt.String(string(req.Messages))) + } +} + +type chatResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + } `json:"usage"` +} + +func extractResponseAttributes(span trace.Span, body []byte) { + var resp chatResponse + if err := json.Unmarshal(body, &resp); err != nil { + return + } + + if resp.ID != "" { + span.SetAttributes(traceai.AttrGenAIResponseID.String(resp.ID)) + } + if resp.Model != "" { + span.SetAttributes(traceai.AttrGenAIResponseModel.String(resp.Model)) + } + if resp.Usage.PromptTokens > 0 { + span.SetAttributes(traceai.AttrGenAIUsageInputTokens.Int(resp.Usage.PromptTokens)) + } + if resp.Usage.CompletionTokens > 0 { + span.SetAttributes(traceai.AttrGenAIUsageOutputTokens.Int(resp.Usage.CompletionTokens)) + } + + if len(resp.Choices) > 0 { + reasons := make([]string, 0, len(resp.Choices)) + var completions []string + for _, c := range resp.Choices { + if c.FinishReason != "" { + reasons = append(reasons, c.FinishReason) + } + if c.Message.Content != "" { + completions = append(completions, c.Message.Content) + } + } + if len(reasons) > 0 { + span.SetAttributes(traceai.AttrGenAIResponseFinishReasons.StringSlice(reasons)) + } + if len(completions) > 0 { + span.SetAttributes(traceai.AttrGenAICompletion.String(strings.Join(completions, "\n"))) + } + } +} diff --git a/go/traceai_openai/options.go b/go/traceai_openai/options.go new file mode 100644 index 0000000..6bc80a1 --- /dev/null +++ b/go/traceai_openai/options.go @@ -0,0 +1,26 @@ +package traceai_openai + +import "go.opentelemetry.io/otel/trace" + +type options struct { + tracerProvider trace.TracerProvider + captureContent bool +} + +func defaultOptions() options { + return options{ + captureContent: true, + } +} + +type Option func(*options) + +// WithTracerProvider sets a specific TracerProvider instead of the global one. +func WithTracerProvider(tp trace.TracerProvider) Option { + return func(o *options) { o.tracerProvider = tp } +} + +// WithContentCapture toggles recording prompt/completion content. Default true. +func WithContentCapture(capture bool) Option { + return func(o *options) { o.captureContent = capture } +} From 574e6606f824a8600ee408d28465c3ac4e0a7f5d Mon Sep 17 00:00:00 2001 From: pmady Date: Fri, 1 May 2026 11:10:30 -0500 Subject: [PATCH 3/6] go: add openai middleware tests Covers chat completion happy path, non-AI endpoint passthrough, error response handling, and content capture toggle. Signed-off-by: pmady --- go/traceai_openai/go.mod | 33 ++++++ go/traceai_openai/go.sum | 59 ++++++++++ go/traceai_openai/openai_test.go | 189 +++++++++++++++++++++++++++++++ 3 files changed, 281 insertions(+) create mode 100644 go/traceai_openai/go.mod create mode 100644 go/traceai_openai/go.sum create mode 100644 go/traceai_openai/openai_test.go diff --git a/go/traceai_openai/go.mod b/go/traceai_openai/go.mod new file mode 100644 index 0000000..1c774d2 --- /dev/null +++ b/go/traceai_openai/go.mod @@ -0,0 +1,33 @@ +module github.com/future-agi/traceAI/go/traceai_openai + +go 1.22.0 + +require ( + github.com/future-agi/traceAI/go/traceai v0.0.0 + go.opentelemetry.io/otel v1.35.0 + go.opentelemetry.io/otel/sdk v1.35.0 + go.opentelemetry.io/otel/trace v1.35.0 +) + +require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect + go.opentelemetry.io/auto/sdk v1.1.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 // indirect + go.opentelemetry.io/otel/metric v1.35.0 // indirect + go.opentelemetry.io/proto/otlp v1.5.0 // indirect + golang.org/x/net v0.35.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.22.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a // indirect + google.golang.org/grpc v1.71.0 // indirect + google.golang.org/protobuf v1.36.5 // indirect +) + +replace github.com/future-agi/traceAI/go/traceai => ../traceai diff --git a/go/traceai_openai/go.sum b/go/traceai_openai/go.sum new file mode 100644 index 0000000..d7f8143 --- /dev/null +++ b/go/traceai_openai/go.sum @@ -0,0 +1,59 @@ +github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8= +github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= +github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 h1:e9Rjr40Z98/clHv5Yg79Is0NtosR5LXRvdr7o/6NwbA= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1/go.mod h1:tIxuGz/9mpox++sgp9fJjHO0+q1X9/UOWd798aAm22M= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ= +go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 h1:1fTNlAIJZGWLP5FVu0fikVry1IsiUnXjf7QFvoNN3Xw= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0/go.mod h1:zjPK58DtkqQFn+YUMbx0M2XV3QgKU0gS9LeGohREyK4= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0 h1:m639+BofXTvcY1q8CGs4ItwQarYtJPOWmVobfM1HpVI= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.35.0/go.mod h1:LjReUci/F4BUyv+y4dwnq3h/26iNOeC3wAIqgvTIZVo= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 h1:xJ2qHD0C1BeYVTLLR9sX12+Qb95kfeD/byKj6Ky1pXg= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0/go.mod h1:u5BF1xyjstDowA1R5QAO9JHzqK+ublenEW/dyqTjBVk= +go.opentelemetry.io/otel/metric v1.35.0 h1:0znxYu2SNyuMSQT4Y9WDWej0VpcsxkuklLa4/siN90M= +go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJGaE8CpgeAU3oE= +go.opentelemetry.io/otel/sdk v1.35.0 h1:iPctf8iprVySXSKJffSS79eOjl9pvxV9ZqOWT0QejKY= +go.opentelemetry.io/otel/sdk v1.35.0/go.mod h1:+ga1bZliga3DxJ3CQGg3updiaAJoNECOgJREo9KHGQg= +go.opentelemetry.io/otel/sdk/metric v1.34.0 h1:5CeK9ujjbFVL5c1PhLuStg1wxA7vQv7ce1EK0Gyvahk= +go.opentelemetry.io/otel/sdk/metric v1.34.0/go.mod h1:jQ/r8Ze28zRKoNRdkjCZxfs6YvBTG1+YIqyFVFYec5w= +go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs= +go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc= +go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU/3i4= +go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8= +golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= +golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= +google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a h1:nwKuGPlUAt+aR+pcrkfFRrTU1BVrSmYyYMxYbUIVHr0= +google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a/go.mod h1:3kWAYMk1I75K4vykHtKt2ycnOgpA6974V7bREqbsenU= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a h1:51aaUVRocpvUOSQKM6Q7VuoaktNIaMCLuhZB6DKksq4= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250218202821-56aae31c358a/go.mod h1:uRxBH1mhmO8PGhU89cMcHaXKZqO+OfakD8QQO0oYwlQ= +google.golang.org/grpc v1.71.0 h1:kF77BGdPTQ4/JZWMlb9VpJ5pa25aqvVqogsxNHHdeBg= +google.golang.org/grpc v1.71.0/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec= +google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM= +google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/go/traceai_openai/openai_test.go b/go/traceai_openai/openai_test.go new file mode 100644 index 0000000..ecf0616 --- /dev/null +++ b/go/traceai_openai/openai_test.go @@ -0,0 +1,189 @@ +package traceai_openai + +import ( + "context" + "io" + "net/http" + "strings" + "testing" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" +) + +func setupTracer() (*tracetest.InMemoryExporter, *sdktrace.TracerProvider) { + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + ) + return exporter, tp +} + +func TestMiddleware_ChatCompletion(t *testing.T) { + exporter, tp := setupTracer() + defer tp.Shutdown(context.Background()) + + mw := Middleware(WithTracerProvider(tp)) + + reqBody := `{"model":"gpt-4","messages":[{"role":"user","content":"hello"}],"temperature":0.7}` + respBody := `{"id":"chatcmpl-abc","model":"gpt-4","choices":[{"finish_reason":"stop","message":{"content":"Hi there!"}}],"usage":{"prompt_tokens":5,"completion_tokens":3}}` + + req, _ := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", io.NopCloser(strings.NewReader(reqBody))) + + mockNext := func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(respBody)), + }, nil + } + + resp, err := mw(req, mockNext) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != 200 { + t.Fatalf("expected status 200, got %d", resp.StatusCode) + } + + spans := exporter.GetSpans() + if len(spans) != 1 { + t.Fatalf("expected 1 span, got %d", len(spans)) + } + + span := spans[0] + if span.Name != "openai.chat" { + t.Errorf("expected span name 'openai.chat', got '%s'", span.Name) + } + if span.Status.Code != codes.Ok { + t.Errorf("expected OK status, got %v", span.Status.Code) + } + + assertAttr(t, span.Attributes, "gen_ai.system", "openai") + assertAttr(t, span.Attributes, "gen_ai.request.model", "gpt-4") + assertAttr(t, span.Attributes, "gen_ai.response.model", "gpt-4") + assertAttr(t, span.Attributes, "gen_ai.response.id", "chatcmpl-abc") + assertIntAttr(t, span.Attributes, "gen_ai.usage.input_tokens", 5) + assertIntAttr(t, span.Attributes, "gen_ai.usage.output_tokens", 3) +} + +func TestMiddleware_NonAIEndpoint(t *testing.T) { + exporter, tp := setupTracer() + defer tp.Shutdown(context.Background()) + + mw := Middleware(WithTracerProvider(tp)) + + req, _ := http.NewRequest("GET", "https://api.openai.com/v1/models", nil) + mockNext := func(req *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + } + + _, err := mw(req, mockNext) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + spans := exporter.GetSpans() + if len(spans) != 0 { + t.Errorf("expected 0 spans for non-AI endpoint, got %d", len(spans)) + } +} + +func TestMiddleware_ErrorResponse(t *testing.T) { + exporter, tp := setupTracer() + defer tp.Shutdown(context.Background()) + + mw := Middleware(WithTracerProvider(tp)) + + reqBody := `{"model":"gpt-4","messages":[{"role":"user","content":"test"}]}` + req, _ := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", io.NopCloser(strings.NewReader(reqBody))) + + mockNext := func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: 429, + Status: "429 Too Many Requests", + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }, nil + } + + resp, err := mw(req, mockNext) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != 429 { + t.Fatalf("expected 429, got %d", resp.StatusCode) + } + + spans := exporter.GetSpans() + if len(spans) != 1 { + t.Fatalf("expected 1 span, got %d", len(spans)) + } + + if spans[0].Status.Code != codes.Error { + t.Errorf("expected Error status for 429, got %v", spans[0].Status.Code) + } +} + +func TestMiddleware_ContentCaptureDisabled(t *testing.T) { + exporter, tp := setupTracer() + defer tp.Shutdown(context.Background()) + + mw := Middleware(WithTracerProvider(tp), WithContentCapture(false)) + + reqBody := `{"model":"gpt-4","messages":[{"role":"user","content":"secret stuff"}]}` + respBody := `{"id":"chatcmpl-abc","model":"gpt-4","choices":[{"finish_reason":"stop","message":{"content":"response"}}],"usage":{"prompt_tokens":5,"completion_tokens":3}}` + + req, _ := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", io.NopCloser(strings.NewReader(reqBody))) + mockNext := func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(respBody)), + }, nil + } + + _, err := mw(req, mockNext) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + spans := exporter.GetSpans() + if len(spans) != 1 { + t.Fatalf("expected 1 span, got %d", len(spans)) + } + + // with content capture off, prompt and completion should not be in attributes + for _, a := range spans[0].Attributes { + if string(a.Key) == "gen_ai.prompt" || string(a.Key) == "gen_ai.completion" { + t.Errorf("content capture disabled but found attribute %s", a.Key) + } + } +} + +func assertAttr(t *testing.T, attrs []attribute.KeyValue, key, want string) { + t.Helper() + for _, a := range attrs { + if string(a.Key) == key { + got := a.Value.AsString() + if got != want { + t.Errorf("attr %s: got %q, want %q", key, got, want) + } + return + } + } + t.Errorf("attr %s not found", key) +} + +func assertIntAttr(t *testing.T, attrs []attribute.KeyValue, key string, want int) { + t.Helper() + for _, a := range attrs { + if string(a.Key) == key { + got := a.Value.AsInt64() + if got != int64(want) { + t.Errorf("attr %s: got %d, want %d", key, got, want) + } + return + } + } + t.Errorf("attr %s not found", key) +} From b86cd33118b93ff2cccaab696c6782417dc20582 Mon Sep 17 00:00:00 2001 From: pmady Date: Wed, 6 May 2026 10:43:01 -0500 Subject: [PATCH 4/6] go: simplify passthrough test Signed-off-by: pmady --- go/traceai_openai/openai_test.go | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/go/traceai_openai/openai_test.go b/go/traceai_openai/openai_test.go index ecf0616..bbb1a92 100644 --- a/go/traceai_openai/openai_test.go +++ b/go/traceai_openai/openai_test.go @@ -73,20 +73,16 @@ func TestMiddleware_NonAIEndpoint(t *testing.T) { defer tp.Shutdown(context.Background()) mw := Middleware(WithTracerProvider(tp)) - req, _ := http.NewRequest("GET", "https://api.openai.com/v1/models", nil) - mockNext := func(req *http.Request) (*http.Response, error) { - return &http.Response{StatusCode: 200, Body: http.NoBody}, nil - } - _, err := mw(req, mockNext) + _, err := mw(req, func(r *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) if err != nil { - t.Fatalf("unexpected error: %v", err) + t.Fatal(err) } - - spans := exporter.GetSpans() - if len(spans) != 0 { - t.Errorf("expected 0 spans for non-AI endpoint, got %d", len(spans)) + if n := len(exporter.GetSpans()); n != 0 { + t.Fatalf("got %d spans, want 0", n) } } @@ -144,7 +140,7 @@ func TestMiddleware_ContentCaptureDisabled(t *testing.T) { _, err := mw(req, mockNext) if err != nil { - t.Fatalf("unexpected error: %v", err) + t.Fatal(err) } spans := exporter.GetSpans() @@ -152,7 +148,6 @@ func TestMiddleware_ContentCaptureDisabled(t *testing.T) { t.Fatalf("expected 1 span, got %d", len(spans)) } - // with content capture off, prompt and completion should not be in attributes for _, a := range spans[0].Attributes { if string(a.Key) == "gen_ai.prompt" || string(a.Key) == "gen_ai.completion" { t.Errorf("content capture disabled but found attribute %s", a.Key) From c7dad53f0d4a3708905e06b3785336c90bfffa10 Mon Sep 17 00:00:00 2001 From: pmady Date: Wed, 6 May 2026 10:43:33 -0500 Subject: [PATCH 5/6] go: readme Signed-off-by: pmady --- go/README.md | 105 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 go/README.md diff --git a/go/README.md b/go/README.md new file mode 100644 index 0000000..59ad7ea --- /dev/null +++ b/go/README.md @@ -0,0 +1,105 @@ +# traceAI Go SDK + +OpenTelemetry instrumentation for AI/LLM frameworks in Go, following the [OTel GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/). + +## Packages + +| Package | Description | +|---------|-------------| +| `traceai` | Core setup — TracerProvider, config, semantic conventions | +| `traceai_openai` | OpenAI Go client instrumentation | + +## Quick Start + +```go +package main + +import ( + "context" + "fmt" + "log" + + "github.com/future-agi/traceAI/go/traceai" + "github.com/future-agi/traceAI/go/traceai_openai" + "github.com/openai/openai-go" + "github.com/openai/openai-go/option" +) + +func main() { + provider, err := traceai.Register( + traceai.WithProjectName("my-go-service"), + ) + if err != nil { + log.Fatal(err) + } + defer provider.Shutdown(context.Background()) + + client := openai.NewClient( + option.WithMiddleware(traceai_openai.Middleware()), + ) + + resp, err := client.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{ + Model: "gpt-4", + Messages: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("What is OpenTelemetry?"), + }, + }) + if err != nil { + log.Fatal(err) + } + fmt.Println(resp.Choices[0].Message.Content) +} +``` + +## Span Attributes + +Spans follow the [OTel GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/) — `gen_ai.system`, `gen_ai.request.model`, `gen_ai.usage.*` tokens, `gen_ai.response.*`, etc. Prompt and completion content is captured by default; disable with `WithContentCapture(false)` for PII-sensitive workloads. + +## Configuration + +### Environment Variables + +| Variable | Description | Default | +|----------|-------------|---------| +| `FI_API_KEY` | Future AGI API key | — | +| `FI_SECRET_KEY` | Future AGI secret key | — | +| `FI_BASE_URL` | OTLP HTTP collector endpoint | `https://api.futureagi.com` | +| `FI_GRPC_URL` | OTLP gRPC collector endpoint | `https://grpc.futureagi.com` | +| `FI_PROJECT_NAME` | Project name for trace grouping | `DEFAULT_PROJECT_NAME` | + +### Options + +```go +traceai.Register(traceai.WithTransport(traceai.TransportGRPC)) +traceai.Register(traceai.WithShutdownTimeout(30 * time.Second)) + +// disable content capture for PII-sensitive workloads +traceai_openai.Middleware(traceai_openai.WithContentCapture(false)) +traceai_openai.Middleware(traceai_openai.WithTracerProvider(myTP)) +``` + +## Custom Backend + +Standard OTLP, point it at any collector: + +```go +provider, _ := traceai.Register( + traceai.WithBaseURL("https://your-otel-collector:4318"), + traceai.WithSetGlobal(true), +) +``` + +## Development + +```bash +cd go/traceai_openai +go test ./... +``` + +## Adding a Framework + +See `traceai_openai/` for the pattern — create `traceai_/`, implement an HTTP middleware or client wrapper, wire up the semconv attributes. + +## License + +Apache 2.0 — same as the traceAI project. From 529dc6ec2938ebc4afc6d7d04ca9e112a6f047e3 Mon Sep 17 00:00:00 2001 From: pmady Date: Mon, 24 Aug 2026 21:50:02 -0500 Subject: [PATCH 6/6] go: address review feedback Sets gen_ai.span.kind so spans classify as LLM calls, scopes WithContentCapture to prompt/completion only so model and usage survive it, and switches the grpc exporter to WithEndpointURL. Adds exporter tests for the ingest path and auth headers. Drops the unused Instrumentor interface and Tracer helper. Signed-off-by: pmady --- go/README.md | 2 +- go/traceai/config.go | 35 ++++++------ go/traceai/config_test.go | 95 ++++++++++++++++++++++++++++++++ go/traceai/instrumentor.go | 17 ------ go/traceai/semconv.go | 31 +++++------ go/traceai_openai/openai.go | 29 +++++----- go/traceai_openai/openai_test.go | 8 +++ 7 files changed, 150 insertions(+), 67 deletions(-) create mode 100644 go/traceai/config_test.go delete mode 100644 go/traceai/instrumentor.go diff --git a/go/README.md b/go/README.md index 59ad7ea..ce00cee 100644 --- a/go/README.md +++ b/go/README.md @@ -53,7 +53,7 @@ func main() { ## Span Attributes -Spans follow the [OTel GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/) — `gen_ai.system`, `gen_ai.request.model`, `gen_ai.usage.*` tokens, `gen_ai.response.*`, etc. Prompt and completion content is captured by default; disable with `WithContentCapture(false)` for PII-sensitive workloads. +Spans follow the [OTel GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/) — `gen_ai.system`, `gen_ai.request.model`, `gen_ai.usage.*` tokens, `gen_ai.response.*`, etc. Prompt and completion content is captured by default; disable with `WithContentCapture(false)` for PII-sensitive workloads. Only `gen_ai.prompt` and `gen_ai.completion` are gated by that toggle — model, token usage, response id and finish reasons are always recorded. ## Configuration diff --git a/go/traceai/config.go b/go/traceai/config.go index 036f2e4..8fe9807 100644 --- a/go/traceai/config.go +++ b/go/traceai/config.go @@ -17,15 +17,15 @@ import ( ) const ( - envAPIKey = "FI_API_KEY" - envSecretKey = "FI_SECRET_KEY" - envBaseURL = "FI_BASE_URL" - envGRPCURL = "FI_GRPC_URL" - envProjectName = "FI_PROJECT_NAME" + envAPIKey = "FI_API_KEY" + envSecretKey = "FI_SECRET_KEY" + envBaseURL = "FI_BASE_URL" + envGRPCURL = "FI_GRPC_URL" + envProjectName = "FI_PROJECT_NAME" - defaultBaseURL = "https://api.futureagi.com" - defaultGRPCURL = "https://grpc.futureagi.com" - defaultProjectName = "DEFAULT_PROJECT_NAME" + defaultBaseURL = "https://api.futureagi.com" + defaultGRPCURL = "https://grpc.futureagi.com" + defaultProjectName = "DEFAULT_PROJECT_NAME" headerAPIKey = "X-Api-Key" headerSecretKey = "X-Secret-Key" @@ -41,14 +41,14 @@ const ( // Config holds the options for initializing a traceAI TracerProvider. type Config struct { - ProjectName string - Transport Transport - BaseURL string - GRPCURL string - APIKey string - SecretKey string - BatchExport bool - SetGlobal bool + ProjectName string + Transport Transport + BaseURL string + GRPCURL string + APIKey string + SecretKey string + BatchExport bool + SetGlobal bool ShutdownTimeout time.Duration } @@ -184,8 +184,9 @@ func newExporter(ctx context.Context, cfg Config) (sdktrace.SpanExporter, error) switch cfg.Transport { case TransportGRPC: + // WithEndpoint wants host:port, GRPCURL has a scheme return otlptracegrpc.New(ctx, - otlptracegrpc.WithEndpoint(cfg.GRPCURL), + otlptracegrpc.WithEndpointURL(cfg.GRPCURL), otlptracegrpc.WithHeaders(headers), ) default: diff --git a/go/traceai/config_test.go b/go/traceai/config_test.go new file mode 100644 index 0000000..b1b4fe8 --- /dev/null +++ b/go/traceai/config_test.go @@ -0,0 +1,95 @@ +package traceai + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestDefaultConfigReadsEnv(t *testing.T) { + t.Setenv(envProjectName, "my-project") + t.Setenv(envBaseURL, "https://example.test") + t.Setenv(envAPIKey, "key") + t.Setenv(envSecretKey, "secret") + + cfg := DefaultConfig() + + if cfg.ProjectName != "my-project" { + t.Errorf("ProjectName = %q, want my-project", cfg.ProjectName) + } + if cfg.BaseURL != "https://example.test" { + t.Errorf("BaseURL = %q, want https://example.test", cfg.BaseURL) + } + if cfg.APIKey != "key" || cfg.SecretKey != "secret" { + t.Errorf("credentials = %q/%q, want key/secret", cfg.APIKey, cfg.SecretKey) + } + if cfg.GRPCURL != defaultGRPCURL { + t.Errorf("GRPCURL = %q, want the default %q", cfg.GRPCURL, defaultGRPCURL) + } +} + +func TestRegisterHTTPWiring(t *testing.T) { + type received struct { + path string + apiKey string + secretKey string + } + got := make(chan received, 1) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case got <- received{r.URL.Path, r.Header.Get(headerAPIKey), r.Header.Get(headerSecretKey)}: + default: + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + provider, err := Register( + WithBaseURL(srv.URL), + WithCredentials("key", "secret"), + WithBatchExport(false), + WithSetGlobal(false), + ) + if err != nil { + t.Fatalf("Register: %v", err) + } + + _, span := provider.TracerProvider().Tracer("test").Start(context.Background(), "span") + span.End() + + if err := provider.Shutdown(context.Background()); err != nil { + t.Fatalf("Shutdown: %v", err) + } + + select { + case r := <-got: + if r.path != "/tracer/v1/traces" { + t.Errorf("path = %q, want /tracer/v1/traces", r.path) + } + if r.apiKey != "key" { + t.Errorf("%s = %q, want key", headerAPIKey, r.apiKey) + } + if r.secretKey != "secret" { + t.Errorf("%s = %q, want secret", headerSecretKey, r.secretKey) + } + case <-time.After(5 * time.Second): + t.Fatal("exporter never reached the server") + } +} + +// default GRPCURL has a scheme, WithEndpoint would choke on it +func TestNewExporterGRPCAcceptsSchemeURL(t *testing.T) { + cfg := DefaultConfig() + cfg.Transport = TransportGRPC + + exporter, err := newExporter(context.Background(), cfg) + if err != nil { + t.Fatalf("newExporter(grpc) with %q: %v", cfg.GRPCURL, err) + } + if err := exporter.Shutdown(context.Background()); err != nil { + t.Errorf("Shutdown: %v", err) + } +} diff --git a/go/traceai/instrumentor.go b/go/traceai/instrumentor.go deleted file mode 100644 index cd0d3b5..0000000 --- a/go/traceai/instrumentor.go +++ /dev/null @@ -1,17 +0,0 @@ -package traceai - -import "go.opentelemetry.io/otel/trace" - -// Instrumentor wraps an AI framework client with OTel tracing. -type Instrumentor interface { - Instrument(tp trace.TracerProvider) error - Uninstrument() error -} - -func Tracer(tp trace.TracerProvider, name string, opts ...trace.TracerOption) trace.Tracer { - if tp == nil { - // fall back to the global provider set by Register() - return trace.NewNoopTracerProvider().Tracer(name) - } - return tp.Tracer(name, opts...) -} diff --git a/go/traceai/semconv.go b/go/traceai/semconv.go index 62a77bc..130608c 100644 --- a/go/traceai/semconv.go +++ b/go/traceai/semconv.go @@ -5,15 +5,14 @@ import "go.opentelemetry.io/otel/attribute" // GenAI semantic convention keys. // See https://opentelemetry.io/docs/specs/semconv/gen-ai/ const ( - AttrGenAISystem = attribute.Key("gen_ai.system") - AttrGenAIRequestModel = attribute.Key("gen_ai.request.model") - AttrGenAIResponseModel = attribute.Key("gen_ai.response.model") - AttrGenAIOperationName = attribute.Key("gen_ai.operation.name") + AttrGenAISystem = attribute.Key("gen_ai.system") + AttrGenAIRequestModel = attribute.Key("gen_ai.request.model") + AttrGenAIResponseModel = attribute.Key("gen_ai.response.model") + AttrGenAIOperationName = attribute.Key("gen_ai.operation.name") - AttrGenAIRequestMaxTokens = attribute.Key("gen_ai.request.max_tokens") - AttrGenAIRequestTemperature = attribute.Key("gen_ai.request.temperature") - AttrGenAIRequestTopP = attribute.Key("gen_ai.request.top_p") - AttrGenAIRequestStopSequences = attribute.Key("gen_ai.request.stop_sequences") + AttrGenAIRequestMaxTokens = attribute.Key("gen_ai.request.max_tokens") + AttrGenAIRequestTemperature = attribute.Key("gen_ai.request.temperature") + AttrGenAIRequestTopP = attribute.Key("gen_ai.request.top_p") AttrGenAIUsageInputTokens = attribute.Key("gen_ai.usage.input_tokens") AttrGenAIUsageOutputTokens = attribute.Key("gen_ai.usage.output_tokens") @@ -23,17 +22,17 @@ const ( AttrGenAIPrompt = attribute.Key("gen_ai.prompt") AttrGenAICompletion = attribute.Key("gen_ai.completion") -) -const ( - GenAISystemOpenAI = "openai" - GenAISystemAnthropic = "anthropic" - GenAISystemCohere = "cohere" - GenAISystemGoogle = "google" + // platform uses this to classify the span as an LLM call + AttrGenAISpanKind = attribute.Key("gen_ai.span.kind") ) +const SpanKindLLM = "LLM" + +const GenAISystemOpenAI = "openai" + const ( OpChat = "chat" - OpCompletion = "completion" - OpEmbedding = "embedding" + OpCompletion = "text_completion" + OpEmbedding = "embeddings" ) diff --git a/go/traceai_openai/openai.go b/go/traceai_openai/openai.go index 2aa407c..0a1b8bb 100644 --- a/go/traceai_openai/openai.go +++ b/go/traceai_openai/openai.go @@ -11,11 +11,11 @@ package traceai_openai import ( + "bytes" "encoding/json" "io" "net/http" "strings" - "time" "github.com/future-agi/traceAI/go/traceai" "go.opentelemetry.io/otel" @@ -60,24 +60,21 @@ func Middleware(opts ...Option) func(req *http.Request, next func(req *http.Requ span.SetAttributes( traceai.AttrGenAISystem.String(traceai.GenAISystemOpenAI), traceai.AttrGenAIOperationName.String(opName), + traceai.AttrGenAISpanKind.String(traceai.SpanKindLLM), ) - if req.Body != nil && cfg.captureContent { + if req.Body != nil { bodyBytes, err := io.ReadAll(req.Body) if err == nil { - req.Body = io.NopCloser(strings.NewReader(string(bodyBytes))) - extractRequestAttributes(span, bodyBytes) + req.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + extractRequestAttributes(span, bodyBytes, cfg.captureContent) } } req = req.WithContext(ctx) - start := time.Now() resp, err := next(req) - span.SetAttributes(attribute.Float64("gen_ai.request.duration_ms", - float64(time.Since(start).Milliseconds()))) - if err != nil { span.RecordError(err) span.SetStatus(codes.Error, err.Error()) @@ -90,12 +87,12 @@ func Middleware(opts ...Option) func(req *http.Request, next func(req *http.Requ return resp, err } - // read response for token counts - if resp != nil && resp.Body != nil && cfg.captureContent { + // model and usage live in the body too, only the text is gated + if resp != nil && resp.Body != nil { bodyBytes, readErr := io.ReadAll(resp.Body) if readErr == nil { - resp.Body = io.NopCloser(strings.NewReader(string(bodyBytes))) - extractResponseAttributes(span, bodyBytes) + resp.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + extractResponseAttributes(span, bodyBytes, cfg.captureContent) } } @@ -125,7 +122,7 @@ type chatRequest struct { TopP *float64 `json:"top_p,omitempty"` } -func extractRequestAttributes(span trace.Span, body []byte) { +func extractRequestAttributes(span trace.Span, body []byte, captureContent bool) { var req chatRequest if err := json.Unmarshal(body, &req); err != nil { return @@ -143,7 +140,7 @@ func extractRequestAttributes(span trace.Span, body []byte) { if req.TopP != nil { span.SetAttributes(traceai.AttrGenAIRequestTopP.Float64(*req.TopP)) } - if len(req.Messages) > 0 { + if captureContent && len(req.Messages) > 0 { span.SetAttributes(traceai.AttrGenAIPrompt.String(string(req.Messages))) } } @@ -163,7 +160,7 @@ type chatResponse struct { } `json:"usage"` } -func extractResponseAttributes(span trace.Span, body []byte) { +func extractResponseAttributes(span trace.Span, body []byte, captureContent bool) { var resp chatResponse if err := json.Unmarshal(body, &resp); err != nil { return @@ -196,7 +193,7 @@ func extractResponseAttributes(span trace.Span, body []byte) { if len(reasons) > 0 { span.SetAttributes(traceai.AttrGenAIResponseFinishReasons.StringSlice(reasons)) } - if len(completions) > 0 { + if captureContent && len(completions) > 0 { span.SetAttributes(traceai.AttrGenAICompletion.String(strings.Join(completions, "\n"))) } } diff --git a/go/traceai_openai/openai_test.go b/go/traceai_openai/openai_test.go index bbb1a92..0209e0d 100644 --- a/go/traceai_openai/openai_test.go +++ b/go/traceai_openai/openai_test.go @@ -61,6 +61,7 @@ func TestMiddleware_ChatCompletion(t *testing.T) { } assertAttr(t, span.Attributes, "gen_ai.system", "openai") + assertAttr(t, span.Attributes, "gen_ai.span.kind", "LLM") assertAttr(t, span.Attributes, "gen_ai.request.model", "gpt-4") assertAttr(t, span.Attributes, "gen_ai.response.model", "gpt-4") assertAttr(t, span.Attributes, "gen_ai.response.id", "chatcmpl-abc") @@ -153,6 +154,13 @@ func TestMiddleware_ContentCaptureDisabled(t *testing.T) { t.Errorf("content capture disabled but found attribute %s", a.Key) } } + + // model and usage aren't content, they stay + assertAttr(t, spans[0].Attributes, "gen_ai.request.model", "gpt-4") + assertAttr(t, spans[0].Attributes, "gen_ai.response.model", "gpt-4") + assertAttr(t, spans[0].Attributes, "gen_ai.response.id", "chatcmpl-abc") + assertIntAttr(t, spans[0].Attributes, "gen_ai.usage.input_tokens", 5) + assertIntAttr(t, spans[0].Attributes, "gen_ai.usage.output_tokens", 3) } func assertAttr(t *testing.T, attrs []attribute.KeyValue, key, want string) {