diff --git a/examples/internal/genai-streaming/go.mod b/examples/internal/genai-streaming/go.mod new file mode 100644 index 0000000..0ce6cf3 --- /dev/null +++ b/examples/internal/genai-streaming/go.mod @@ -0,0 +1,39 @@ +module example.com/genai-streaming + +go 1.25.0 + +replace github.com/braintrustdata/braintrust-sdk-go => ../../.. + +require ( + github.com/braintrustdata/braintrust-sdk-go v0.0.0-00010101000000-000000000000 + go.opentelemetry.io/otel/sdk v1.43.0 + google.golang.org/genai v1.52.1 +) + +require ( + cloud.google.com/go v0.123.0 // indirect + cloud.google.com/go/auth v0.17.0 // indirect + cloud.google.com/go/compute/metadata v0.9.0 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/felixge/httpsnoop v1.0.4 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/google/go-cmp v0.7.0 // indirect + github.com/google/s2a-go v0.1.9 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect + github.com/googleapis/gax-go/v2 v2.15.0 // indirect + github.com/gorilla/websocket v1.5.3 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect + go.opentelemetry.io/otel v1.43.0 // indirect + go.opentelemetry.io/otel/metric v1.43.0 // indirect + go.opentelemetry.io/otel/trace v1.43.0 // indirect + golang.org/x/crypto v0.45.0 // indirect + golang.org/x/net v0.47.0 // indirect + golang.org/x/sys v0.42.0 // indirect + golang.org/x/text v0.31.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 // indirect + google.golang.org/grpc v1.76.0 // indirect + google.golang.org/protobuf v1.36.10 // indirect +) diff --git a/examples/internal/genai-streaming/go.sum b/examples/internal/genai-streaming/go.sum new file mode 100644 index 0000000..21c0cbe --- /dev/null +++ b/examples/internal/genai-streaming/go.sum @@ -0,0 +1,75 @@ +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/auth v0.17.0 h1:74yCm7hCj2rUyyAocqnFzsAYXgJhrG26XCFimrc/Kz4= +cloud.google.com/go/auth v0.17.0/go.mod h1:6wv/t5/6rOPAX4fJiRjKkJCvswLwdet7G8+UGXt7nCQ= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= +github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/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/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= +github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= +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/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= +github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= +github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81vgd/bo= +github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg= +go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= +go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0= +go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= +go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= +go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= +go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= +go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= +go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= +go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +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/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= +golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= +golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= +golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= +golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= +golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= +golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM= +golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM= +gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= +gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= +google.golang.org/genai v1.52.1 h1:dYoljKtLDXMiBdVaClSJ/ZPwZ7j1N0lGjMhwOKOQUlk= +google.golang.org/genai v1.52.1/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 h1:M1rk8KBnUsBDg1oPGHNCxG4vc1f49epmTO7xscSajMk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk= +google.golang.org/grpc v1.76.0 h1:UnVkv1+uMLYXoIz6o7chp59WfQUYA2ex/BXQ9rHZu7A= +google.golang.org/grpc v1.76.0/go.mod h1:Ju12QI8M6iQJtbcsV+awF5a4hfJMLi4X0JLo94ULZ6c= +google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE= +google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/dnaeon/go-vcr.v3 v3.2.0 h1:Rltp0Vf+Aq0u4rQXgmXgtgoRDStTnFN83cWgSGSoRzM= +gopkg.in/dnaeon/go-vcr.v3 v3.2.0/go.mod h1:2IMOnnlx9I6u9x+YBsM3tAMx6AlOxnJ0pWxQAzZ79Ag= +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/examples/internal/genai-streaming/main.go b/examples/internal/genai-streaming/main.go new file mode 100644 index 0000000..03414af --- /dev/null +++ b/examples/internal/genai-streaming/main.go @@ -0,0 +1,119 @@ +// Demonstrates that Gemini streaming calls (streamGenerateContent) are now +// fully instrumented. Before this fix, only non-streaming generateContent +// calls produced spans -- streaming calls silently passed through without +// any tracing. +// +// The example makes one non-streaming call and one streaming call, then +// inspects the captured spans to show that both produce complete trace data +// including output, token metrics, and time_to_first_token. +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "os" + + "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "google.golang.org/genai" + + tracegenai "github.com/braintrustdata/braintrust-sdk-go/trace/contrib/genai" +) + +func main() { + apiKey := os.Getenv("GOOGLE_API_KEY") + if apiKey == "" { + log.Fatal("GOOGLE_API_KEY is required") + } + + // Set up an in-memory exporter so we can inspect spans locally. + exporter := tracetest.NewInMemoryExporter() + tp := trace.NewTracerProvider( + trace.WithSyncer(exporter), + ) + defer tp.Shutdown(context.Background()) //nolint:errcheck + + // Create Gemini client with tracing. + client, err := genai.NewClient(context.Background(), &genai.ClientConfig{ + HTTPClient: tracegenai.Client(tracegenai.WithTracerProvider(tp)), + APIKey: apiKey, + Backend: genai.BackendGeminiAPI, + }) + if err != nil { + log.Fatal(err) + } + + ctx := context.Background() + + // --- Non-streaming call --- + fmt.Println("=== Non-streaming call ===") + resp, err := client.Models.GenerateContent(ctx, "gemini-2.0-flash", + genai.Text("What is 2+2? Answer with just the number."), nil) + if err != nil { + log.Fatal(err) + } + fmt.Printf("Response: %s\n", resp.Text()) + + // --- Streaming call (was previously uninstrumented) --- + fmt.Println("\n=== Streaming call ===") + iter := client.Models.GenerateContentStream(ctx, "gemini-2.0-flash", + genai.Text("Count from 1 to 3, one number per line."), nil) + fmt.Print("Response: ") + for chunk, err := range iter { + if err != nil { + log.Fatal(err) + } + fmt.Print(chunk.Text()) + } + fmt.Println() + + // --- Inspect spans --- + fmt.Println("\n=== Captured spans ===") + spans := exporter.GetSpans() + fmt.Printf("Total spans: %d (before this fix, streaming calls produced 0 spans)\n\n", len(spans)) + + for i, span := range spans { + fmt.Printf("Span %d: %s\n", i+1, span.Name) + + metrics := jsonAttr(span, "braintrust.metrics") + metadata := jsonAttr(span, "braintrust.metadata") + + fmt.Printf(" provider: %v\n", metadata["provider"]) + fmt.Printf(" model: %v\n", metadata["model"]) + fmt.Printf(" prompt_tokens: %v\n", metrics["prompt_tokens"]) + fmt.Printf(" completion_tokens: %v\n", metrics["completion_tokens"]) + fmt.Printf(" time_to_first_token: %v seconds\n", metrics["time_to_first_token"]) + + output := jsonAttr(span, "braintrust.output_json") + if candidates, ok := output["candidates"].([]any); ok && len(candidates) > 0 { + if c, ok := candidates[0].(map[string]any); ok { + if content, ok := c["content"].(map[string]any); ok { + if parts, ok := content["parts"].([]any); ok && len(parts) > 0 { + if p, ok := parts[0].(map[string]any); ok { + text := fmt.Sprintf("%v", p["text"]) + if len(text) > 80 { + text = text[:80] + "..." + } + fmt.Printf(" output text: %q\n", text) + } + } + } + } + } + fmt.Println() + } +} + +func jsonAttr(span tracetest.SpanStub, key string) map[string]any { + for _, attr := range span.Attributes { + if string(attr.Key) == key { + var m map[string]any + if err := json.Unmarshal([]byte(attr.Value.AsString()), &m); err == nil { + return m + } + } + } + return nil +} diff --git a/examples/internal/genai/go.mod b/examples/internal/genai/go.mod index 039668f..eb3d220 100644 --- a/examples/internal/genai/go.mod +++ b/examples/internal/genai/go.mod @@ -7,13 +7,14 @@ replace github.com/braintrustdata/braintrust-sdk-go => ../../.. require ( github.com/braintrustdata/braintrust-sdk-go v0.0.12 go.opentelemetry.io/otel v1.38.0 - google.golang.org/genai v1.30.0 + go.opentelemetry.io/otel/sdk v1.38.0 + google.golang.org/genai v1.41.0 ) require ( - cloud.google.com/go v0.116.0 // indirect - cloud.google.com/go/auth v0.14.0 // indirect - cloud.google.com/go/compute/metadata v0.6.0 // indirect + cloud.google.com/go v0.123.0 // indirect + cloud.google.com/go/auth v0.17.0 // indirect + cloud.google.com/go/compute/metadata v0.9.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.3 // indirect @@ -21,25 +22,24 @@ require ( github.com/google/go-cmp v0.7.0 // indirect github.com/google/s2a-go v0.1.9 // indirect github.com/google/uuid v1.6.0 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect - github.com/googleapis/gax-go/v2 v2.14.1 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect + github.com/googleapis/gax-go/v2 v2.15.0 // indirect github.com/gorilla/websocket v1.5.3 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect - go.opentelemetry.io/auto/sdk v1.1.0 // indirect - go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 // indirect go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.36.0 // indirect go.opentelemetry.io/otel/metric v1.38.0 // indirect - go.opentelemetry.io/otel/sdk v1.36.0 // indirect go.opentelemetry.io/otel/trace v1.38.0 // indirect go.opentelemetry.io/proto/otlp v1.5.0 // indirect - golang.org/x/crypto v0.41.0 // indirect - golang.org/x/net v0.43.0 // indirect - golang.org/x/sys v0.35.0 // indirect - golang.org/x/text v0.28.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 + golang.org/x/crypto v0.45.0 // indirect + golang.org/x/net v0.47.0 // indirect + golang.org/x/sys v0.38.0 // indirect + golang.org/x/text v0.31.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 // indirect + google.golang.org/grpc v1.76.0 // indirect + google.golang.org/protobuf v1.36.10 // indirect ) diff --git a/examples/internal/genai/go.sum b/examples/internal/genai/go.sum index 03e508e..871a6b4 100644 --- a/examples/internal/genai/go.sum +++ b/examples/internal/genai/go.sum @@ -1,9 +1,9 @@ -cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE= -cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U= -cloud.google.com/go/auth v0.14.0 h1:A5C4dKV/Spdvxcl0ggWwWEzzP7AZMJSEIgrkngwhGYM= -cloud.google.com/go/auth v0.14.0/go.mod h1:CYsoRL1PdiDuqeQpZE0bP2pnPrGqFcOkI0nldEQis+A= -cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I= -cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg= +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/auth v0.17.0 h1:74yCm7hCj2rUyyAocqnFzsAYXgJhrG26XCFimrc/Kz4= +cloud.google.com/go/auth v0.17.0/go.mod h1:6wv/t5/6rOPAX4fJiRjKkJCvswLwdet7G8+UGXt7nCQ= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= 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.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= @@ -23,10 +23,10 @@ github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= 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/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw= -github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA= -github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q= -github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA= +github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= +github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= +github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81vgd/bo= +github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 h1:e9Rjr40Z98/clHv5Yg79Is0NtosR5LXRvdr7o/6NwbA= @@ -35,10 +35,10 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -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/contrib/instrumentation/net/http/otelhttp v0.61.0 h1:F7Jx+6hwnZ41NSFTO5q4LYDtJRXBf2PD0rNBkeB/lus= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0/go.mod h1:UHB22Z8QsdRDrnAtX4PntOl36ajSxcdUMt1sF7Y6E7Q= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg= go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8= go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM= go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 h1:1fTNlAIJZGWLP5FVu0fikVry1IsiUnXjf7QFvoNN3Xw= @@ -49,35 +49,39 @@ go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.36.0 h1:G8Xec/SgZQricwW go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.36.0/go.mod h1:PD57idA/AiFD5aqoxGxCvT/ILJPeHy3MjqU/NS7KogY= go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA= go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI= -go.opentelemetry.io/otel/sdk v1.36.0 h1:b6SYIuLRs88ztox4EyrvRti80uXIFy+Sqzoh9kFULbs= -go.opentelemetry.io/otel/sdk v1.36.0/go.mod h1:+lC+mTgD+MUWfjJubi2vvXWcVxyr9rmlshZni72pXeY= -go.opentelemetry.io/otel/sdk/metric v1.36.0 h1:r0ntwwGosWGaa0CrSt8cuNuTcccMXERFwHX4dThiPis= -go.opentelemetry.io/otel/sdk/metric v1.36.0/go.mod h1:qTNOhFDfKRwX0yXOqJYegL5WRaW376QbB7P4Pb0qva4= +go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E= +go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg= +go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM= +go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA= go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE= go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs= 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/crypto v0.41.0 h1:WKYxWedPGCTVVl5+WHSSrOBT0O8lx32+zxmHxijgXp4= -golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sUc= -golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE= -golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg= -golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw= -golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= -golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= -golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= -golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng= -golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU= -google.golang.org/genai v1.30.0 h1:7021aneIvl24nEBLbtQFEWleHsMbjzpcQvkT4WcJ1dc= -google.golang.org/genai v1.30.0/go.mod h1:7pAilaICJlQBonjKKJNhftDFv3SREhZcTe9F6nRcjbg= -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= +golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= +golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= +golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= +golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= +golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= +golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= +golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM= +golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM= +gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= +gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= +google.golang.org/genai v1.41.0 h1:ayXl75LjTmqTu0y94yr96d17gIb4zF8gWVzX2TgioEY= +google.golang.org/genai v1.41.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= +google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f h1:OiFuztEyBivVKDvguQJYWq1yDcfAHIID/FVrPR4oiI0= +google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f/go.mod h1:kprOiu9Tr0JYyD6DORrc4Hfyk3RFXqkQ3ctHEum3ZbM= +google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 h1:M1rk8KBnUsBDg1oPGHNCxG4vc1f49epmTO7xscSajMk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk= +google.golang.org/grpc v1.76.0 h1:UnVkv1+uMLYXoIz6o7chp59WfQUYA2ex/BXQ9rHZu7A= +google.golang.org/grpc v1.76.0/go.mod h1:Ju12QI8M6iQJtbcsV+awF5a4hfJMLi4X0JLo94ULZ6c= +google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE= +google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/dnaeon/go-vcr.v3 v3.2.0 h1:Rltp0Vf+Aq0u4rQXgmXgtgoRDStTnFN83cWgSGSoRzM= +gopkg.in/dnaeon/go-vcr.v3 v3.2.0/go.mod h1:2IMOnnlx9I6u9x+YBsM3tAMx6AlOxnJ0pWxQAzZ79Ag= 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/examples/internal/genai/main.go b/examples/internal/genai/main.go index 9f1f8c6..db8d6ce 100644 --- a/examples/internal/genai/main.go +++ b/examples/internal/genai/main.go @@ -38,7 +38,7 @@ func (g *GeminiBot) basicText(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("What is the capital of France?"), nil, ) @@ -59,7 +59,7 @@ func (g *GeminiBot) withSystemInstruction(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("Explain quantum computing."), &genai.GenerateContentConfig{ SystemInstruction: genai.NewContentFromText("You are a helpful science educator. Be concise and clear.", ""), @@ -96,7 +96,7 @@ func (g *GeminiBot) multiTurnConversation(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", history, &genai.GenerateContentConfig{ Temperature: genai.Ptr[float32](0.5), @@ -119,7 +119,7 @@ func (g *GeminiBot) streaming(ctx context.Context) error { iter := g.client.Models.GenerateContentStream( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("Count from 1 to 5 slowly, with one number per line."), &genai.GenerateContentConfig{ Temperature: genai.Ptr[float32](0.9), @@ -168,7 +168,7 @@ func (g *GeminiBot) functionCalling(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("What's the weather like in Tokyo?"), &genai.GenerateContentConfig{ Tools: []*genai.Tool{ @@ -210,7 +210,7 @@ func (g *GeminiBot) safetySettings(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("Tell me about internet safety."), &genai.GenerateContentConfig{ SafetySettings: []*genai.SafetySetting{ @@ -279,7 +279,7 @@ func (g *GeminiBot) jsonMode(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("Create a profile for Albert Einstein"), &genai.GenerateContentConfig{ ResponseMIMEType: "application/json", @@ -316,7 +316,7 @@ func (g *GeminiBot) multimodal(ctx context.Context) error { resp, err := g.client.Models.GenerateContent( ctx, - "gemini-2.0-flash-exp", + "gemini-2.0-flash", genai.Text("Describe a beautiful sunset over mountains."), nil, ) @@ -334,7 +334,8 @@ func main() { // Initialize braintrust tracing with a specific project tp := trace.NewTracerProvider() - defer tp.Shutdown(context.Background()) + defer tp.Shutdown(context.Background()) //nolint:errcheck + otel.SetTracerProvider(tp) bt, err := braintrust.New(tp, braintrust.WithProject("go-sdk-examples"), diff --git a/trace/contrib/genai/generatecontent.go b/trace/contrib/genai/generatecontent.go index c68a1a7..738847d 100644 --- a/trace/contrib/genai/generatecontent.go +++ b/trace/contrib/genai/generatecontent.go @@ -3,6 +3,7 @@ package genai // this file parses the generateContent API. import ( + "bufio" "context" "encoding/json" "io" @@ -20,12 +21,13 @@ type generateContentTracer struct { streaming bool metadata map[string]any model string + startTime time.Time } -func newGenerateContentTracer(cfg *config, model string) *generateContentTracer { +func newGenerateContentTracer(cfg *config, model string, streaming bool) *generateContentTracer { return &generateContentTracer{ cfg: cfg, - streaming: false, + streaming: streaming, model: model, metadata: map[string]any{ "provider": "gemini", @@ -34,6 +36,7 @@ func newGenerateContentTracer(cfg *config, model string) *generateContentTracer } func (gt *generateContentTracer) StartSpan(ctx context.Context, t time.Time, request io.Reader) (context.Context, trace.Span, error) { + gt.startTime = t ctx, span := gt.cfg.tracer().Start( ctx, "generate_content", @@ -122,11 +125,173 @@ func (gt *generateContentTracer) StartSpan(ctx context.Context, t time.Time, req } func (gt *generateContentTracer) TagSpan(span trace.Span, body io.Reader) error { - // For now, handle non-streaming responses - // Streaming will be handled separately + if gt.streaming { + return gt.parseStreamingResponse(span, body) + } return gt.parseResponse(span, body) } +func (gt *generateContentTracer) parseStreamingResponse(span trace.Span, body io.Reader) error { + scanner := bufio.NewScanner(body) + var allResults []map[string]any + var timeToFirstToken time.Duration + + for scanner.Scan() { + line := scanner.Text() + + if !strings.HasPrefix(line, "data: ") { + continue + } + + line = strings.TrimPrefix(line, "data: ") + if line == "[DONE]" { + break + } + + if timeToFirstToken == 0 { + timeToFirstToken = time.Since(gt.startTime) + } + + var chunk map[string]any + if err := json.Unmarshal([]byte(line), &chunk); err != nil { + return err + } + + allResults = append(allResults, chunk) + } + + // Aggregate chunks into a single response + output := gt.postprocessStreamingResults(allResults) + if output != nil { + if err := internal.SetJSONAttr(span, "braintrust.output_json", output); err != nil { + return err + } + } + + // Collect usage from the last chunk (Gemini includes usage in the final chunk) + metrics := make(map[string]any) + for i := len(allResults) - 1; i >= 0; i-- { + if usageMetadata, ok := allResults[i]["usageMetadata"].(map[string]any); ok { + for k, v := range parseUsageTokens(usageMetadata) { + metrics[k] = v + } + break + } + } + + // Extract model version from any chunk + for _, chunk := range allResults { + if modelVersion, ok := chunk["modelVersion"].(string); ok { + gt.metadata["model"] = modelVersion + break + } + } + if err := internal.SetJSONAttr(span, "braintrust.metadata", gt.metadata); err != nil { + return err + } + + metrics["time_to_first_token"] = timeToFirstToken.Seconds() + if err := internal.SetJSONAttr(span, "braintrust.metrics", metrics); err != nil { + return err + } + + return scanner.Err() +} + +// postprocessStreamingResults aggregates streaming chunks into a single response +// matching the non-streaming generateContent response format. +func (gt *generateContentTracer) postprocessStreamingResults(allResults []map[string]any) map[string]any { + if len(allResults) == 0 { + return nil + } + + // Aggregate text parts from all candidates across chunks + var textParts []string + var finishReason any + var role string + + for _, chunk := range allResults { + candidates, ok := chunk["candidates"].([]any) + if !ok || len(candidates) == 0 { + continue + } + candidate, ok := candidates[0].(map[string]any) + if !ok { + continue + } + + if fr, ok := candidate["finishReason"]; ok && fr != nil { + finishReason = fr + } + + content, ok := candidate["content"].(map[string]any) + if !ok { + continue + } + + if r, ok := content["role"].(string); ok && role == "" { + role = r + } + + parts, ok := content["parts"].([]any) + if !ok { + continue + } + + for _, p := range parts { + part, ok := p.(map[string]any) + if !ok { + continue + } + if text, ok := part["text"].(string); ok { + textParts = append(textParts, text) + } + } + } + + // Build aggregated response in Gemini format + result := map[string]any{ + "candidates": []any{ + map[string]any{ + "content": map[string]any{ + "parts": []any{ + map[string]any{ + "text": strings.Join(textParts, ""), + }, + }, + "role": role, + }, + }, + }, + } + + if finishReason != nil { + if candidates, ok := result["candidates"].([]any); ok && len(candidates) > 0 { + if c, ok := candidates[0].(map[string]any); ok { + c["finishReason"] = finishReason + } + } + } + + // Include usage from the last chunk + for i := len(allResults) - 1; i >= 0; i-- { + if usage, ok := allResults[i]["usageMetadata"]; ok { + result["usageMetadata"] = usage + break + } + } + + // Include model version + for _, chunk := range allResults { + if mv, ok := chunk["modelVersion"]; ok { + result["modelVersion"] = mv + break + } + } + + return result +} + func (gt *generateContentTracer) parseResponse(span trace.Span, body io.Reader) error { var raw map[string]interface{} err := json.NewDecoder(body).Decode(&raw) @@ -153,13 +318,17 @@ func (gt *generateContentTracer) handleResponse(span trace.Span, raw map[string] return err } - // Parse usage metadata (token counts) + // Parse usage metadata (token counts) and time_to_first_token + metrics := make(map[string]any) if usageMetadata, ok := raw["usageMetadata"].(map[string]any); ok { - metrics := parseUsageTokens(usageMetadata) - if err := internal.SetJSONAttr(span, "braintrust.metrics", metrics); err != nil { - return err + for k, v := range parseUsageTokens(usageMetadata) { + metrics[k] = v } } + metrics["time_to_first_token"] = time.Since(gt.startTime).Seconds() + if err := internal.SetJSONAttr(span, "braintrust.metrics", metrics); err != nil { + return err + } return nil } diff --git a/trace/contrib/genai/testdata/cassettes/TestStreamingGenerateContent.yaml b/trace/contrib/genai/testdata/cassettes/TestStreamingGenerateContent.yaml new file mode 100644 index 0000000..675d284 --- /dev/null +++ b/trace/contrib/genai/testdata/cassettes/TestStreamingGenerateContent.yaml @@ -0,0 +1,51 @@ +--- +version: 2 +interactions: + - id: 0 + request: + proto: HTTP/1.1 + proto_major: 1 + proto_minor: 1 + content_length: 109 + transfer_encoding: [] + trailer: {} + host: generativelanguage.googleapis.com + remote_addr: "" + request_uri: "" + body: | + {"contents":[{"parts":[{"text":"Count from 1 to 3. Output only the numbers."}],"role":"user"}]} + form: {} + headers: + Content-Type: + - application/json + User-Agent: + - google-genai-sdk/1.23.0 gl-go/go1.23.12 + X-Goog-Api-Client: + - google-genai-sdk/1.23.0 gl-go/go1.23.12 + X-Server-Timeout: + - "30" + url: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-exp:streamGenerateContent?alt=sse + method: POST + response: + proto: HTTP/2.0 + proto_major: 2 + proto_minor: 0 + transfer_encoding: [] + trailer: {} + content_length: -1 + uncompressed: false + body: "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"1\"}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"gemini-2.0-flash-exp\"}\r\n\r\ndata: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\" 2\"}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"gemini-2.0-flash-exp\"}\r\n\r\ndata: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\" 3\\n\"}],\"role\":\"model\"},\"index\":0,\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":12,\"candidatesTokenCount\":8,\"totalTokenCount\":20},\"modelVersion\":\"gemini-2.0-flash-exp\"}\r\n\r\n" + headers: + Content-Type: + - text/event-stream + Date: + - Fri, 07 Nov 2025 17:16:55 GMT + Server: + - scaffolding on HTTPServer2 + Vary: + - Origin + - X-Origin + - Referer + status: 200 OK + code: 200 + duration: 623.688792ms diff --git a/trace/contrib/genai/tracegenai.go b/trace/contrib/genai/tracegenai.go index ba518f4..8bee459 100644 --- a/trace/contrib/genai/tracegenai.go +++ b/trace/contrib/genai/tracegenai.go @@ -141,20 +141,27 @@ func (rt *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) { // genaiRouter maps Gemini API paths to their corresponding tracers. func genaiRouter(cfg *config, path string) internal.MiddlewareTracer { - // Match both Gemini API and Vertex AI paths - // Gemini API: /v1beta/models/{model}/generateContent + // Match both Gemini API and Vertex AI paths for generateContent and streamGenerateContent + // Gemini API: /v1beta/models/{model}:generateContent or :streamGenerateContent // Vertex AI: /v1/projects/{project}/locations/{location}/publishers/google/models/{model}:generateContent if containsGenerateContent(path) { model := extractModelFromPath(path) - return newGenerateContentTracer(cfg, model) + streaming := isStreamingPath(path) + return newGenerateContentTracer(cfg, model, streaming) } return nil } -// containsGenerateContent checks if the path is for a generateContent endpoint +// containsGenerateContent checks if the path is for a generateContent or streamGenerateContent endpoint func containsGenerateContent(path string) bool { - return strings.Contains(path, "/generateContent") || - strings.Contains(path, ":generateContent") + return strings.Contains(path, "generateContent") || + strings.Contains(path, "GenerateContent") +} + +// isStreamingPath checks if the path is for the streaming endpoint +func isStreamingPath(path string) bool { + return strings.Contains(path, "streamGenerateContent") || + strings.Contains(path, "StreamGenerateContent") } // extractModelFromPath extracts the model name from the URL path @@ -170,9 +177,15 @@ func extractModelFromPath(path string) string { // Get the model part (after /models/) modelPart := parts[1] - // Remove :generateContent or /generateContent suffix - modelPart = strings.TrimSuffix(modelPart, ":generateContent") - modelPart = strings.TrimSuffix(modelPart, "/generateContent") + // Remove the endpoint suffix (handles both streaming and non-streaming) + for _, suffix := range []string{ + ":streamGenerateContent", + "/streamGenerateContent", + ":generateContent", + "/generateContent", + } { + modelPart = strings.TrimSuffix(modelPart, suffix) + } return modelPart } diff --git a/trace/contrib/genai/tracegenai_test.go b/trace/contrib/genai/tracegenai_test.go index 90b2210..c23b18a 100644 --- a/trace/contrib/genai/tracegenai_test.go +++ b/trace/contrib/genai/tracegenai_test.go @@ -91,11 +91,65 @@ func TestBasicGenerateContent(t *testing.T) { output := ts.Output() require.NotNil(output) - // Verify metrics (token counts) + // Verify metrics (token counts + time_to_first_token) metrics := ts.Metrics() assert.Greater(metrics["prompt_tokens"], float64(0)) assert.Greater(metrics["completion_tokens"], float64(0)) assert.Greater(metrics["tokens"], float64(0)) + assert.Greater(metrics["time_to_first_token"], float64(0)) +} + +func TestStreamingGenerateContent(t *testing.T) { + client, exporter := setUpTest(t) + + assert := assert.New(t) + require := require.New(t) + + // Make a streaming generateContent request + timer := oteltest.NewTimer() + iter := client.Models.GenerateContentStream( + context.Background(), + "gemini-2.0-flash-exp", + genai.Text("Count from 1 to 3. Output only the numbers."), + nil, + ) + + var fullText string + for resp, err := range iter { + require.NoError(err) + fullText += resp.Text() + } + timeRange := timer.Tick() + + assert.Contains(fullText, "1") + assert.Contains(fullText, "2") + assert.Contains(fullText, "3") + + // Verify span was created + ts := exporter.FlushOne() + ts.AssertInTimeRange(timeRange) + ts.AssertNameIs("generate_content") + assert.Equal(codes.Unset, ts.Status().Code) + + // Verify metadata + metadata := ts.Metadata() + assert.Equal("gemini", metadata["provider"]) + assert.Equal("gemini-2.0-flash-exp", metadata["model"]) + + // Verify input + input := ts.Input() + require.NotNil(input) + + // Verify output was reconstructed from stream + output := ts.Output() + require.NotNil(output) + + // Verify metrics (token counts + time_to_first_token) + metrics := ts.Metrics() + assert.Greater(metrics["prompt_tokens"], float64(0)) + assert.Greater(metrics["completion_tokens"], float64(0)) + assert.Greater(metrics["tokens"], float64(0)) + assert.Greater(metrics["time_to_first_token"], float64(0)) } func TestParseUsageTokens(t *testing.T) { @@ -148,6 +202,65 @@ func TestParseUsageTokens(t *testing.T) { }) } +func TestContainsGenerateContent(t *testing.T) { + tests := []struct { + path string + matches bool + }{ + // Non-streaming + {"/v1beta/models/gemini-2.0-flash/generateContent", true}, + {"/v1beta/models/gemini-2.0-flash:generateContent", true}, + {"/v1/projects/p/locations/l/publishers/google/models/gemini-2.0-flash:generateContent", true}, + // Streaming + {"/v1beta/models/gemini-2.0-flash:streamGenerateContent", true}, + {"/v1beta/models/gemini-2.0-flash/streamGenerateContent", true}, + {"/v1/projects/p/locations/l/publishers/google/models/gemini-2.0-flash:streamGenerateContent", true}, + // Non-matching + {"/v1beta/models/gemini-2.0-flash/embedContent", false}, + {"/v1beta/models", false}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + assert.Equal(t, tt.matches, containsGenerateContent(tt.path)) + }) + } +} + +func TestExtractModelFromPath(t *testing.T) { + tests := []struct { + path string + model string + }{ + {"/v1beta/models/gemini-2.0-flash:generateContent", "gemini-2.0-flash"}, + {"/v1beta/models/gemini-2.0-flash:streamGenerateContent", "gemini-2.0-flash"}, + {"/v1beta/models/gemini-2.0-flash/generateContent", "gemini-2.0-flash"}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + assert.Equal(t, tt.model, extractModelFromPath(tt.path)) + }) + } +} + +func TestIsStreamingPath(t *testing.T) { + tests := []struct { + path string + streaming bool + }{ + {"/v1beta/models/gemini-2.0-flash:generateContent", false}, + {"/v1beta/models/gemini-2.0-flash:streamGenerateContent", true}, + {"/v1beta/models/gemini-2.0-flash/streamGenerateContent", true}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + assert.Equal(t, tt.streaming, isStreamingPath(tt.path)) + }) + } +} + func TestCamelToSnake(t *testing.T) { tests := []struct { input string