perf: parallelize EmbedBatch using goroutines and errgroup

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/qdrant/go-client v1.15.2
 	github.com/spf13/cobra v1.10.2
 	golang.org/x/oauth2 v0.27.0
+	golang.org/x/sync v0.18.0
 	google.golang.org/api v0.186.0
 	google.golang.org/grpc v1.71.0-dev
 	gopkg.in/yaml.v3 v3.0.1
@@ -41,7 +42,6 @@ require (
 	go.opentelemetry.io/otel/trace v1.38.0 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/net v0.33.0 // indirect
-	golang.org/x/sync v0.18.0 // indirect
 	golang.org/x/sys v0.38.0 // indirect
 	golang.org/x/text v0.31.0 // indirect
 	golang.org/x/time v0.14.0 // indirect

internal/integrations/ai/embedder.go

@@ -15,9 +15,13 @@ import (
 	"time"
 
 	"github.com/google/generative-ai-go/genai"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/api/option"
 )
 
+// maxBatchConcurrency limits the number of concurrent embedding requests in EmbedBatch.
+const maxBatchConcurrency = 10
+
 // Embedder generates embeddings using Gemini or OpenAI.
 type Embedder struct {
 	provider    Provider
@@ -166,19 +170,30 @@ func (e *Embedder) embedOpenAI(ctx context.Context, text string) ([]float32, err
 	})
 }
 
-// EmbedBatch generates embeddings for multiple texts.
+// EmbedBatch generates embeddings for multiple texts in parallel.
+// It uses up to maxBatchConcurrency concurrent requests to the embedding API.
 func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
 	if len(texts) == 0 {
 		return nil, fmt.Errorf("texts cannot be empty")
 	}
 
 	embeddings := make([][]float32, len(texts))
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(maxBatchConcurrency)
+
 	for i, text := range texts {
-		embedding, err := e.Embed(ctx, text)
-		if err != nil {
-			return nil, fmt.Errorf("failed to embed text %d: %w", i, err)
-		}
-		embeddings[i] = embedding
+		g.Go(func() error {
+			embedding, err := e.Embed(ctx, text)
+			if err != nil {
+				return fmt.Errorf("failed to embed text %d: %w", i, err)
+			}
+			embeddings[i] = embedding
+			return nil
+		})
+	}
+
+	if err := g.Wait(); err != nil {
+		return nil, err
 	}
 
 	return embeddings, nil

internal/integrations/ai/embedder_test.go

@@ -6,8 +6,15 @@
 package ai
 
 import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 )
 
 func TestInferEmbeddingDimensions(t *testing.T) {
@@ -89,6 +96,142 @@ func TestIsLikelyGeminiEmbeddingModel(t *testing.T) {
 	}
 }
 
+func TestEmbedBatch_EmptyInput(t *testing.T) {
+	srv, _ := statusServer([]int{200}, func(_ int) []byte { return embeddingOKBody() })
+	defer srv.Close()
+
+	e := newTestEmbedder(srv.URL)
+	_, err := e.EmbedBatch(context.Background(), []string{})
+	if err == nil {
+		t.Fatal("expected error for empty input, got nil")
+	}
+}
+
+func TestEmbedBatch_ResultsPreserveOrder(t *testing.T) {
+	// The handler parses the input text ("text-N") and encodes N+1 as the first
+	// embedding value. This lets us deterministically verify that results[i]
+	// corresponds to texts[i] regardless of goroutine scheduling order.
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var req struct {
+			Input string `json:"input"`
+		}
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			http.Error(w, "bad request", http.StatusBadRequest)
+			return
+		}
+		var idx int
+		if _, err := fmt.Sscanf(req.Input, "text-%d", &idx); err != nil {
+			http.Error(w, "invalid input format", http.StatusBadRequest)
+			return
+		}
+
+		type embItem struct {
+			Embedding []float64 `json:"embedding"`
+		}
+		type embResp struct {
+			Data []embItem `json:"data"`
+		}
+		b, _ := json.Marshal(embResp{Data: []embItem{{Embedding: []float64{float64(idx + 1), 0.0, 0.0}}}})
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(200)
+		_, _ = w.Write(b)
+	}))
+	defer srv.Close()
+
+	e := newTestEmbedder(srv.URL)
+	e.retryConfig = fastRetry
+
+	texts := make([]string, 20)
+	for i := range texts {
+		texts[i] = fmt.Sprintf("text-%d", i)
+	}
+
+	results, err := e.EmbedBatch(context.Background(), texts)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(results) != len(texts) {
+		t.Fatalf("expected %d results, got %d", len(texts), len(results))
+	}
+	for i, emb := range results {
+		if len(emb) == 0 {
+			t.Errorf("result[%d] is empty", i)
+			continue
+		}
+		// texts[i] = "text-i", so the handler encodes i+1 as emb[0].
+		if got, want := emb[0], float32(i+1); got != want {
+			t.Errorf("result[%d][0] = %v, want %v (order not preserved)", i, got, want)
+		}
+	}
+}
+
+func TestEmbedBatch_PropagatesError(t *testing.T) {
+	// First request succeeds, subsequent ones fail with a non-retryable 400.
+	var counter atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := counter.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		if n == 1 {
+			w.WriteHeader(200)
+			_, _ = w.Write(embeddingOKBody())
+		} else {
+			w.WriteHeader(400)
+			_, _ = w.Write([]byte(`{"error":{"message":"bad request"}}`))
+		}
+	}))
+	defer srv.Close()
+
+	e := newTestEmbedder(srv.URL)
+	e.retryConfig = fastRetry
+
+	_, err := e.EmbedBatch(context.Background(), []string{"a", "b", "c"})
+	if err == nil {
+		t.Fatal("expected error from failed embedding, got nil")
+	}
+}
+
+func TestEmbedBatch_ConcurrencyLimit(t *testing.T) {
+	// Verify that EmbedBatch honours maxBatchConcurrency: concurrent in-flight
+	// requests should never exceed the limit.
+	var (
+		inFlight    atomic.Int32
+		maxObserved atomic.Int32
+	)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cur := inFlight.Add(1)
+		defer inFlight.Add(-1)
+		// Update max observed — simple CAS loop.
+		for {
+			old := maxObserved.Load()
+			if cur <= old || maxObserved.CompareAndSwap(old, cur) {
+				break
+			}
+		}
+		// Hold the request briefly so multiple goroutines overlap in-flight,
+		// making the concurrency cap observable.
+		time.Sleep(5 * time.Millisecond)
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(200)
+		_, _ = w.Write(embeddingOKBody())
+	}))
+	defer srv.Close()
+
+	e := newTestEmbedder(srv.URL)
+	e.retryConfig = fastRetry
+
+	texts := make([]string, maxBatchConcurrency*3)
+	for i := range texts {
+		texts[i] = fmt.Sprintf("text-%d", i)
+	}
+
+	if _, err := e.EmbedBatch(context.Background(), texts); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if got := maxObserved.Load(); got > int32(maxBatchConcurrency) {
+		t.Errorf("max concurrent requests = %d, want <= %d", got, maxBatchConcurrency)
+	}
+}
+
 func TestNewEmbedderRejectsLegacyGeminiModels(t *testing.T) {
 	// Fake a Gemini API key so provider resolution picks Gemini.
 	t.Setenv("GEMINI_API_KEY", "fake-key-for-test")