diff --git a/pkg/rag/embed/embed.go b/pkg/rag/embed/embed.go index 312480b417..4ee9cbff33 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,13 @@ 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, 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, err } slog.DebugContext(ctx, "Embedding generated", @@ -76,19 +64,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, 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, 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, error) { if len(texts) == 0 { - return [][]float64{}, nil + return [][]float64{}, 0, nil } // Check if provider supports batch embeddings. @@ -102,19 +90,21 @@ func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, "text_count", len(texts)) embeddings := make([][]float64, len(texts)) + var totalTokens int64 for i, text := range texts { - embedding, err := e.Embed(ctx, text) + embedding, tokens, err := e.Embed(ctx, text) + totalTokens += tokens if err != nil { - return nil, 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 } - return embeddings, 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, 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(), @@ -125,6 +115,7 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid // Pre-allocate results embeddings := make([][]float64, totalTexts) var mu sync.Mutex + var totalTokens int64 // Create errgroup with concurrency limit g, ctx := errgroup.WithContext(ctx) @@ -154,13 +145,9 @@ 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 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 +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, err + return nil, totalTokens, err } slog.DebugContext(ctx, "Batch embedding completed", @@ -181,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, 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..ab1db7b7a9 --- /dev/null +++ b/pkg/rag/embed/embed_test.go @@ -0,0 +1,66 @@ +package embed + +import ( + "context" + "errors" + "testing" + + "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 +} + +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) + _ 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(t.Context(), 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 5aec611a86..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), @@ -241,7 +242,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 @@ -275,13 +276,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", @@ -314,7 +315,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S "num_results", len(results)) } - return results, nil + return results, usage, nil } } @@ -326,6 +327,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 } @@ -342,10 +344,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) @@ -359,7 +362,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S select { case result = <-resultsChan: case <-ctx.Done(): - return nil, ctx.Err() + return nil, types.Usage{}, ctx.Err() } if result.err != nil { @@ -367,9 +370,11 @@ 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, "strategy", result.name, @@ -386,7 +391,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) @@ -394,7 +399,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", @@ -428,11 +433,7 @@ 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, nil + return fusedResults, usage, nil } // Helper to get strategy names for logging @@ -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)), ) 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/manager_test.go b/pkg/rag/manager_test.go index f8ddefe100..48599f82ca 100644 --- a/pkg/rag/manager_test.go +++ b/pkg/rag/manager_test.go @@ -12,8 +12,100 @@ import ( "github.com/docker/docker-agent/pkg/rag/database" "github.com/docker/docker-agent/pkg/rag/strategy" + "github.com/docker/docker-agent/pkg/rag/types" ) +func TestGetAbsolutePaths_WithBasePath(t *testing.T) { + t.Parallel() + result := GetAbsolutePaths("/base", []string{"relative/file.go", "/absolute/file.go"}) + assert.Equal(t, []string{"/base/relative/file.go", "/absolute/file.go"}, result) +} + +func TestGetAbsolutePaths_EmptyBasePath(t *testing.T) { + t.Parallel() + // When basePath is empty (OCI/URL sources), relative paths should be + // resolved against the current working directory instead of producing + // broken paths like "relative/file.go". + cwd, err := os.Getwd() + require.NoError(t, err) + + result := GetAbsolutePaths("", []string{"relative/file.go", "/absolute/file.go"}) + + assert.Equal(t, filepath.Join(cwd, "relative", "file.go"), result[0]) + assert.Equal(t, "/absolute/file.go", result[1]) +} + +func TestGetAbsolutePaths_NilInput(t *testing.T) { + t.Parallel() + 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(t.Context(), "test-rag", cfg, events) + require.NoError(t, err) + + _, usage, err := mgr.Query(t.Context(), "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") +} + // blockingStrategy deliberately ignores ctx and blocks Initialize and Query // until release is closed. type blockingStrategy struct { @@ -35,9 +127,9 @@ func (s *blockingStrategy) Initialize(context.Context, []string, strategy.Chunki return nil } -func (s *blockingStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, error) { +func (s *blockingStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, types.Usage, error) { s.block() - return nil, nil + return nil, types.Usage{}, nil } func (s *blockingStrategy) CheckAndReindexChangedFiles(context.Context, []string, strategy.ChunkingConfig) error { @@ -105,7 +197,7 @@ func TestQueryReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { } resCh := make(chan queryResult, 1) go func() { - results, err := m.Query(ctx, "some query") + results, _, err := m.Query(ctx, "some query") resCh <- queryResult{results: results, err: err} }() @@ -121,29 +213,3 @@ func TestQueryReturnsOnContextCancellationWithBlockedStrategy(t *testing.T) { 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"}) - assert.Equal(t, []string{"/base/relative/file.go", "/absolute/file.go"}, result) -} - -func TestGetAbsolutePaths_EmptyBasePath(t *testing.T) { - t.Parallel() - // When basePath is empty (OCI/URL sources), relative paths should be - // resolved against the current working directory instead of producing - // broken paths like "relative/file.go". - cwd, err := os.Getwd() - require.NoError(t, err) - - result := GetAbsolutePaths("", []string{"relative/file.go", "/absolute/file.go"}) - - assert.Equal(t, filepath.Join(cwd, "relative", "file.go"), result[0]) - assert.Equal(t, "/absolute/file.go", result[1]) -} - -func TestGetAbsolutePaths_NilInput(t *testing.T) { - t.Parallel() - result := GetAbsolutePaths("/base", nil) - assert.Nil(t, result) -} 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..82629e6786 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) + cost := s.calculateCost(ctx, tokens) + usage := types.Usage{TotalTokens: tokens, Cost: cost, ModelID: s.modelID.String()} + if err != nil { - return nil, fmt.Errorf("failed to generate query embedding: %w", err) + return nil, usage, fmt.Errorf("failed to generate query embedding: %w", err) } results, err := s.db.SearchSimilarVectors(ctx, queryEmbedding, numResults) if err != nil { - return nil, 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 @@ -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,7 +602,12 @@ 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) + + // 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)) } diff --git a/pkg/rag/types/types.go b/pkg/rag/types/types.go index ce315a463f..4090d172b4 100644 --- a/pkg/rag/types/types.go +++ b/pkg/rag/types/types.go @@ -9,6 +9,24 @@ type Document struct { ChunkIndex int // Position of this chunk within the source document (0-based) } +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 const ( diff --git a/pkg/rag/types/types_test.go b/pkg/rag/types/types_test.go new file mode 100644 index 0000000000..7a21e89511 --- /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 36833402b6..8ca720492a 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,19 @@ 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 { + 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)) for _, r := range results { out = append(out, queryResult{ diff --git a/pkg/tools/builtin/rag/rag_test.go b/pkg/tools/builtin/rag/rag_test.go index b8f6ec1e5a..69e2d2996d 100644 --- a/pkg/tools/builtin/rag/rag_test.go +++ b/pkg/tools/builtin/rag/rag_test.go @@ -2,11 +2,17 @@ package rag import ( "cmp" + "context" "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" ) func TestRAGTool_ToolName(t *testing.T) { @@ -75,3 +81,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(t.Context(), "test-rag", cfg, events) + require.NoError(t, err) + + tool := &ToolSet{ + manager: mgr, + toolName: "test-rag", + } + + res, err := tool.handleQueryRAG(t.Context(), queryRAGArgs{Query: "test"}) + require.NoError(t, err) + assert.NotNil(t, res) +}