feat(cloudflare): add Braintrust tracing for Cloudflare Workers AI - #238
Atharva Mhaske (atharvamhaske) wants to merge 3 commits into
Conversation
…I Go SDK Signed-off-by: atharvamhaske <atharvamhaske76@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 235339902b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| if _, ok := raw["image"]; ok { | ||
| return taskMultimodalEmbedding | ||
| } | ||
| return taskTextEmbeddings |
There was a problem hiding this comment.
Because cloudflareRouter creates an aiRunTracer for every /ai/run/ URL, any Workers AI operation whose JSON omits both messages and image takes this default branch. Those calls are recorded as text embeddings and the subsequent input/output handling looks only for text and {data, shape}, so other task shapes produce mislabeled and incomplete traces. Restrict this default to recognized embedding bodies or add classifiers for the remaining Workers AI task schemas.
Useful? React with 👍 / 👎.
| if shape, ok := result["shape"]; ok { | ||
| summary["shape"] = shape | ||
| } | ||
| return internal.SetJSONAttr(span, "braintrust.output_json", summary) |
There was a problem hiding this comment.
This embeddings path only writes braintrust.output_json and never records the usage object returned by Workers AI. Consequently every text- or multimodal-embedding trace loses prompt/total-token accounting, unlike the text-generation path above, making token-based cost and experiment summaries inaccurate. Parse result["usage"] here and set braintrust.metrics as well.
Useful? React with 👍 / 👎.
|
can you make sure the instrumentation follows this skill? https://github.com/braintrustdata/braintrust-spec/blob/main/skills/instrumentation-spec/SKILL.md |
yes sure, adding this on my workflows |
…output Response-side classification (instead of guessing from the ambiguous request shape) covers all Workers AI task types; text generation now emits OpenAI-format choices/tool_calls and OpenAI-nested metadata.tools per the Braintrust instrumentation spec. Signed-off-by: atharvamhaske <atharvamhaske76@gmail.com>
…kers-ai-tracing # Conflicts: # internal/genorchestrion/genorchestrion_test.go # trace/contrib/all/all.go # trace/contrib/testdata/orchestrion/go.mod # trace/contrib/testdata/orchestrion/go.sum # trace/contrib/testdata/orchestrion/orchestrion_test.go
|
Abhijeet Prasad (@AbhiPrasad) i have pushed some changes after using skills and added screenshots, can be reviewed and test on your end ! |
2 participants
Summary
trace/contrib/cloudflare, covering text generation (chat, tool calls), text/multimodal embeddings, and every other Workers AI task type throughclient.AI.Run.{"image":[...]}and bare{"text":"..."}each map to two different task types). The tracer classifies from the response instead of the request, which is always unambiguous:{data,shape}is an embeddings result, ausagekey is a text generation result, anything else is logged as Cloudflare's own response shape.choicesarray,tool_callswith string-encodedarguments) and available tools convert to the OpenAI-nestedmetadata.toolsshape, per the Braintrust instrumentation spec.cloudflare.NewClientat compile time. No manualWithMiddlewarecall is needed.examples/internal/cloudflare, now including tool-calling and text classification.Verification
github.com/cloudflare/cloudflare-go/v6. The current version isv7.10.0. Every type the audit cited exists under thev7path.@cf/meta/llama-3.1-8b-instruct-fast(text generation, tool calls),@cf/baai/bge-base-en-v1.5(text embeddings), and@cf/huggingface/distilbert-sst-2-int8(text classification) confirmed the response shapes above, including that Cloudflare sets"response": null(not omitted) on tool-calling-only replies, so classification keys onusageinstead.{data,shape}response shape as text embeddings, so it's covered by the same code path and test.braintrust_attachmentreferences yet; oversized binary fields are summarized to{type, length}so a span never carries megabytes of raw bytes, but that's a stopgap, not a real attachment upload. Left as a follow-up.Changes
trace/contrib/cloudflare/:tracecloudflare.go,airun.go(response-based classification, OpenAI-format conversion),orchestrion.go,orchestrion.yml, tests, cassettes.examples/internal/cloudflare/: example, now covering tool-calling and text classification.examples/go.mod,go.work,scripts/nested_modules.txt: registered the new module.trace/contrib/all/: registered cloudflare in the meta-module and regeneratedorchestrion.yml.trace/contrib/orchestrion_test.go: added cloudflare to the fixture's replace-directive map and import list.trace/contrib/testdata/orchestrion/: addedTestCloudflare. It confirms orchestrion instruments a plaincloudflare.NewClient(...)call with no manualWithMiddleware.internal/genorchestrion/genorchestrion_test.go: raised the expected aspect count from 18 to 19.Notes
make cipasses locally, including bothTestOrchestrionInjectionvariants.Screenshots
1. text generation
2. embeddings
Closes #236
Note on task type coverage
This PR does not add 14 example calls for the 14 Workers AI task types. The tracer code handles all task types. It sorts each response into one of 3 groups: embeddings, text generation, or a fallback group for everything else. The fallback group logs Cloudflare's own response shape as-is.
The example script shows one live call per group, plus a tool-calling call. That gives 4 calls total: text generation, tool calls, embeddings, and text classification (the fallback group).
The other 9 task types (translation, summarization, image tasks, speech tasks) do not have example calls yet. Most of them need image or audio input, and the instrumentation spec does not yet define an attachment format for that input. This PR does not add those attachments.