From 5b11d2b19a497ddafeeb80f9c1c48d9bad399600 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Sat, 11 Jul 2026 19:42:46 +0530 Subject: [PATCH] fix(rag): track and emit query embedding usage to active session --- pkg/rag/embed/embed.go | 52 ++++++++++++++------------------ pkg/rag/manager.go | 22 ++++++++------ pkg/rag/manager_rerank_test.go | 11 ++++--- pkg/rag/strategy/bm25.go | 10 +++--- pkg/rag/strategy/types.go | 3 +- pkg/rag/strategy/vector_store.go | 26 ++++++++-------- pkg/rag/types/types.go | 12 ++++++++ pkg/tools/builtin/rag/rag.go | 7 ++++- 8 files changed, 78 insertions(+), 65 deletions(-) diff --git a/pkg/rag/embed/embed.go b/pkg/rag/embed/embed.go index 312480b417..6537bc4543 100644 --- a/pkg/rag/embed/embed.go +++ b/pkg/rag/embed/embed.go @@ -14,9 +14,8 @@ import ( // Embedder generates vector embeddings for text type Embedder struct { provider provider.Provider - usageHandler func(ctx context.Context, tokens int64, cost float64) // Callback to emit usage events - batchSize int // Batch size for API calls - maxConcurrency int // Maximum concurrent embedding batch requests + batchSize int // Batch size for API calls + maxConcurrency int // Maximum concurrent embedding batch requests } // Option is a functional option for configuring the Embedder @@ -51,24 +50,14 @@ func New(p provider.Provider, opts ...Option) *Embedder { return e } -// SetUsageHandler sets a callback to be called after each embedding with usage info -func (e *Embedder) SetUsageHandler(handler func(ctx context.Context, tokens int64, cost float64)) { - e.usageHandler = handler -} // 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) { // Try to use the provider's embedding API if it implements EmbeddingProvider. if embeddingProvider, ok := e.provider.(provider.EmbeddingProvider); ok { result, err := embeddingProvider.CreateEmbedding(ctx, text) if err != nil { - return nil, err - } - - // Emit usage event immediately - if e.usageHandler != nil { - e.usageHandler(ctx, result.TotalTokens, result.Cost) + return nil, 0, 0, err } slog.DebugContext(ctx, "Embedding generated", @@ -76,19 +65,19 @@ func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, error) { "tokens", result.TotalTokens, "cost", result.Cost) - return result.Embedding, nil + return result.Embedding, result.TotalTokens, result.Cost, nil } // Provider does not support embeddings via the standard interface; fail fast. - return nil, fmt.Errorf("provider %s does not support embeddings", e.provider.ID()) + return nil, 0, 0, fmt.Errorf("provider %s does not support embeddings", e.provider.ID()) } // EmbedBatch generates embeddings for multiple texts using intelligent batching // If the provider supports batch embeddings, it will use parallel batch API calls // Otherwise, it falls back to sequential processing -func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, error) { +func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, int64, float64, error) { if len(texts) == 0 { - return [][]float64{}, nil + return [][]float64{}, 0, 0, nil } // Check if provider supports batch embeddings. @@ -102,19 +91,23 @@ func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, "text_count", len(texts)) embeddings := make([][]float64, len(texts)) + var totalTokens int64 + var totalCost float64 for i, text := range texts { - embedding, err := e.Embed(ctx, text) + embedding, tokens, cost, err := e.Embed(ctx, text) if err != nil { - return nil, fmt.Errorf("failed to embed text %d: %w", i, err) + return nil, 0, 0, fmt.Errorf("failed to embed text %d: %w", i, err) } embeddings[i] = embedding + totalTokens += tokens + totalCost += cost } - return embeddings, nil + return embeddings, totalTokens, totalCost, nil } // embedBatchOptimized processes texts in optimized batches with parallel API calls -func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, error) { +func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, int64, float64, error) { totalTexts := len(texts) slog.DebugContext(ctx, "Starting optimized batch embedding", "provider", e.provider.ID(), @@ -125,6 +118,8 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid // Pre-allocate results embeddings := make([][]float64, totalTexts) var mu sync.Mutex + var totalTokens int64 + var totalCost float64 // Create errgroup with concurrency limit g, ctx := errgroup.WithContext(ctx) @@ -154,13 +149,10 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid // Store results (mutex protects slice writes) mu.Lock() copy(embeddings[start:end], result.Embeddings) + totalTokens += result.TotalTokens + totalCost += result.Cost mu.Unlock() - // Emit usage event (handler should be thread-safe) - if e.usageHandler != nil { - e.usageHandler(ctx, result.TotalTokens, result.Cost) - } - slog.DebugContext(ctx, "Batch completed", "batch", batchNum, "embeddings", len(result.Embeddings), @@ -173,7 +165,7 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid // Wait for all batches and return first error if any if err := g.Wait(); err != nil { - return nil, err + return nil, 0, 0, err } slog.DebugContext(ctx, "Batch embedding completed", @@ -181,5 +173,5 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid "total_embeddings", len(embeddings), "batches_processed", (totalTexts+e.batchSize-1)/e.batchSize) - return embeddings, nil + return embeddings, totalTokens, totalCost, nil } diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index ec5088b2ff..3648c03e94 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -236,7 +236,7 @@ func (m *Manager) Initialize(ctx context.Context) (err error) { // Query searches for relevant documents using all configured strategies // If multiple strategies are configured, results are combined using the fusion strategy -func (m *Manager) Query(ctx context.Context, query string) (results []database.SearchResult, err error) { +func (m *Manager) Query(ctx context.Context, query string) (results []database.SearchResult, usage types.Usage, err error) { // Start a `retrieval {rag_name}` span per the OTel GenAI semconv. // The query text itself is sensitive so we never capture it on the // span here — content capture is gated by a separate environment @@ -270,13 +270,13 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S // `:=`) so the deferred span closure sees the live values // even if a future change replaces the explicit // `return X, Y` form below with a bare `return`. - results, err = strategyImpl.Query(ctx, query, strategyCfg.Limit, strategyCfg.Threshold) + results, usage, err = strategyImpl.Query(ctx, query, strategyCfg.Limit, strategyCfg.Threshold) if err != nil { slog.ErrorContext(ctx, "[RAG Manager] Strategy query failed", "rag_name", m.name, "strategy", strategyName, "error", err) - return nil, err + return nil, types.Usage{}, err } slog.DebugContext(ctx, "[RAG Manager] Single strategy results", @@ -309,7 +309,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "num_results", len(results)) } - return results, nil + return results, usage, nil } } @@ -321,6 +321,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S type strategyResult struct { name string results []database.SearchResult + usage types.Usage err error } @@ -337,10 +338,11 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "strategy_threshold", strategyCfg.Threshold) go func(name string, strategyImpl strategy.Strategy, cfg strategy.Config) { - results, err := strategyImpl.Query(ctx, query, cfg.Limit, cfg.Threshold) + results, usage, err := strategyImpl.Query(ctx, query, cfg.Limit, cfg.Threshold) resultsChan <- strategyResult{ name: name, results: results, + usage: usage, err: err, } }(strategyName, strategyImpl, strategyCfg) @@ -356,8 +358,10 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "rag_name", m.name, "strategy", result.name, "error", result.err) - return nil, fmt.Errorf("strategy %s failed: %w", result.name, result.err) + return nil, types.Usage{}, fmt.Errorf("strategy %s failed: %w", result.name, result.err) } + + usage.Add(result.usage) slog.DebugContext(ctx, "[RAG Manager] Strategy returned results", "rag_name", m.name, @@ -375,7 +379,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S // Safety check: fusion should never be nil with multiple strategies if m.fusion == nil { - return nil, errors.New("fusion strategy is nil but multiple strategies are configured (this is a bug)") + return nil, types.Usage{}, errors.New("fusion strategy is nil but multiple strategies are configured (this is a bug)") } fusedResults, err := m.fusion.Fuse(strategyResults) @@ -383,7 +387,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S slog.ErrorContext(ctx, "[RAG Manager] Fusion failed", "rag_name", m.name, "error", err) - return nil, fmt.Errorf("failed to fuse results: %w", err) + return nil, types.Usage{}, fmt.Errorf("failed to fuse results: %w", err) } slog.DebugContext(ctx, "[RAG Manager] Fusion complete", @@ -421,7 +425,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S // For queries during agent execution, usage should be added to agent's session // This requires passing session context through the RAG tool - return fusedResults, nil + return fusedResults, usage, nil } // Helper to get strategy names for logging diff --git a/pkg/rag/manager_rerank_test.go b/pkg/rag/manager_rerank_test.go index 66ba3b1b50..60aca86a7c 100644 --- a/pkg/rag/manager_rerank_test.go +++ b/pkg/rag/manager_rerank_test.go @@ -13,6 +13,7 @@ import ( "github.com/docker/docker-agent/pkg/modelerrors" "github.com/docker/docker-agent/pkg/rag/database" "github.com/docker/docker-agent/pkg/rag/strategy" + "github.com/docker/docker-agent/pkg/rag/types" ) // failingReranker counts calls and always fails with a fixed error. @@ -35,8 +36,8 @@ func (s *staticStrategy) Initialize(context.Context, []string, strategy.Chunking return nil } -func (s *staticStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, error) { - return s.results, nil +func (s *staticStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, types.Usage, error) { + return s.results, types.Usage{}, nil } func (s *staticStrategy) CheckAndReindexChangedFiles(context.Context, []string, strategy.ChunkingConfig) error { @@ -83,7 +84,7 @@ func TestQueryDisablesRerankerAfterNonRetryableError(t *testing.T) { m, reranker := newRerankTestManager(t, rerankErr) for range 3 { - results, err := m.Query(t.Context(), "some query") + results, _, err := m.Query(t.Context(), "some query") require.NoError(t, err, "rerank failures must not fail the query") assert.Len(t, results, 2, "original results are returned as fallback") } @@ -101,7 +102,7 @@ func TestQueryKeepsRerankerOnTransientError(t *testing.T) { m, reranker := newRerankTestManager(t, rerankErr) for range 3 { - results, err := m.Query(t.Context(), "some query") + results, _, err := m.Query(t.Context(), "some query") require.NoError(t, err) assert.Len(t, results, 2) } @@ -114,7 +115,7 @@ func TestQueryKeepsRerankerOnContextCancellation(t *testing.T) { t.Parallel() m, reranker := newRerankTestManager(t, context.Canceled) - results, err := m.Query(t.Context(), "some query") + results, _, err := m.Query(t.Context(), "some query") require.NoError(t, err) assert.Len(t, results, 2) assert.Equal(t, int64(1), reranker.calls.Load()) diff --git a/pkg/rag/strategy/bm25.go b/pkg/rag/strategy/bm25.go index 77cd5c2471..3b3fbd7240 100644 --- a/pkg/rag/strategy/bm25.go +++ b/pkg/rag/strategy/bm25.go @@ -271,21 +271,21 @@ func (s *BM25Strategy) Initialize(ctx context.Context, docPaths []string, chunki } // Query searches for relevant documents using BM25 scoring -func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error) { +func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) { // Tokenize query queryTerms := s.tokenize(query) if len(queryTerms) == 0 { - return nil, errors.New("query contains no valid terms") + return nil, types.Usage{}, errors.New("query contains no valid terms") } // Get all documents allDocs, err := s.getAllDocuments(ctx) if err != nil { - return nil, fmt.Errorf("failed to retrieve documents: %w", err) + return nil, types.Usage{}, fmt.Errorf("failed to retrieve documents: %w", err) } if len(allDocs) == 0 { - return []database.SearchResult{}, nil + return []database.SearchResult{}, types.Usage{}, nil } // Pre-tokenize all documents once: build term frequency maps and lengths. @@ -333,7 +333,7 @@ func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int, scores = scores[:numResults] } - return scores, nil + return scores, types.Usage{}, nil } // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed diff --git a/pkg/rag/strategy/types.go b/pkg/rag/strategy/types.go index 8617dfb27f..1a2e9f1742 100644 --- a/pkg/rag/strategy/types.go +++ b/pkg/rag/strategy/types.go @@ -4,6 +4,7 @@ import ( "context" "github.com/docker/docker-agent/pkg/rag/database" + "github.com/docker/docker-agent/pkg/rag/types" ) // Strategy defines the interface for different retrieval strategies. @@ -14,7 +15,7 @@ type Strategy interface { // Query searches for relevant documents using the strategy's retrieval method. // numResults is the maximum number of candidates to retrieve (before fusion). - Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error) + Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed. CheckAndReindexChangedFiles(ctx context.Context, docPaths []string, chunking ChunkingConfig) error diff --git a/pkg/rag/strategy/vector_store.go b/pkg/rag/strategy/vector_store.go index 03f1dc8b6f..77e96558bd 100644 --- a/pkg/rag/strategy/vector_store.go +++ b/pkg/rag/strategy/vector_store.go @@ -147,14 +147,6 @@ func NewVectorStore(cfg VectorStoreConfig) *VectorStore { embeddingConcurrency: cfg.EmbeddingConcurrency, fileIndexConcurrency: cfg.FileIndexConcurrency, } - - // Set usage handler to calculate cost from models.dev and emit events with CUMULATIVE totals - // This matches how chat completions calculate cost in runtime.go - cfg.Embedder.SetUsageHandler(func(ctx context.Context, tokens int64, _ float64) { - cost := s.calculateCost(ctx, tokens) - s.recordUsage(tokens, cost) - }) - return s } @@ -385,15 +377,18 @@ func (s *VectorStore) Initialize(ctx context.Context, docPaths []string, chunkin } // Query searches for relevant documents using vector similarity -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) if err != nil { - return nil, fmt.Errorf("failed to generate query embedding: %w", err) + return nil, types.Usage{}, fmt.Errorf("failed to generate query embedding: %w", err) } + cost := s.calculateCost(ctx, tokens) + usage := types.Usage{TotalTokens: tokens, Cost: cost} + results, err := s.db.SearchSimilarVectors(ctx, queryEmbedding, numResults) if err != nil { - return nil, fmt.Errorf("failed to search: %w", err) + return nil, types.Usage{}, fmt.Errorf("failed to search: %w", err) } // Convert internal result type to public SearchResult type @@ -407,7 +402,7 @@ func (s *VectorStore) Query(ctx context.Context, query string, numResults int, t } } - return filtered, nil + return filtered, usage, nil } // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed @@ -607,11 +602,14 @@ func (s *VectorStore) indexFile(ctx context.Context, filePath string) error { "path", filePath, "chunk_count", len(validChunks)) - embeddings, err := s.embedder.EmbedBatch(ctx, chunkContents) + embeddings, batchTokens, _, err := s.embedder.EmbedBatch(ctx, chunkContents) if err != nil { return fmt.Errorf("failed to generate embeddings: %w", classifyModelCallError(err)) } + batchCost := s.calculateCost(ctx, batchTokens) + s.recordUsage(batchTokens, batchCost) + if len(embeddings) != len(validChunks) { return fmt.Errorf("embedding count mismatch: got %d embeddings for %d chunks", len(embeddings), len(validChunks)) } diff --git a/pkg/rag/types/types.go b/pkg/rag/types/types.go index ce315a463f..8be0fedcc5 100644 --- a/pkg/rag/types/types.go +++ b/pkg/rag/types/types.go @@ -9,6 +9,18 @@ type Document struct { ChunkIndex int // Position of this chunk within the source document (0-based) } +// Usage tracks token usage and cost for RAG operations. +type Usage struct { + TotalTokens int64 + Cost float64 +} + +// Add combines two usage structs. +func (u *Usage) Add(other Usage) { + u.TotalTokens += other.TotalTokens + u.Cost += other.Cost +} + type EventTye string const ( diff --git a/pkg/tools/builtin/rag/rag.go b/pkg/tools/builtin/rag/rag.go index 36833402b6..9554840a2d 100644 --- a/pkg/tools/builtin/rag/rag.go +++ b/pkg/tools/builtin/rag/rag.go @@ -13,6 +13,7 @@ import ( "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/rag" ragtypes "github.com/docker/docker-agent/pkg/rag/types" + "github.com/docker/docker-agent/pkg/telemetry" "github.com/docker/docker-agent/pkg/tools" ) @@ -172,11 +173,15 @@ func (t *ToolSet) handleQueryRAG(ctx context.Context, args queryRAGArgs) (*tools return nil, errors.New("query cannot be empty") } - results, err := t.manager.Query(ctx, args.Query) + results, usage, err := t.manager.Query(ctx, args.Query) if err != nil { return nil, fmt.Errorf("RAG query failed: %w", err) } + if usage.TotalTokens > 0 || usage.Cost > 0 { + telemetry.RecordTokenUsage(ctx, t.toolName, usage.TotalTokens, 0, usage.Cost) + } + out := make([]queryResult, 0, len(results)) for _, r := range results { out = append(out, queryResult{