gdmcode 0.1.0__py3-none-any.whl

src/models/client.py

@@ -0,0 +1,715 @@
+"""GdmClient — structured-output-aware wrapper around openai.OpenAI.
+
+Provides retry with exponential backoff, structured parsing via
+client.beta.chat.completions.parse(), and provider-specific configuration.
+"""
+from __future__ import annotations
+
+import io
+import json
+import logging
+import random
+import time
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass, field, replace as _dc_replace
+from typing import Any, TypeVar
+
+import openai
+from pydantic import BaseModel
+
+from src.config import GdmConfig
+from src.exceptions import ApiError, ApiRateLimitError, ConfigError, SchemaError
+from src.models.definitions import PROVIDER_BASE_URLS, Provider
+
+__all__ = [
+    "GdmClient",
+    "BaseModelT",
+    "LocalProviderConfig",
+    "LocalModelCapabilities",
+    "OllamaProvider",
+    "VLLMProvider",
+    "_validate_local_base_url",
+    "should_use_batch",
+    "OpenAIBatchBackend",
+    "GeminiBatchBackend",
+    "GrokSequentialBackend",
+    "BatchClient",
+    "_map_status",
+]
+
+log = logging.getLogger(__name__)
+
+BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
+
+_MAX_RETRIES: int = 3
+_BACKOFF_BASE: float = 2.0
+_PARSE_MAX_RETRIES: int = 2
+_SCHEMA_CORRECTION_PROMPT: str = (
+    "Your previous response did not conform to the required JSON schema. "
+    "Please try again and ensure your response matches the schema exactly."
+)
+_RETRYABLE_ERRORS: tuple[type[Exception], ...] = (
+    openai.RateLimitError,
+    openai.InternalServerError,
+    openai.APIConnectionError,
+    openai.APITimeoutError,
+)
+_FINISH_REASON_TRUNCATED: str = "length"
+_FINISH_REASON_FILTERED: str = "content_filter"
+
+# ---------------------------------------------------------------------------
+# Local provider support
+# ---------------------------------------------------------------------------
+
+_ALLOWED_LOCAL_HOSTS: frozenset[str] = frozenset({"localhost", "127.0.0.1", "::1"})
+_LOCAL_PROVIDERS: frozenset[str] = frozenset({"ollama", "vllm"})
+
+
+@dataclass
+class LocalProviderConfig:
+    """Configuration for a local (non-cloud) model provider."""
+
+    provider: str           # "ollama" or "vllm"
+    base_url: str
+    model: str
+    local_only: bool = True
+    allow_cloud_fallback: bool = False
+    allow_external: bool = False
+
+
+@dataclass
+class LocalModelCapabilities:
+    """Capabilities reported by a local model provider."""
+
+    supports_tools: bool = False
+    supports_json_mode: bool = False
+    context_window: int = 4096
+    streaming: bool = True
+
+
+def _validate_local_base_url(url: str, allow_external: bool = False) -> None:
+    """Raise ConfigError if url is non-local and allow_external is False."""
+    from urllib.parse import urlparse
+    parsed = urlparse(url)
+    if not allow_external and parsed.hostname not in _ALLOWED_LOCAL_HOSTS:
+        raise ConfigError(
+            f"Local provider base_url '{url}' resolves to non-local host "
+            f"'{parsed.hostname}'. Set allow_external=true to override."
+        )
+
+
+class OllamaProvider:
+    """Native Ollama /api/chat client (NOT OpenAI-compatible)."""
+
+    def __init__(self, config: LocalProviderConfig) -> None:
+        _validate_local_base_url(config.base_url, config.allow_external)
+        self.config = config
+        self._capabilities: LocalModelCapabilities | None = None
+        import urllib.request
+        self._http = urllib.request  # stdlib only — no requests dep
+
+    def complete(self, messages: list[dict], *, tools: list | None = None, **kw: Any) -> dict:
+        """POST to /api/chat and return adapted OpenAI-shaped dict."""
+        import json
+
+        if tools is not None:
+            caps = self.probe_capabilities()
+            if not caps.supports_tools:
+                raise ConfigError(
+                    f"Model '{self.config.model}' does not support tool calls "
+                    "(supports_tools=False). Remove tools or use a capable model."
+                )
+
+        payload: dict[str, Any] = {
+            "model": self.config.model,
+            "messages": messages,
+            "stream": False,
+        }
+        if tools is not None:
+            payload["tools"] = tools
+
+        url = f"{self.config.base_url.rstrip('/')}/api/chat"
+        data = json.dumps(payload).encode("utf-8")
+        req = self._http.Request(
+            url,
+            data=data,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        with self._http.urlopen(req, timeout=60) as resp:
+            raw: dict = json.loads(resp.read())
+        return self.to_openai_response(raw)
+
+    def to_openai_response(self, raw: dict) -> dict:
+        """Adapt Ollama /api/chat response to OpenAI ChatCompletion shape.
+
+        raw["message"]["content"]    -> choices[0].message.content
+        raw["message"]["tool_calls"] -> choices[0].message.tool_calls (if present)
+        """
+        message = raw.get("message", {})
+        content = message.get("content", "")
+        tool_calls = message.get("tool_calls")
+
+        msg: dict[str, Any] = {"role": "assistant", "content": content}
+        if tool_calls is not None:
+            msg["tool_calls"] = tool_calls
+
+        return {
+            "id": f"ollama-{raw.get('created_at', 'response')}",
+            "object": "chat.completion",
+            "model": raw.get("model", self.config.model),
+            "choices": [
+                {
+                    "index": 0,
+                    "message": msg,
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": raw.get("prompt_eval_count", 0),
+                "completion_tokens": raw.get("eval_count", 0),
+                "total_tokens": (
+                    raw.get("prompt_eval_count", 0) + raw.get("eval_count", 0)
+                ),
+            },
+        }
+
+    def probe_capabilities(self) -> LocalModelCapabilities:
+        """Lazy-probe /api/show for model capabilities. Cached per instance."""
+        if self._capabilities is not None:
+            return self._capabilities
+
+        import json
+
+        try:
+            url = f"{self.config.base_url.rstrip('/')}/api/show"
+            data = json.dumps({"name": self.config.model}).encode("utf-8")
+            req = self._http.Request(
+                url,
+                data=data,
+                headers={"Content-Type": "application/json"},
+                method="POST",
+            )
+            with self._http.urlopen(req, timeout=10) as resp:
+                show_data: dict = json.loads(resp.read())
+
+            modelinfo = show_data.get("modelinfo", {})
+            context_window = (
+                modelinfo.get("llama.context_length")
+                or modelinfo.get("context_length")
+                or 4096
+            )
+            self._capabilities = LocalModelCapabilities(
+                supports_tools=bool(show_data.get("supports_tools", False)),
+                supports_json_mode=bool(show_data.get("supports_json_mode", False)),
+                context_window=int(context_window),
+                streaming=True,
+            )
+        except Exception:
+            self._capabilities = LocalModelCapabilities()
+
+        return self._capabilities
+
+
+class VLLMProvider:
+    """vLLM provider -- OpenAI-compatible, reuses openai.OpenAI with custom base_url."""
+
+    def __init__(self, config: LocalProviderConfig) -> None:
+        _validate_local_base_url(config.base_url, config.allow_external)
+        self.config = config
+
+    def get_client(self) -> "openai.OpenAI":
+        """Return openai.OpenAI(base_url=..., api_key='not-needed')."""
+        return openai.OpenAI(
+            base_url=self.config.base_url,
+            api_key="not-needed",
+        )
+
+
+
+class GdmClient:
+    """OpenAI-compatible client with retry, structured output, and provider abstraction.
+
+    Wraps openai.OpenAI with:
+    - Structured output via client.beta.chat.completions.parse() with Pydantic models
+    - Exponential backoff retry (3 retries, 2x backoff) on rate limit / server errors
+    - Provider-specific configuration (base URL, API key)
+    - Sync interface only (async in Phase 3)
+    """
+
+    def __init__(self, cfg: GdmConfig) -> None:
+        """Initialise the client using provider config from GdmConfig.
+
+        Raises:
+            ConfigError: if cfg.provider is not a recognised provider.
+        """
+        if cfg.provider not in PROVIDER_BASE_URLS:
+            raise ConfigError(
+                f"Unknown provider '{cfg.provider}'. "
+                f"Valid providers: {list(PROVIDER_BASE_URLS)}"
+            )
+        self._client = openai.OpenAI(
+            api_key=cfg.api_key,
+            base_url=PROVIDER_BASE_URLS[cfg.provider],
+        )
+        self._cfg = cfg
+
+    @classmethod
+    def for_provider(cls, provider: str, cfg: GdmConfig) -> "GdmClient":
+        """Create a GdmClient configured for *provider* using keys from *cfg*.
+
+        Used to switch providers at runtime (e.g. API fallback) without
+        reloading config. The returned client has a modified cfg snapshot
+        reflecting the fallback provider/key.
+
+        Raises:
+            ConfigError: if no API key is available for *provider*, or if
+                         *provider* is not a recognised provider.
+        """
+        if provider not in PROVIDER_BASE_URLS:
+            raise ConfigError(
+                f"Unknown provider '{provider}'. "
+                f"Valid providers: {list(PROVIDER_BASE_URLS)}"
+            )
+        key_map: dict[str, str | None] = {
+            Provider.GROK: cfg.xai_api_key,
+            Provider.GEMINI: cfg.gemini_api_key,
+            Provider.CODEX: cfg.openai_api_key,
+        }
+        api_key = key_map.get(provider)
+        if not api_key:
+            raise ConfigError(
+                f"No API key configured for fallback provider '{provider}'. "
+                f"Set XAI_API_KEY / GEMINI_API_KEY / OPENAI_API_KEY as appropriate."
+            )
+        instance = cls.__new__(cls)
+        instance._client = openai.OpenAI(
+            api_key=api_key,
+            base_url=PROVIDER_BASE_URLS[provider],
+        )
+        instance._cfg = _dc_replace(cfg, provider=provider, api_key=api_key)
+        return instance
+
+    @classmethod
+    def for_proxy(cls, proxy_url: str, proxy_token: str) -> "GdmClient":
+        """Route all LLM calls through a proxy server.
+
+        Used when direct API access is geo-restricted (e.g. Grok unavailable
+        in DRC). The proxy must expose an OpenAI-compatible /v1 endpoint.
+        ``proxy_token`` is sent as the Bearer auth header; HTTPS is strongly
+        recommended so the token is protected in transit.
+        """
+        if not proxy_url.startswith("https://") and not proxy_url.startswith("http://localhost"):
+            log.warning(
+                "Proxy URL '%s' does not use HTTPS — token will be sent unencrypted",
+                proxy_url,
+            )
+        instance = cls.__new__(cls)
+        instance._client = openai.OpenAI(
+            api_key=proxy_token,
+            base_url=proxy_url.rstrip("/"),
+        )
+        # _cfg = None is safe: _dispatch() reads it via getattr(self._cfg, ..., None)
+        instance._cfg = None
+        return instance
+
+    def _call_with_retry(self, fn: Callable[[], Any]) -> Any:
+        """Execute fn with exponential backoff retry on transient API errors.
+
+        Retries up to _MAX_RETRIES times on rate-limit, server, connection,
+        and timeout errors. Non-transient APIErrors are re-raised immediately
+        as ApiError.
+        """
+        last_exc: BaseException = RuntimeError("unreachable sentinel")
+        for attempt in range(_MAX_RETRIES + 1):
+            try:
+                return fn()
+            except _RETRYABLE_ERRORS as exc:
+                last_exc = exc
+                if attempt < _MAX_RETRIES:
+                    wait = _BACKOFF_BASE ** attempt
+                    log.warning(
+                        "Retryable API error (%s), backing off %.0fs (retry %d/%d)",
+                        type(exc).__name__,
+                        wait,
+                        attempt + 1,
+                        _MAX_RETRIES,
+                    )
+                    time.sleep(wait)
+            except openai.APIError as exc:
+                raise ApiError(
+                    str(exc), status_code=getattr(exc, "status_code", None)
+                ) from exc
+        status = getattr(last_exc, "status_code", None)
+        if isinstance(last_exc, openai.RateLimitError):
+            raise ApiRateLimitError(
+                f"Rate limit exceeded after {_MAX_RETRIES} retries: {last_exc}",
+                status_code=status,
+            ) from last_exc
+        raise ApiError(
+            f"API call failed after {_MAX_RETRIES} retries: {last_exc}",
+            status_code=status,
+        ) from last_exc
+
+    def complete(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        model: str,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str = "auto",
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> Any:
+        """Make a non-structured completion request. Used by the agent loop.
+
+        Returns an openai.types.chat.ChatCompletion object.
+        """
+        _tools = tools if tools is not None else openai.NOT_GIVEN
+        _tool_choice: Any = tool_choice if tools is not None else openai.NOT_GIVEN
+        _max_tokens: Any = max_tokens if max_tokens is not None else openai.NOT_GIVEN
+        _temperature: Any = temperature if temperature is not None else openai.NOT_GIVEN
+
+        def _call() -> Any:
+            return self._client.chat.completions.create(
+                model=model,
+                messages=messages,  # type: ignore[arg-type]  # SDK type is narrower
+                tools=_tools,
+                tool_choice=_tool_choice,
+                max_tokens=_max_tokens,
+                temperature=_temperature,
+            )
+
+        return self._call_with_retry(_call)
+
+    def parse(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        model: str,
+        response_format: type[BaseModelT],
+        max_tokens: int | None = None,
+    ) -> BaseModelT:
+        """Make a structured output request via client.beta.chat.completions.parse().
+
+        Returns the parsed Pydantic model instance.
+        Retries up to _PARSE_MAX_RETRIES times on schema parse failure with an
+        explicit correction prompt.
+
+        Raises:
+            ApiError: on API failure, refusal, or content filter.
+            SchemaError: on repeated schema parse failure or response truncation.
+        """
+        msgs: list[dict[str, Any]] = list(messages)
+        _max_tokens: Any = max_tokens if max_tokens is not None else openai.NOT_GIVEN
+
+        for attempt in range(_PARSE_MAX_RETRIES + 1):
+            def _call(m: list[dict[str, Any]] = msgs) -> Any:
+                return self._client.beta.chat.completions.parse(
+                    model=model,
+                    messages=m,  # type: ignore[arg-type]
+                    response_format=response_format,
+                    max_tokens=_max_tokens,
+                )
+
+            result = self._call_with_retry(_call)
+            if not result.choices:
+                raise ApiError("API returned no completion choices")
+
+            choice = result.choices[0]
+            refusal = getattr(choice.message, "refusal", None)
+            if refusal:
+                raise ApiError(f"Model refused request: {refusal}")
+            if choice.finish_reason == _FINISH_REASON_TRUNCATED:
+                raise SchemaError(
+                    f"Response truncated (max_tokens too low) for {response_format.__name__}"
+                )
+            if choice.finish_reason == _FINISH_REASON_FILTERED:
+                raise ApiError("Response blocked by content policy")
+            if choice.message.parsed is not None:
+                return choice.message.parsed  # type: ignore[return-value]
+
+            if attempt < _PARSE_MAX_RETRIES:
+                log.warning(
+                    "Schema parse failed (attempt %d/%d), retrying with correction",
+                    attempt + 1,
+                    _PARSE_MAX_RETRIES + 1,
+                )
+                msgs = msgs + [
+                    {"role": "assistant", "content": choice.message.content or ""},
+                    {"role": "user", "content": _SCHEMA_CORRECTION_PROMPT},
+                ]
+
+        raise SchemaError(
+            f"Failed to parse {response_format.__name__} "
+            f"after {_PARSE_MAX_RETRIES + 1} attempts"
+        )
+
+    def _dispatch(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        model: str,
+        tools: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Route request to a local provider if configured, else return None.
+
+        Returns the provider response when a local provider handles the
+        request, or None to signal the caller should use the cloud path.
+
+        1. If a local provider is configured and reachable, use it.
+        2. If local_only=True and local is unreachable -> raise ConfigError.
+        3. Otherwise fall back to cloud with a WARNING log and return None.
+
+        Raises:
+            ConfigError: if local_only=True and local provider is unreachable.
+            ConfigError: if local provider rejects the request (e.g. no tools).
+        """
+        local_providers = getattr(self._cfg, "local_providers", None) or []
+        if not local_providers:
+            return None
+
+        local_cfg: LocalProviderConfig = local_providers[0]
+
+        try:
+            if local_cfg.provider == "ollama":
+                provider = OllamaProvider(local_cfg)
+                return provider.complete(messages, tools=tools)
+            elif local_cfg.provider == "vllm":
+                vllm_provider = VLLMProvider(local_cfg)
+                client = vllm_provider.get_client()
+                _tools: Any = tools if tools is not None else openai.NOT_GIVEN
+                return client.chat.completions.create(
+                    model=model,
+                    messages=messages,  # type: ignore[arg-type]
+                    tools=_tools,
+                )
+        except ConfigError:
+            raise  # Re-raise config errors (e.g. tool support rejection)
+        except Exception as exc:
+            if local_cfg.local_only:
+                raise ConfigError(
+                    "Local provider unreachable and local_only=true"
+                    " -- refusing cloud fallback"
+                ) from exc
+            log.warning(
+                "Local provider unreachable (%s), falling back to cloud: %s",
+                local_cfg.provider,
+                exc,
+            )
+            return None
+
+
+    def complete_text(
+        self,
+        prompt: str,
+        *,
+        model: str,
+        max_tokens: int = 200,
+        system: str | None = None,
+    ) -> str:
+        """Convenience wrapper for simple text completion. Returns response text.
+
+        Raises:
+            ApiError: if the API returns no choices.
+        """
+        messages: list[dict[str, Any]] = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+        result = self.complete(messages, model=model, max_tokens=max_tokens)
+        if not result.choices:
+            raise ApiError("API returned no completion choices")
+        return result.choices[0].message.content or ""
+
+
+# ---------------------------------------------------------------------------
+# Batch API
+# ---------------------------------------------------------------------------
+
+_BATCH_DISCOUNT_RATE: dict[str, float] = {"gemini": 0.50, "codex": 0.50}
+_BATCH_MIN_SAVINGS_USD: float = 0.01
+
+
+def should_use_batch(requests: list[dict], provider: str) -> bool:
+    """Return True when batch submission is worthwhile for this provider/request set."""
+    if provider == "grok":
+        return False
+    if len(requests) < 5:
+        return False
+    estimated_tokens = sum(len(json.dumps(r)) // 4 for r in requests)
+    savings = estimated_tokens / 1_000_000 * _BATCH_DISCOUNT_RATE.get(provider, 0)
+    return savings > _BATCH_MIN_SAVINGS_USD
+
+
+class OpenAIBatchBackend:
+    """OpenAI Batch API backend — uploads JSONL, creates batch, polls, retrieves."""
+
+    def __init__(self, client: Any) -> None:
+        self._client = client
+
+    def submit(self, requests: list[dict]) -> str:
+        buf = io.BytesIO("\n".join(json.dumps(r) for r in requests).encode())
+        f = self._client.files.create(file=("batch.jsonl", buf), purpose="batch")
+        batch = self._client.batches.create(
+            input_file_id=f.id,
+            endpoint="/v1/chat/completions",
+            completion_window="24h",
+        )
+        return batch.id
+
+    def poll(self, batch_id: str) -> str:
+        return self._client.batches.retrieve(batch_id).status
+
+    def retrieve(self, batch_id: str) -> list[dict]:
+        batch = self._client.batches.retrieve(batch_id)
+        content = self._client.files.content(batch.output_file_id).text
+        return [json.loads(line) for line in content.splitlines() if line.strip()]
+
+
+class GeminiBatchBackend:
+    """Gemini batchGenerateContent backend."""
+
+    def __init__(self, api_key: str, base_url: str = "https://generativelanguage.googleapis.com") -> None:
+        self._api_key = api_key
+        self._base_url = base_url
+
+    def submit(self, requests: list[dict]) -> str:
+        import urllib.request as _ur
+        payload = json.dumps({"requests": requests}).encode()
+        req = _ur.Request(
+            f"{self._base_url}/v1beta/models:batchGenerateContent?key={self._api_key}",
+            data=payload,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        with _ur.urlopen(req, timeout=30) as r:
+            body = json.loads(r.read())
+        return body.get("name", str(uuid.uuid4()))
+
+    def poll(self, batch_id: str) -> str:
+        import urllib.request as _ur
+        url = f"{self._base_url}/v1beta/{batch_id}?key={self._api_key}"
+        with _ur.urlopen(url, timeout=30) as r:
+            body = json.loads(r.read())
+        return "completed" if body.get("done") else "in_progress"
+
+    def retrieve(self, batch_id: str) -> list[dict]:
+        import urllib.request as _ur
+        url = f"{self._base_url}/v1beta/{batch_id}/responses?key={self._api_key}"
+        with _ur.urlopen(url, timeout=30) as r:
+            return json.loads(r.read()).get("responses", [])
+
+
+class GrokSequentialBackend:
+    """Grok has no batch API — sequential calls with jittered exponential backoff."""
+
+    def __init__(self, client: Any) -> None:
+        self._client = client
+        self._results: dict[str, list[dict]] = {}
+
+    def submit(self, requests: list[dict]) -> str:
+        results = []
+        for i, req in enumerate(requests):
+            for attempt in range(5):
+                try:
+                    result = self._client.chat.completions.create(**req)
+                    results.append(
+                        result.model_dump() if hasattr(result, "model_dump") else dict(result)
+                    )
+                    break
+                except Exception as e:
+                    if "rate" in str(e).lower():
+                        wait = (2 ** attempt) + random.uniform(0, 1)
+                        log.warning("Rate limit on request %d, backing off %.1fs", i, wait)
+                        time.sleep(wait)
+                    else:
+                        raise
+        batch_id = f"grok-seq-{uuid.uuid4().hex[:8]}"
+        self._results[batch_id] = results
+        return batch_id
+
+    def poll(self, batch_id: str) -> str:
+        return "completed"
+
+    def retrieve(self, batch_id: str) -> list[dict]:
+        return self._results.get(batch_id, [])
+
+
+class BatchClient:
+    """Provider-agnostic batch submission client with DB persistence."""
+
+    def __init__(self, provider: str, backend: Any) -> None:
+        self.provider = provider
+        self._backend = backend
+        self._db_conn: Any = None  # set externally if crash-recovery needed
+
+    def submit(self, requests: list[dict], job_type: str = "index") -> str:
+        """Redact secrets, submit batch, persist to DB. Returns job UUID."""
+        import datetime
+        from src.security import redact
+        safe = [json.loads(redact(json.dumps(r))) for r in requests]
+        log.info("Submitting batch of %d requests (%s/%s)", len(safe), self.provider, job_type)
+        batch_id = self._backend.submit(safe)
+        job_id = str(uuid.uuid4())
+        if self._db_conn is not None:
+            now = datetime.datetime.utcnow().isoformat()
+            self._db_conn.execute(
+                "INSERT INTO batch_jobs "
+                "(id, provider, batch_id, job_type, status, request_count, created_at, submitted_at) "
+                "VALUES (?, ?, ?, ?, 'submitted', ?, ?, ?)",
+                (job_id, self.provider, batch_id, job_type, len(safe), now, now),
+            )
+            self._db_conn.commit()
+        log.info("Batch submitted: job_id=%s batch_id=%s", job_id, batch_id)
+        return job_id
+
+    def poll(self, job_id: str) -> str:
+        """Poll provider status, update DB, return internal status string."""
+        import datetime
+        if self._db_conn is None:
+            return "unknown"
+        row = self._db_conn.execute(
+            "SELECT batch_id, status FROM batch_jobs WHERE id=?", (job_id,)
+        ).fetchone()
+        if not row:
+            return "unknown"
+        batch_id, current_status = row[0], row[1]
+        if current_status in ("completed", "failed", "cancelled"):
+            return current_status
+        provider_status = self._backend.poll(batch_id)
+        status = _map_status(provider_status, self.provider)
+        updates: dict[str, str] = {"status": status}
+        if status in ("completed", "failed"):
+            updates["completed_at"] = datetime.datetime.utcnow().isoformat()
+        for k, v in updates.items():
+            self._db_conn.execute(f"UPDATE batch_jobs SET {k}=? WHERE id=?", (v, job_id))
+        self._db_conn.commit()
+        return status
+
+    def retrieve(self, job_id: str) -> list[dict]:
+        """Retrieve completed batch results."""
+        if self._db_conn is not None:
+            row = self._db_conn.execute(
+                "SELECT batch_id FROM batch_jobs WHERE id=?", (job_id,)
+            ).fetchone()
+            if row:
+                return self._backend.retrieve(row[0])
+        return []
+
+
+def _map_status(provider_status: str, provider: str) -> str:
+    """Normalise provider-specific status strings to internal status."""
+    _completed = {"completed", "succeeded", "done"}
+    _failed = {"failed", "error", "expired"}
+    _cancelled = {"cancelled", "canceled"}
+    if provider_status in _completed:
+        return "completed"
+    if provider_status in _failed:
+        return "failed"
+    if provider_status in _cancelled:
+        return "cancelled"
+    return "polling"
+