-
Couldn't load subscription status.
- Fork 2.1k
feat(tracing): custom tracing processor which injects tenant_id #5918
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
88ed676 to
383f7e2
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Greptile Overview
Greptile Summary
This PR introduces per-tenant usage tracking for Braintrust traces by injecting tenant_id into trace metadata. The implementation uses a custom TenantContextTracingProcessor that adds tenant context when traces start, and updates @traced decorators across LLM invocation points to include tenant_id metadata.
Key changes:
- Created
TenantContextTracingProcessorthat injects tenant_id from context variables into trace metadata - Fixed decorator metadata to use lambda functions, ensuring tenant_id is captured at runtime rather than at module import time (this addresses the issue from the parent commit)
- Upgraded braintrust SDK from 0.2.6 to 0.3.5 to resolve processor-related bugs
- Applied tenant_id tracking to
invoke llm,stream llm, andclarifier stream and processtraces
The PR correctly addresses the timing issue mentioned by the author - using metadata=lambda: defers the contextvar evaluation until function execution, which is the proper approach for capturing dynamic runtime context.
Confidence Score: 5/5
- Safe to merge - properly implements tenant tracking with correct timing semantics
- The PR correctly fixes the timing issue from the parent commit by using lambda functions to defer contextvar evaluation. The custom processor implementation is straightforward and follows expected patterns. The braintrust SDK upgrade addresses a known bug. All changes are isolated to tracing infrastructure with no breaking changes to business logic.
- No files require special attention
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| backend/onyx/tracing/braintrust_tracing.py | 4/5 | Adds TenantContextTracingProcessor to inject tenant_id into trace metadata; implementation is sound but could handle None case more explicitly |
| backend/onyx/llm/interfaces.py | 5/5 | Updated @traced decorator with lambda metadata to correctly capture tenant_id at runtime; fixes timing issue from parent commit |
| backend/onyx/agents/agent_search/shared_graph_utils/llm.py | 5/5 | Updated @traced decorator with lambda metadata for runtime tenant_id capture; properly defers evaluation |
| backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py | 5/5 | Fixed @traced decorator to use lambda for tenant_id metadata; correctly defers contextvar evaluation until function execution |
| backend/requirements/default.txt | 5/5 | Upgrades braintrust SDK from 0.2.6 to 0.3.5 to fix processor-related bug |
Sequence Diagram
sequenceDiagram
participant App as Application Code
participant Dec as @traced Decorator
participant Proc as TenantContextTracingProcessor
participant Ctx as CURRENT_TENANT_ID_CONTEXTVAR
participant BT as Braintrust SDK
Note over App,BT: Setup Phase (module import)
App->>Dec: Define @traced(metadata=lambda: {...})
Note over Dec: Lambda NOT evaluated yet
Note over App,BT: Runtime Phase (function call)
App->>Dec: Call decorated function
Dec->>Dec: Evaluate metadata lambda
Dec->>Ctx: get() tenant_id
Ctx-->>Dec: Return current tenant_id
Dec->>Proc: on_trace_start(trace)
Proc->>Ctx: get() tenant_id
Ctx-->>Proc: Return current tenant_id
Proc->>Proc: Set trace.metadata["tenant_id"]
Proc->>BT: super().on_trace_start(trace)
BT-->>Proc: Trace started
Proc-->>Dec: Continue
Dec->>App: Execute function
App-->>Dec: Return result
Dec->>BT: Log trace with metadata
BT-->>Dec: Trace logged
Dec-->>App: Return result
5 files reviewed, no comments
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
2 issues found across 5 files
Prompt for AI agents (all 2 issues)
Understand the root cause of the following 2 issues and fix them.
<file name="backend/onyx/llm/interfaces.py">
<violation number="1" location="backend/onyx/llm/interfaces.py:94">
Evaluating CURRENT_TENANT_ID_CONTEXTVAR.get() in the decorator runs at import time, so each span ends up with tenant_id stuck at the ContextVar’s default (often None) instead of the active request’s tenant. Move the lookup to runtime so it reads the ContextVar per invocation.</violation>
</file>
<file name="backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py">
<violation number="1" location="backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py:692">
`traced.metadata` expects a mapping (e.g., with a `tenant_id` key). Passing the raw context-var string loses the key/value structure, so the trace won’t include the tenant metadata. Wrap the value in a dict.</violation>
</file>
React with 👍 or 👎 to teach cubic. Mention @cubic-dev-ai to give feedback, ask questions, or re-run the review.
Description
Includes
tenant_idalong with braintrust traces which allows for per-tenant usage tracking.I'm not entirely sure why the why the custom processor doesn't capture the
@traceddecorator. AFAICT, it seems like it should -- it does look like it uses the samespan.log()machinery, but couldn't get those callsites to even register the custom processor. Defining the metadata along with the decorator seems non-ideal but also not too burdensome. There is also a non-zero chance there is a bug/quirk in the braintrust sdk that is being exposed.Also upgrades braintrust sdk to latest as that appeared to fix a bug relevant to the
TenantContextTracingProcessor.How Has This Been Tested?
Tested on existing
invoke llm,stream llmandfast_chat_turntraces and confirmed their metadata appears in Braintrust,Additional Options
Summary by cubic
Add a custom Braintrust tracing processor that injects tenant_id into trace metadata and updates LLM trace decorators to include tenant_id, enabling per-tenant usage tracking. Also upgrade Braintrust SDK to v0.3.5 to resolve a processor issue.
New Features
Dependencies