Skip to content

Conversation

@trial2onyx
Copy link
Collaborator

@trial2onyx trial2onyx commented Oct 24, 2025

Description

Includes tenant_id along 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 @traced decorator. AFAICT, it seems like it should -- it does look like it uses the same span.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 llm and fast_chat_turn traces and confirmed their metadata appears in Braintrust,

Screenshot 2025-10-24 at 3 01 33 PM

Additional Options

  • Override Linear Check

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

    • TenantContextTracingProcessor injects tenant_id on trace start.
    • Added tenant_id from CURRENT_TENANT_ID_CONTEXTVAR to @Traced for invoke llm, stream llm, and clarifier.
    • Switched set_trace_processors to use the new processor.
  • Dependencies

    • braintrust[openai-agents]: 0.2.6 → 0.3.5

@trial2onyx trial2onyx requested a review from a team as a code owner October 24, 2025 22:05
@vercel
Copy link

vercel bot commented Oct 24, 2025

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Preview Comments Updated (UTC)
internal-search Ready Ready Preview Comment Oct 25, 2025 1:52am

greptile-apps[bot]

This comment was marked as outdated.

@trial2onyx trial2onyx force-pushed the jamison/braintrust-processor branch from 88ed676 to 383f7e2 Compare October 24, 2025 22:11
@trial2onyx
Copy link
Collaborator Author

@greptileai

Copy link
Contributor

@greptile-apps greptile-apps bot left a 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 TenantContextTracingProcessor that 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, and clarifier stream and process traces

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
Loading

5 files reviewed, no comments

Edit Code Review Agent Settings | Greptile

Copy link
Contributor

@cubic-dev-ai cubic-dev-ai bot left a 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants