Add adaptive embedding throughput shaping for 429 throttling

BenjaminMichaelis · BenjaminMichaelis · commit 273d0cbe10d9 · 2026-05-15T22:42:10.000-07:00
- Downshift embedding batch size on repeated 429s by recursively splitting batches
- Reuse successful smaller batch size for subsequent requests in the same run
- Fail clearly when batch size 1 still receives sustained 429 throttling
- Add sequential request pacing with configurable min inter-request delay
- Add configurable MaxEmbeddingBatchSize and MinInterRequestDelayMs options
- Harden Retry-After parsing for retry-after, retry-after-ms, x-ms-retry-after-ms, and message hints
- Update configuration comments and default appsettings values
diff --git a/EssentialCSharp.Chat.Shared/Extensions/ServiceCollectionExtensions.cs b/EssentialCSharp.Chat.Shared/Extensions/ServiceCollectionExtensions.cs
@@ -100,8 +100,8 @@ public static IServiceCollection AddAzureOpenAIServices(
         // Configure AI options from configuration
         services.Configure<AIOptions>(configuration.GetSection("AIOptions"));
 
-        // Configure retry options from configuration section
-        // Environment variables like EmbeddingRetry:MaxRetries will override defaults
+        // Configure retry options from configuration section.
+        // Environment variables can override via AIOptions__EmbeddingRetry__*.
         services.AddOptions<EmbeddingRetryOptions>()
             .Bind(configuration.GetSection(EmbeddingRetryOptions.SectionPath))
             .ValidateDataAnnotations()
diff --git a/EssentialCSharp.Chat.Shared/Models/EmbeddingRetryOptions.cs b/EssentialCSharp.Chat.Shared/Models/EmbeddingRetryOptions.cs
@@ -34,6 +34,20 @@ public sealed class EmbeddingRetryOptions
     [Range(1, 600000)]
     public int MaxDelayMs { get; set; } = 60000;
 
+    /// <summary>
+    /// Maximum embedding request payload size sent per API call.
+    /// The service may adaptively downshift below this value when throttled.
+    /// </summary>
+    [Range(1, 2048)]
+    public int MaxEmbeddingBatchSize { get; set; } = 2048;
+
+    /// <summary>
+    /// Minimum delay between embedding API requests in milliseconds.
+    /// This adds request pacing to reduce sustained rate-limit pressure.
+    /// </summary>
+    [Range(0, 600000)]
+    public int MinInterRequestDelayMs { get; set; } = 250;
+
     /// <summary>
     /// Exponential backoff multiplier. Each retry delay is multiplied by this value.
     /// For example, with baseDelay=1000ms and multiplier=2.0:
@@ -74,6 +88,15 @@ public void Validate()
         if (BaseDelayMs > MaxDelayMs)
             throw new InvalidOperationException("BaseDelayMs must be less than or equal to MaxDelayMs.");
 
+        if (MaxEmbeddingBatchSize <= 0)
+            throw new InvalidOperationException("MaxEmbeddingBatchSize must be positive.");
+
+        if (MaxEmbeddingBatchSize > 2048)
+            throw new InvalidOperationException("MaxEmbeddingBatchSize cannot exceed Azure embedding API limit (2048).");
+
+        if (MinInterRequestDelayMs < 0)
+            throw new InvalidOperationException("MinInterRequestDelayMs must be non-negative.");
+
         if (BackoffMultiplier < 1.0)
             throw new InvalidOperationException("BackoffMultiplier must be >= 1.0.");
 
diff --git a/EssentialCSharp.Chat.Shared/Services/EmbeddingService.cs b/EssentialCSharp.Chat.Shared/Services/EmbeddingService.cs
@@ -6,6 +6,7 @@
 using Microsoft.Extensions.VectorData;
 using Npgsql;
 using System.ClientModel;
+using System.ClientModel.Primitives;
 using System.Globalization;
 
 namespace EssentialCSharp.Chat.Common.Services;
@@ -32,9 +33,12 @@ public partial class EmbeddingService(
 
     private readonly EmbeddingRetryOptions _retryOptions = ValidateRetryOptions(retryOptions?.Value ?? new EmbeddingRetryOptions());
     private readonly ILogger<EmbeddingService>? _logger = logger;
+    private static readonly SemaphoreSlim _embeddingRequestLock = new(1, 1);
+    private DateTimeOffset _lastEmbeddingRequestStartedUtc = DateTimeOffset.MinValue;
 
     // Only allow simple identifiers: letters, digits, and underscores, starting with a letter or underscore.
     private static readonly Regex _safeIdentifierRegex = new(@"^[a-zA-Z_][a-zA-Z0-9_]*$", RegexOptions.Compiled);
+    private static readonly Regex _retryAfterSecondsRegex = new(@"retry\s+after\s+(?<seconds>\d+)\s+seconds?", RegexOptions.Compiled | RegexOptions.IgnoreCase);
 
     private static EmbeddingRetryOptions ValidateRetryOptions(EmbeddingRetryOptions options)
     {
@@ -96,6 +100,10 @@ private static bool IsTransientStatusCode(int statusCode) =>
         return ex.InnerException is null ? null : TryGetStatusCode(ex.InnerException);
     }
 
+    private static bool IsRateLimitError(Exception ex) =>
+        TryGetStatusCode(ex) == 429
+        || ex.Message.Contains("RateLimitReached", StringComparison.OrdinalIgnoreCase);
+
     /// <summary>
     /// Extracts the Retry-After delay from known exception types if present.
     /// Returns null if the header is not present or invalid.
@@ -105,19 +113,61 @@ private static bool IsTransientStatusCode(int statusCode) =>
         if (ex is ClientResultException clientResultException)
         {
             var rawResponse = clientResultException.GetRawResponse();
-            var headerValue = rawResponse?.Headers.TryGetValue("retry-after", out var value) == true
-                ? value
-                : null;
-            if (TryParseRetryAfterValue(headerValue, out var retryAfter))
+            if (TryGetHeaderValue(rawResponse, "x-ms-retry-after-ms", out var msHeaderValue)
+                && TryParseMilliseconds(msHeaderValue, out var msRetryAfter))
+            {
+                return msRetryAfter;
+            }
+
+            if (TryGetHeaderValue(rawResponse, "retry-after-ms", out var retryAfterMsHeaderValue)
+                && TryParseMilliseconds(retryAfterMsHeaderValue, out var retryAfterMs))
+            {
+                return retryAfterMs;
+            }
+
+            if (TryGetHeaderValue(rawResponse, "retry-after", out var headerValue)
+                && TryParseRetryAfterValue(headerValue, out var retryAfter))
+            {
                 return retryAfter;
+            }
         }
 
-        if (ex is HttpRequestException)
-            return null;
+        if (TryParseRetryAfterValue(ex.Message, out var messageRetryAfter))
+            return messageRetryAfter;
 
         return ex.InnerException is null ? null : ExtractRetryAfter(ex.InnerException);
     }
 
+    private static bool TryGetHeaderValue(PipelineResponse? response, string headerName, out string? headerValue)
+    {
+        headerValue = null;
+        if (response is null)
+            return false;
+
+        if (response.Headers.TryGetValue(headerName, out var value) && !string.IsNullOrWhiteSpace(value))
+        {
+            headerValue = value;
+            return true;
+        }
+
+        return false;
+    }
+
+    private static bool TryParseMilliseconds(string? value, out TimeSpan retryAfter)
+    {
+        retryAfter = default;
+        if (string.IsNullOrWhiteSpace(value))
+            return false;
+
+        if (int.TryParse(value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var msValue) && msValue >= 0)
+        {
+            retryAfter = TimeSpan.FromMilliseconds(msValue);
+            return true;
+        }
+
+        return false;
+    }
+
     private static bool TryParseRetryAfterValue(string? headerValue, out TimeSpan retryAfter)
     {
         retryAfter = default;
@@ -140,6 +190,15 @@ private static bool TryParseRetryAfterValue(string? headerValue, out TimeSpan re
             }
         }
 
+        var secondsMatch = _retryAfterSecondsRegex.Match(headerValue);
+        if (secondsMatch.Success
+            && int.TryParse(secondsMatch.Groups["seconds"].Value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var extractedSeconds)
+            && extractedSeconds >= 0)
+        {
+            retryAfter = TimeSpan.FromSeconds(extractedSeconds);
+            return true;
+        }
+
         return false;
     }
 
@@ -165,6 +224,31 @@ private TimeSpan ClampRetryDelay(TimeSpan delay) =>
             ? TimeSpan.FromMilliseconds(_retryOptions.MaxDelayMs)
             : delay;
 
+    private async Task<T> ExecuteEmbeddingRequestWithPacingAsync<T>(
+        Func<CancellationToken, Task<T>> embeddingRequest,
+        CancellationToken cancellationToken)
+    {
+        await _embeddingRequestLock.WaitAsync(cancellationToken);
+        try
+        {
+            var minimumSpacing = TimeSpan.FromMilliseconds(_retryOptions.MinInterRequestDelayMs);
+            if (minimumSpacing > TimeSpan.Zero && _lastEmbeddingRequestStartedUtc != DateTimeOffset.MinValue)
+            {
+                var elapsed = DateTimeOffset.UtcNow - _lastEmbeddingRequestStartedUtc;
+                var remainingDelay = minimumSpacing - elapsed;
+                if (remainingDelay > TimeSpan.Zero)
+                    await Task.Delay(remainingDelay, cancellationToken);
+            }
+
+            _lastEmbeddingRequestStartedUtc = DateTimeOffset.UtcNow;
+            return await embeddingRequest(cancellationToken);
+        }
+        finally
+        {
+            _embeddingRequestLock.Release();
+        }
+    }
+
     /// <summary>
     /// Wraps an async operation with retry logic for transient failures.
     /// </summary>
@@ -237,7 +321,9 @@ private async Task<T> ExecuteWithRetryAsync<T>(
     public async Task<ReadOnlyMemory<float>> GenerateEmbeddingAsync(string text, CancellationToken cancellationToken = default)
     {
         var embedding = await ExecuteWithRetryAsync(
-            async ct => await embeddingGenerator.GenerateAsync(text, cancellationToken: ct),
+            async ct => await ExecuteEmbeddingRequestWithPacingAsync(
+                async pacingCt => await embeddingGenerator.GenerateAsync(text, cancellationToken: pacingCt),
+                ct),
             "GenerateEmbedding",
             cancellationToken);
         return embedding.Vector;
@@ -287,44 +373,89 @@ public async Task GenerateBookContentEmbeddingsAndUploadToVectorStore(
 
         // ── Step 2 & 3: Batch-embed and immediately upsert each batch ─────────────────
         // Azure OpenAI supports at most EmbeddingBatchSize inputs per GenerateAsync call.
-        // bookContents is streamed in fixed-size batches without full upfront materialization,
-        // keeping peak memory bounded to one batch of chunk objects and their embeddings at a time.
+        // The effective request size starts at min(EmbeddingBatchSize, MaxEmbeddingBatchSize)
+        // and adaptively downshifts on 429 throttling responses.
+        // bookContents is streamed in batches without full upfront materialization,
+        // keeping peak memory bounded to one batch (or adaptive split) at a time.
         // The staging-swap (Step 3) is safe because it only runs after all batches have
         // been successfully upserted.
-        var buffer = new List<BookContentChunk>(EmbeddingBatchSize);
+        var configuredMaxBatchSize = Math.Clamp(_retryOptions.MaxEmbeddingBatchSize, 1, EmbeddingBatchSize);
+        var adaptiveBatchSize = configuredMaxBatchSize;
+        var buffer = new List<BookContentChunk>(configuredMaxBatchSize);
         int totalCount = 0;
 
-        async Task EmbedAndUpsertBatchAsync()
+        async Task EmbedAndUpsertExactBatchAsync(IReadOnlyList<BookContentChunk> batch)
         {
             var batchEmbeddings = await ExecuteWithRetryAsync(
-                async ct => await embeddingGenerator.GenerateAsync(
-                    buffer.Select(c => c.ChunkText), cancellationToken: ct),
-                $"GenerateBatchEmbeddings(size={buffer.Count})",
+                async ct => await ExecuteEmbeddingRequestWithPacingAsync(
+                    async pacingCt => await embeddingGenerator.GenerateAsync(
+                        batch.Select(c => c.ChunkText), cancellationToken: pacingCt),
+                    ct),
+                $"GenerateBatchEmbeddings(size={batch.Count})",
                 cancellationToken);
 
-            if (batchEmbeddings.Count != buffer.Count)
+            if (batchEmbeddings.Count != batch.Count)
                 throw new InvalidOperationException(
-                    $"Embedding count mismatch: expected {buffer.Count}, got {batchEmbeddings.Count}.");
+                    $"Embedding count mismatch: expected {batch.Count}, got {batchEmbeddings.Count}.");
+
+            for (int i = 0; i < batch.Count; i++)
+                batch[i].TextEmbedding = batchEmbeddings[i].Vector;
 
-            for (int i = 0; i < buffer.Count; i++)
-                buffer[i].TextEmbedding = batchEmbeddings[i].Vector;
+            await staging.UpsertAsync(batch, cancellationToken);
+            totalCount += batch.Count;
+        }
+
+        async Task EmbedAndUpsertBatchAdaptiveAsync(IReadOnlyList<BookContentChunk> batch)
+        {
+            try
+            {
+                await EmbedAndUpsertExactBatchAsync(batch);
+            }
+            catch (Exception ex) when (IsRateLimitError(ex) && batch.Count > 1)
+            {
+                var splitSize = Math.Max(1, batch.Count / 2);
+                if (adaptiveBatchSize > splitSize)
+                {
+                    var previousAdaptiveBatchSize = adaptiveBatchSize;
+                    adaptiveBatchSize = splitSize;
+                    if (_logger is not null)
+                    {
+                        LogEmbeddingBatchDownshift(_logger, previousAdaptiveBatchSize, adaptiveBatchSize, _retryOptions.MaxRetries + 1);
+                    }
+                }
 
-            await staging.UpsertAsync(buffer, cancellationToken);
-            totalCount += buffer.Count;
-            buffer.Clear();
+                var first = batch.Take(splitSize).ToArray();
+                var second = batch.Skip(splitSize).ToArray();
+                await EmbedAndUpsertBatchAdaptiveAsync(first);
+                await EmbedAndUpsertBatchAdaptiveAsync(second);
+            }
+            catch (Exception ex) when (IsRateLimitError(ex) && batch.Count == 1)
+            {
+                throw new InvalidOperationException(
+                    $"Embedding request failed with repeated 429 rate limits even at batch size 1 after {_retryOptions.MaxRetries + 1} attempts.",
+                    ex);
+            }
         }
 
         try
         {
             foreach (var chunk in bookContents)
             {
                 buffer.Add(chunk);
-                if (buffer.Count == EmbeddingBatchSize)
-                    await EmbedAndUpsertBatchAsync();
+                if (buffer.Count >= adaptiveBatchSize)
+                {
+                    var batchToProcess = buffer.ToArray();
+                    buffer.Clear();
+                    await EmbedAndUpsertBatchAdaptiveAsync(batchToProcess);
+                }
             }
 
             if (buffer.Count > 0)
-                await EmbedAndUpsertBatchAsync();
+            {
+                var batchToProcess = buffer.ToArray();
+                buffer.Clear();
+                await EmbedAndUpsertBatchAdaptiveAsync(batchToProcess);
+            }
 
             Console.WriteLine($"Uploaded {totalCount} chunks to staging collection '{stagingName}'.");
         }
@@ -417,4 +548,14 @@ private static partial void LogEmbeddingRetryAttemptsExhausted(
         int attemptCount,
         string lastError,
         int? statusCode);
+
+    [LoggerMessage(
+        EventId = 12004,
+        Level = LogLevel.Warning,
+        Message = "Embedding batch downshift triggered after throttling. PreviousBatchSize={PreviousBatchSize}, NewBatchSize={NewBatchSize}, RetryAttemptsPerRequest={RetryAttemptsPerRequest}")]
+    private static partial void LogEmbeddingBatchDownshift(
+        ILogger logger,
+        int previousBatchSize,
+        int newBatchSize,
+        int retryAttemptsPerRequest);
 }
diff --git a/EssentialCSharp.Web/appsettings.json b/EssentialCSharp.Web/appsettings.json
@@ -25,6 +25,8 @@
       "MaxRetries": 5,
       "BaseDelayMs": 1000,
       "MaxDelayMs": 60000,
+      "MaxEmbeddingBatchSize": 2048,
+      "MinInterRequestDelayMs": 250,
       "BackoffMultiplier": 2.0,
       "MaxJitterFraction": 0.2
     },
