fix: throttle Azure OpenAI embedding calls and back off on rate limits

The data loader previously fanned out all 4 contribution types
("Spoken", "Written", "Corrections", "Petitions") concurrently in a
TaskGroup, with each type spawning one task per page. With nothing
limiting concurrent embed_batch calls, the loader could trivially
exceed Azure's TPM quota (especially the default S0 tier at 350k TPM).
The retry decorator stop_after_attempt(3) had no wait, so it ignored
the Retry-After header.

Changes:
- Add an aiolimiter + semaphore around every embeddings.create call,
  configurable via OPENAI_EMBED_RATE_PER_SECOND, OPENAI_EMBED_MAX_CONCURRENCY,
  OPENAI_EMBED_BATCH_SIZE env vars. Defaults are tuned for the S0 tier.
- Switch tenacity to wait_random_exponential with jitter, retry only on
  RateLimitError / APITimeoutError / APIConnectionError, up to 8 attempts.
- Pass max_retries=6 to AsyncAzureOpenAI so the SDK honours Retry-After
  before tenacity takes over for longer waits.
- Reduce default batch_size from 100 to 32 (~10k tokens/call) to fit
  the per-minute quota.
- Run load_all_contributions sequentially; pages within each type still
  parallelise.

Author: jonahanton

parliament_mcp/openai_helpers.py

@@ -1,22 +1,59 @@
+import asyncio
 import logging
+import os
 from itertools import batched
 
 import httpx
+import openai
+from aiolimiter import AsyncLimiter
 from openai import AsyncAzureOpenAI
-from tenacity import retry, stop_after_attempt
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 from parliament_mcp.settings import ParliamentMCPSettings
 
 logger = logging.getLogger(__name__)
 
 
+def _env_float(name: str, default: float) -> float:
+    try:
+        return float(os.environ.get(name, default))
+    except (TypeError, ValueError):
+        return default
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, default))
+    except (TypeError, ValueError):
+        return default
+
+
+# Throttles for Azure OpenAI embedding calls. Defaults are conservative enough
+# for the S0 tier (350k TPM / 2,100 RPM); raise via env vars when on a larger
+# quota. The TPM cap, not RPM, is the binding constraint at default tier.
+_EMBED_RATE_PER_SECOND = _env_float("OPENAI_EMBED_RATE_PER_SECOND", 1.5)
+_EMBED_MAX_CONCURRENCY = _env_int("OPENAI_EMBED_MAX_CONCURRENCY", 2)
+_EMBED_DEFAULT_BATCH_SIZE = _env_int("OPENAI_EMBED_BATCH_SIZE", 32)
+
+_embed_limiter = AsyncLimiter(max_rate=_EMBED_RATE_PER_SECOND, time_period=1.0)
+_embed_semaphore = asyncio.Semaphore(_EMBED_MAX_CONCURRENCY)
+
+
 def get_openai_client(settings: ParliamentMCPSettings) -> AsyncAzureOpenAI:
     """Get an async Azure OpenAI client."""
     return AsyncAzureOpenAI(
         api_key=settings.AZURE_OPENAI_API_KEY,
         api_version=settings.AZURE_OPENAI_API_VERSION,
         azure_endpoint=settings.AZURE_OPENAI_ENDPOINT,
-        http_client=httpx.AsyncClient(timeout=30.0),
+        # The SDK honours Retry-After on 429s; let it ride out short windows
+        # before tenacity takes over for longer waits.
+        max_retries=6,
+        http_client=httpx.AsyncClient(timeout=60.0),
     )
 
 
@@ -27,21 +64,27 @@ async def embed_single(
     dimensions: int = 1024,
 ) -> list[float]:
     """Generate a single embedding for a text using Azure OpenAI."""
-    response = await client.embeddings.create(
-        input=text,
-        model=model,
-        dimensions=dimensions,
-    )
+    async with _embed_semaphore, _embed_limiter:
+        response = await client.embeddings.create(
+            input=text,
+            model=model,
+            dimensions=dimensions,
+        )
     return response.data[0].embedding
 
 
-@retry(stop=stop_after_attempt(3))
+@retry(
+    retry=retry_if_exception_type((openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError)),
+    wait=wait_random_exponential(min=2, max=60),
+    stop=stop_after_attempt(8),
+    reraise=True,
+)
 async def embed_batch(
     client: AsyncAzureOpenAI,
     texts: list[str],
     model: str,
     dimensions: int = 1024,
-    batch_size: int = 100,
+    batch_size: int = _EMBED_DEFAULT_BATCH_SIZE,
 ) -> list[list[float]]:
     """Generate embeddings for a list of texts using Azure OpenAI.
 
@@ -50,7 +93,7 @@ async def embed_batch(
         texts: List of texts to embed
         model: Deployment name for the embedding model
         dimensions: Number of dimensions for the embeddings (default 1024)
-        batch_size: Number of texts to process in each API call
+        batch_size: Number of texts per API call (default tuned for S0 tier)
 
     Returns:
         List of embedding vectors
@@ -59,15 +102,20 @@ async def embed_batch(
 
     for i, batch in enumerate(batched(texts, batch_size)):
         try:
-            response = await client.embeddings.create(
-                input=batch,
-                model=model,
-                dimensions=dimensions,
-            )
+            async with _embed_semaphore, _embed_limiter:
+                response = await client.embeddings.create(
+                    input=batch,
+                    model=model,
+                    dimensions=dimensions,
+                )
 
             batch_embeddings = [item.embedding for item in response.data]
             all_embeddings.extend(batch_embeddings)
 
+        except openai.RateLimitError:
+            # Let tenacity retry the whole batch with backoff.
+            logger.warning("Azure OpenAI rate limit hit on batch %d; backing off", i // batch_size + 1)
+            raise
         except Exception:
             logger.exception("Error generating embeddings for batch %d", i // batch_size + 1)
             raise

parliament_mcp/qdrant_data_loaders.py

@@ -234,12 +234,15 @@ async def load_all_contributions(
         from_date: str = "2020-01-01",
         to_date: str = "2020-02-10",
     ) -> None:
-        """Load all contribution types concurrently."""
+        """Load all contribution types sequentially.
+
+        Sequential rather than concurrent to keep total in-flight embedding
+        calls bounded; the per-type loader still parallelises pages internally.
+        """
         contribution_types = ["Spoken", "Written", "Corrections", "Petitions"]
         with self.progress_context():
-            async with asyncio.TaskGroup() as tg:
-                for contrib_type in contribution_types:
-                    tg.create_task(self.load_contributions_by_type(contrib_type, from_date, to_date))
+            for contrib_type in contribution_types:
+                await self.load_contributions_by_type(contrib_type, from_date, to_date)
 
     async def load_contributions_by_type(
         self,