Make batches, transcription, and audio-output chat SDK-reachable

CHANGELOG.d/openai-sdk-batches-audio-modalities.md

@@ -0,0 +1 @@
+Made four operations reachable with the unmodified stock OpenAI SDK (`api_key` + `base_url` only). Added a standard `POST/GET /v1/batches` and `POST /v1/batches/{id}/cancel` surface (`client.batches.create/retrieve/list/cancel()`) that wraps the existing `submit_batch`/`LocalBatchBackend`/`PgLlmBatchBackend` machinery, sourced from a file already uploaded through `/v1/files` with `purpose=batch`; a completed batch's real, already-computed results are reshaped into an OpenAI-shaped output file downloadable through the same `client.files.content()` call the SDK already supports. Gave `POST /v1/audio/transcriptions` the same multipart-detection treatment `/v1/files` already had, so `client.audio.transcriptions.create(file=..., model=...)` no longer 415s. Let Chat Completions accept `modalities: ["text","audio"]` + `audio: {voice, format}` (the SDK-native, and only SDK-reachable, way to request spoken-audio chat output) by routing it to the same single-agent `capability="audio"` passthrough `/v1/audio/generations` already used, instead of failing closed on every non-text modality. `GET /v1/videos` (list) stays intentionally unimplemented — see the code comment above `GET /v1/videos/{id}` for why a genuine list needs a different (per-job, provider-fanned) implementation than retrieve.

README.md

@@ -272,6 +272,14 @@ is read from a **KV config store**, never `os.getenv`.
   Submit via `POST /api/v1/batch_routing_jobs`, poll
   `GET /api/v1/batch_routing_jobs/{id}`, retrieve
   `POST /api/v1/batch_routing_jobs/{id}/results` (which records usage + cost).
+- **Standard OpenAI Batch API.** `POST /v1/batches`, `GET /v1/batches/{id}`,
+  `GET /v1/batches`, and `POST /v1/batches/{id}/cancel` wrap the same batch
+  machinery above in the real `client.batches.create/retrieve/list/cancel()`
+  shape the stock OpenAI SDK expects: `input_file_id` is a file already
+  uploaded through `POST /v1/files` with `purpose=batch` (an OpenAI Batch
+  input-line JSONL), and a completed batch's `output_file_id` downloads
+  through the same `client.files.content()` call. Only
+  `endpoint: "/v1/chat/completions"` is wired today.
 - **Batch embeddings.** Bulk, latency-tolerant embedding work (e.g. naruon's
   email-import backfill) submits to `POST /v1/batch/embeddings`
   (`{model, input|inputs:[...], endpoint, metadata|attribution}`) and polls

contextual_orchestrator/batch_routing.py

@@ -190,6 +190,7 @@ class BatchRequest:
     attribution: Dict[str, Any] = field(default_factory=dict)
     mode: str = "auto"
     zdr_only: bool = False
+    parameters: Dict[str, Any] = field(default_factory=dict)
 
     def to_jsonl_line(self, endpoint: str = "/v1/chat/completions") -> Dict[str, Any]:
         """Render this request as an OpenAI Batch API JSONL line."""
@@ -199,7 +200,11 @@ def to_jsonl_line(self, endpoint: str = "/v1/chat/completions") -> Dict[str, Any
             "url": endpoint,
             # ``zdr_only`` is a contextual-orchestrator selection policy, not
             # an upstream provider request field.
-            "body": {"model": self.model, "messages": self.messages},
+            "body": {
+                **self.parameters,
+                "model": self.model,
+                "messages": self.messages,
+            },
         }
 
 
@@ -324,13 +329,17 @@ def retrieve(self, job: BatchJob) -> List[BatchResultItem]:
         """Retrieve completed results for a batch job."""
         ...
 
+    def cancel(self, job: BatchJob) -> Dict[str, Any]:
+        """Request backend cancellation without inventing a terminal state."""
+        ...
+
 
 class LocalBatchBackend:
     """In-process batch backend that runs each request via an injected runner.
 
     Preserves the mock/local path: no external service, no Postgres. The runner
-    is any callable ``(messages, mode, model) -> {"answer": str, "mode": str}`` — the
-    orchestrator's own ``complete`` fits directly. Results are computed eagerly
+    is any callable ``(messages, mode, model[, parameters]) -> result``; the fourth
+    argument is used only when an SDK batch line carries provider options. Results are computed eagerly
     on submit and returned verbatim on retrieve, so the batch lifecycle is fully
     observable in tests.
     """
@@ -368,7 +377,12 @@ def run(request: BatchRequest) -> BatchResultItem:
                 else nullcontext()
             )
             with context:
-                result = self._runner(request.messages, request.mode, request.model)
+                if request.parameters:
+                    result = self._runner(
+                        request.messages, request.mode, request.model, request.parameters
+                    )
+                else:
+                    result = self._runner(request.messages, request.mode, request.model)
             answer = result.get("answer", "")
             # The runner (typically orchestrator.complete()) has no top-level
             # "usage" key -- real provider usage is nested per-step inside
@@ -430,6 +444,10 @@ def retrieve(self, job: BatchJob) -> List[BatchResultItem]:
         """Return the results computed at submit time."""
         return self._results.get(job.job_id, [])
 
+    def cancel(self, job: BatchJob) -> Dict[str, Any]:
+        """Report the already-terminal local job; no cancellation is fabricated."""
+        return {"accepted": False, **self.poll(job)}
+
 
 class PgLlmBatchBackend:
     """Batch backend that submits to **pg-llm-batch** and retrieves results.
@@ -581,6 +599,25 @@ async def _download() -> Dict[str, Any]:
             )
         return items
 
+    def cancel(self, job: BatchJob) -> Dict[str, Any]:
+        """Ask pg-llm-batch to cancel and return its authoritative status."""
+        cancel_batch_job = getattr(self._client, "cancel_batch_job", None)
+        if not callable(cancel_batch_job):
+            return {"accepted": False, **self.poll(job)}
+
+        async def _cancel() -> Dict[str, Any]:
+            return await cancel_batch_job(job.job_id, self._endpoint_alias)
+
+        result = self._run(_cancel())
+        status = result.get("status")
+        accepted = result.get("accepted") is True or status in {"cancelling", "cancelled"}
+        return {
+            "accepted": accepted,
+            "job_id": job.job_id,
+            "status": status,
+            "is_complete": status in {"completed", "failed", "expired", "cancelled"},
+        }
+
 
 def _extract_answer(body: Dict[str, Any]) -> str:
     choices = body.get("choices") or []

contextual_orchestrator/cost_router.py

@@ -176,7 +176,11 @@ def __init__(
                 start_embedding_job(recovered_job)
 
     def _run_local_batch(
-        self, messages: List[Dict[str, str]], mode: str, model: str
+        self,
+        messages: List[Dict[str, str]],
+        mode: str,
+        model: str,
+        parameters: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """Run one local item while retaining completed endpoint-race usage."""
         context = {
@@ -186,7 +190,21 @@ def _run_local_batch(
         }
         token = self._race_usage_context.set(context)
         try:
-            result = self.orchestrator.complete(messages, mode=mode, model_name=model)
+            if parameters:
+                proxied = self.orchestrator.proxy_completion(
+                    {**parameters, "model": model, "messages": messages},
+                    endpoint="chat/completions",
+                    single_agent=True,
+                )
+                choices = proxied.get("choices") or []
+                message = choices[0].get("message", {}) if choices else {}
+                result = {
+                    "answer": message.get("content", ""),
+                    "mode": "route",
+                    "trace": [{"usage": proxied.get("usage", {})}],
+                }
+            else:
+                result = self.orchestrator.complete(messages, mode=mode, model_name=model)
         finally:
             self._race_usage_context.reset(token)
         race_usage = []
@@ -1030,6 +1048,11 @@ def poll_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict[str
         job = self._require_job(job_id, owner_id=owner_id)
         return self.batch_backend.poll(job)
 
+    def cancel_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict[str, Any]:
+        """Request cancellation from the backend that owns the batch."""
+        job = self._require_job(job_id, owner_id=owner_id)
+        return self.batch_backend.cancel(job)
+
     def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict[str, Any]:
         """Retrieve results for a batch owned by ``owner_id`` and record usage.
 
@@ -1241,7 +1264,9 @@ def _batch_usage_record_id(job_id: str, custom_id: str, kind: str, index: int) -
     def _batch_item_usage_valid(item: BatchResultItem) -> bool:
         """True when a batch result item's provider-reported usage is trustworthy."""
         return (
-            item.prompt_tokens >= 0
+            type(item.prompt_tokens) is int
+            and item.prompt_tokens >= 0
+            and type(item.completion_tokens) is int
             and item.completion_tokens >= 0
             and (
                 item.usage_valid is True

contextual_orchestrator/file_registry.py

@@ -3,11 +3,18 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, replace
+import base64
 import hashlib
 import json
+import time
 from typing import Any
 import uuid
 
+# Sentinel ``agent_id``/``provider_file_id`` for a gateway-generated file (for
+# example batch output JSONL) that has no upstream provider replica at all —
+# its bytes are stored locally instead. See ``register_local``/``is_local``.
+_LOCAL_FILE_AGENT_ID = "__local__"
+
 
 class FileContractError(RuntimeError):
     """Raised when a provider response cannot become a gateway file resource."""
@@ -49,6 +56,19 @@ def __init__(self, job_registry: Any) -> None:
         self._owners = job_registry.mapping(
             "file_owners", decode=lambda raw: FileOwner(**raw)
         )
+        # base64 text (JSON-mapping values must be JSON-serializable), keyed
+        # by gateway file id. Only populated for locally registered files.
+        self._content = job_registry.mapping("local_file_content")
+        self._retention_seconds = job_registry.retention_seconds
+
+    def _purge_expired_local_content(self) -> None:
+        """Remove orphaned or expired gateway-generated file bodies."""
+        cutoff = int(time.time()) - self._retention_seconds
+        for gateway_file_id in list(self._content):
+            owner = self._owners.get(gateway_file_id)
+            created_at = owner.document.get("created_at") if owner is not None else None
+            if owner is None or type(created_at) is not int or created_at < cutoff:
+                self._content.pop(gateway_file_id, None)
 
     def register(
         self,
@@ -100,6 +120,58 @@ def register_replicas(
         self._owners[gateway_file_id] = owner
         return self.public_response(provider_result, owner)
 
+    def register_local(
+        self,
+        content: bytes,
+        *,
+        filename: str,
+        purpose: str,
+        owner_id: str,
+    ) -> dict[str, Any]:
+        """Register gateway-computed content (no upstream provider replica).
+
+        Used for artifacts the gateway builds itself — batch-output JSONL is
+        the first caller — rather than receiving from an upload. There is no
+        provider file to bind, so ``agent_id``/``provider_file_id`` are the
+        ``_LOCAL_FILE_AGENT_ID`` sentinel and the bytes are served straight
+        back out of ``self._content`` (see ``is_local``/``local_content``).
+        """
+        self._purge_expired_local_content()
+        if not owner_id:
+            raise FileContractError("file owner is unavailable")
+        gateway_file_id = f"file_{uuid.uuid4().hex}"
+        document = {
+            "id": gateway_file_id,
+            "object": "file",
+            "bytes": len(content),
+            "created_at": int(time.time()),
+            "filename": filename,
+            "purpose": purpose,
+            "status": "processed",
+        }
+        owner = FileOwner(
+            gateway_file_id=gateway_file_id,
+            provider_file_id=gateway_file_id,
+            agent_id=_LOCAL_FILE_AGENT_ID,
+            owner_id=owner_id,
+            agent_affinity_key=_LOCAL_FILE_AGENT_ID,
+            document=document,
+            replicas={},
+        )
+        self._owners[gateway_file_id] = owner
+        self._content[gateway_file_id] = base64.b64encode(content).decode("ascii")
+        return dict(document)
+
+    @staticmethod
+    def is_local(owner: FileOwner) -> bool:
+        """True when ``owner`` is gateway-generated content, not provider-backed."""
+        return owner.agent_id == _LOCAL_FILE_AGENT_ID
+
+    def local_content(self, gateway_file_id: str) -> bytes:
+        """Return the raw bytes registered for a locally generated file."""
+        self._purge_expired_local_content()
+        return base64.b64decode(self._content[gateway_file_id])
+
     @staticmethod
     def public_response(document: dict[str, Any], owner: FileOwner) -> dict[str, Any]:
         """Return a provider document with only the gateway identity exposed."""
@@ -129,6 +201,7 @@ def delete(self, gateway_file_id: str, owner_id: str) -> FileOwner:
         """Remove and return a principal-owned binding after provider deletion."""
         owner = self.owner(gateway_file_id, owner_id)
         del self._owners[gateway_file_id]
+        self._content.pop(gateway_file_id, None)
         return owner
 
     def retain_replicas(

contextual_orchestrator/orchestrator.py

@@ -2594,12 +2594,17 @@ def proxy_upload(
     ) -> dict[str, Any]:
         """Stream a seekable multipart body to one validated provider account."""
         if agent.base_url.startswith("mock://"):
+            # created_at/status are required by the real OpenAI FileObject
+            # shape (client.files.create() return type) -- omitting them left
+            # the mock transport unusable by a genuine SDK client.
             return {
                 "id": f"provider_file_{uuid.uuid4().hex}",
                 "object": "file",
                 "bytes": content_length,
+                "created_at": int(time.time()),
                 "purpose": "user_data",
                 "filename": "upload",
+                "status": "processed",
             }
         api_key = _provider_credential(agent)  # pragma: no cover
         headers = {  # pragma: no cover