From caace32bc6a1197a57e7d8e8dbaa465992d2be72 Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Fri, 3 Apr 2026 10:31:42 -0400 Subject: [PATCH 1/2] Instrument Gemini streaming (streamGenerateContent) calls The genai integration was not tracing streaming API calls at all. The URL router could not match streamGenerateContent endpoints, and the response parser only handled single-JSON responses. This adds full streaming support: SSE chunk parsing, response aggregation, time_to_first_token metrics, and proper URL routing for both streaming and non-streaming Gemini endpoints. Also adds time_to_first_token to non-streaming responses. Fixes #50 Co-Authored-By: Claude Opus 4.6 (1M context) --- examples/internal/genai/go.mod | 34 ++-- examples/internal/genai/go.sum | 80 ++++---- examples/internal/genai/main.go | 19 +- trace/contrib/genai/generatecontent.go | 185 +++++++++++++++++- .../TestStreamingGenerateContent.yaml | 51 +++++ trace/contrib/genai/tracegenai.go | 31 ++- trace/contrib/genai/tracegenai_test.go | 115 ++++++++++- 7 files changed, 433 insertions(+), 82 deletions(-) create mode 100644 trace/contrib/genai/testdata/cassettes/TestStreamingGenerateContent.yaml 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 From ba38905b1af7b91a19aca7dd9d2e8b9a4c90d22c Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Fri, 3 Apr 2026 11:35:31 -0400 Subject: [PATCH 2/2] Add genai-streaming example demonstrating new streaming instrumentation Shows that both non-streaming and streaming Gemini calls now produce complete spans with output, token metrics, and time_to_first_token. Co-Authored-By: Claude Opus 4.6 (1M context) --- examples/internal/genai-streaming/go.mod | 39 +++++++ examples/internal/genai-streaming/go.sum | 75 ++++++++++++++ examples/internal/genai-streaming/main.go | 119 ++++++++++++++++++++++ 3 files changed, 233 insertions(+) create mode 100644 examples/internal/genai-streaming/go.mod create mode 100644 examples/internal/genai-streaming/go.sum create mode 100644 examples/internal/genai-streaming/main.go 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 +}