Skip to content

Commit df45781

Browse files
markuswustenbergmarkuswustenberg
authored
Add OTel tracing (#39)
Fixes #36
1 parent 5df97d9 commit df45781

File tree

5 files changed

+162
-19
lines changed

5 files changed

+162
-19
lines changed

‎chat_complete.go‎

Lines changed: 69 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,14 @@ import (
55
"encoding/json"
66
"fmt"
77
"log/slog"
8+
"sort"
89
"strings"
910

1011
"github.com/openai/openai-go"
12+
"go.opentelemetry.io/otel"
13+
"go.opentelemetry.io/otel/attribute"
14+
"go.opentelemetry.io/otel/codes"
15+
"go.opentelemetry.io/otel/trace"
1116
"maragu.dev/gai"
1217
)
1318

@@ -22,6 +27,7 @@ type ChatCompleter struct {
2227
Client openai.Client
2328
log *slog.Logger
2429
model ChatCompleteModel
30+
tracer trace.Tracer
2531
}
2632

2733
type NewChatCompleterOptions struct {
@@ -33,15 +39,29 @@ func (c *Client) NewChatCompleter(opts NewChatCompleterOptions) *ChatCompleter {
3339
Client: c.Client,
3440
log: c.log,
3541
model: opts.Model,
42+
tracer: otel.Tracer("maragu.dev/gai-openai"),
3643
}
3744
}
3845

3946
// ChatComplete satisfies [gai.ChatCompleter].
4047
func (c *ChatCompleter) ChatComplete(ctx context.Context, req gai.ChatCompleteRequest) (gai.ChatCompleteResponse, error) {
48+
ctx, span := c.tracer.Start(ctx, "openai.chat_complete",
49+
trace.WithSpanKind(trace.SpanKindClient),
50+
trace.WithAttributes(
51+
attribute.String("ai.model", string(c.model)),
52+
attribute.Int("ai.message_count", len(req.Messages)),
53+
),
54+
)
55+
defer span.End()
56+
4157
var messages []openai.ChatCompletionMessageParamUnion
4258

4359
if req.System != nil {
4460
messages = append(messages, openai.SystemMessage(*req.System))
61+
span.SetAttributes(
62+
attribute.Bool("ai.has_system_prompt", true),
63+
attribute.String("ai.system_prompt", *req.System),
64+
)
4565
}
4666

4767
for _, m := range req.Messages {
@@ -132,6 +152,7 @@ func (c *ChatCompleter) ChatComplete(ctx context.Context, req gai.ChatCompleteRe
132152
}
133153

134154
var tools []openai.ChatCompletionToolParam
155+
var toolNames []string
135156
for _, tool := range req.Tools {
136157
tools = append(tools, openai.ChatCompletionToolParam{
137158
Function: openai.FunctionDefinitionParam{
@@ -143,21 +164,33 @@ func (c *ChatCompleter) ChatComplete(ctx context.Context, req gai.ChatCompleteRe
143164
},
144165
},
145166
})
167+
toolNames = append(toolNames, tool.Name)
146168
}
169+
sort.Strings(toolNames)
170+
span.SetAttributes(
171+
attribute.Int("ai.tool_count", len(tools)),
172+
attribute.StringSlice("ai.tools", toolNames),
173+
)
147174

148175
params := openai.ChatCompletionNewParams{
149176
Messages: messages,
150177
Model: openai.ChatModel(c.model),
151178
Tools: tools,
179+
StreamOptions: openai.ChatCompletionStreamOptionsParam{
180+
IncludeUsage: openai.Bool(true),
181+
},
152182
}
153183

154184
if req.Temperature != nil {
155185
params.Temperature = openai.Opt(req.Temperature.Float64())
186+
span.SetAttributes(attribute.Float64("ai.temperature", req.Temperature.Float64()))
156187
}
157188

158189
stream := c.Client.Chat.Completions.NewStreaming(ctx, params)
159190

160-
return gai.NewChatCompleteResponse(func(yield func(gai.MessagePart, error) bool) {
191+
meta := &gai.ChatCompleteResponseMetadata{}
192+
193+
res := gai.NewChatCompleteResponse(func(yield func(gai.MessagePart, error) bool) {
161194
defer func() {
162195
if err := stream.Close(); err != nil {
163196
c.log.Info("Error closing stream", "error", err)
@@ -169,33 +202,54 @@ func (c *ChatCompleter) ChatComplete(ctx context.Context, req gai.ChatCompleteRe
169202
chunk := stream.Current()
170203
acc.AddChunk(chunk)
171204

172-
if _, ok := acc.JustFinishedContent(); ok {
173-
break
174-
}
205+
if _, ok := acc.JustFinishedContent(); !ok {
206+
if toolCall, ok := acc.JustFinishedToolCall(); ok {
207+
if !yield(gai.ToolCallPart(toolCall.ID, toolCall.Name, json.RawMessage(toolCall.Arguments)), nil) {
208+
return
209+
}
210+
continue
211+
}
175212

176-
if toolCall, ok := acc.JustFinishedToolCall(); ok {
177-
if !yield(gai.ToolCallPart(toolCall.ID, toolCall.Name, json.RawMessage(toolCall.Arguments)), nil) {
213+
if refusal, ok := acc.JustFinishedRefusal(); ok {
214+
err := fmt.Errorf("refusal: %v", refusal)
215+
span.RecordError(err)
216+
span.SetStatus(codes.Error, "model refused request")
217+
yield(gai.MessagePart{}, err)
178218
return
179219
}
180-
continue
220+
221+
if len(chunk.Choices) > 0 {
222+
if !yield(gai.TextMessagePart(chunk.Choices[0].Delta.Content), nil) {
223+
return
224+
}
225+
}
181226
}
182227

183-
if refusal, ok := acc.JustFinishedRefusal(); ok {
184-
yield(gai.MessagePart{}, fmt.Errorf("refusal: %v", refusal))
185-
return
228+
if chunk.Usage.PromptTokens == 0 {
229+
continue
186230
}
187231

188-
if len(chunk.Choices) > 0 {
189-
if !yield(gai.TextMessagePart(chunk.Choices[0].Delta.Content), nil) {
190-
return
191-
}
232+
meta.Usage = gai.ChatCompleteResponseUsage{
233+
PromptTokens: int(chunk.Usage.PromptTokens),
234+
CompletionTokens: int(chunk.Usage.CompletionTokens),
192235
}
236+
span.SetAttributes(
237+
attribute.Int("ai.prompt_tokens", int(chunk.Usage.PromptTokens)),
238+
attribute.Int("ai.completion_tokens", int(chunk.Usage.CompletionTokens)),
239+
attribute.Int("ai.total_tokens", int(chunk.Usage.TotalTokens)),
240+
)
193241
}
194242

195243
if err := stream.Err(); err != nil {
244+
span.RecordError(err)
245+
span.SetStatus(codes.Error, "stream error")
196246
yield(gai.MessagePart{}, err)
197247
}
198-
}), nil
248+
})
249+
250+
res.Meta = meta
251+
252+
return res, nil
199253
}
200254

201255
// normalizeToolSchemaProperties recursively normalizes schema properties for OpenAI compatibility

‎chat_complete_test.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,38 @@ func TestChatCompleter_ChatComplete(t *testing.T) {
231231

232232
is.Equal(t, "Bonjour ! Comment puis-je vous aider aujourd'hui ?", output)
233233
})
234+
235+
t.Run("tracks token usage", func(t *testing.T) {
236+
cc := newChatCompleter(t)
237+
238+
req := gai.ChatCompleteRequest{
239+
Messages: []gai.Message{
240+
gai.NewUserTextMessage("Hi!"),
241+
},
242+
Temperature: gai.Ptr(gai.Temperature(0)),
243+
}
244+
245+
res, err := cc.ChatComplete(t.Context(), req)
246+
is.NotError(t, err)
247+
248+
// Consume the response to ensure token usage is populated
249+
var output string
250+
for part, err := range res.Parts() {
251+
is.NotError(t, err)
252+
if part.Type == gai.MessagePartTypeText {
253+
output += part.Text()
254+
}
255+
}
256+
257+
// Check that we got a response
258+
is.True(t, len(output) > 0, "should have response text")
259+
260+
// Check token usage in Meta.Usage
261+
is.NotNil(t, res.Meta, "should have metadata")
262+
t.Log(res.Meta.Usage.PromptTokens, res.Meta.Usage.CompletionTokens)
263+
is.True(t, res.Meta.Usage.PromptTokens > 0, "should have prompt tokens")
264+
is.True(t, res.Meta.Usage.CompletionTokens > 0, "should have completion tokens")
265+
})
234266
}
235267

236268
func newChatCompleter(t *testing.T) *openai.ChatCompleter {

‎embed.go‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@ import (
55
"log/slog"
66

77
"github.com/openai/openai-go"
8+
"go.opentelemetry.io/otel"
9+
"go.opentelemetry.io/otel/attribute"
10+
"go.opentelemetry.io/otel/codes"
11+
"go.opentelemetry.io/otel/trace"
812
"maragu.dev/errors"
913
"maragu.dev/gai"
1014
)
@@ -21,6 +25,7 @@ type Embedder struct {
2125
dimensions int
2226
log *slog.Logger
2327
model EmbedModel
28+
tracer trace.Tracer
2429
}
2530

2631
type NewEmbedderOptions struct {
@@ -49,12 +54,23 @@ func (c *Client) NewEmbedder(opts NewEmbedderOptions) *Embedder {
4954
dimensions: opts.Dimensions,
5055
log: c.log,
5156
model: opts.Model,
57+
tracer: otel.Tracer("maragu.dev/gai-openai"),
5258
}
5359
}
5460

5561
// Embed satisfies [gai.Embedder].
5662
func (e *Embedder) Embed(ctx context.Context, req gai.EmbedRequest) (gai.EmbedResponse[float64], error) {
63+
ctx, span := e.tracer.Start(ctx, "openai.embed",
64+
trace.WithSpanKind(trace.SpanKindClient),
65+
trace.WithAttributes(
66+
attribute.String("ai.model", string(e.model)),
67+
attribute.Int("ai.dimensions", e.dimensions),
68+
),
69+
)
70+
defer span.End()
71+
5772
v := gai.ReadAllString(req.Input)
73+
span.SetAttributes(attribute.Int("ai.input_length", len(v)))
5874

5975
res, err := e.Client.Embeddings.New(ctx, openai.EmbeddingNewParams{
6076
Input: openai.EmbeddingNewParamsInputUnion{OfString: openai.Opt(v)},
@@ -63,10 +79,23 @@ func (e *Embedder) Embed(ctx context.Context, req gai.EmbedRequest) (gai.EmbedRe
6379
Dimensions: openai.Opt(int64(e.dimensions)),
6480
})
6581
if err != nil {
82+
span.RecordError(err)
83+
span.SetStatus(codes.Error, "embedding request failed")
6684
return gai.EmbedResponse[float64]{}, errors.Wrap(err, "error embedding")
6785
}
6886
if len(res.Data) == 0 {
69-
return gai.EmbedResponse[float64]{}, errors.New("no embeddings returned")
87+
err := errors.New("no embeddings returned")
88+
span.RecordError(err)
89+
span.SetStatus(codes.Error, "no embeddings in response")
90+
return gai.EmbedResponse[float64]{}, err
91+
}
92+
93+
// Record token usage if available
94+
if res.Usage.PromptTokens > 0 {
95+
span.SetAttributes(
96+
attribute.Int("ai.prompt_tokens", int(res.Usage.PromptTokens)),
97+
attribute.Int("ai.total_tokens", int(res.Usage.TotalTokens)),
98+
)
7099
}
71100

72101
return gai.EmbedResponse[float64]{

‎go.mod‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ go 1.24
44

55
require (
66
github.com/openai/openai-go v1.12.0
7+
go.opentelemetry.io/otel v1.38.0
8+
go.opentelemetry.io/otel/trace v1.38.0
79
maragu.dev/env v0.2.0
810
maragu.dev/errors v0.3.0
911
maragu.dev/gai v0.0.0-20250826105131-21f642fac70d
@@ -13,12 +15,16 @@ require (
1315
require (
1416
github.com/bahlo/generic-list-go v0.2.0 // indirect
1517
github.com/buger/jsonparser v1.1.1 // indirect
18+
github.com/go-logr/logr v1.4.3 // indirect
19+
github.com/go-logr/stdr v1.2.2 // indirect
1620
github.com/invopop/jsonschema v0.13.0 // indirect
1721
github.com/mailru/easyjson v0.7.7 // indirect
1822
github.com/tidwall/gjson v1.14.4 // indirect
1923
github.com/tidwall/match v1.1.1 // indirect
2024
github.com/tidwall/pretty v1.2.1 // indirect
2125
github.com/tidwall/sjson v1.2.5 // indirect
2226
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
27+
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
28+
go.opentelemetry.io/otel/metric v1.38.0 // indirect
2329
gopkg.in/yaml.v3 v3.0.1 // indirect
2430
)

‎go.sum‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,30 @@ github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMU
44
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
55
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
66
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
7+
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
8+
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
9+
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
10+
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
11+
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
12+
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
13+
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
714
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
815
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
916
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
17+
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
18+
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
19+
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
20+
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
1021
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
1122
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
1223
github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0=
1324
github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y=
1425
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
1526
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
16-
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
17-
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
27+
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
28+
github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
29+
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
30+
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
1831
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
1932
github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
2033
github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
@@ -27,8 +40,17 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
2740
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
2841
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
2942
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
30-
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
43+
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
44+
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
45+
go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
46+
go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
47+
go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
48+
go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
49+
go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
50+
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
3151
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
52+
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
53+
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
3254
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
3355
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
3456
maragu.dev/env v0.2.0 h1:nQKitDEB65ArZsh6E7vxzodOqY9bxEVFdBg+tskS1ys=

0 commit comments

Comments
 (0)