Learn3 · The AI Layer

3 · The AI Layer

Everything model-related lives behind one interface in src/ai/. The pipeline never imports a vendor SDK.

The interface

# src/ai/base.py
class AIProvider(ABC):
    @abstractmethod
    async def complete(self, messages, options) -> CompletionResult: ...
    @abstractmethod
    async def stream(self, messages, options) -> AsyncIterator[str]: ...
    @abstractmethod
    async def embed(self, texts) -> list[list[float]]: ...
    @abstractmethod
    async def health_check(self) -> bool: ...
 
    @property
    @abstractmethod
    def supports_function_calling(self) -> bool: ...
    @property
    @abstractmethod
    def context_window(self) -> int: ...

Plus vendor-neutral dataclasses: Message, CompletionOptions, CompletionResult (which carries token counts, model and finish reason), and AIProviderConfig.

ProviderNotes
openaiGPT models
anthropicClaude models
geminiDefault. Unified google-genai SDK
ollamaLocal — no key, no data leaves the machine
lmstudioLocal — OpenAI-compatible endpoint
customAny OpenAI-compatible URL

Adding one is a file in src/ai/providers/ plus a line in registry.py. Nothing in the pipeline changes.

Why local providers matter. A retail customer may not be allowed to send schema or row samples to a third party. Ollama and LM Studio make the same product deployable fully on-premise. That is a business answer to an architecture question, which is the kind interviewers remember.

The registry

# src/ai/registry.py
class AIProviderFactory:
    _providers = {"openai": OpenAIProvider, "anthropic": AnthropicProvider, ...}
 
    @classmethod
    def create(cls, provider_type: str, config) -> AIProvider:
        ...

Strategy pattern behind a factory. /v1/meta exposes supported_types() so clients can discover capability at runtime rather than hardcoding a list.

A migration worth understanding

The Gemini provider moved from the deprecated google-generativeai to google-genai. The interesting reason is not deprecation:

The old SDK configured credentials through a module-level genai.configure() call. Retail Mind can hold several Gemini configs at once (different tenants, different keys), and a global call meant whichever provider was constructed last silently set the API key for every other one. The new SDK scopes credentials to a Client instance.

Global state and multi-tenancy do not mix. That is the transferable lesson, and it generalises to connection pools, logging contexts and feature flags.

Prompts

Jinja2 templates in src/ai/prompts/, not f-strings in Python:

sql_generation.j2       schema DDL + dialect + row cap → SQL
answer_synthesis.j2     question + columns + rows → prose
schema_compression.j2   verbose schema → compact DDL
insight_summary.j2
anomaly_detection.j2
forecast_interpretation.j2

Why templates? Prompts change far more often than code. Keeping them as files means a prompt change is a reviewable diff that does not touch Python, and the eval suite can measure whether it helped.

The SQL prompt is rendered with the dialect, the compressed schema, and the row cap:

system_prompt = template.render(
    dialect=ctx.dialect,
    schema_ddl=ctx.compressed_schema or "",
    max_rows=ctx.options.get("max_rows", 1000),
)

Retries — and the failure taxonomy

src/ai/retry.py is the single choke point for every provider call. Its job is to answer one question: is this failure worth trying again?

async def with_retry(operation, *, max_attempts=4, base_delay=2.0, label="ai_call"):
    for attempt in range(1, max_attempts + 1):
        try:
            return await operation()
        except ProviderQuotaExhausted:
            raise
        except Exception as e:
            if is_daily_quota(e):
                raise ProviderQuotaExhausted(str(e)) from e
            if attempt >= max_attempts or not is_transient(e):
                raise
            delay = suggested_delay(e) or min(base_delay * 2 ** (attempt - 1), MAX_DELAY)
            delay += random.uniform(0, delay * 0.25)   # jitter
            await asyncio.sleep(delay)

Three classes:

ClassExamplesAction
Transient5xx, connection reset, timeout, per-minute rate limitRetry with exponential backoff + jitter, honouring Retry-After
Terminal401, invalid key, permission deniedFail immediately — retrying cannot succeed
Quota exhaustedDaily cap, insufficient_quotaFail immediately as typed ProviderQuotaExhausted

Why jitter

Backoff is jittered so callers rate-limited together do not all wake at the same instant and trigger the limit again — the thundering-herd problem. One line, and a classic interview question.

The expensive lesson

A 429 usually means “wait and retry”. But Gemini’s free tier caps generate_content at 20 requests per day, and still returns retryDelay: 50s — which is the wait until the next per-minute window, not until the daily counter resets.

Honouring it meant four backoffs (~180 s) per call, every call, for the rest of the day. A 23-case eval run spent three hours retrying a quota that could not come back, and reported 26% accuracy when 17 cases never reached the model.

Per-minute and per-day 429s carry identical prose and both carry a RetryInfo. Only the quota ID distinguishes them:

_DAILY_QUOTA_MARKERS = (
    "perdayper",        # GenerateRequestsPerDayPerProjectPerModel-FreeTier
    "requestsperday",
    "dailylimit",
    "dailyquota",
    "insufficientquota",  # OpenAI: out of credit
)
 
def is_daily_quota(error) -> bool:
    message = str(error).lower()
    if not any(m in message for m in ("429", "resource_exhausted", "quota")):
        return False
    squashed = re.sub(r"[^a-z]", "", message)   # survive casing, hyphens, underscores
    return any(marker in squashed for marker in _DAILY_QUOTA_MARKERS)
⚠️

Match on the machine-readable identifier, not the prose. Vendor error messages are unstable and often identical across genuinely different conditions. The quotaId is structured and stable. This generalises well beyond LLM APIs.

The full story is in Bugs That Shipped. tests/test_retry.py pins it using verbatim response bodies for both 429 kinds.

Surfacing a dead quota

ProviderQuotaExhausted is a typed error so callers can act on it:

except ProviderQuotaExhausted as e:
    raise HTTPException(
        status_code=503,
        detail=QUOTA_MESSAGE,
        headers={"Retry-After": str(_seconds_until_utc_midnight())},
    ) from e

503, not 500. The deployment is healthy; the upstream allowance is spent. Retry-After points at the next UTC midnight, when the daily counter rolls over. Streaming clients get an SSE error carrying code: "provider_quota_exhausted" so the UI can explain rather than just fail.

Cost control

Ways this codebase avoids paying a model:

  1. Cache first — a hit costs nothing (stage 1).
  2. Only two AI calls per query; the other seven stages are code.
  3. Chart generation is deterministic — no call at all.
  4. Intent is keyword-based — no call.
  5. Schema is compressed — fewer prompt tokens per request.
  6. Row caps — bounded data into the answer prompt.
  7. Small model by default — gemini-3.5-flash; the 23-case eval cost $0.0023.

ctx.tokens_used accumulates across both calls and is written to query_logs, so per-query cost is attributable after the fact.

“How would you cut cost 10×?” Cache harder (longer TTL, semantic near-miss matching), route by difficulty — small model for simple aggregations, large only when validation fails — and compress the schema more aggressively. Then measure each with the eval suite rather than guessing.