Skip to content

fix(rag): correctly attribute and track query embedding usage for active sessions#3603

Open
Piyush0049 wants to merge 1 commit into
docker:mainfrom
Piyush0049:fix/rag-usage-tracking
Open

fix(rag): correctly attribute and track query embedding usage for active sessions#3603
Piyush0049 wants to merge 1 commit into
docker:mainfrom
Piyush0049:fix/rag-usage-tracking

Conversation

@Piyush0049

Copy link
Copy Markdown
Contributor

Description

This PR resolves a long-standing TODO in the RAG Manager to properly track, aggregate, and emit embedding token usage for agent search queries, ensuring they are accurately billed and attributed to the active session telemetry.

Previously, VectorStore relied on a hidden global usageHandler callback attached to the Embedder. This blindly fired during any embedding generation and logged the tokens as Indexing usage. As a result, tokens consumed by active RAG queries during a chat session were miscategorized as background file-indexing operations and completely missed by the session token tracker.

Changes Made

  • Refactored Embedder for Purity: Removed the usageHandler callback side-effect. Embed() and EmbedBatch() now explicitly return (tokens, cost) alongside the vectors.
  • Separated Query vs Indexing Attribution: VectorStore now manually triggers recordUsage() only during indexFile(). Query tokens are cleanly returned up the stack.
  • Updated Strategy Interfaces: Strategy.Query() was updated to return a new types.Usage struct, allowing Manager.Query() to cleanly aggregate usage across parallel search strategies (e.g., in a hybrid RAG configuration).
  • Wired Telemetry: Hooked the RAG tool handler (handleQueryRAG) to capture the aggregated usage and emit it via telemetry.RecordTokenUsage(ctx, ...). Because the tool execution context inherently carries the telemetry.Client, this seamlessly solves the missing session attribution.

Testing

  • Verified go build ./... passes without errors.
  • Executed task test ./pkg/rag/... to ensure all strategy mocks, implementations, and fallback batches aggregate token counts safely without regressions.
  • Confirmed high-concurrency batching loops are correctly protected by mutex locks for token accumulation.

@Piyush0049 Piyush0049 requested a review from a team as a code owner July 11, 2026 14:20
@aheritier aheritier added area/rag For work/issues that have to do with the RAG features area/tools For features/issues/fixes related to the usage of built-in and MCP tools kind/fix PR fixes a bug (maps to fix:). Use on PRs only. labels Jul 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/rag For work/issues that have to do with the RAG features area/tools For features/issues/fixes related to the usage of built-in and MCP tools kind/fix PR fixes a bug (maps to fix:). Use on PRs only.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants