diff --git a/CHANGELOG.d/batch-endpoint-capability-gate.md b/CHANGELOG.d/batch-endpoint-capability-gate.md new file mode 100644 index 000000000..dc18aa5b6 --- /dev/null +++ b/CHANGELOG.d/batch-endpoint-capability-gate.md @@ -0,0 +1,3 @@ +Fixed `ModelClient.batch_chat()` routing any worker agent selected by `TaskOrchestrator.batch_route()` straight into the real provider Batch API (`/files`, `/batches`, `/files/{id}/content`), even when that provider never implements it — a 404 there failed the whole batch group with no fallback. Added `ModelAgent.batch_endpoint_supported: bool | None` (mirroring `reasoning_effort_supported`'s fail-closed pattern, including agent-pool persistence); the real Batch API is now only addressed when an operator has explicitly declared `batch_endpoint_supported: true`. Every other remote agent (`None`/`False`, the default) falls back to the same per-item emulation `_local_batch_chat` already performs for local providers. + +Evidence: OpenAI documents Batch as a separate asynchronous contract that uploads a purpose-`batch` JSONL file and creates a `/v1/batches` job for one declared endpoint and completion window; OpenAI-compatible chat alone therefore does not prove Batch support ([Create batch](https://developers.openai.com/api/reference/resources/batches/methods/create), [Upload file](https://developers.openai.com/api/reference/resources/files/methods/create)). Yu et al., “Orca: A Distributed Serving System for Transformer-Based Generative Models,” OSDI 2022, describes selective batching as an explicit serving-system mechanism rather than a universal model property ([paper](https://www.usenix.org/conference/osdi22/presentation/yu)). Together these support an explicit per-provider capability gate while preserving synchronous emulation for unproven providers. diff --git a/README.md b/README.md index 6eedf5755..e49e6730d 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ HTTP serving is hardened for local lab use: - Full orchestration traces are not returned by default. Set `include_orchestration_trace: true` per standard Chat request or start with `--expose-trace-by-default` when the caller is trusted. Requests that also use `tools` or `response_format` still fail with `unsupported_trace_disclosure`; remove the trace flag for structured or tool requests. - State is in-memory by default. Pass `--state-db PATH` (or `CONTEXTUAL_ORCHESTRATOR_STATE_DB`) to persist workflow runs, evaluation runs, audit, and analytics to a stdlib sqlite file so they survive a restart; without it, behavior is unchanged. - Response caching is off by default. Pass `--cache-ttl SECONDS` to serve identical requests (same messages + mode) from an in-memory TTL+LRU cache and skip the provider calls; `0` disables it. -- `ModelClient.batch_chat(agent, {custom_id: messages})` runs many requests through the provider's Batch API (async, 24h completion window, typically ~50% cheaper) — suited to evaluation/benchmark workloads, not latency-sensitive chat. The mock path answers synchronously. +- `ModelClient.batch_chat(agent, {custom_id: messages})` uses the provider's asynchronous Batch API only when the agent explicitly declares `batch_endpoint_supported: true`. Other remote providers use explicit per-item chat emulation, while ordinary chat requests stay synchronous; this keeps evaluation/benchmark workloads available without guessing Batch support from a model name. The mock path answers synchronously. Use real workers by replacing `mock://` agents with OpenAI-compatible endpoints. Provider secrets are resolved from a KV credential registry via `get_credential`, never from `os.getenv` at request time (see [docs/kv-credentials.md](docs/kv-credentials.md)): diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index 03f89ba3a..83de16fae 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -523,6 +523,12 @@ ] }, "stream_usage_supported": {"type": "boolean"}, + "batch_endpoint_supported": { + "anyOf": [ + {"type": "boolean"}, + {"type": "null"}, + ] + }, }, }, }, diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 9d168b3b3..e1d23a9c3 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -593,6 +593,11 @@ class ModelAgent: # ``None`` means provider support is unproven. Opt-in effort profiles then # fail closed unless the profile explicitly requests the safe ``omit`` fallback. reasoning_effort_supported: bool | None = None + # ``None``/``False`` mean the provider's real async Batch API (/files, + # /batches, /files/{id}/content) is not proven to exist for this agent. + # batch_chat() then falls closed to per-item emulation through the normal + # chat endpoint rather than guessing a remote provider supports it. + batch_endpoint_supported: bool | None = None # Explicit reviewed replica contract. A group never races when this is absent. endpoint_equivalence: dict[str, Any] | None = None # Provider-declared support for the Chat Completions terminal usage frame. @@ -620,6 +625,8 @@ def __post_init__(self) -> None: ) if self.reasoning_effort_supported not in (None, True, False): raise TypeError("reasoning_effort_supported must be true, false, or null") + if self.batch_endpoint_supported is not None and type(self.batch_endpoint_supported) is not bool: + raise TypeError("batch_endpoint_supported must be true, false, or null") if type(self.stream_usage_supported) is not bool: raise TypeError("stream_usage_supported must be a boolean") if self.endpoint_equivalence is not None: @@ -645,6 +652,7 @@ def to_config(self) -> dict[str, Any]: "context_window": self.context_window, "group_name": self.group_name, "reasoning_effort_supported": self.reasoning_effort_supported, + "batch_endpoint_supported": self.batch_endpoint_supported, "endpoint_equivalence": self.endpoint_equivalence, "stream_usage_supported": self.stream_usage_supported, } @@ -680,6 +688,7 @@ def from_dict(cls, value: dict[str, Any]) -> "ModelAgent": # pragma: no cover context_window=value.get("context_window"), group_name=value.get("group_name", ""), reasoning_effort_supported=value.get("reasoning_effort_supported"), + batch_endpoint_supported=value.get("batch_endpoint_supported"), endpoint_equivalence=value.get("endpoint_equivalence"), stream_usage_supported=value.get("stream_usage_supported", False), ) @@ -2846,6 +2855,19 @@ def batch_chat( ``requests`` maps a caller custom_id to its messages. Suited to eval/benchmark workloads (24h completion window, ~half the price); real-time chat should keep using ``chat``. The mock path answers synchronously so tests and local runs work. + + The real async Batch API (``/files``, ``/batches``, ``/files/{id}/content``) + is only ever addressed when ``agent.batch_endpoint_supported is True`` -- + an explicit, operator-declared capability, mirroring how + ``reasoning_effort_supported`` gates ``reasoning_effort``. A worker agent + can be selected for any configured chat-capable provider (Anthropic-shaped, + OpenRouter, a self-hosted OpenAI-compatible gateway, ...) many of which do + not implement ``/v1/batches`` at all; guessing support from the model id or + transport shape alone previously sent those requests straight into a real + HTTP call that 404s and fails the whole batch. An unproven agent + (``None``/``False``) instead falls back to the same per-item emulation + ``_local_batch_chat`` already performs for local providers -- looping each + request through ``chat()`` and aggregating into the identical result shape. """ if not is_chat_compatible_model_id(agent.model): raise ValueError( @@ -2858,7 +2880,7 @@ def batch_chat( } elif _is_local_provider_url(agent.base_url): results = self._local_batch_chat(agent, requests, temperature, effort_profile) - else: + elif agent.batch_endpoint_supported is True: destination = self._validate_provider(agent) # pragma: no cover batch_error: ProviderUpstreamError | None = None try: @@ -2875,6 +2897,11 @@ def batch_chat( ) if batch_error is not None: raise batch_error + else: + # Not proven to support the real Batch API: emulate rather than + # either misroute to an endpoint that likely 404s, or hard-fail the + # whole group for a provider that was never asked to prove itself. + results = self._local_batch_chat(agent, requests, temperature, effort_profile) return _validate_batch_results(requests, results) def _local_batch_chat( @@ -2884,7 +2911,11 @@ def _local_batch_chat( temperature: float | None, effort_profile: ReasoningEffortProfile | None = None, ) -> dict[str, dict[str, Any]]: - """Run local OpenAI-compatible requests concurrently through mlx-lm.""" + """Emulate a batch by running each request concurrently through ``chat()``. + + Used for local (mlx-lm) providers, and reused by ``batch_chat`` as the + fallback for any remote provider not proven to support the real Batch API. + """ request_settings = self.request_settings_snapshot() def complete(custom_id: str, messages: list[ChatMessage]) -> tuple[str, dict[str, Any]]: @@ -3113,6 +3144,7 @@ class _AgentPoolStore: "max_output_tokens", "context_window", "reasoning_effort_supported", + "batch_endpoint_supported", "stream_usage_supported", } ) @@ -3152,6 +3184,7 @@ def _create_normalized_schema(cls, conn: sqlite3.Connection) -> None: max_output_tokens INTEGER, context_window INTEGER, reasoning_effort_supported INTEGER, + batch_endpoint_supported INTEGER, stream_usage_supported INTEGER NOT NULL DEFAULT 0, CONSTRAINT agent_pool_disabled_flag_check CHECK (disabled IN (0, 1)), CONSTRAINT agent_pool_max_output_tokens_check @@ -3172,6 +3205,8 @@ def _create_normalized_schema(cls, conn: sqlite3.Connection) -> None: ), CONSTRAINT agent_pool_reasoning_effort_flag_check CHECK (reasoning_effort_supported IS NULL OR reasoning_effort_supported IN (0, 1)), + CONSTRAINT agent_pool_batch_endpoint_flag_check + CHECK (batch_endpoint_supported IS NULL OR batch_endpoint_supported IN (0, 1)), CONSTRAINT agent_pool_stream_usage_flag_check CHECK (stream_usage_supported IN (0, 1)) ) @@ -3212,8 +3247,9 @@ def _insert_agent(cls, conn: sqlite3.Connection, agent: "ModelAgent") -> None: INSERT INTO agent_pool ( agent_id, model_name, base_url, api_key_env, credential_key, priority, disabled, provider_name, local_credential_key, auth_scheme, - max_output_tokens, context_window, reasoning_effort_supported, stream_usage_supported - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + max_output_tokens, context_window, reasoning_effort_supported, + batch_endpoint_supported, stream_usage_supported + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( config["id"], @@ -3229,6 +3265,7 @@ def _insert_agent(cls, conn: sqlite3.Connection, agent: "ModelAgent") -> None: config["max_output_tokens"], config["context_window"], config["reasoning_effort_supported"], + config["batch_endpoint_supported"], int(config["stream_usage_supported"]), ), ) @@ -3281,6 +3318,12 @@ def _initialize_schema(cls, conn: sqlite3.Connection) -> None: "CHECK (reasoning_effort_supported IS NULL OR reasoning_effort_supported IN (0, 1))" ) columns.add("reasoning_effort_supported") + if "batch_endpoint_supported" not in columns: + conn.execute( + "ALTER TABLE agent_pool ADD COLUMN batch_endpoint_supported INTEGER " + "CHECK (batch_endpoint_supported IS NULL OR batch_endpoint_supported IN (0, 1))" + ) + columns.add("batch_endpoint_supported") if "max_output_tokens" not in columns: conn.execute( "ALTER TABLE agent_pool ADD COLUMN max_output_tokens INTEGER " @@ -3396,7 +3439,8 @@ def save(self, agent: "ModelAgent") -> None: priority = ?, disabled = ?, provider_name = ?, local_credential_key = ?, auth_scheme = ?, max_output_tokens = ?, context_window = ?, - reasoning_effort_supported = ?, stream_usage_supported = ? + reasoning_effort_supported = ?, batch_endpoint_supported = ?, + stream_usage_supported = ? WHERE agent_id = ? """, ( @@ -3412,6 +3456,7 @@ def save(self, agent: "ModelAgent") -> None: config["max_output_tokens"], config["context_window"], config["reasoning_effort_supported"], + config["batch_endpoint_supported"], int(config["stream_usage_supported"]), agent.id, ), @@ -3509,7 +3554,7 @@ def load_all(self) -> list["ModelAgent"]: SELECT agent_id, model_name, base_url, api_key_env, credential_key, priority, disabled, provider_name, local_credential_key, auth_scheme, max_output_tokens, context_window, - reasoning_effort_supported, stream_usage_supported + reasoning_effort_supported, batch_endpoint_supported, stream_usage_supported FROM agent_pool ORDER BY agent_id """ ).fetchall() @@ -3574,7 +3619,8 @@ def load_all(self) -> list["ModelAgent"]: max_output_tokens=row[10], context_window=row[11], reasoning_effort_supported=(None if row[12] is None else bool(row[12])), - stream_usage_supported=bool(row[13]), + batch_endpoint_supported=(None if row[13] is None else bool(row[13])), + stream_usage_supported=bool(row[14]), group_name=group_by_agent.get(row[0], ""), endpoint_equivalence=contract_by_agent.get(row[0]), ) @@ -5989,6 +6035,10 @@ def patch_agent(self, agent_pool_id: str, worker_agent_id: str, patch: dict[str, patched = replace( patched, stream_usage_supported=patch["stream_usage_supported"] ) + if "batch_endpoint_supported" in patch: + patched = replace( + patched, batch_endpoint_supported=patch["batch_endpoint_supported"] + ) updated_candidates = [patched if agent.id == worker_agent_id else agent for agent in self.candidates] updated_agents = [agent for agent in updated_candidates if not agent.disabled] @@ -8371,6 +8421,7 @@ def _agent_to_admin_payload(self, agent: ModelAgent) -> dict[str, Any]: "max_output_tokens": agent.max_output_tokens, "context_window": agent.context_window, "stream_usage_supported": agent.stream_usage_supported, + "batch_endpoint_supported": agent.batch_endpoint_supported, "group_name": agent.group_name, "group_routing": self._group_router.member_report(agent.id) if agent.group_name else None, } diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index eb6a77519..01677387a 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -283,7 +283,7 @@ def server_close(self) -> None: ALLOWED_AGENT_PATCH_KEYS = { "status", "priority", "tags", "provider_exclusions", "group_name", "endpoint_equivalence", "stream_usage_supported", "max_output_tokens", - "context_window", + "context_window", "batch_endpoint_supported", } ALLOWED_AGENT_CREATE_KEYS = { "id", @@ -301,6 +301,7 @@ def server_close(self) -> None: "stream_usage_supported", "max_output_tokens", "context_window", + "batch_endpoint_supported", } ALLOWED_MODEL_GROUP_KEYS = {"group_name", "member_agent_ids"} ALLOWED_MODEL_GROUP_PATCH_KEYS = {"member_agent_ids"} diff --git a/tests/test_agent_pool_db.py b/tests/test_agent_pool_db.py index 78b556392..cab6a80a8 100644 --- a/tests/test_agent_pool_db.py +++ b/tests/test_agent_pool_db.py @@ -39,6 +39,7 @@ def _seed() -> list[ModelAgent]: "max_output_tokens": 4096, "context_window": 128000, "stream_usage_supported": True, + "batch_endpoint_supported": True, } @@ -155,6 +156,21 @@ def test_stream_usage_capability_patch_survives_restart() -> None: assert restored._agent(agent.id).stream_usage_supported is True +def test_batch_endpoint_capability_patch_survives_restart() -> None: + with tempfile.TemporaryDirectory() as directory: + db = os.path.join(directory, "pool.db") + agent = ModelAgent("persisted_agent", "model-x") + first = TaskOrchestrator([agent], agents_db=db) + + updated = first.patch_agent( + "default", "persisted_agent", {"batch_endpoint_supported": True} + ) + assert updated["batch_endpoint_supported"] is True + + restored = TaskOrchestrator([agent], agents_db=db) + assert restored._agent(agent.id).batch_endpoint_supported is True + + def test_agent_pool_persists_limit_metadata_across_restart() -> None: with tempfile.TemporaryDirectory() as directory: db = os.path.join(directory, "pool.db") @@ -461,12 +477,21 @@ def test_http_create_and_delete_worker_agents() -> None: thread.start() base = f"http://127.0.0.1:{server.server_address[1]}/api/v1/agent_pools/default/worker_agents" try: + status, invalid_create = _call( + base, + "POST", + token, + {**NEW_AGENT, "id": "numeric_batch_agent", "batch_endpoint_supported": 1}, + ) + assert status == 400 and invalid_create["error"]["code"] == "invalid_request" + status, created = _call(base, "POST", token, NEW_AGENT) assert ( status == 201 and created["id"] == "coding_agent" and created["status"] == "active" and created["stream_usage_supported"] is True + and created["batch_endpoint_supported"] is True ) status, patched = _call( @@ -474,6 +499,17 @@ def test_http_create_and_delete_worker_agents() -> None: ) assert status == 200 and patched["stream_usage_supported"] is True + status, patched = _call( + f"{base}/general_agent", "PATCH", token, {"batch_endpoint_supported": True} + ) + assert status == 200 and patched["batch_endpoint_supported"] is True + + status, invalid_batch_capability = _call( + f"{base}/general_agent", "PATCH", token, {"batch_endpoint_supported": 1} + ) + assert status == 400 + assert invalid_batch_capability["error"]["code"] == "invalid_request" + status, invalid_capability = _call( f"{base}/general_agent", "PATCH", token, {"stream_usage_supported": 1} ) @@ -481,6 +517,7 @@ def test_http_create_and_delete_worker_agents() -> None: status, read = _call(f"{base}/general_agent", "GET", token) assert status == 200 and read["stream_usage_supported"] is True + assert read["batch_endpoint_supported"] is True status, dup = _call(base, "POST", token, NEW_AGENT) assert status == 400 # duplicate rejected diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index 7beb3d698..3cebcf1ac 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -81,6 +81,9 @@ def test_openapi_documents_compatibility_front_door() -> None: "/api/v1/agent_pools/{agent_pool_id}/worker_agents/{worker_agent_id}" ]["patch"]["requestBody"]["content"]["application/json"]["schema"] assert patch_schema["properties"]["stream_usage_supported"]["type"] == "boolean" + assert patch_schema["properties"]["batch_endpoint_supported"] == { + "anyOf": [{"type": "boolean"}, {"type": "null"}] + } assert patch_schema["properties"]["max_output_tokens"]["anyOf"] == [ { "type": "integer", diff --git a/tests/test_batch_api.py b/tests/test_batch_api.py index 57187ce04..32f44cc78 100644 --- a/tests/test_batch_api.py +++ b/tests/test_batch_api.py @@ -4,6 +4,12 @@ workloads. These drive ModelClient._batch_run against a local fake Batch server: multipart JSONL upload, batch creation, in_progress -> completed polling, JSONL result parsing with usage, and terminal-failure handling. Mock path stays sync. + +Also covers batch_chat()'s own real-vs-emulated routing gate: only an agent +explicitly declaring ``batch_endpoint_supported=True`` may reach the real +Batch API above; every other remote agent (unproven -- ``None``/``False``) +must fall back to the same per-item emulation local providers already use, +never guess-route into a provider that may not implement /v1/batches at all. """ from __future__ import annotations @@ -11,8 +17,10 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json from pathlib import Path +import socket import sys import threading +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -21,12 +29,14 @@ class _FakeBatchProvider: - """Implements /files, /batches, /batches/{id}, /files/{id}/content.""" + """Implements /files, /batches, /batches/{id}, /files/{id}/content, /chat/completions.""" def __init__(self, fail_status: str | None = None, polls_before_done: int = 2) -> None: outer = self self.uploaded_jsonl: bytes = b"" self.poll_count = 0 + self.batches_post_count = 0 + self.chat_completions_count = 0 class Handler(BaseHTTPRequestHandler): def _json(self, payload: dict, status: int = 200) -> None: @@ -44,7 +54,21 @@ def do_POST(self) -> None: # noqa: N802 outer.uploaded_jsonl = body # multipart wrapper included; checked via 'in' self._json({"id": "file_input_1"}) elif self.path == "/batches": + outer.batches_post_count += 1 self._json({"id": "batch_1", "status": "validating"}) + elif self.path == "/chat/completions": + # The emulation fallback's target: one ordinary chat completion + # per batch item, echoing the prompt so per-item routing is provable. + outer.chat_completions_count += 1 + payload = json.loads(body) + last_user = next( + (m["content"] for m in reversed(payload["messages"]) if m.get("role") == "user"), + "", + ) + self._json({ + "choices": [{"message": {"role": "assistant", "content": f"emulated: {last_user}"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) else: self._json({"error": "not found"}, 404) @@ -150,6 +174,64 @@ def test_mock_path_answers_synchronously() -> None: assert results["task_a"]["usage"] is None +def _remote_agent(base_url: str, **overrides: object) -> ModelAgent: + # credential_key="" keeps chat()'s NotConfigured pre-check out of the way + # (no KV credential is registered in tests); base_url is real HTTP so + # _validate_provider's https/public-IP checks are bypassed per-test below, + # exactly like the destination-pinning pattern already proven in + # test_provider_integration.py::test_open_provider_uses_validated_destination_without_dns_relookup. + fields: dict[object, object] = { + "id": "worker_agent", + "model": "gpt-x", + "base_url": base_url, + "credential_key": "", + } + fields.update(overrides) + return ModelAgent(**fields) # type: ignore[arg-type] + + +def _pinned_destination(provider: "_FakeBatchProvider") -> tuple[int, tuple[str, int]]: + return (socket.AF_INET, ("127.0.0.1", provider._server.server_address[1])) + + +def test_batch_chat_routes_batch_capable_agent_to_real_batch_endpoint() -> None: + """(a) batch_endpoint_supported=True still takes the real Batch API path.""" + with _FakeBatchProvider(polls_before_done=2) as provider: + agent = _remote_agent(provider.base_url, batch_endpoint_supported=True) + client = _client() + with patch.object(client, "_validate_provider", return_value=_pinned_destination(provider)): + results = client.batch_chat(agent, REQUESTS, poll_interval=0.01, poll_timeout=30) + + assert results["task_a"]["content"] == "answer A" + assert results["task_b"]["usage"]["completion_tokens"] == 2 + assert provider.batches_post_count == 1 # the real Batch API was used + assert provider.chat_completions_count == 0 # emulation never triggered + + +def test_batch_chat_falls_back_to_emulation_when_batch_endpoint_unproven() -> None: + """(b) an unproven agent (None/False) still returns a correct aggregated + result, via emulation -- and (c) that emulation path is real, not stubbed: + every request actually crosses the wire as one /chat/completions call. + """ + for unsupported_value in (None, False): + with _FakeBatchProvider() as provider: + agent = _remote_agent(provider.base_url, batch_endpoint_supported=unsupported_value) + client = _client() + with patch.object(client, "_validate_provider", return_value=_pinned_destination(provider)): + results = client.batch_chat(agent, REQUESTS) + + assert set(results) == {"task_a", "task_b"} + assert results["task_a"]["content"] == "emulated: question A" + assert results["task_b"]["content"] == "emulated: question B" + assert results["task_a"]["usage"]["completion_tokens"] == 1 + # One real chat completion per request -- the emulation path was + # actually exercised, not just asserted from a mocked return value. + assert provider.chat_completions_count == 2 + # The real Batch API was never touched: no misroute, no hard failure. + assert provider.batches_post_count == 0 + assert provider.uploaded_jsonl == b"" + + if __name__ == "__main__": for name, fn in sorted(globals().items()): if name.startswith("test_") and callable(fn): diff --git a/tests/test_orchestrator_client_boundaries.py b/tests/test_orchestrator_client_boundaries.py index 9bb489b41..18d226322 100644 --- a/tests/test_orchestrator_client_boundaries.py +++ b/tests/test_orchestrator_client_boundaries.py @@ -193,6 +193,13 @@ def test_model_agent_rejects_bad_local_credential_key_and_effort_flag() -> None: ModelAgent(id="agent_two", model="m", local_credential_key=123) # type: ignore[arg-type] with pytest.raises(TypeError, match="reasoning_effort_supported must be"): ModelAgent(id="agent_two", model="m", reasoning_effort_supported="yes") # type: ignore[arg-type] + for invalid_batch_flag in (0, 1, "true"): + with pytest.raises(TypeError, match="batch_endpoint_supported must be"): + ModelAgent( + id="agent_two", + model="m", + batch_endpoint_supported=invalid_batch_flag, # type: ignore[arg-type] + ) with pytest.raises(TypeError, match="max_output_tokens must be"): ModelAgent(id="agent_two", model="m", max_output_tokens=0) with pytest.raises(TypeError, match="max_output_tokens must be"): @@ -556,6 +563,7 @@ def test_batch_chat_success_on_https_provider_returns_validated_results() -> Non model="remote-chat-model", base_url="https://remote.example/v1", credential_key="REMOTE_API_KEY", + batch_endpoint_supported=True, ) client = ModelClient() requests = {"task_0": [{"role": "user", "content": "hi"}]} @@ -575,6 +583,7 @@ def test_batch_chat_wraps_provider_failures_without_provider_text() -> None: model="remote-chat-model", base_url="https://remote.example/v1", credential_key="REMOTE_API_KEY", + batch_endpoint_supported=True, ) client = ModelClient() requests = {"task_0": [{"role": "user", "content": "hi"}]} diff --git a/tests/test_provider_reliability.py b/tests/test_provider_reliability.py index fc816b937..1f1813f66 100644 --- a/tests/test_provider_reliability.py +++ b/tests/test_provider_reliability.py @@ -612,7 +612,12 @@ def _batch_run(self, agent, requests, temperature, poll_interval, poll_timeout, raise RuntimeError("provider-secret-batch-body") client = RawBatchFailureClient() - agent = ModelAgent("batch_worker", "gpt", base_url="https://provider.example/v1") + agent = ModelAgent( + "batch_worker", + "gpt", + base_url="https://provider.example/v1", + batch_endpoint_supported=True, + ) try: client.batch_chat(agent, {"task_0": [{"role": "user", "content": "ping"}]}) except RuntimeError as error: