diff --git a/provider_adapters/anthropic.py b/provider_adapters/anthropic.py index 4bb2c70..c89b03d 100644 --- a/provider_adapters/anthropic.py +++ b/provider_adapters/anthropic.py @@ -8,7 +8,10 @@ from provider_adapters.common import ( AsyncCallProviderHook, + StreamAcc, auth_token, + drive_http_sse, + ignore_delta, json_args, text_from_content, _classify_status, @@ -117,6 +120,142 @@ def _parse_anthropic_response(data: dict, status: int, latency: int) -> dict: } +def _anthropic_request(request: dict, token: str, extra_headers: dict) -> tuple[str, dict, dict]: + messages, system = _openai_messages_to_anthropic( + request.get("messages") or []) + body: dict[str, Any] = { + "model": (request.get("offer") or {}).get("wire_model_id") + or request["served_model_id"], + "messages": messages, + "max_tokens": request.get("max_tokens") or 4096, + } + if system: + body["system"] = system + if request.get("temperature") is not None: + body["temperature"] = request["temperature"] + tools = _openai_tools_to_anthropic(request.get("tools")) + if tools: + body["tools"] = tools + headers = { + "Content-Type": "application/json", + "x-api-key": token, + "anthropic-version": "2023-06-01", + **extra_headers, + } + url = ( + (request.get("base_url") or "https://api.anthropic.com/v1").rstrip("/") + + "/messages" + ) + return url, body, headers + + +async def stream_anthropic( + request: dict, + emit, + *, + env_get: Callable[[str], str | None] | None = None, + timeout_s: float = 30.0, + extra_headers: dict[str, str] | None = None, + client: Any = None, +) -> dict: + """Native Anthropic streaming backend for api_kind='anthropic'.""" + import httpx + + token, error = auth_token(request, env_get or os.environ.get, "ANTHROPIC_API_KEY") + if error is not None: + return error + url, body, headers = _anthropic_request(request, token, dict(extra_headers or {})) + body["stream"] = True + timeout = (request.get("timeout_ms") or int(timeout_s * 1000)) / 1000.0 + if client is None: + client = httpx.AsyncClient() + + async def _on_event(ev: dict, acc: StreamAcc) -> "dict | None": + etype = ev.get("type") + if etype == "message_start": + msg = ev.get("message") or {} + acc.raw_model = msg.get("model") or acc.raw_model + acc.usage.update(msg.get("usage") or {}) + elif etype == "content_block_start": + idx = ev.get("index", 0) + block = ev.get("content_block") or {} + if block.get("type") == "tool_use": + acc.saw_output = True + if acc.tool_calls is None: + acc.tool_calls = {} + acc.tool_calls[idx] = { + "id": block.get("id"), + "type": "function", + "function": { + "name": block.get("name") or "", + "arguments": ( + json.dumps(block.get("input")) + if block.get("input") else "" + ), + }, + } + elif etype == "content_block_delta": + idx = ev.get("index", 0) + delta = ev.get("delta") or {} + if delta.get("type") == "text_delta" and delta.get("text"): + acc.saw_output = True + acc.emitted = True + acc.text_parts.append(delta["text"]) + await emit(delta["text"]) + elif delta.get("type") == "input_json_delta": + acc.saw_output = True + if acc.tool_calls is None: + acc.tool_calls = {} + tc = acc.tool_calls.setdefault(idx, { + "id": None, + "type": "function", + "function": {"name": "", "arguments": ""}, + }) + if delta.get("partial_json"): + tc["function"]["arguments"] += delta["partial_json"] + elif str(delta.get("type") or "").endswith("_delta"): + acc.saw_output = True + elif etype == "message_delta": + delta = ev.get("delta") or {} + acc.finish_reason = delta.get("stop_reason") or acc.finish_reason + acc.usage.update(ev.get("usage") or {}) + elif etype == "error": + e = ev.get("error") or {} + msg = e.get("message") or str(e or ev) + return _err(_classify_status(acc.status, msg), acc.status, + acc.latency(), msg[:500]) + return None + + acc, error = await drive_http_sse( + client=client, url=url, body=body, headers=headers, timeout=timeout, + request=request, on_event=_on_event) + if error is not None: + return error + + by_index = acc.tool_calls or {} + tool_calls = [by_index[i] for i in sorted(by_index)] or None + text = "".join(acc.text_parts) + if not text.strip() and not tool_calls: + return _err("bad_response", 200, acc.latency(), "empty assistant content") + usage = acc.usage + return { + "ok": True, + "latency_ms": acc.latency(), + "response": { + "text": text, + "tool_calls": tool_calls, + "finish_reason": acc.finish_reason, + "tokens_in": usage.get("input_tokens"), + "tokens_out": usage.get("output_tokens"), + "tokens_total": ( + (usage.get("input_tokens") or 0) + (usage.get("output_tokens") or 0) + if usage else None + ), + "raw_model": acc.raw_model, + }, + } + + def make_anthropic_async_call_provider( env_get: Callable[[str], str | None] | None = None, timeout_s: float = 30.0, @@ -133,31 +272,11 @@ async def call(request: dict) -> dict: token, error = auth_token(request, _env_get, "ANTHROPIC_API_KEY") if error is not None: return error - messages, system = _openai_messages_to_anthropic( - request.get("messages") or []) - body: dict[str, Any] = { - "model": (request.get("offer") or {}).get("wire_model_id") - or request["served_model_id"], - "messages": messages, - "max_tokens": request.get("max_tokens") or 4096, - } - if system: - body["system"] = system - if request.get("temperature") is not None: - body["temperature"] = request["temperature"] - tools = _openai_tools_to_anthropic(request.get("tools")) - if tools: - body["tools"] = tools - headers = { - "Content-Type": "application/json", - "x-api-key": token, - "anthropic-version": "2023-06-01", - **_extra, - } - url = ( - (request.get("base_url") or "https://api.anthropic.com/v1").rstrip("/") - + "/messages" - ) + if request.get("first_token_timeout_ms") is not None: + return await stream_anthropic( + request, ignore_delta, env_get=_env_get, timeout_s=timeout_s, + extra_headers=_extra, client=client) + url, body, headers = _anthropic_request(request, token, _extra) timeout = (request.get("timeout_ms") or int(timeout_s * 1000)) / 1000.0 t0 = time.monotonic() try: diff --git a/provider_adapters/bedrock.py b/provider_adapters/bedrock.py index 9b8d889..c142d19 100644 --- a/provider_adapters/bedrock.py +++ b/provider_adapters/bedrock.py @@ -9,6 +9,10 @@ from provider_adapters.common import ( AsyncCallProviderHook, + before_first_output, + first_token_timeout_err, + first_token_timeout_s, + ignore_delta, json_args, text_from_content, _elapsed_ms, @@ -170,6 +174,27 @@ def _parse_bedrock_response(data: dict, latency: int) -> dict: } +def _bedrock_request(request: dict) -> dict[str, Any]: + messages, system = _openai_messages_to_bedrock(request.get("messages") or []) + body: dict[str, Any] = { + "modelId": _bedrock_model_id(request), + "messages": messages, + } + if system: + body["system"] = system + inference = {} + if request.get("max_tokens") is not None: + inference["maxTokens"] = request["max_tokens"] + if request.get("temperature") is not None: + inference["temperature"] = request["temperature"] + if inference: + body["inferenceConfig"] = inference + tools = _openai_tools_to_bedrock(request.get("tools")) + if tools: + body["toolConfig"] = tools + return body + + def _error_code(exc: Exception) -> str: response = getattr(exc, "response", None) if isinstance(response, dict): @@ -203,6 +228,162 @@ def _classify_bedrock_error(exc: Exception) -> str: return "unknown" +def _next_stream_event(events): + try: + return next(events) + except StopIteration: + return None + + +def _bedrock_client(region: str, timeout_s: float): + import boto3 + from botocore.config import Config + return boto3.client( + "bedrock-runtime", + region_name=region, + config=Config(read_timeout=timeout_s, connect_timeout=min(timeout_s, 10)), + ) + + +async def stream_bedrock( + request: dict, + emit, + *, + env_get: Callable[[str], str | None] | None = None, + timeout_s: float = 30.0, + client: Any = None, + client_factory: Callable[[str], Any] | None = None, +) -> dict: + """Native Bedrock ConverseStream backend for api_kind='bedrock'.""" + api_kind = request.get("api_kind") + if api_kind != "bedrock": + return _err("unsupported_api_kind", 0, 0, + f"api_kind={api_kind!r} not supported by Bedrock backend") + + _env_get = env_get or os.environ.get + region = _aws_region(request, _env_get) + bedrock = client or (client_factory(region) if client_factory + else _bedrock_client(region, timeout_s)) + body = _bedrock_request(request) + t0 = time.monotonic() + saw_output = False + emitted = False + text_parts: list[str] = [] + tool_calls_acc: dict[int, dict] = {} + finish_reason = None + usage: dict = {} + first_timeout_s = first_token_timeout_s(request) + + def _latency() -> int: + return _elapsed_ms(t0) + + def _saw_output() -> bool: + return saw_output + + def _timeout_err() -> dict: + return first_token_timeout_err(first_timeout_s, _latency()) + + def _tool_call(idx: int) -> dict: + return tool_calls_acc.setdefault(idx, { + "id": None, + "type": "function", + "function": {"name": "", "arguments": ""}, + }) + + try: + try: + data = await before_first_output( + asyncio.to_thread(bedrock.converse_stream, **body), + first_timeout_s, t0, _saw_output) + except (asyncio.TimeoutError, TimeoutError): + return _timeout_err() + events = iter((data or {}).get("stream") or []) + while True: + try: + event = await before_first_output( + asyncio.to_thread(_next_stream_event, events), + first_timeout_s, t0, _saw_output) + except (asyncio.TimeoutError, TimeoutError): + if not saw_output: + return _timeout_err() + raise + if event is None: + break + if "messageStop" in event: + finish_reason = (event["messageStop"] or {}).get("stopReason") \ + or finish_reason + continue + if "metadata" in event: + usage = (event["metadata"] or {}).get("usage") or usage + continue + if "contentBlockStart" in event: + payload = event["contentBlockStart"] or {} + idx = payload.get("contentBlockIndex", 0) + start = payload.get("start") or {} + tool = start.get("toolUse") + if isinstance(tool, dict): + saw_output = True + acc = _tool_call(idx) + acc["id"] = tool.get("toolUseId") or acc["id"] + if tool.get("name"): + acc["function"]["name"] = tool["name"] + continue + if "contentBlockDelta" in event: + payload = event["contentBlockDelta"] or {} + idx = payload.get("contentBlockIndex", 0) + delta = payload.get("delta") or {} + if delta.get("text"): + saw_output = True + emitted = True + text_parts.append(delta["text"]) + await emit(delta["text"]) + tool = delta.get("toolUse") + if isinstance(tool, dict): + saw_output = True + acc = _tool_call(idx) + if tool.get("toolUseId"): + acc["id"] = tool["toolUseId"] + if tool.get("name"): + acc["function"]["name"] = tool["name"] + if tool.get("input") is not None: + inp = tool["input"] + acc["function"]["arguments"] += ( + inp if isinstance(inp, str) else json.dumps(inp) + ) + if delta.get("reasoningContent"): + saw_output = True + continue + error_key = next((k for k in event if k.endswith("Exception")), None) + if error_key: + return _err("server_error", 0, _latency(), + str(event.get(error_key) or event)[:500]) + except Exception as exc: # noqa: BLE001 + if emitted: + partial = "".join(text_parts) + return _err("stream_interrupted", 0, _latency(), + f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})") + return _err(_classify_bedrock_error(exc), 0, _latency(), str(exc)[:500]) + + tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None + text = "".join(text_parts) + if not text.strip() and not tool_calls: + return _err("bad_response", 200, _latency(), "empty assistant content") + return { + "ok": True, + "latency_ms": _latency(), + "response": { + "text": text, + "tool_calls": tool_calls, + "finish_reason": finish_reason, + "tokens_in": usage.get("inputTokens"), + "tokens_out": usage.get("outputTokens"), + "tokens_total": usage.get("totalTokens"), + "tokens_cached": usage.get("cacheReadInputTokens"), + "raw_model": body.get("modelId"), + }, + } + + def make_bedrock_async_call_provider( env_get: Callable[[str], str | None] | None = None, timeout_s: float = 30.0, @@ -217,37 +398,19 @@ def _client(region: str): return client if client_factory is not None: return client_factory(region) - import boto3 - from botocore.config import Config - return boto3.client( - "bedrock-runtime", - region_name=region, - config=Config(read_timeout=timeout_s, connect_timeout=min(timeout_s, 10)), - ) + return _bedrock_client(region, timeout_s) async def call(request: dict) -> dict: api_kind = request.get("api_kind") if api_kind != "bedrock": return _err("unsupported_api_kind", 0, 0, f"api_kind={api_kind!r} not supported by Bedrock backend") + if request.get("first_token_timeout_ms") is not None: + return await stream_bedrock( + request, ignore_delta, env_get=_env_get, timeout_s=timeout_s, + client=client, client_factory=client_factory) - messages, system = _openai_messages_to_bedrock(request.get("messages") or []) - body: dict[str, Any] = { - "modelId": _bedrock_model_id(request), - "messages": messages, - } - if system: - body["system"] = system - inference = {} - if request.get("max_tokens") is not None: - inference["maxTokens"] = request["max_tokens"] - if request.get("temperature") is not None: - inference["temperature"] = request["temperature"] - if inference: - body["inferenceConfig"] = inference - tools = _openai_tools_to_bedrock(request.get("tools")) - if tools: - body["toolConfig"] = tools + body = _bedrock_request(request) region = _aws_region(request, _env_get) bedrock = _client(region) diff --git a/provider_adapters/common.py b/provider_adapters/common.py index 1f4d1c3..4200995 100644 --- a/provider_adapters/common.py +++ b/provider_adapters/common.py @@ -4,6 +4,8 @@ import json import asyncio import time +from contextlib import AsyncExitStack +from dataclasses import dataclass, field from typing import Awaitable, Callable, Any CallProviderHook = Callable[[dict], dict] @@ -188,6 +190,106 @@ def first_token_timeout_err(timeout_s: float, latency_ms: int) -> dict: f"first token timed out after {int(timeout_s * 1000)}ms") +async def ignore_delta(_delta: str) -> None: + """A no-op emit for driving a streaming backend as a NON-streaming call + (the first-token deadline path): the deltas are discarded, only the + aggregated final response is returned.""" + return None + + +@dataclass +class StreamAcc: + """Mutable accumulator an SSE event handler folds a native stream into, then + the adapter reads to build its complete-response dict. `tool_calls` is left + handler-owned (a dict keyed by content index, or a list) so each provider + accumulates in its own shape.""" + t0: float + status: int = 0 + text_parts: list = field(default_factory=list) + tool_calls: Any = None + finish_reason: Any = None + usage: dict = field(default_factory=dict) + raw_model: Any = None + saw_output: bool = False + emitted: bool = False + + def latency(self) -> int: + return elapsed_ms(self.t0) + + +async def drive_http_sse( + *, + client: Any, + url: str, + body: dict, + headers: dict, + timeout: float, + request: dict, + on_event: Callable[[dict, "StreamAcc"], Awaitable["dict | None"]], +) -> "tuple[StreamAcc | None, dict | None]": + """Drive an httpx Server-Sent-Events stream with the first-token deadline, + shared by the native httpx adapters (Anthropic, Gemini). Opens the stream + (guarded by the first-token timeout), maps a non-2xx to a classified error, + then feeds each decoded `data:` JSON object to `on_event(ev, acc)` — which + mutates `acc` and returns an error dict to abort or None to continue. Returns + `(acc, None)` on a clean end or `(None, error_dict)`; the caller finalizes the + accumulator into its provider-specific response shape. Bedrock is NOT a + consumer (its transport is the boto3 event stream, not httpx SSE).""" + acc = StreamAcc(t0=time.monotonic()) + timeout_s = first_token_timeout_s(request) + + def _saw() -> bool: + return acc.saw_output + + try: + async with AsyncExitStack() as stack: + try: + resp = await before_first_output( + stack.enter_async_context( + client.stream("POST", url, json=body, headers=headers, + timeout=timeout)), + timeout_s, acc.t0, _saw) + except (asyncio.TimeoutError, TimeoutError): + return None, first_token_timeout_err(timeout_s, acc.latency()) + acc.status = resp.status_code + if not (200 <= resp.status_code < 300): + raw = (await resp.aread()).decode("utf-8", "replace")[:500] + return None, err(classify_status(resp.status_code, raw), + resp.status_code, acc.latency(), raw) + + lines = resp.aiter_lines().__aiter__() + while True: + try: + line = await before_first_output( + lines.__anext__(), timeout_s, acc.t0, _saw) + except StopAsyncIteration: + break + except (asyncio.TimeoutError, TimeoutError): + if not acc.saw_output: + return None, first_token_timeout_err(timeout_s, acc.latency()) + raise + if not line or not line.startswith("data:"): + continue + data = line[len("data:"):].strip() + if data == "[DONE]": + break + try: + ev = json.loads(data) + except ValueError: + continue + aborted = await on_event(ev, acc) + if aborted is not None: + return None, aborted + except Exception as exc: # noqa: BLE001 + if acc.emitted: + partial = "".join(acc.text_parts) + return None, err("stream_interrupted", 0, acc.latency(), + f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})") + return None, err("network_error", 0, acc.latency(), + f"{type(exc).__name__}: {exc}") + return acc, None + + # Back-compat names used by older tests/scripts that import through # llm_router_host's re-export layer. _cached_tokens = cached_tokens diff --git a/provider_adapters/google.py b/provider_adapters/google.py index c45c8c5..aeef227 100644 --- a/provider_adapters/google.py +++ b/provider_adapters/google.py @@ -9,7 +9,10 @@ from provider_adapters.common import ( AsyncCallProviderHook, + StreamAcc, auth_token, + drive_http_sse, + ignore_delta, json_args, text_from_content, _classify_status, @@ -144,6 +147,121 @@ def _parse_gemini_response(data: dict, status: int, latency: int) -> dict: } +def _google_request( + request: dict, + token: str, + extra_headers: dict, + *, + stream: bool = False, +) -> tuple[str, dict, dict]: + model = (request.get("offer") or {}).get("wire_model_id") \ + or request["served_model_id"] + model_path = model if str(model).startswith("models/") else f"models/{model}" + contents, system = _openai_messages_to_gemini(request.get("messages") or []) + body: dict[str, Any] = {"contents": contents} + if system: + body["systemInstruction"] = system + generation = {} + if request.get("max_tokens") is not None: + generation["maxOutputTokens"] = request["max_tokens"] + if request.get("temperature") is not None: + generation["temperature"] = request["temperature"] + if generation: + body["generationConfig"] = generation + tools = _openai_tools_to_gemini(request.get("tools")) + if tools: + body["tools"] = tools + base = ( + request.get("base_url") + or "https://generativelanguage.googleapis.com/v1beta" + ).rstrip("/") + suffix = "streamGenerateContent?alt=sse" if stream else "generateContent" + # The API key rides the x-goog-api-key header, never the URL query — + # a `?key=` URL leaks into error_message/logs/traces on + # timeout/network errors (§3: keys are never logged or echoed). + url = f"{base}/{quote(model_path, safe='/')}:{suffix}" + headers = { + "Content-Type": "application/json", + "x-goog-api-key": token, + **extra_headers, + } + return url, body, headers + + +async def stream_google( + request: dict, + emit, + *, + env_get: Callable[[str], str | None] | None = None, + timeout_s: float = 30.0, + extra_headers: dict[str, str] | None = None, + client: Any = None, +) -> dict: + """Native Gemini streamGenerateContent backend for api_kind='google'.""" + import httpx + + token, error = auth_token(request, env_get or os.environ.get, "GEMINI_API_KEY") + if error is not None: + return error + url, body, headers = _google_request( + request, token, dict(extra_headers or {}), stream=True) + timeout = (request.get("timeout_ms") or int(timeout_s * 1000)) / 1000.0 + if client is None: + client = httpx.AsyncClient() + + async def _on_event(chunk: dict, acc: StreamAcc) -> "dict | None": + acc.raw_model = chunk.get("modelVersion") or acc.raw_model + if chunk.get("usageMetadata"): + acc.usage = chunk["usageMetadata"] + for cand in chunk.get("candidates") or []: + acc.finish_reason = cand.get("finishReason") or acc.finish_reason + for part in ((cand.get("content") or {}).get("parts") or []): + if part.get("text"): + acc.saw_output = True + acc.emitted = True + acc.text_parts.append(part["text"]) + await emit(part["text"]) + if part.get("functionCall"): + acc.saw_output = True + if acc.tool_calls is None: + acc.tool_calls = [] + fc = part["functionCall"] + acc.tool_calls.append({ + "id": fc.get("id") or fc.get("name"), + "type": "function", + "function": { + "name": fc.get("name") or "", + "arguments": json.dumps(fc.get("args") or {}), + }, + }) + return None + + acc, error = await drive_http_sse( + client=client, url=url, body=body, headers=headers, timeout=timeout, + request=request, on_event=_on_event) + if error is not None: + return error + + text = "".join(acc.text_parts) + tool_calls = acc.tool_calls or None + if not text.strip() and not tool_calls: + return _err("bad_response", 200, acc.latency(), "empty assistant content") + usage = acc.usage + return { + "ok": True, + "latency_ms": acc.latency(), + "response": { + "text": text, + "tool_calls": tool_calls, + "finish_reason": acc.finish_reason, + "tokens_in": usage.get("promptTokenCount"), + "tokens_out": usage.get("candidatesTokenCount"), + "tokens_total": usage.get("totalTokenCount"), + "raw_model": acc.raw_model, + }, + } + + def make_google_async_call_provider( env_get: Callable[[str], str | None] | None = None, timeout_s: float = 30.0, @@ -160,36 +278,11 @@ async def call(request: dict) -> dict: token, error = auth_token(request, _env_get, "GEMINI_API_KEY") if error is not None: return error - model = (request.get("offer") or {}).get("wire_model_id") \ - or request["served_model_id"] - model_path = model if str(model).startswith("models/") else f"models/{model}" - contents, system = _openai_messages_to_gemini(request.get("messages") or []) - body: dict[str, Any] = {"contents": contents} - if system: - body["systemInstruction"] = system - generation = {} - if request.get("max_tokens") is not None: - generation["maxOutputTokens"] = request["max_tokens"] - if request.get("temperature") is not None: - generation["temperature"] = request["temperature"] - if generation: - body["generationConfig"] = generation - tools = _openai_tools_to_gemini(request.get("tools")) - if tools: - body["tools"] = tools - base = ( - request.get("base_url") - or "https://generativelanguage.googleapis.com/v1beta" - ).rstrip("/") - # The API key rides the x-goog-api-key header, never the URL query — - # a `?key=` URL leaks into error_message/logs/traces on - # timeout/network errors (§3: keys are never logged or echoed). - url = f"{base}/{quote(model_path, safe='/')}:generateContent" - headers = { - "Content-Type": "application/json", - "x-goog-api-key": token, - **_extra, - } + if request.get("first_token_timeout_ms") is not None: + return await stream_google( + request, ignore_delta, env_get=_env_get, timeout_s=timeout_s, + extra_headers=_extra, client=client) + url, body, headers = _google_request(request, token, _extra) timeout = (request.get("timeout_ms") or int(timeout_s * 1000)) / 1000.0 t0 = time.monotonic() try: diff --git a/providers.py b/providers.py index 24412f9..db5449f 100644 --- a/providers.py +++ b/providers.py @@ -27,6 +27,7 @@ from __future__ import annotations import os +import functools from dataclasses import dataclass, field from typing import Any, Callable @@ -51,6 +52,7 @@ class Provider: # openai_compatible backend (no dedicated adapter). api_kind: "str | None" = None adapter: "Callable[..., Any] | None" = None + stream_adapter: "Callable[..., Any] | None" = None # KNOBS — operator-tunable settings.SCHEMA entries (keyed bare; namespaced # `.` at registration). Declared next to the provider. knobs: "dict[str, dict]" = field(default_factory=dict) @@ -83,16 +85,31 @@ def _anthropic_adapter(timeout_s): return make_anthropic_async_call_provider(timeout_s=timeout_s) +def _anthropic_stream_adapter(timeout_s): + from provider_adapters.anthropic import stream_anthropic + return functools.partial(stream_anthropic, timeout_s=timeout_s) + + def _bedrock_adapter(timeout_s): from provider_adapters.bedrock import make_bedrock_async_call_provider return make_bedrock_async_call_provider(timeout_s=timeout_s) +def _bedrock_stream_adapter(timeout_s): + from provider_adapters.bedrock import stream_bedrock + return functools.partial(stream_bedrock, timeout_s=timeout_s) + + def _google_adapter(timeout_s): from provider_adapters.google import make_google_async_call_provider return make_google_async_call_provider(timeout_s=timeout_s) +def _google_stream_adapter(timeout_s): + from provider_adapters.google import stream_google + return functools.partial(stream_google, timeout_s=timeout_s) + + def _codex_source(catalog, env_get): # codex's source just needs the codex provider id from the loaded catalog; # the OBSERVE/bind coupling (the source watching the backend's quota traffic) @@ -180,6 +197,7 @@ def _present(provider_id): enabled=lambda c: _has(c, lambda pid, p: p.get("source") == "bedrock" or str(pid).startswith("bedrock")), api_kind="bedrock", adapter=_bedrock_adapter, + stream_adapter=_bedrock_stream_adapter, ), # Direct first-party providers: no marketplace catalog of their own, but a # price source feeding host_store.provider_prices from each official pricing @@ -187,8 +205,10 @@ def _present(provider_id): Provider("openai", source=_official_price_source("openai"), enabled=_present("openai")), Provider("anthropic", api_kind="anthropic", adapter=_anthropic_adapter, + stream_adapter=_anthropic_stream_adapter, source=_official_price_source("anthropic"), enabled=_present("anthropic")), Provider("google", api_kind="google", adapter=_google_adapter, + stream_adapter=_google_stream_adapter, source=_official_price_source("google"), enabled=_present("google")), # codex: its source is built like any provider's (via _codex_source); only the # ADAPTER and the observe/bind coupling (the source watching its own backend's @@ -253,6 +273,14 @@ def native_adapter_handlers(timeout_s: float) -> "dict[str, Any]": if not p.special and p.api_kind and p.adapter} +def native_streaming_adapter_handlers(timeout_s: float) -> "dict[str, Any]": + """api_kind -> true streaming backend for native providers that support it + (codex remains wired separately in serve.py because it needs `observe`).""" + return {p.api_kind: p.stream_adapter(timeout_s) + for p in PROVIDERS + if not p.special and p.api_kind and p.stream_adapter} + + def _price_multiplier_knob(provider_id: str) -> dict: # Default 1.0 (no nudge) for every provider: a routing preference is an # operator decision, set + persisted from the Config tab, not hardcoded here diff --git a/serve.py b/serve.py index a4545da..49a18dc 100644 --- a/serve.py +++ b/serve.py @@ -151,15 +151,15 @@ def main() -> None: make_streaming_dispatcher, stream_codex, stream_openai_compatible, - stream_unsupported_api_kind, ) + _native_streaming = providers.native_streaming_adapter_handlers(args.timeout_s) streaming_call = make_streaming_dispatcher( default=functools.partial(stream_openai_compatible, timeout_s=args.timeout_s, provider_rules=provider_rules), - # the native adapters don't stream natively (fall back to unsupported); - # codex has a real streaming twin that also feeds `observe`. - handlers={**{ak: stream_unsupported_api_kind for ak in _native}, + # Native providers and Codex have real streaming twins; Codex also feeds + # `observe` for quota/scarcity pricing. + handlers={**_native_streaming, "openai_codex": functools.partial(stream_codex, auth=codex_auth, observe=observe)}, diff --git a/tests/test_native_providers.py b/tests/test_native_providers.py index 0848729..097df29 100644 --- a/tests/test_native_providers.py +++ b/tests/test_native_providers.py @@ -14,8 +14,11 @@ sys.path.insert(0, str(ROOT)) from provider_adapters.anthropic import make_anthropic_async_call_provider # noqa: E402 +from provider_adapters.anthropic import stream_anthropic # noqa: E402 from provider_adapters.bedrock import make_bedrock_async_call_provider # noqa: E402 +from provider_adapters.bedrock import stream_bedrock # noqa: E402 from provider_adapters.google import make_google_async_call_provider # noqa: E402 +from provider_adapters.google import stream_google # noqa: E402 class FakeResponse: @@ -43,12 +46,68 @@ async def post(self, url, json=None, headers=None, timeout=None): return self.response -class FakeBedrockClient: +class FakeStreamResponse: + def __init__(self, status_code, lines=None, body=b"", open_delay=0, line_delay=0): + self.status_code = status_code + self._lines = lines or [] + self._body = body + self._open_delay = open_delay + self._line_delay = line_delay + self.headers = {} + + async def __aenter__(self): + if self._open_delay: + import asyncio + await asyncio.sleep(self._open_delay) + return self + + async def __aexit__(self, *exc): + return False + + async def aiter_lines(self): + if self._line_delay: + import asyncio + await asyncio.sleep(self._line_delay) + for line in self._lines: + yield line + + async def aread(self): + return self._body + + +class FakeStreamClient: def __init__(self, response): self.response = response self.requests = [] + def stream(self, method, url, json=None, headers=None, timeout=None): + self.requests.append({ + "method": method, + "url": url, + "json": json, + "headers": headers, + "timeout": timeout, + }) + return self.response + + +class FakeBedrockClient: + def __init__(self, response, stream_delay=0): + self.response = response + self.stream_delay = stream_delay + self.requests = [] + self.methods = [] + def converse(self, **kwargs): + self.methods.append("converse") + self.requests.append(kwargs) + return self.response + + def converse_stream(self, **kwargs): + if self.stream_delay: + import time + time.sleep(self.stream_delay) + self.methods.append("converse_stream") self.requests.append(kwargs) return self.response @@ -67,6 +126,10 @@ def converse(self, **kwargs): } +def _sse(data: dict) -> str: + return "data: " + json.dumps(data) + + @pytest.mark.asyncio async def test_anthropic_native_handler_translates_openai_shape(): client = FakeClient(FakeResponse(200, { @@ -238,6 +301,218 @@ async def test_bedrock_native_handler_translates_openai_shape(): assert req["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"]["required"] == ["id"] +@pytest.mark.asyncio +async def test_anthropic_native_stream_emits_and_aggregates(): + deltas = [] + client = FakeStreamClient(FakeStreamResponse(200, [ + _sse({"type": "message_start", "message": { + "model": "claude-sonnet-4-6", + "usage": {"input_tokens": 11}, + }}), + _sse({"type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": "Use the tool."}}), + _sse({"type": "content_block_start", "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", + "name": "lookup", "input": {}}}), + _sse({"type": "content_block_delta", "index": 1, + "delta": {"type": "input_json_delta", + "partial_json": '{"id": "abc"}'}}), + _sse({"type": "message_delta", "delta": {"stop_reason": "tool_use"}, + "usage": {"output_tokens": 7}}), + ])) + + async def emit(delta): + deltas.append(delta) + + result = await stream_anthropic({ + "provider_id": "anthropic", + "api_kind": "anthropic", + "auth_env": "ANTHROPIC_TEST_KEY", + "served_model_id": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + }, emit, env_get={"ANTHROPIC_TEST_KEY": "anth-key"}.get, client=client) + + assert deltas == ["Use the tool."] + assert result["ok"] is True + assert result["response"]["text"] == "Use the tool." + assert result["response"]["tokens_total"] == 18 + assert result["response"]["tool_calls"][0]["function"] == { + "name": "lookup", + "arguments": '{"id": "abc"}', + } + assert client.requests[0]["json"]["stream"] is True + + +@pytest.mark.asyncio +async def test_google_native_stream_emits_and_aggregates(): + deltas = [] + client = FakeStreamClient(FakeStreamResponse(200, [ + _sse({"modelVersion": "gemini-3.1-pro-preview", + "candidates": [{"content": {"parts": [{"text": "Done."}]}}]}), + _sse({"candidates": [{"finishReason": "STOP", + "content": {"parts": [ + {"functionCall": {"name": "lookup", "args": {"id": "abc"}}}, + ]}}], + "usageMetadata": { + "promptTokenCount": 13, + "candidatesTokenCount": 5, + "totalTokenCount": 18, + }}), + ])) + + async def emit(delta): + deltas.append(delta) + + result = await stream_google({ + "provider_id": "gemini", + "api_kind": "google", + "auth_env": "GEMINI_TEST_KEY", + "served_model_id": "gemini-3.1-pro-preview", + "messages": [{"role": "user", "content": "hi"}], + }, emit, env_get={"GEMINI_TEST_KEY": "gem-key"}.get, client=client) + + assert deltas == ["Done."] + assert result["ok"] is True + assert result["response"]["text"] == "Done." + assert result["response"]["raw_model"] == "gemini-3.1-pro-preview" + assert result["response"]["tokens_total"] == 18 + assert result["response"]["tool_calls"][0]["function"] == { + "name": "lookup", + "arguments": '{"id": "abc"}', + } + assert client.requests[0]["url"].endswith( + "/models/gemini-3.1-pro-preview:streamGenerateContent?alt=sse") + + +@pytest.mark.asyncio +async def test_bedrock_native_stream_emits_and_aggregates(): + deltas = [] + client = FakeBedrockClient({"stream": [ + {"contentBlockDelta": {"contentBlockIndex": 0, + "delta": {"text": "Use the tool."}}}, + {"contentBlockStart": {"contentBlockIndex": 1, + "start": {"toolUse": { + "toolUseId": "toolu_1", + "name": "lookup", + }}}}, + {"contentBlockDelta": {"contentBlockIndex": 1, + "delta": {"toolUse": { + "input": '{"id": "abc"}', + }}}}, + {"messageStop": {"stopReason": "tool_use"}}, + {"metadata": {"usage": {"inputTokens": 11, "outputTokens": 7, + "totalTokens": 18}}}, + ]}) + + async def emit(delta): + deltas.append(delta) + + result = await stream_bedrock({ + "provider_id": "bedrock", + "api_kind": "bedrock", + "served_model_id": "us.anthropic.claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + }, emit, env_get={"AWS_REGION": "us-east-1"}.get, client=client) + + assert deltas == ["Use the tool."] + assert result["ok"] is True + assert result["response"]["text"] == "Use the tool." + assert result["response"]["raw_model"] == "us.anthropic.claude-sonnet-4-6" + assert result["response"]["tokens_total"] == 18 + assert result["response"]["tool_calls"][0]["function"] == { + "name": "lookup", + "arguments": '{"id": "abc"}', + } + assert client.methods == ["converse_stream"] + + +@pytest.mark.asyncio +async def test_native_call_providers_use_streaming_when_first_token_timeout_present(): + anthropic = FakeStreamClient(FakeStreamResponse(200, [ + _sse({"type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": "a"}}), + ])) + call = make_anthropic_async_call_provider( + env_get={"ANTHROPIC_API_KEY": "k"}.get, client=anthropic) + result = await call({ + "served_model_id": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 100, + }) + assert result["ok"] is True + assert anthropic.requests[0]["json"]["stream"] is True + + google = FakeStreamClient(FakeStreamResponse(200, [ + _sse({"candidates": [{"content": {"parts": [{"text": "g"}]}}]}), + ])) + call = make_google_async_call_provider( + env_get={"GEMINI_API_KEY": "k"}.get, client=google) + result = await call({ + "served_model_id": "gemini-3.1-pro-preview", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 100, + }) + assert result["ok"] is True + assert ":streamGenerateContent?alt=sse" in google.requests[0]["url"] + + bedrock = FakeBedrockClient({"stream": [ + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "b"}}}, + ]}) + call = make_bedrock_async_call_provider( + env_get={"AWS_REGION": "us-east-1"}.get, client=bedrock) + result = await call({ + "api_kind": "bedrock", + "served_model_id": "us.anthropic.claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 100, + }) + assert result["ok"] is True + assert bedrock.methods == ["converse_stream"] + + +@pytest.mark.asyncio +async def test_native_streams_enforce_first_token_timeout_before_output(): + async def emit(_delta): + raise AssertionError("no text should be emitted before timeout") + + anthropic = await stream_anthropic({ + "served_model_id": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 10, + }, emit, env_get={"ANTHROPIC_API_KEY": "k"}.get, + client=FakeStreamClient(FakeStreamResponse( + 200, + [_sse({"type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": "late"}})], + open_delay=0.05, + ))) + assert anthropic["error_kind"] == "timeout" + + google = await stream_google({ + "served_model_id": "gemini-3.1-pro-preview", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 10, + }, emit, env_get={"GEMINI_API_KEY": "k"}.get, + client=FakeStreamClient(FakeStreamResponse( + 200, + [_sse({"candidates": [{"content": {"parts": [{"text": "late"}]}}]})], + line_delay=0.05, + ))) + assert google["error_kind"] == "timeout" + + bedrock = await stream_bedrock({ + "api_kind": "bedrock", + "served_model_id": "us.anthropic.claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hi"}], + "first_token_timeout_ms": 10, + }, emit, env_get={"AWS_REGION": "us-east-1"}.get, + client=FakeBedrockClient({"stream": [ + {"contentBlockDelta": {"contentBlockIndex": 0, + "delta": {"text": "late"}}}, + ]}, stream_delay=0.05)) + assert bedrock["error_kind"] == "timeout" + + @pytest.mark.asyncio async def test_native_handlers_require_provider_credentials(): call = make_google_async_call_provider(env_get={}.get, client=FakeClient(None)) diff --git a/tests/test_providers.py b/tests/test_providers.py index 29a8152..7a13a43 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -17,7 +17,7 @@ def test_every_provider_declares_at_least_one_aspect(): # composition, not inheritance: a provider supplies only the aspects it has, # but it must contribute *something* (a source, a wire adapter, or knobs). for p in providers.PROVIDERS: - assert p.source or p.adapter or p.knobs or p.special, p.id + assert p.source or p.adapter or p.stream_adapter or p.knobs or p.special, p.id def test_knob_schema_is_namespaced_and_grouped(): @@ -65,6 +65,8 @@ def test_enabled_predicates_gate_on_the_catalog(): def test_native_api_kinds_declared_and_codex_is_the_one_exception(): native = {p.api_kind for p in providers.PROVIDERS if p.api_kind and not p.special} assert {"anthropic", "bedrock", "google"} <= native + streaming = providers.native_streaming_adapter_handlers(timeout_s=1) + assert {"anthropic", "bedrock", "google"} <= set(streaming) codex = next(p for p in providers.PROVIDERS if p.id == "codex") # codex is special only in its WIRE: its adapter + the observe/bind coupling # live in serve.py, so it declares no api_kind/adapter here and never lands in @@ -72,3 +74,41 @@ def test_native_api_kinds_declared_and_codex_is_the_one_exception(): assert codex.special and codex.api_kind is None and codex.adapter is None cat = {"providers": {"oai": {"api_kind": "openai_codex"}}, "models": {}} assert [s.name for s in providers.build_source_registry(cat)] == ["codex"] + + +def test_native_streaming_handlers_bind_the_real_backends(): + # serve.py wires the streaming dispatcher from these handlers: a native + # api_kind must reach its REAL stream backend, not the old + # stream_unsupported_api_kind fallback (#63). + from provider_adapters import anthropic, bedrock, google + handlers = providers.native_streaming_adapter_handlers(timeout_s=1) + assert handlers["anthropic"].func is anthropic.stream_anthropic + assert handlers["google"].func is google.stream_google + assert handlers["bedrock"].func is bedrock.stream_bedrock + + +def test_streaming_dispatcher_routes_by_api_kind(): + # the dispatcher selects the per-api_kind streaming backend and falls back to + # the default for openai_compatible — the wiring serve.py relies on. + import asyncio + from streaming import make_streaming_dispatcher + + seen = {} + + async def _default(req, emit): + seen["h"] = "default" + return {"ok": True} + + async def _anthropic(req, emit): + seen["h"] = "anthropic" + return {"ok": True} + + async def _emit(_d): + return None + + dispatch = make_streaming_dispatcher(default=_default, + handlers={"anthropic": _anthropic}) + asyncio.run(dispatch({"api_kind": "anthropic"}, _emit)) + assert seen["h"] == "anthropic" + asyncio.run(dispatch({"api_kind": "openai_compatible"}, _emit)) + assert seen["h"] == "default"