From 5b11d2b19a497ddafeeb80f9c1c48d9bad399600 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Sat, 11 Jul 2026 19:42:46 +0530 Subject: [PATCH 1/8] 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{ From 570d7fee72b809f84e368cd96bfc34cbf71672ec Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Mon, 13 Jul 2026 14:02:38 +0530 Subject: [PATCH 2/8] style: fix gci lint formatting issues --- pkg/rag/embed/embed.go | 1 - pkg/rag/manager.go | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/pkg/rag/embed/embed.go b/pkg/rag/embed/embed.go index 6537bc4543..440171b9d9 100644 --- a/pkg/rag/embed/embed.go +++ b/pkg/rag/embed/embed.go @@ -50,7 +50,6 @@ func New(p provider.Provider, opts ...Option) *Embedder { return e } - // Embed generates an embedding for a single text 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. diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index 3648c03e94..163c5ad8da 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -360,7 +360,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "error", 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", From e08a8db9ca0aa3bbf9a973bae822a8f6e920c8be Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Thu, 16 Jul 2026 17:53:24 +0530 Subject: [PATCH 3/8] fix: resolve telemetry model dimension pollution and capture partial 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 --- pkg/rag/embed/embed.go | 30 ++++++------- pkg/rag/embed/embed_test.go | 63 +++++++++++++++++++++++++++ pkg/rag/manager.go | 3 -- pkg/rag/manager_test.go | 69 ++++++++++++++++++++++++++++++ pkg/rag/strategy/vector_store.go | 22 +++++----- pkg/rag/types/types.go | 8 +++- pkg/rag/types/types_test.go | 71 +++++++++++++++++++++++++++++++ pkg/tools/builtin/rag/rag.go | 6 ++- pkg/tools/builtin/rag/rag_test.go | 67 +++++++++++++++++++++++++++++ 9 files changed, 307 insertions(+), 32 deletions(-) create mode 100644 pkg/rag/embed/embed_test.go create mode 100644 pkg/rag/types/types_test.go diff --git a/pkg/rag/embed/embed.go b/pkg/rag/embed/embed.go index 440171b9d9..4ee9cbff33 100644 --- a/pkg/rag/embed/embed.go +++ b/pkg/rag/embed/embed.go @@ -51,12 +51,12 @@ func New(p provider.Provider, opts ...Option) *Embedder { } // Embed generates an embedding for a single text -func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, int64, float64, error) { +func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, int64, 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, 0, 0, err + return nil, 0, err } slog.DebugContext(ctx, "Embedding generated", @@ -64,19 +64,19 @@ func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, int64, fl "tokens", result.TotalTokens, "cost", result.Cost) - return result.Embedding, result.TotalTokens, result.Cost, nil + return result.Embedding, result.TotalTokens, nil } // Provider does not support embeddings via the standard interface; fail fast. - return nil, 0, 0, fmt.Errorf("provider %s does not support embeddings", e.provider.ID()) + return nil, 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, int64, float64, error) { +func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, int64, error) { if len(texts) == 0 { - return [][]float64{}, 0, 0, nil + return [][]float64{}, 0, nil } // Check if provider supports batch embeddings. @@ -91,22 +91,20 @@ func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, embeddings := make([][]float64, len(texts)) var totalTokens int64 - var totalCost float64 for i, text := range texts { - embedding, tokens, cost, err := e.Embed(ctx, text) + embedding, tokens, err := e.Embed(ctx, text) + totalTokens += tokens if err != nil { - return nil, 0, 0, fmt.Errorf("failed to embed text %d: %w", i, err) + return nil, totalTokens, fmt.Errorf("failed to embed text %d: %w", i, err) } embeddings[i] = embedding - totalTokens += tokens - totalCost += cost } - return embeddings, totalTokens, totalCost, nil + return embeddings, totalTokens, nil } // embedBatchOptimized processes texts in optimized batches with parallel API calls -func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, int64, float64, error) { +func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, int64, error) { totalTexts := len(texts) slog.DebugContext(ctx, "Starting optimized batch embedding", "provider", e.provider.ID(), @@ -118,7 +116,6 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid embeddings := make([][]float64, totalTexts) var mu sync.Mutex var totalTokens int64 - var totalCost float64 // Create errgroup with concurrency limit g, ctx := errgroup.WithContext(ctx) @@ -149,7 +146,6 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid mu.Lock() copy(embeddings[start:end], result.Embeddings) totalTokens += result.TotalTokens - totalCost += result.Cost mu.Unlock() slog.DebugContext(ctx, "Batch completed", @@ -164,7 +160,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, 0, 0, err + return nil, totalTokens, err } slog.DebugContext(ctx, "Batch embedding completed", @@ -172,5 +168,5 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid "total_embeddings", len(embeddings), "batches_processed", (totalTexts+e.batchSize-1)/e.batchSize) - return embeddings, totalTokens, totalCost, nil + return embeddings, totalTokens, nil } diff --git a/pkg/rag/embed/embed_test.go b/pkg/rag/embed/embed_test.go new file mode 100644 index 0000000000..60c4c68f12 --- /dev/null +++ b/pkg/rag/embed/embed_test.go @@ -0,0 +1,63 @@ +package embed + +import ( + "context" + "errors" + "testing" + + "github.com/docker/docker-agent/pkg/model/provider" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockEmbeddingProvider fakes a provider that fails on the 3rd embedding +type mockEmbeddingProvider struct { + provider.Provider + callCount int +} + +func (m *mockEmbeddingProvider) ID() modelsdev.ID { + return modelsdev.ID{} +} + +func (m *mockEmbeddingProvider) BaseConfig() base.Config { + return base.Config{} +} + + +func (m *mockEmbeddingProvider) CreateEmbedding(ctx context.Context, text string) (*base.EmbeddingResult, error) { + m.callCount++ + if m.callCount == 3 { + return nil, errors.New("simulated failure") + } + return &base.EmbeddingResult{ + Embedding: []float64{0.1, 0.2, 0.3}, + TotalTokens: 10, + Cost: 0.001, + }, nil +} + +// verify we implement EmbeddingProvider +var _ provider.EmbeddingProvider = (*mockEmbeddingProvider)(nil) +var _ provider.Provider = (*mockEmbeddingProvider)(nil) + +func TestEmbedBatch_PartialUsageOnError(t *testing.T) { + mockProv := &mockEmbeddingProvider{} + embedder := New(mockProv, WithBatchSize(10), WithMaxConcurrency(1)) + + // Provide 4 texts. The 3rd will fail. + texts := []string{"text1", "text2", "text3", "text4"} + + embeddings, tokens, err := embedder.EmbedBatch(context.Background(), texts) + + require.Error(t, err) + assert.Contains(t, err.Error(), "simulated failure") + + // EmbedBatch should return nil for embeddings since it failed + assert.Nil(t, embeddings) + + // Tokens should be 20 (10 from text1 + 10 from text2) + assert.Equal(t, int64(20), tokens) +} diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index 163c5ad8da..c73fd0a0a1 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -421,9 +421,6 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "num_results", len(fusedResults)) } - // TODO: Track and emit query embedding usage - // For queries during agent execution, usage should be added to agent's session - // This requires passing session context through the RAG tool return fusedResults, usage, nil } diff --git a/pkg/rag/manager_test.go b/pkg/rag/manager_test.go index a89e419355..0440ae7771 100644 --- a/pkg/rag/manager_test.go +++ b/pkg/rag/manager_test.go @@ -1,10 +1,14 @@ package rag import ( + "context" "os" "path/filepath" "testing" + "github.com/docker/docker-agent/pkg/rag/database" + "github.com/docker/docker-agent/pkg/rag/strategy" + "github.com/docker/docker-agent/pkg/rag/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -34,3 +38,68 @@ func TestGetAbsolutePaths_NilInput(t *testing.T) { result := GetAbsolutePaths("/base", nil) assert.Nil(t, result) } + +type mockStrategy struct { + usage types.Usage +} + +func (m *mockStrategy) Initialize(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) { + return nil, m.usage, nil +} + +func (m *mockStrategy) CheckAndReindexChangedFiles(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) StartFileWatcher(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) Close() error { + return nil +} + +func TestManager_Query_UsageAggregation(t *testing.T) { + events := make(chan types.Event) + defer close(events) + + stratA := &mockStrategy{ + usage: types.Usage{ + TotalTokens: 10, + Cost: 0.1, + ModelID: "model-a", + }, + } + stratB := &mockStrategy{ + usage: types.Usage{ + TotalTokens: 20, + Cost: 0.2, + ModelID: "model-b", + }, + } + + cfg := Config{ + StrategyConfigs: []strategy.Config{ + {Name: "stratA", Strategy: stratA}, + {Name: "stratB", Strategy: stratB}, + }, + Results: ResultsConfig{ + Limit: 10, + }, + } + + mgr, err := New(context.Background(), "test-rag", cfg, events) + require.NoError(t, err) + + _, usage, err := mgr.Query(context.Background(), "test query") + require.NoError(t, err) + + assert.Equal(t, int64(30), usage.TotalTokens) + assert.InDelta(t, 0.3, usage.Cost, 0.0001) + assert.Contains(t, usage.ModelID, "model-a") + assert.Contains(t, usage.ModelID, "model-b") +} diff --git a/pkg/rag/strategy/vector_store.go b/pkg/rag/strategy/vector_store.go index 77e96558bd..14bce9c6df 100644 --- a/pkg/rag/strategy/vector_store.go +++ b/pkg/rag/strategy/vector_store.go @@ -378,17 +378,17 @@ 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, types.Usage, error) { - queryEmbedding, tokens, _, err := s.embedder.Embed(ctx, query) + queryEmbedding, tokens, err := s.embedder.Embed(ctx, query) + cost := s.calculateCost(ctx, tokens) + usage := types.Usage{TotalTokens: tokens, Cost: cost, ModelID: s.modelID.String()} + if err != nil { - return nil, types.Usage{}, fmt.Errorf("failed to generate query embedding: %w", err) + return nil, 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, types.Usage{}, fmt.Errorf("failed to search: %w", err) + return nil, usage, fmt.Errorf("failed to search: %w", err) } // Convert internal result type to public SearchResult type @@ -602,14 +602,16 @@ func (s *VectorStore) indexFile(ctx context.Context, filePath string) error { "path", filePath, "chunk_count", len(validChunks)) - embeddings, batchTokens, _, err := s.embedder.EmbedBatch(ctx, chunkContents) + embeddings, batchTokens, err := s.embedder.EmbedBatch(ctx, chunkContents) + + // Record usage even on partial failure (tokens consumed before error) + batchCost := s.calculateCost(ctx, batchTokens) + s.recordUsage(batchTokens, batchCost) + 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 8be0fedcc5..4090d172b4 100644 --- a/pkg/rag/types/types.go +++ b/pkg/rag/types/types.go @@ -9,16 +9,22 @@ 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 + ModelID string } // Add combines two usage structs. func (u *Usage) Add(other Usage) { u.TotalTokens += other.TotalTokens u.Cost += other.Cost + if u.ModelID == "" { + u.ModelID = other.ModelID + } else if other.ModelID != "" && u.ModelID != other.ModelID { + // Just take the first if there are multiple, or concatenate + u.ModelID = u.ModelID + "," + other.ModelID + } } type EventTye string diff --git a/pkg/rag/types/types_test.go b/pkg/rag/types/types_test.go new file mode 100644 index 0000000000..470fab4981 --- /dev/null +++ b/pkg/rag/types/types_test.go @@ -0,0 +1,71 @@ +package types + +import ( + "testing" +) + +func TestUsage_Add(t *testing.T) { + tests := []struct { + name string + base Usage + other Usage + expected Usage + }{ + { + name: "Both empty", + base: Usage{}, + other: Usage{}, + expected: Usage{}, + }, + { + name: "Base empty, other has values", + base: Usage{}, + other: Usage{TotalTokens: 10, Cost: 0.5, ModelID: "model-a"}, + expected: Usage{TotalTokens: 10, Cost: 0.5, ModelID: "model-a"}, + }, + { + name: "Base has values, other empty", + base: Usage{TotalTokens: 5, Cost: 0.2, ModelID: "model-a"}, + other: Usage{}, + expected: Usage{TotalTokens: 5, Cost: 0.2, ModelID: "model-a"}, + }, + { + name: "Same model ID", + base: Usage{TotalTokens: 5, Cost: 0.2, ModelID: "model-a"}, + other: Usage{TotalTokens: 10, Cost: 0.5, ModelID: "model-a"}, + expected: Usage{TotalTokens: 15, Cost: 0.7, ModelID: "model-a"}, + }, + { + name: "Different model IDs", + base: Usage{TotalTokens: 5, Cost: 0.2, ModelID: "model-a"}, + other: Usage{TotalTokens: 10, Cost: 0.5, ModelID: "model-b"}, + expected: Usage{TotalTokens: 15, Cost: 0.7, ModelID: "model-a,model-b"}, + }, + { + name: "Multiple concatenations", + base: Usage{TotalTokens: 5, Cost: 0.2, ModelID: "model-a,model-b"}, + other: Usage{TotalTokens: 10, Cost: 0.5, ModelID: "model-c"}, + expected: Usage{TotalTokens: 15, Cost: 0.7, ModelID: "model-a,model-b,model-c"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + u := tt.base + u.Add(tt.other) + + if u.TotalTokens != tt.expected.TotalTokens { + t.Errorf("expected TotalTokens %d, got %d", tt.expected.TotalTokens, u.TotalTokens) + } + + // Float comparison with small epsilon since we're just adding constants + if u.Cost != tt.expected.Cost { + t.Errorf("expected Cost %f, got %f", tt.expected.Cost, u.Cost) + } + + if u.ModelID != tt.expected.ModelID { + t.Errorf("expected ModelID %q, got %q", tt.expected.ModelID, u.ModelID) + } + }) + } +} diff --git a/pkg/tools/builtin/rag/rag.go b/pkg/tools/builtin/rag/rag.go index 9554840a2d..8ca720492a 100644 --- a/pkg/tools/builtin/rag/rag.go +++ b/pkg/tools/builtin/rag/rag.go @@ -179,7 +179,11 @@ func (t *ToolSet) handleQueryRAG(ctx context.Context, args queryRAGArgs) (*tools } if usage.TotalTokens > 0 || usage.Cost > 0 { - telemetry.RecordTokenUsage(ctx, t.toolName, usage.TotalTokens, 0, usage.Cost) + model := usage.ModelID + if model == "" { + model = t.toolName // Fallback to tool name for strategies that do not specify a model ID (e.g. BM25) + } + telemetry.RecordTokenUsage(ctx, model, usage.TotalTokens, 0, usage.Cost) } out := make([]queryResult, 0, len(results)) diff --git a/pkg/tools/builtin/rag/rag_test.go b/pkg/tools/builtin/rag/rag_test.go index b8f6ec1e5a..bd97bdf724 100644 --- a/pkg/tools/builtin/rag/rag_test.go +++ b/pkg/tools/builtin/rag/rag_test.go @@ -2,9 +2,14 @@ package rag import ( "cmp" + "context" "slices" "testing" + "github.com/docker/docker-agent/pkg/rag" + "github.com/docker/docker-agent/pkg/rag/database" + "github.com/docker/docker-agent/pkg/rag/strategy" + ragtypes "github.com/docker/docker-agent/pkg/rag/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -75,3 +80,65 @@ func TestRAGTool_SortResults(t *testing.T) { assert.Equal(t, "a.txt", results[2].SourcePath) assert.Equal(t, "c.txt", results[3].SourcePath) } + +type mockStrategy struct { + usage ragtypes.Usage +} + +func (m *mockStrategy) Initialize(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, ragtypes.Usage, error) { + return nil, m.usage, nil +} + +func (m *mockStrategy) CheckAndReindexChangedFiles(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) StartFileWatcher(ctx context.Context, docPaths []string, chunking strategy.ChunkingConfig) error { + return nil +} + +func (m *mockStrategy) Close() error { + return nil +} + +func TestRAGTool_HandleQuery_Telemetry(t *testing.T) { + // This test asserts handleQueryRAG runs and doesn't panic. + // Since telemetry is global, we can't easily assert on the emitted event here without + // exposing internal test utilities, but we can verify it doesn't crash on non-zero usage. + + events := make(chan ragtypes.Event) + defer close(events) + + stratA := &mockStrategy{ + usage: ragtypes.Usage{ + TotalTokens: 10, + Cost: 0.1, + ModelID: "test-model", + }, + } + + cfg := rag.Config{ + StrategyConfigs: []strategy.Config{ + {Name: "stratA", Strategy: stratA}, + }, + Results: rag.ResultsConfig{ + Limit: 10, + }, + } + + mgr, err := rag.New(context.Background(), "test-rag", cfg, events) + require.NoError(t, err) + + tool := &ToolSet{ + manager: mgr, + toolName: "test-rag", + } + + res, err := tool.handleQueryRAG(context.Background(), queryRAGArgs{Query: "test"}) + require.NoError(t, err) + assert.NotNil(t, res) +} From 2c75a809a0bbec6e7f2573291f6b33e754fea785 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Thu, 16 Jul 2026 22:17:10 +0530 Subject: [PATCH 4/8] fix: reconcile RAG telemetry with new strategy query signature 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. --- pkg/rag/manager.go | 10 +++- pkg/rag/manager_test.go | 109 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 2 deletions(-) diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index c73fd0a0a1..29bffdf282 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -348,10 +348,16 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S }(strategyName, strategyImpl, strategyCfg) } - // Collect results from all strategies + // Collect results from all strategies. On cancellation the buffered + // channel lets abandoned strategy goroutines send and exit. strategyResults := make(map[string][]database.SearchResult) for range len(m.strategies) { - result := <-resultsChan + var result strategyResult + select { + case result = <-resultsChan: + case <-ctx.Done(): + return nil, types.Usage{}, ctx.Err() + } if result.err != nil { slog.ErrorContext(ctx, "[RAG Manager] Strategy query failed", diff --git a/pkg/rag/manager_test.go b/pkg/rag/manager_test.go index 0440ae7771..5041806f43 100644 --- a/pkg/rag/manager_test.go +++ b/pkg/rag/manager_test.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" "testing" + "time" "github.com/docker/docker-agent/pkg/rag/database" "github.com/docker/docker-agent/pkg/rag/strategy" @@ -103,3 +104,111 @@ func TestManager_Query_UsageAggregation(t *testing.T) { assert.Contains(t, usage.ModelID, "model-a") assert.Contains(t, usage.ModelID, "model-b") } + +// blockingStrategy deliberately ignores ctx and blocks Initialize and Query +// until release is closed. +type blockingStrategy struct { + started chan struct{} + release chan struct{} +} + +func newBlockingStrategy(release chan struct{}) *blockingStrategy { + return &blockingStrategy{started: make(chan struct{}), release: release} +} + +func (s *blockingStrategy) block() { + close(s.started) + <-s.release +} + +func (s *blockingStrategy) Initialize(context.Context, []string, strategy.ChunkingConfig) error { + s.block() + return nil +} + +func (s *blockingStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, types.Usage, error) { + s.block() + return nil, types.Usage{}, nil +} + +func (s *blockingStrategy) CheckAndReindexChangedFiles(context.Context, []string, strategy.ChunkingConfig) error { + return nil +} + +func (s *blockingStrategy) StartFileWatcher(context.Context, []string, strategy.ChunkingConfig) error { + return nil +} + +func (s *blockingStrategy) Close() error { return nil } + +func TestInitializeReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { + t.Parallel() + release := make(chan struct{}) + t.Cleanup(func() { close(release) }) + blocked := newBlockingStrategy(release) + + cfg := Config{ + StrategyConfigs: []strategy.Config{{Name: "blocked", Strategy: blocked}}, + } + m, err := New(t.Context(), "test", cfg, nil) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + errCh := make(chan error, 1) + go func() { errCh <- m.Initialize(ctx) }() + + <-blocked.started + cancel() + + select { + case err := <-errCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(2 * time.Second): + t.Fatal("Initialize did not return after context cancellation") + } +} + +func TestQueryReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { + t.Parallel() + release := make(chan struct{}) + t.Cleanup(func() { close(release) }) + first := newBlockingStrategy(release) + second := newBlockingStrategy(release) + + // Two strategies so Query takes the multi-strategy fan-in path. + cfg := Config{ + StrategyConfigs: []strategy.Config{ + {Name: "first", Strategy: first, Limit: 5}, + {Name: "second", Strategy: second, Limit: 5}, + }, + } + m, err := New(t.Context(), "test", cfg, nil) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + type queryResult struct { + results []database.SearchResult + err error + } + resCh := make(chan queryResult, 1) + go func() { + results, _, err := m.Query(ctx, "some query") + resCh <- queryResult{results: results, err: err} + }() + + <-first.started + <-second.started + cancel() + + select { + case res := <-resCh: + require.ErrorIs(t, res.err, context.Canceled) + assert.Nil(t, res.results) + case <-time.After(2 * time.Second): + t.Fatal("Query did not return after context cancellation") + } +} From 8cf06c0c9f6fc7af2a29bb572ebf8e5e5029395d Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Fri, 17 Jul 2026 17:14:10 +0530 Subject: [PATCH 5/8] fix: resolve duplicate mock strategy definitions after merge conflict --- pkg/rag/manager_test.go | 111 ---------------------------------------- 1 file changed, 111 deletions(-) diff --git a/pkg/rag/manager_test.go b/pkg/rag/manager_test.go index 37ede7d9c2..5041806f43 100644 --- a/pkg/rag/manager_test.go +++ b/pkg/rag/manager_test.go @@ -12,119 +12,8 @@ import ( "github.com/docker/docker-agent/pkg/rag/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/docker/docker-agent/pkg/rag/database" - "github.com/docker/docker-agent/pkg/rag/strategy" ) -// blockingStrategy deliberately ignores ctx and blocks Initialize and Query -// until release is closed. -type blockingStrategy struct { - started chan struct{} - release chan struct{} -} - -func newBlockingStrategy(release chan struct{}) *blockingStrategy { - return &blockingStrategy{started: make(chan struct{}), release: release} -} - -func (s *blockingStrategy) block() { - close(s.started) - <-s.release -} - -func (s *blockingStrategy) Initialize(context.Context, []string, strategy.ChunkingConfig) error { - s.block() - return nil -} - -func (s *blockingStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, error) { - s.block() - return nil, nil -} - -func (s *blockingStrategy) CheckAndReindexChangedFiles(context.Context, []string, strategy.ChunkingConfig) error { - return nil -} - -func (s *blockingStrategy) StartFileWatcher(context.Context, []string, strategy.ChunkingConfig) error { - return nil -} - -func (s *blockingStrategy) Close() error { return nil } - -func TestInitializeReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { - t.Parallel() - release := make(chan struct{}) - t.Cleanup(func() { close(release) }) - blocked := newBlockingStrategy(release) - - cfg := Config{ - StrategyConfigs: []strategy.Config{{Name: "blocked", Strategy: blocked}}, - } - m, err := New(t.Context(), "test", cfg, nil) - require.NoError(t, err) - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - errCh := make(chan error, 1) - go func() { errCh <- m.Initialize(ctx) }() - - <-blocked.started - cancel() - - select { - case err := <-errCh: - require.ErrorIs(t, err, context.Canceled) - case <-time.After(2 * time.Second): - t.Fatal("Initialize did not return after context cancellation") - } -} - -func TestQueryReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { - t.Parallel() - release := make(chan struct{}) - t.Cleanup(func() { close(release) }) - first := newBlockingStrategy(release) - second := newBlockingStrategy(release) - - // Two strategies so Query takes the multi-strategy fan-in path. - cfg := Config{ - StrategyConfigs: []strategy.Config{ - {Name: "first", Strategy: first, Limit: 5}, - {Name: "second", Strategy: second, Limit: 5}, - }, - } - m, err := New(t.Context(), "test", cfg, nil) - require.NoError(t, err) - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - type queryResult struct { - results []database.SearchResult - err error - } - resCh := make(chan queryResult, 1) - go func() { - results, err := m.Query(ctx, "some query") - resCh <- queryResult{results: results, err: err} - }() - - <-first.started - <-second.started - cancel() - - select { - case res := <-resCh: - require.ErrorIs(t, res.err, context.Canceled) - assert.Nil(t, res.results) - case <-time.After(2 * time.Second): - t.Fatal("Query did not return after context cancellation") - } -} - func TestGetAbsolutePaths_WithBasePath(t *testing.T) { t.Parallel() result := GetAbsolutePaths("/base", []string{"relative/file.go", "/absolute/file.go"}) From a70cde34cc3b785140a83dd3382941c70f96c430 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Fri, 17 Jul 2026 17:41:25 +0530 Subject: [PATCH 6/8] style: fix gci, forbidigo, and embeddedstructfieldcheck linting issues --- pkg/rag/embed/embed_test.go | 10 ++++++---- pkg/rag/manager.go | 1 - pkg/rag/manager_test.go | 9 +++++---- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/pkg/rag/embed/embed_test.go b/pkg/rag/embed/embed_test.go index 60c4c68f12..d0bbe53f83 100644 --- a/pkg/rag/embed/embed_test.go +++ b/pkg/rag/embed/embed_test.go @@ -5,16 +5,18 @@ import ( "errors" "testing" - "github.com/docker/docker-agent/pkg/model/provider" - "github.com/docker/docker-agent/pkg/modelsdev" - "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/model/provider" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/modelsdev" ) // mockEmbeddingProvider fakes a provider that fails on the 3rd embedding type mockEmbeddingProvider struct { provider.Provider + callCount int } @@ -50,7 +52,7 @@ func TestEmbedBatch_PartialUsageOnError(t *testing.T) { // Provide 4 texts. The 3rd will fail. texts := []string{"text1", "text2", "text3", "text4"} - embeddings, tokens, err := embedder.EmbedBatch(context.Background(), texts) + embeddings, tokens, err := embedder.EmbedBatch(t.Context(), texts) require.Error(t, err) assert.Contains(t, err.Error(), "simulated failure") diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index dacaec04e0..e3c538489c 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -432,7 +432,6 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "num_results", len(fusedResults)) } - return fusedResults, usage, nil } diff --git a/pkg/rag/manager_test.go b/pkg/rag/manager_test.go index 5041806f43..48599f82ca 100644 --- a/pkg/rag/manager_test.go +++ b/pkg/rag/manager_test.go @@ -7,11 +7,12 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/docker/docker-agent/pkg/rag/database" "github.com/docker/docker-agent/pkg/rag/strategy" "github.com/docker/docker-agent/pkg/rag/types" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) func TestGetAbsolutePaths_WithBasePath(t *testing.T) { @@ -93,10 +94,10 @@ func TestManager_Query_UsageAggregation(t *testing.T) { }, } - mgr, err := New(context.Background(), "test-rag", cfg, events) + mgr, err := New(t.Context(), "test-rag", cfg, events) require.NoError(t, err) - _, usage, err := mgr.Query(context.Background(), "test query") + _, usage, err := mgr.Query(t.Context(), "test query") require.NoError(t, err) assert.Equal(t, int64(30), usage.TotalTokens) From 0a96189c4c9a533a944cd2fe29d4766e06079da6 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Fri, 17 Jul 2026 18:12:18 +0530 Subject: [PATCH 7/8] style: fix remaining formatting and forbidigo linting issues --- pkg/rag/embed/embed_test.go | 5 ++--- pkg/rag/strategy/vector_store.go | 2 +- pkg/rag/types/types_test.go | 6 +++--- pkg/tools/builtin/rag/rag_test.go | 11 ++++++----- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/pkg/rag/embed/embed_test.go b/pkg/rag/embed/embed_test.go index d0bbe53f83..351ddb5c8b 100644 --- a/pkg/rag/embed/embed_test.go +++ b/pkg/rag/embed/embed_test.go @@ -28,7 +28,6 @@ func (m *mockEmbeddingProvider) BaseConfig() base.Config { return base.Config{} } - func (m *mockEmbeddingProvider) CreateEmbedding(ctx context.Context, text string) (*base.EmbeddingResult, error) { m.callCount++ if m.callCount == 3 { @@ -56,10 +55,10 @@ func TestEmbedBatch_PartialUsageOnError(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "simulated failure") - + // EmbedBatch should return nil for embeddings since it failed assert.Nil(t, embeddings) - + // Tokens should be 20 (10 from text1 + 10 from text2) assert.Equal(t, int64(20), tokens) } diff --git a/pkg/rag/strategy/vector_store.go b/pkg/rag/strategy/vector_store.go index 14bce9c6df..82629e6786 100644 --- a/pkg/rag/strategy/vector_store.go +++ b/pkg/rag/strategy/vector_store.go @@ -603,7 +603,7 @@ func (s *VectorStore) indexFile(ctx context.Context, filePath string) error { "chunk_count", len(validChunks)) embeddings, batchTokens, err := s.embedder.EmbedBatch(ctx, chunkContents) - + // Record usage even on partial failure (tokens consumed before error) batchCost := s.calculateCost(ctx, batchTokens) s.recordUsage(batchTokens, batchCost) diff --git a/pkg/rag/types/types_test.go b/pkg/rag/types/types_test.go index 470fab4981..7a21e89511 100644 --- a/pkg/rag/types/types_test.go +++ b/pkg/rag/types/types_test.go @@ -53,16 +53,16 @@ func TestUsage_Add(t *testing.T) { t.Run(tt.name, func(t *testing.T) { u := tt.base u.Add(tt.other) - + if u.TotalTokens != tt.expected.TotalTokens { t.Errorf("expected TotalTokens %d, got %d", tt.expected.TotalTokens, u.TotalTokens) } - + // Float comparison with small epsilon since we're just adding constants if u.Cost != tt.expected.Cost { t.Errorf("expected Cost %f, got %f", tt.expected.Cost, u.Cost) } - + if u.ModelID != tt.expected.ModelID { t.Errorf("expected ModelID %q, got %q", tt.expected.ModelID, u.ModelID) } diff --git a/pkg/tools/builtin/rag/rag_test.go b/pkg/tools/builtin/rag/rag_test.go index bd97bdf724..69e2d2996d 100644 --- a/pkg/tools/builtin/rag/rag_test.go +++ b/pkg/tools/builtin/rag/rag_test.go @@ -6,12 +6,13 @@ import ( "slices" "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/docker/docker-agent/pkg/rag" "github.com/docker/docker-agent/pkg/rag/database" "github.com/docker/docker-agent/pkg/rag/strategy" ragtypes "github.com/docker/docker-agent/pkg/rag/types" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) func TestRAGTool_ToolName(t *testing.T) { @@ -107,7 +108,7 @@ func (m *mockStrategy) Close() error { func TestRAGTool_HandleQuery_Telemetry(t *testing.T) { // This test asserts handleQueryRAG runs and doesn't panic. - // Since telemetry is global, we can't easily assert on the emitted event here without + // Since telemetry is global, we can't easily assert on the emitted event here without // exposing internal test utilities, but we can verify it doesn't crash on non-zero usage. events := make(chan ragtypes.Event) @@ -130,7 +131,7 @@ func TestRAGTool_HandleQuery_Telemetry(t *testing.T) { }, } - mgr, err := rag.New(context.Background(), "test-rag", cfg, events) + mgr, err := rag.New(t.Context(), "test-rag", cfg, events) require.NoError(t, err) tool := &ToolSet{ @@ -138,7 +139,7 @@ func TestRAGTool_HandleQuery_Telemetry(t *testing.T) { toolName: "test-rag", } - res, err := tool.handleQueryRAG(context.Background(), queryRAGArgs{Query: "test"}) + res, err := tool.handleQueryRAG(t.Context(), queryRAGArgs{Query: "test"}) require.NoError(t, err) assert.NotNil(t, res) } From 2bae6c7d04d2158df8eda9908a0cc5f46cb08607 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Fri, 17 Jul 2026 19:02:38 +0530 Subject: [PATCH 8/8] style: fix gofumpt formatting issues --- pkg/rag/embed/embed_test.go | 6 ++++-- pkg/rag/manager.go | 9 ++++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/pkg/rag/embed/embed_test.go b/pkg/rag/embed/embed_test.go index 351ddb5c8b..ab1db7b7a9 100644 --- a/pkg/rag/embed/embed_test.go +++ b/pkg/rag/embed/embed_test.go @@ -41,8 +41,10 @@ func (m *mockEmbeddingProvider) CreateEmbedding(ctx context.Context, text string } // verify we implement EmbeddingProvider -var _ provider.EmbeddingProvider = (*mockEmbeddingProvider)(nil) -var _ provider.Provider = (*mockEmbeddingProvider)(nil) +var ( + _ provider.EmbeddingProvider = (*mockEmbeddingProvider)(nil) + _ provider.Provider = (*mockEmbeddingProvider)(nil) +) func TestEmbedBatch_PartialUsageOnError(t *testing.T) { mockProv := &mockEmbeddingProvider{} diff --git a/pkg/rag/manager.go b/pkg/rag/manager.go index e3c538489c..81cd7ef495 100644 --- a/pkg/rag/manager.go +++ b/pkg/rag/manager.go @@ -154,7 +154,8 @@ func New(_ context.Context, name string, config Config, strategyEvents <-chan ty // Strategies are initialized in parallel for better performance func (m *Manager) Initialize(ctx context.Context) (err error) { tracer := otel.Tracer("github.com/docker/docker-agent/pkg/rag") - ctx, span := tracer.Start(ctx, "rag.initialize", + ctx, span := tracer.Start( + ctx, "rag.initialize", trace.WithSpanKind(trace.SpanKindInternal), trace.WithAttributes( attribute.String(genai.AttrDataSourceID, m.name), @@ -482,7 +483,8 @@ func (m *Manager) rerank(ctx context.Context, query string, results []database.S // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed func (m *Manager) CheckAndReindexChangedFiles(ctx context.Context) (err error) { tracer := otel.Tracer("github.com/docker/docker-agent/pkg/rag") - ctx, span := tracer.Start(ctx, "rag.reindex", + ctx, span := tracer.Start( + ctx, "rag.reindex", trace.WithSpanKind(trace.SpanKindInternal), trace.WithAttributes(attribute.String(genai.AttrDataSourceID, m.name)), ) @@ -506,7 +508,8 @@ func (m *Manager) CheckAndReindexChangedFiles(ctx context.Context) (err error) { // StartFileWatcher starts monitoring files and directories for changes func (m *Manager) StartFileWatcher(ctx context.Context) (err error) { tracer := otel.Tracer("github.com/docker/docker-agent/pkg/rag") - ctx, span := tracer.Start(ctx, "rag.file_watcher.start", + ctx, span := tracer.Start( + ctx, "rag.file_watcher.start", trace.WithSpanKind(trace.SpanKindInternal), trace.WithAttributes(attribute.String(genai.AttrDataSourceID, m.name)), )