✨ feat: Add rate limit to Gemini API requests

Author: nmdra

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/pgvector/pgvector-go v0.3.0
 	github.com/redis/go-redis/v9 v9.11.0
 	github.com/stretchr/testify v1.10.0
+	golang.org/x/time v0.11.0
 	google.golang.org/genai v1.15.0
 )
 
@@ -45,7 +46,6 @@ require (
 	golang.org/x/sync v0.14.0 // indirect
 	golang.org/x/sys v0.33.0 // indirect
 	golang.org/x/text v0.25.0 // indirect
-	golang.org/x/time v0.11.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
 	google.golang.org/grpc v1.66.2 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect

internal/embed/embed.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 
 	_ "github.com/joho/godotenv/autoload"
+	"golang.org/x/time/rate"
 	"google.golang.org/genai"
 )
 
@@ -15,8 +16,9 @@ type Embedder interface {
 }
 
 type GeminiEmbedder struct {
-	client *genai.Client
-	logger *slog.Logger
+	client  *genai.Client
+	logger  *slog.Logger
+	limiter *rate.Limiter
 }
 
 func NewGeminiEmbedder(ctx context.Context, logger *slog.Logger, apiKey string) (*GeminiEmbedder, error) {
@@ -35,13 +37,20 @@ func NewGeminiEmbedder(ctx context.Context, logger *slog.Logger, apiKey string)
 	}
 
 	return &GeminiEmbedder{
-		client: client,
-		logger: logger,
+		client:  client,
+		logger:  logger,
+		limiter: rate.NewLimiter(rate.Limit(5), 2), // 5 requests/sec, with up to 2 sent instantly in a burst
 	}, nil
 }
 
 func (g *GeminiEmbedder) Embed(ctx context.Context, input string) ([]float32, error) {
 	dim := int32(768)
+
+	if err := g.limiter.Wait(ctx); err != nil {
+		g.logger.Warn("rate limiter blocked request", "error", err)
+		return nil, fmt.Errorf("rate limit exceeded: %w", err)
+	}
+
 	contents := []*genai.Content{
 		genai.NewContentFromText(input, genai.RoleUser),
 	}