fix(rag): correctly attribute and track query embedding usage for active sessions#3603
fix(rag): correctly attribute and track query embedding usage for active sessions#3603Piyush0049 wants to merge 5 commits into
Conversation
aheritier
added
area/rag
Sayt-0
left a comment
Sayt-0
left a comment
There was a problem hiding this comment.
Thanks for tackling this long-standing TODO. The direction is right: removing the hidden usageHandler side effect in favor of explicit return values fixes a real accounting bug (query embedding tokens were credited to indexing and never reached session telemetry). Build and tests pass locally, and the concurrent accumulation in embedBatchOptimized is correctly done under the existing mutex.
A few points need to be addressed before merge:
| # | Finding | Severity | Location |
|---|---|---|---|
| 1 | The TODO comment this PR resolves is still present | blocking | pkg/rag/manager.go:424-426 |
| 2 | Tool name passed as the model parameter of RecordTokenUsage, polluting the model dimension in telemetry |
blocking | pkg/tools/builtin/rag/rag.go:182 |
| 3 | The cost returned by Embed/EmbedBatch is ignored by every caller, which recomputes it via calculateCost |
suggestion | pkg/rag/embed/embed.go, pkg/rag/strategy/vector_store.go |
| 4 | No test covers the new behavior (usage aggregation, telemetry emission) | strongly recommended | pkg/rag/ |
| 5 | Partial usage is discarded on error paths, a small accounting regression vs the callback approach | note | pkg/rag/strategy/vector_store.go, pkg/rag/embed/embed.go |
Details in the inline comments. None of this questions the overall approach; once 1 and 2 are addressed and some test coverage is added, this is good to go.
| // For queries during agent execution, usage should be added to agent's session | ||
| // This requires passing session context through the RAG tool |
There was a problem hiding this comment.
This TODO is resolved by this PR but the comment is left in place (the TODO starts one line above, at line 424). It should be removed, or rewritten to describe the new behavior: usage is returned to the caller and emitted by the RAG tool handler. Keeping it as-is will mislead the next reader into thinking query usage is still untracked.
| } | ||
|
|
||
| if usage.TotalTokens > 0 || usage.Cost > 0 { | ||
| telemetry.RecordTokenUsage(ctx, t.toolName, usage.TotalTokens, 0, usage.Cost) |
There was a problem hiding this comment.
RecordTokenUsage(ctx, model string, ...) expects a model name as its first argument and feeds TokenEvent.ModelName. Passing t.toolName here means the tool name (rag, or any custom toolset name) shows up as a model in analytics, which breaks per-model token aggregation.
The embedding model is known to VectorStore (s.modelID). One option: add a Model string field to types.Usage, set it in VectorStore.Query, and use it here (falling back to t.toolName when empty, e.g. for BM25-only configs which report no usage anyway). If mixing embedding models per query is a concern with hybrid configs, emitting one event per strategy usage would also work.
| // Embed generates an embedding for a single text | ||
| // Emits usage event immediately via handler if set | ||
| func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, error) { | ||
| func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, int64, float64, error) { |
There was a problem hiding this comment.
([]float64, int64, float64, error) carries two anonymous numeric returns, which is easy to misread or swap at call sites. Since this PR introduces types.Usage, returning ([]float64, types.Usage, error) (and the batch equivalent) would be clearer. pkg/rag/types has no dependency back on this package, so no import cycle.
Related: the returned cost is currently dead weight. Both call sites (VectorStore.Query, indexFile) discard it and recompute cost via calculateCost (models.dev pricing). Either the provider-reported cost should be trusted and used, or it should be dropped from the return values to keep a single source of truth. The pre-PR comment on the removed handler ("matches how chat completions calculate cost in runtime.go") suggests calculateCost is the intended source, in which case dropping the cost return is the consistent choice.
| return nil, types.Usage{}, fmt.Errorf("strategy %s failed: %w", result.name, result.err) | ||
| } | ||
|
|
||
| usage.Add(result.usage) |
There was a problem hiding this comment.
The core claim of this PR (usage is aggregated across strategies and surfaced to the caller) has no test. Suggested minimal coverage:
- a multi-strategy
Manager.Querytest assertingusageequals the sum of per-strategy usages (astaticStrategyreturning non-zerotypes.Usageis enough), - a
handleQueryRAGtest assertingRecordTokenUsagefires with the aggregated values (and does not fire when usage is zero), - optionally a
Usage.Addunit test.
| func (s *VectorStore) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error) { | ||
| queryEmbedding, err := s.embedder.Embed(ctx, query) | ||
| func (s *VectorStore) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) { | ||
| queryEmbedding, tokens, _, err := s.embedder.Embed(ctx, query) |
There was a problem hiding this comment.
Two small accounting regressions vs the callback approach, acceptable if intentional but worth stating:
- In
Query, ifSearchSimilarVectorsfails after a successfulEmbed, the tokens already consumed are returned astypes.Usage{}and never recorded anywhere. - In
indexFile,recordUsagenow fires once per file after the wholeEmbedBatchcompletes. On mid-batch failure, tokens consumed by successful batches are not recorded (previously the per-batch handler recorded them incrementally). This also coarsens TUI usage-event granularity from per-batch to per-file on large files.
Returning the partial usage alongside the error, or documenting that failed operations are not accounted, would settle this.
…usage in RAG * Add ModelID to types.Usage and implement aggregation in Usage.Add() * Refactor EmbedBatch to accumulate partial tokens on mid-batch error * Update VectorStore to always record partial usage and set correct ModelID * Update handleQueryRAG to use usage.ModelID for RecordTokenUsage telemetry * Remove unused cost return values from Embedder interfaces * Add unit tests for Usage.Add, Manager.Query usage aggregation, and EmbedBatch partial usage
|
@Sayt-0 Thank you for the review, I have pushed an update addressing all your feedback.
Please do let me know if any other changes are required. |
Updates the usage tracking implementation to supply the newly required context cancellation callback to strategy.Query added in upstream main, resolving the CI signature mismatch errors.
3 participants
Description
This PR resolves a long-standing TODO in the RAG
Managerto 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,
VectorStorerelied on a hidden globalusageHandlercallback attached to theEmbedder. 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
Embedderfor Purity: Removed theusageHandlercallback side-effect.Embed()andEmbedBatch()now explicitly return(tokens, cost)alongside the vectors.VectorStorenow manually triggersrecordUsage()only duringindexFile(). Query tokens are cleanly returned up the stack.Strategy.Query()was updated to return a newtypes.Usagestruct, allowingManager.Query()to cleanly aggregate usage across parallel search strategies (e.g., in a hybrid RAG configuration).handleQueryRAG) to capture the aggregated usage and emit it viatelemetry.RecordTokenUsage(ctx, ...). Because the tool execution context inherently carries thetelemetry.Client, this seamlessly solves the missing session attribution.Testing
go build ./...passes without errors.task test ./pkg/rag/...to ensure all strategy mocks, implementations, and fallback batches aggregate token counts safely without regressions.