feat(weaviate): add Braintrust tracing for Weaviate generative search - #241
Atharva Mhaske (atharvamhaske) wants to merge 2 commits into
Conversation
Traces only GraphQL Get queries using WithGenerativeSearch (server-side RAG), classified by the literal generate(...) field the SDK always emits. Plain vector search and CRUD produce no spans, since they aren't a generative-AI execution surface. Weaviate has no request middleware hook, so tracing wraps the client's *http.Client like the genai integration does. Closes braintrustdata#234 Signed-off-by: atharvamhaske <atharvamhaske76@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cc7c6c90ec
ℹ️ 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 !strings.Contains(raw.Query, "generate(") { | ||
| return ctx, trace.SpanFromContext(ctx), nil |
There was a problem hiding this comment.
Return an independent no-op span for untraced queries
For a plain vector/CRUD GraphQL query issued with an active parent span in its context, SpanFromContext(ctx) returns that caller-owned span rather than a no-op span. The shared internal.Middleware subsequently calls TagSpan, records transport errors, and calls End on the returned span, so an intentionally untraced query can prematurely end and mutate the application's enclosing span. This is not covered by the current no-trace test because it uses context.Background().
Useful? React with 👍 / 👎.
|
Abhijeet Prasad (@AbhiPrasad) i have added cassettes and screenshots on actual bt dashboard |
MultiClassGet bundles several Get queries into one GraphQL request, so data.Get can come back with more than one class key. The tracer only kept whichever class Go's map iteration landed on first, silently dropping the rest. metadata.class_name is now metadata.class_names (all classes present), and output is a class-keyed map when there is more than one class. Signed-off-by: atharvamhaske <atharvamhaske76@gmail.com>
Summary
trace/contrib/weaviate. Closes [bot] Instrument Weaviate (2,169,002 weekly downloads) #234.client.GraphQL().Get().WithGenerativeSearch(...)calls are traced. Plain vector search and CRUD produce no spans.*http.Clientthroughweaviate.Config.ConnectionClient, the same approach thegenaiintegration uses.MultiClassGet: a single request can return more than one class, sometadata.class_nameslists every class and the output covers all of them.metadata.modelis not set. Weaviate's generative search response never names the underlying LLM model.examples/internal/weaviate.Note on token counts
The dashboard can show a token count and cost estimate on a
weaviate.graphql.generatespan. That comes from the Braintrust UI's own tokenizer, not from Weaviate. The tracer only setsmetrics.start,metrics.end, andmetrics.duration.Changes
trace/contrib/weaviate/: new module (tracewaeviate.go,generativesearch.go,orchestrion.go,orchestrion.yml, tests, cassettes).examples/internal/weaviate/: new example.go.work,scripts/nested_modules.txt: registered the new module.go.work'sgodirective moved from1.25.0to1.26for a transitiveweaviate-go-clientdependency.trace/contrib/all/: registered weaviate in the meta-module and regeneratedorchestrion.yml.trace/contrib/orchestrion_test.go,trace/contrib/testdata/orchestrion/: addedTestWeaviateand registered it in the fixture.internal/genorchestrion/genorchestrion_test.go: raised the expected aspect count from 20 to 21.trace/contrib/README.md: added a Weaviate section.Notes
make cipasses locally. All tests run against a real local Weaviate instance and a real OpenAI key, recorded to VCR cassettes. No mocks.Screenshots