From 64ea3279a6881f91af4a0216f0102319644ea7ab Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Tue, 8 Sep 2026 11:30:03 -0600 Subject: [PATCH 1/8] LLM Gateway Module V1 --- assemblyai/__init__.py | 56 ++ assemblyai/llm_gateway/__init__.py | 65 ++ assemblyai/llm_gateway/v1/__init__.py | 56 ++ assemblyai/llm_gateway/v1/api.py | 191 ++++++ assemblyai/llm_gateway/v1/async_api.py | 154 +++++ assemblyai/llm_gateway/v1/async_client.py | 341 ++++++++++ assemblyai/llm_gateway/v1/client.py | 303 +++++++++ assemblyai/llm_gateway/v1/models.py | 179 +++++ assemblyai/llm_gateway/v1/params.py | 20 + assemblyai/llm_gateway/v1/stream.py | 90 +++ assemblyai/types.py | 30 + tests/unit/test_llm_gateway.py | 759 ++++++++++++++++++++++ tests/unit/test_llm_gateway_async.py | 480 ++++++++++++++ 13 files changed, 2724 insertions(+) create mode 100644 assemblyai/llm_gateway/__init__.py create mode 100644 assemblyai/llm_gateway/v1/__init__.py create mode 100644 assemblyai/llm_gateway/v1/api.py create mode 100644 assemblyai/llm_gateway/v1/async_api.py create mode 100644 assemblyai/llm_gateway/v1/async_client.py create mode 100644 assemblyai/llm_gateway/v1/client.py create mode 100644 assemblyai/llm_gateway/v1/models.py create mode 100644 assemblyai/llm_gateway/v1/params.py create mode 100644 assemblyai/llm_gateway/v1/stream.py create mode 100644 tests/unit/test_llm_gateway.py create mode 100644 tests/unit/test_llm_gateway_async.py diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index 982f0ceb..0125ac39 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -1,6 +1,34 @@ from .__version__ import __version__ from .async_client import AsyncClient from .client import Client +from .llm_gateway.v1 import ( + AsyncLLMGateway, + AsyncLLMGatewayStream, + LLMGateway, + LLMGatewayCacheCreation, + LLMGatewayChatCompletion, + LLMGatewayChoice, + LLMGatewayChunkChoice, + LLMGatewayChunkDelta, + LLMGatewayCompletionChunk, + LLMGatewayCompletionTokensDetails, + LLMGatewayDefaultParameters, + LLMGatewayFunction, + LLMGatewayMessageParam, + LLMGatewayMessageParamUnion, + LLMGatewayModel, + LLMGatewayModelList, + LLMGatewayPricing, + LLMGatewayPricingData, + LLMGatewayPromptTokensDetails, + LLMGatewayResponseMessage, + LLMGatewayStream, + LLMGatewayStreamAccumulator, + LLMGatewayToolCall, + LLMGatewayTopProvider, + LLMGatewayUnderstandingResponse, + LLMGatewayUsage, +) from .prerecorded.v2 import AsyncTranscriber, AsyncTranscript from .sync.v1 import AsyncSyncTranscriber, SyncTranscriber from .transcriber import Transcriber, Transcript, TranscriptGroup @@ -25,6 +53,7 @@ LanguageDetectionOptions, ListTranscriptParameters, ListTranscriptResponse, + LLMGatewayError, PageDetails, Paragraph, PIIRedactedAudioMethod, @@ -96,6 +125,8 @@ def __getattr__(name: str): # types "AssemblyAIError", "AsyncClient", + "AsyncLLMGateway", + "AsyncLLMGatewayStream", "AsyncSyncTranscriber", "AsyncTranscriber", "AsyncTranscript", @@ -119,6 +150,31 @@ def __getattr__(name: str): "LanguageDetectionOptions", "ListTranscriptParameters", "ListTranscriptResponse", + "LLMGateway", + "LLMGatewayCacheCreation", + "LLMGatewayChatCompletion", + "LLMGatewayChoice", + "LLMGatewayChunkChoice", + "LLMGatewayChunkDelta", + "LLMGatewayCompletionChunk", + "LLMGatewayCompletionTokensDetails", + "LLMGatewayDefaultParameters", + "LLMGatewayError", + "LLMGatewayFunction", + "LLMGatewayMessageParam", + "LLMGatewayMessageParamUnion", + "LLMGatewayModel", + "LLMGatewayModelList", + "LLMGatewayPricing", + "LLMGatewayPricingData", + "LLMGatewayPromptTokensDetails", + "LLMGatewayResponseMessage", + "LLMGatewayStream", + "LLMGatewayStreamAccumulator", + "LLMGatewayToolCall", + "LLMGatewayTopProvider", + "LLMGatewayUnderstandingResponse", + "LLMGatewayUsage", "PageDetails", "Sentence", "Sentiment", diff --git a/assemblyai/llm_gateway/__init__.py b/assemblyai/llm_gateway/__init__.py new file mode 100644 index 00000000..f8cb8fad --- /dev/null +++ b/assemblyai/llm_gateway/__init__.py @@ -0,0 +1,65 @@ +"""Re-exports of the LLM Gateway client surface. + +The canonical location is ``assemblyai.llm_gateway.v1``, mirroring +``assemblyai/sync/v1``. +""" + +from .v1 import ( # noqa: F401 + AsyncLLMGateway, + AsyncLLMGatewayStream, + LLMGateway, + LLMGatewayCacheCreation, + LLMGatewayChatCompletion, + LLMGatewayChoice, + LLMGatewayChunkChoice, + LLMGatewayChunkDelta, + LLMGatewayCompletionChunk, + LLMGatewayCompletionTokensDetails, + LLMGatewayDefaultParameters, + LLMGatewayError, + LLMGatewayFunction, + LLMGatewayMessageParam, + LLMGatewayMessageParamUnion, + LLMGatewayModel, + LLMGatewayModelList, + LLMGatewayPricing, + LLMGatewayPricingData, + LLMGatewayPromptTokensDetails, + LLMGatewayResponseMessage, + LLMGatewayStream, + LLMGatewayStreamAccumulator, + LLMGatewayToolCall, + LLMGatewayTopProvider, + LLMGatewayUnderstandingResponse, + LLMGatewayUsage, +) + +__all__ = [ + "AsyncLLMGateway", + "AsyncLLMGatewayStream", + "LLMGateway", + "LLMGatewayCacheCreation", + "LLMGatewayChatCompletion", + "LLMGatewayChoice", + "LLMGatewayChunkChoice", + "LLMGatewayChunkDelta", + "LLMGatewayCompletionChunk", + "LLMGatewayCompletionTokensDetails", + "LLMGatewayDefaultParameters", + "LLMGatewayError", + "LLMGatewayFunction", + "LLMGatewayMessageParam", + "LLMGatewayMessageParamUnion", + "LLMGatewayModel", + "LLMGatewayModelList", + "LLMGatewayPricing", + "LLMGatewayPricingData", + "LLMGatewayPromptTokensDetails", + "LLMGatewayResponseMessage", + "LLMGatewayStream", + "LLMGatewayStreamAccumulator", + "LLMGatewayToolCall", + "LLMGatewayTopProvider", + "LLMGatewayUnderstandingResponse", + "LLMGatewayUsage", +] diff --git a/assemblyai/llm_gateway/v1/__init__.py b/assemblyai/llm_gateway/v1/__init__.py new file mode 100644 index 00000000..4c3a22b8 --- /dev/null +++ b/assemblyai/llm_gateway/v1/__init__.py @@ -0,0 +1,56 @@ +from ...types import LLMGatewayError +from .async_client import AsyncLLMGateway +from .client import LLMGateway +from .models import ( + LLMGatewayCacheCreation, + LLMGatewayChatCompletion, + LLMGatewayChoice, + LLMGatewayChunkChoice, + LLMGatewayChunkDelta, + LLMGatewayCompletionChunk, + LLMGatewayCompletionTokensDetails, + LLMGatewayDefaultParameters, + LLMGatewayFunction, + LLMGatewayModel, + LLMGatewayModelList, + LLMGatewayPricing, + LLMGatewayPricingData, + LLMGatewayPromptTokensDetails, + LLMGatewayResponseMessage, + LLMGatewayToolCall, + LLMGatewayTopProvider, + LLMGatewayUnderstandingResponse, + LLMGatewayUsage, +) +from .params import LLMGatewayMessageParam, LLMGatewayMessageParamUnion +from .stream import AsyncLLMGatewayStream, LLMGatewayStream, LLMGatewayStreamAccumulator + +__all__ = [ + "AsyncLLMGateway", + "AsyncLLMGatewayStream", + "LLMGateway", + "LLMGatewayCacheCreation", + "LLMGatewayChatCompletion", + "LLMGatewayChoice", + "LLMGatewayChunkChoice", + "LLMGatewayChunkDelta", + "LLMGatewayCompletionChunk", + "LLMGatewayCompletionTokensDetails", + "LLMGatewayDefaultParameters", + "LLMGatewayError", + "LLMGatewayFunction", + "LLMGatewayMessageParam", + "LLMGatewayMessageParamUnion", + "LLMGatewayModel", + "LLMGatewayModelList", + "LLMGatewayPricing", + "LLMGatewayPricingData", + "LLMGatewayPromptTokensDetails", + "LLMGatewayResponseMessage", + "LLMGatewayStream", + "LLMGatewayStreamAccumulator", + "LLMGatewayToolCall", + "LLMGatewayTopProvider", + "LLMGatewayUnderstandingResponse", + "LLMGatewayUsage", +] diff --git a/assemblyai/llm_gateway/v1/api.py b/assemblyai/llm_gateway/v1/api.py new file mode 100644 index 00000000..cd9c0fc0 --- /dev/null +++ b/assemblyai/llm_gateway/v1/api.py @@ -0,0 +1,191 @@ +import json +from typing import Any, Dict, Iterator, List, Optional + +import httpx + +from ... import types +from . import models + +ENDPOINT_MODELS = "/v1/models" +ENDPOINT_CHAT_COMPLETIONS = "/v1/chat/completions" +ENDPOINT_UNDERSTANDING = "/v1/understanding" +ENDPOINT_UNDERSTANDING_VALIDATE = "/v1/understanding/validate" + + +def _error_from_response(response: httpx.Response) -> types.LLMGatewayError: + """ + Builds an `LLMGatewayError` from a non-200 response. + + The service uses two envelopes: business errors + (`{"message", "code", "request_id", "metadata": {"errors": [...]}}`, from + route handlers) and 401s raised by auth middleware before a request + reaches a handler (`{"error", "status", "request_id"}`). Both are + detected here. A missing or unparsable body falls back to the raw + response text or a generic message. + """ + request_id: Optional[str] = None + message: Optional[str] = None + errors: Optional[List[str]] = None + + print(response.text) + + try: + body = response.json() + if isinstance(body, dict): + request_id = body.get("request_id") + if "message" in body: + message = body.get("message") + metadata = body.get("metadata") + if isinstance(metadata, dict) and isinstance(metadata.get("errors"), list): + errors = metadata["errors"] + elif "error" in body: + message = body.get("error") + except Exception: + message = response.text or None + + if not message: + message = f"LLM Gateway request failed with status {response.status_code}" + + return types.LLMGatewayError( + message, + status_code=response.status_code, + request_id=request_id, + errors=errors, + ) + + +def list_models( + client: httpx.Client, + *, + base_url: str, + timeout: float, +) -> models.LLMGatewayModelList: + """ + Lists the models available through the LLM Gateway. + + Args: + client: the HTTP client (carries the `authorization` header, though + this endpoint does not require one). + base_url: the LLM Gateway base URL, e.g. `https://llm-gateway.assemblyai.com`. + timeout: per-request timeout in seconds. + + Raises: `LLMGatewayError` on any non-200 response. + """ + response = client.get( + base_url.rstrip("/") + ENDPOINT_MODELS, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayModelList.parse_obj(response.json()) + + +def create_completion( + client: httpx.Client, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> models.LLMGatewayChatCompletion: + """ + Posts a non-streaming chat completion request. + + Raises: `LLMGatewayError` on any non-200 response. + """ + + response = client.post( + base_url.rstrip("/") + ENDPOINT_CHAT_COMPLETIONS, + json=body, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayChatCompletion.parse_obj(response.json()) + + +def stream_completion( + client: httpx.Client, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> Iterator[models.LLMGatewayCompletionChunk]: + """ + Posts a streaming chat completion request and yields parsed chunks. + + Chunk parsing is verified against OpenAI-routed models only (`data: {...}` + SSE lines terminated by `data: [DONE]`); other providers' streaming + framing was not verified and may not parse into `LLMGatewayCompletionChunk`. + + Raises: `LLMGatewayError` if the request fails before any bytes stream. + """ + with client.stream( + "POST", + base_url.rstrip("/") + ENDPOINT_CHAT_COMPLETIONS, + json=body, + timeout=timeout, + ) as response: + if response.status_code != httpx.codes.OK: + response.read() + raise _error_from_response(response) + + for line in response.iter_lines(): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :].strip() + if payload == "[DONE]": + break + yield models.LLMGatewayCompletionChunk.parse_obj(json.loads(payload)) + + +def create_understanding( + client: httpx.Client, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> models.LLMGatewayUnderstandingResponse: + """ + Runs speech understanding features against an existing transcript. + + Raises: `LLMGatewayError` on any non-200 response. + """ + response = client.post( + base_url.rstrip("/") + ENDPOINT_UNDERSTANDING, + json=body, + timeout=timeout, + ) + + from pprint import pprint + pprint (response.json()) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayUnderstandingResponse.parse_obj(response.json()) + + +def validate_understanding( + client: httpx.Client, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> None: + """ + Validates a speech understanding request without running it. + + Raises: `LLMGatewayError` (with `.errors` populated) if the request is invalid. + """ + response = client.post( + base_url.rstrip("/") + ENDPOINT_UNDERSTANDING_VALIDATE, + json=body, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) diff --git a/assemblyai/llm_gateway/v1/async_api.py b/assemblyai/llm_gateway/v1/async_api.py new file mode 100644 index 00000000..3088ccb2 --- /dev/null +++ b/assemblyai/llm_gateway/v1/async_api.py @@ -0,0 +1,154 @@ +"""The asyncio counterpart of `api.py`. + +Calls the same endpoints as their sync twins and raises the same +`LLMGatewayError` through `api._error_from_response`. +""" + +import json +from typing import Any, AsyncIterator, Dict + +import httpx + +from . import models +from .api import ( + ENDPOINT_CHAT_COMPLETIONS, + ENDPOINT_MODELS, + ENDPOINT_UNDERSTANDING, + ENDPOINT_UNDERSTANDING_VALIDATE, + _error_from_response, +) + +__all__ = [ + "list_models", + "create_completion", + "stream_completion", + "create_understanding", + "validate_understanding", +] + + +async def list_models( + client: httpx.AsyncClient, + *, + base_url: str, + timeout: float, +) -> models.LLMGatewayModelList: + """ + Lists the models available through the LLM Gateway. + + Raises: `LLMGatewayError` on any non-200 response. + """ + response = await client.get( + base_url.rstrip("/") + ENDPOINT_MODELS, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayModelList.parse_obj(response.json()) + + +async def create_completion( + client: httpx.AsyncClient, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> models.LLMGatewayChatCompletion: + """ + Posts a non-streaming chat completion request. + + Raises: `LLMGatewayError` on any non-200 response. + """ + response = await client.post( + base_url.rstrip("/") + ENDPOINT_CHAT_COMPLETIONS, + json=body, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayChatCompletion.parse_obj(response.json()) + + +async def stream_completion( + client: httpx.AsyncClient, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> AsyncIterator[models.LLMGatewayCompletionChunk]: + """ + Posts a streaming chat completion request and yields parsed chunks. + + Chunk parsing is verified against OpenAI-routed models only (`data: {...}` + SSE lines terminated by `data: [DONE]`); other providers' streaming + framing was not verified and may not parse into `LLMGatewayCompletionChunk`. + + Raises: `LLMGatewayError` if the request fails before any bytes stream. + """ + async with client.stream( + "POST", + base_url.rstrip("/") + ENDPOINT_CHAT_COMPLETIONS, + json=body, + timeout=timeout, + ) as response: + if response.status_code != httpx.codes.OK: + await response.aread() + raise _error_from_response(response) + + async for line in response.aiter_lines(): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :].strip() + if payload == "[DONE]": + break + yield models.LLMGatewayCompletionChunk.parse_obj(json.loads(payload)) + + +async def create_understanding( + client: httpx.AsyncClient, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> models.LLMGatewayUnderstandingResponse: + """ + Runs speech understanding features against an existing transcript. + + Raises: `LLMGatewayError` on any non-200 response. + """ + response = await client.post( + base_url.rstrip("/") + ENDPOINT_UNDERSTANDING, + json=body, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return models.LLMGatewayUnderstandingResponse.parse_obj(response.json()) + + +async def validate_understanding( + client: httpx.AsyncClient, + *, + base_url: str, + timeout: float, + body: Dict[str, Any], +) -> None: + """ + Validates a speech understanding request without running it. + + Raises: `LLMGatewayError` (with `.errors` populated) if the request is invalid. + """ + response = await client.post( + base_url.rstrip("/") + ENDPOINT_UNDERSTANDING_VALIDATE, + json=body, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) diff --git a/assemblyai/llm_gateway/v1/async_client.py b/assemblyai/llm_gateway/v1/async_client.py new file mode 100644 index 00000000..aca155ca --- /dev/null +++ b/assemblyai/llm_gateway/v1/async_client.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import json +from types import TracebackType +from typing import Any, Callable, Dict, List, Optional, Type, Union + +from typing_extensions import Self + +from ... import async_client as _async_client +from . import async_api, models +from .params import LLMGatewayMessageParamUnion +from .stream import AsyncLLMGatewayStream + + +class AsyncModelsResource: + """The `models` resource of `AsyncLLMGateway`.""" + + def __init__(self, gateway: "AsyncLLMGateway") -> None: + self._gateway = gateway + + async def list(self) -> models.LLMGatewayModelList: + """ + Lists the models available through the LLM Gateway. + + Raises: `LLMGatewayError` if the request fails. + """ + settings = self._gateway.client.settings + + return await async_api.list_models( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + ) + + +class AsyncCompletionsResource: + """The `chat.completions` resource of `AsyncLLMGateway`.""" + + def __init__(self, gateway: "AsyncLLMGateway") -> None: + self._gateway = gateway + + async def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParamUnion], + stream: bool = False, + **kwargs: Any, + ) -> Union[models.LLMGatewayChatCompletion, AsyncLLMGatewayStream]: + """ + Creates a chat completion. + + `messages` are plain dicts (see `LLMGatewayMessageParam`), e.g. + `[{"role": "user", "content": "..."}]`. Extra keyword arguments — + `max_tokens`, `temperature`, `tools`, `fallbacks`, `fallback_config`, + `zero_data_retention`, `transcript_id`, `reasoning`, and any other + field the API accepts — are forwarded directly into the request body. + `fallbacks` is a list of dicts, each with the API's expected shape, + e.g. `[{"model": "gpt-4o", "messages": [...]}]`. + + With `stream=True`, resolves to an `AsyncLLMGatewayStream` — iterate it + for `LLMGatewayCompletionChunk`s exactly like the raw stream it + replaces (`async for chunk in await gateway.chat.completions.create(..., stream=True)`), + then call `.get_final_message()` for the assembled text/usage — + instead of a single `LLMGatewayChatCompletion`. Chunk parsing is + verified for OpenAI-routed models only. + + Raises: `LLMGatewayError` if the request fails (for a stream, only if + it fails before any bytes are received — a failure mid-stream + just truncates the iterator). + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = { + "model": model, + "messages": messages, + "stream": stream, + **kwargs, + } + + if stream: + return AsyncLLMGatewayStream( + async_api.stream_completion( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + ) + + return await async_api.create_completion( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + async def run_tools( + self, + *, + model: str, + messages: List[LLMGatewayMessageParamUnion], + tools: List[Dict[str, Any]], + functions: Dict[str, Callable[..., Any]], + max_rounds: int = 10, + **kwargs: Any, + ) -> models.LLMGatewayChatCompletion | AsyncLLMGatewayStream: + """ + Runs the tool-calling loop to completion: calls the model, invokes any + requested tools from `functions`, feeds their results back, and repeats + until the model responds without requesting a tool call. + + `messages` is mutated in place with every tool round-trip and the final + assistant reply, so it's a complete transcript once this returns. Tool + round-trips are replayed as `{"role": "assistant", "tool_calls": [...]}` + followed by one `{"role": "tool", "tool_call_id": ..., "content": ...}` + per call, per https://www.assemblyai.com/docs/llm-gateway/tool-calling. + + `functions` maps tool names (as declared in `tools`) to the Python + callables that implement them; each is called with the model's parsed + arguments as keyword arguments, and its return value is sent back to + the model (JSON-encoded, unless it's already a string). + + Raises: + ValueError: if `stream=True` is passed (unsupported), or the model + requests a tool with no matching entry in `functions`. + RuntimeError: if the loop doesn't finish within `max_rounds`. + LLMGatewayError: if any underlying `create()` call fails. + """ + if kwargs.get("stream"): + raise ValueError("run_tools() does not support stream=True") + + for _ in range(max_rounds): + completion = await self.create( + model=model, messages=messages, tools=tools, **kwargs + ) + message = completion.choices[0].message + + if not message.tool_calls: + messages.append({"role": "assistant", "content": message.content}) + return completion + + messages.append( + { + "role": "assistant", + "content": message.content, + "tool_calls": [ + { + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + if isinstance(tool_call.function.arguments, str) + else json.dumps(tool_call.function.arguments), + }, + } + for tool_call in message.tool_calls + ], + } + ) + + for tool_call in message.tool_calls: + function = functions.get(tool_call.function.name) + if function is None: + raise ValueError( + f"No function registered for tool call {tool_call.function.name!r}" + ) + + arguments = tool_call.function.arguments + if isinstance(arguments, str): + arguments = json.loads(arguments) + + result = function(**arguments) + output = result if isinstance(result, str) else json.dumps(result) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": output, + "name": tool_call.function.name, + } + ) + + raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") + + +class AsyncChatResource: + """The `chat` resource of `AsyncLLMGateway`.""" + + def __init__(self, gateway: "AsyncLLMGateway") -> None: + self.completions = AsyncCompletionsResource(gateway) + + +class AsyncUnderstandingResource: + """The `understanding` resource of `AsyncLLMGateway`.""" + + def __init__(self, gateway: "AsyncLLMGateway") -> None: + self._gateway = gateway + + async def create( + self, + *, + transcript_id: str, + request: Dict[str, Any], + **kwargs: Any, + ) -> models.LLMGatewayUnderstandingResponse: + """ + Runs speech understanding features against an existing transcript. + + `request` is the `speech_understanding.request` feature config, e.g. + `{"speaker_identification": {"speaker_type": "name"}}` — a plain + dict, since the server itself has no fixed schema for it. + + Raises: `LLMGatewayError` if the request fails. + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = { + "transcript_id": transcript_id, + "speech_understanding": {"request": request}, + **kwargs, + } + + return await async_api.create_understanding( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + async def validate( + self, + *, + request: Dict[str, Any], + transcript_id: Optional[str] = None, + **kwargs: Any, + ) -> None: + """ + Validates a speech understanding request without running it. + + Returns `None` if the request is valid. Note this still requires + the same authorization/balance as a real `create()` call — it is + not a free dry run. + + Raises: `LLMGatewayError` (with `.errors` populated with the + validation messages) if the request is invalid. + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = {"speech_understanding": {"request": request}, **kwargs} + if transcript_id is not None: + body["transcript_id"] = transcript_id + + await async_api.validate_understanding( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + +class AsyncLLMGateway: + """ + The asyncio counterpart of `LLMGateway`. + + The gateway owns an HTTP connection pool. Close it with `aclose()`, or + use the gateway as an async context manager. + + Example: + ```python + import asyncio + import assemblyai as aai + + aai.settings.api_key = "your-key" + + async def main(): + async with aai.AsyncLLMGateway() as gateway: + models = await gateway.models.list() + for model in models.data: + print(model.id) + + completion = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], + ) + print(completion.choices[0].message.content) + + asyncio.run(main()) + ``` + """ + + def __init__( + self, + *, + client: Optional[_async_client.AsyncClient] = None, + api_key: Optional[str] = None, + ) -> None: + """ + Creates an `AsyncLLMGateway`. + + Args: + client: The `AsyncClient` to use. If `None`, the gateway creates + one from the global `settings` and closes it on `aclose()`. + Pass a client to share one pool between gateways. + api_key: The API key to authenticate with. The gateway builds its + own `AsyncClient` from it and closes that client on + `aclose()`. Given alongside `client`, it takes precedence: the + gateway builds and owns a client made from a copy of that + client's settings with the key replaced, and the given client + is left untouched and stays the caller's to close. + """ + self._owns_client = client is None or api_key is not None + self._client = _async_client._resolve_client(client, api_key) + self.models = AsyncModelsResource(self) + self.chat = AsyncChatResource(self) + self.understanding = AsyncUnderstandingResource(self) + + @property + def client(self) -> _async_client.AsyncClient: + """The `AsyncClient` this gateway sends requests with.""" + + return self._client + + async def aclose(self) -> None: + """ + Closes the HTTP connection pool. + + Leaves a client that was passed in alone. Its creator closes it. + """ + if self._owns_client: + await self._client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await self.aclose() diff --git a/assemblyai/llm_gateway/v1/client.py b/assemblyai/llm_gateway/v1/client.py new file mode 100644 index 00000000..4817e0cc --- /dev/null +++ b/assemblyai/llm_gateway/v1/client.py @@ -0,0 +1,303 @@ +import json +from typing import Any, Callable, Dict, List, Optional, Union + +from ... import client as _client +from . import api, models +from .params import LLMGatewayMessageParamUnion +from .stream import LLMGatewayStream + + +class ModelsResource: + """The `models` resource of `LLMGateway`.""" + + def __init__(self, gateway: "LLMGateway") -> None: + self._gateway = gateway + + def list(self) -> models.LLMGatewayModelList: + """ + Lists the models available through the LLM Gateway. + + Raises: `LLMGatewayError` if the request fails. + """ + settings = self._gateway.client.settings + + return api.list_models( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + ) + + +class CompletionsResource: + """The `chat.completions` resource of `LLMGateway`.""" + + def __init__(self, gateway: "LLMGateway") -> None: + self._gateway = gateway + + def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParamUnion], + stream: bool = False, + **kwargs: Any, + ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: + """ + Creates a chat completion. + + `messages` are plain dicts (see `LLMGatewayMessageParam`), e.g. + `[{"role": "user", "content": "..."}]`. Extra keyword arguments — + `max_tokens`, `temperature`, `tools`, `fallbacks`, `fallback_config`, + `zero_data_retention`, `transcript_id`, `reasoning`, and any other + field the API accepts — are forwarded directly into the request body. + `fallbacks` is a list of dicts, each with the API's expected shape, + e.g. `[{"model": "gpt-4o", "messages": [...]}]`. + + With `stream=True`, returns a `LLMGatewayStream` — iterate it for + `LLMGatewayCompletionChunk`s exactly like the raw stream it replaces, + then call `.get_final_message()` for the assembled text/usage — instead + of a single `LLMGatewayChatCompletion`. Chunk parsing is verified for + OpenAI-routed models only. + + Raises: `LLMGatewayError` if the request fails (for a stream, only if + it fails before any bytes are received — a failure mid-stream + just truncates the iterator). + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = { + "model": model, + "messages": messages, + "stream": stream, + **kwargs, + } + + if stream: + return LLMGatewayStream( + api.stream_completion( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + ) + + return api.create_completion( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + def run_tools( + self, + *, + model: str, + messages: List[LLMGatewayMessageParamUnion], + tools: List[Dict[str, Any]], + functions: Dict[str, Callable[..., Any]], + max_rounds: int = 10, + **kwargs: Any, + ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: + """ + Runs the tool-calling loop to completion: calls the model, invokes any + requested tools from `functions`, feeds their results back, and repeats + until the model responds without requesting a tool call. + + `messages` is mutated in place with every tool round-trip and the final + assistant reply, so it's a complete transcript once this returns. Tool + round-trips are replayed as `{"role": "assistant", "tool_calls": [...]}` + followed by one `{"role": "tool", "tool_call_id": ..., "content": ...}` + per call, per https://www.assemblyai.com/docs/llm-gateway/tool-calling. + + `functions` maps tool names (as declared in `tools`) to the Python + callables that implement them; each is called with the model's parsed + arguments as keyword arguments, and its return value is sent back to + the model (JSON-encoded, unless it's already a string). + + Raises: + ValueError: if `stream=True` is passed (unsupported), or the model + requests a tool with no matching entry in `functions`. + RuntimeError: if the loop doesn't finish within `max_rounds`. + LLMGatewayError: if any underlying `create()` call fails. + """ + if kwargs.get("stream"): + raise ValueError("run_tools() does not support stream=True") + + for _ in range(max_rounds): + completion = self.create( + model=model, messages=messages, tools=tools, **kwargs + ) + message = completion.choices[0].message + + if not message.tool_calls: + messages.append({"role": "assistant", "content": message.content}) + return completion + + messages.append( + { + "role": "assistant", + "content": message.content, + "tool_calls": [ + { + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + if isinstance(tool_call.function.arguments, str) + else json.dumps(tool_call.function.arguments), + }, + } + for tool_call in message.tool_calls + ], + } + ) + + for tool_call in message.tool_calls: + function = functions.get(tool_call.function.name) + if function is None: + raise ValueError( + f"No function registered for tool call {tool_call.function.name!r}" + ) + + arguments = tool_call.function.arguments + if isinstance(arguments, str): + arguments = json.loads(arguments) + + result = function(**arguments) + output = result if isinstance(result, str) else json.dumps(result) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": output, + "name": tool_call.function.name, + } + ) + + raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") + + +class ChatResource: + """The `chat` resource of `LLMGateway`.""" + + def __init__(self, gateway: "LLMGateway") -> None: + self.completions = CompletionsResource(gateway) + + +class UnderstandingResource: + """The `understanding` resource of `LLMGateway`.""" + + def __init__(self, gateway: "LLMGateway") -> None: + self._gateway = gateway + + def create( + self, + *, + transcript_id: str, + request: Dict[str, Any], + **kwargs: Any, + ) -> models.LLMGatewayUnderstandingResponse: + """ + Runs speech understanding features against an existing transcript. + + `request` is the `speech_understanding.request` feature config, e.g. + `{"speaker_identification": {"speaker_type": "name"}}` — a plain + dict, since the server itself has no fixed schema for it. + + Raises: `LLMGatewayError` if the request fails. + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = { + "transcript_id": transcript_id, + "speech_understanding": {"request": request}, + **kwargs, + } + + return api.create_understanding( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + def validate( + self, + *, + request: Dict[str, Any], + transcript_id: Optional[str] = None, + **kwargs: Any, + ) -> None: + """ + Validates a speech understanding request without running it. + + Returns `None` if the request is valid. Note this still requires + the same authorization/balance as a real `create()` call — it is + not a free dry run. + + Raises: `LLMGatewayError` (with `.errors` populated with the + validation messages) if the request is invalid. + """ + settings = self._gateway.client.settings + body: Dict[str, Any] = {"speech_understanding": {"request": request}, **kwargs} + if transcript_id is not None: + body["transcript_id"] = transcript_id + + api.validate_understanding( + self._gateway.client.http_client, + base_url=settings.llm_gateway_base_url, + timeout=settings.llm_gateway_http_timeout, + body=body, + ) + + +class LLMGateway: + """ + Client for AssemblyAI's LLM Gateway. + + Example: + ```python + import assemblyai as aai + + aai.settings.api_key = "your-key" + + gateway = aai.LLMGateway() + for model in gateway.models.list().data: + print(model.id) + + completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], + ) + print(completion.choices[0].message.content) + ``` + """ + + def __init__( + self, + *, + client: Optional[_client.Client] = None, + api_key: Optional[str] = None, + ) -> None: + """ + Creates an `LLMGateway`. + + Args: + client: The HTTP client to use. Defaults to the shared default client. + api_key: The API key to authenticate with. Builds a `Client` for this + gateway. Given alongside `client`, it takes precedence: the + gateway builds its own client from a copy of that client's + settings with the key replaced, and the given client is left + untouched. + """ + self._client = _client._resolve_client(client, api_key) + self.models = ModelsResource(self) + self.chat = ChatResource(self) + self.understanding = UnderstandingResource(self) + + @property + def client(self) -> _client.Client: + """The `Client` this gateway sends requests with.""" + + return self._client diff --git a/assemblyai/llm_gateway/v1/models.py b/assemblyai/llm_gateway/v1/models.py new file mode 100644 index 00000000..5a5c640b --- /dev/null +++ b/assemblyai/llm_gateway/v1/models.py @@ -0,0 +1,179 @@ +from typing import Any, Dict, List, Optional + +from ...types import BaseModel, ConfigDict, Field, pydantic_v2 + + +class LLMGatewayPricingData(BaseModel): + completions: float + prompt: float + input_cache_read: Optional[float] = None + input_cache_write: Optional[float] = None + input_cache_write_1h: Optional[float] = None + + +class LLMGatewayPricing(BaseModel): + us: Optional[LLMGatewayPricingData] = None + eu: Optional[LLMGatewayPricingData] = None + global_: LLMGatewayPricingData = Field(alias="global") + regional_increase_percent: Optional[float] = None + + if pydantic_v2: + model_config = ConfigDict(populate_by_name=True) + else: + + class Config: + allow_population_by_field_name = True + + +class LLMGatewayDefaultParameters(BaseModel): + temperature: Optional[float] = None + top_p: Optional[float] = None + frequency_penalty: Optional[int] = None + + +class LLMGatewayTopProvider(BaseModel): + is_moderated: bool + context_length: int + max_completion_tokens: int + + +class LLMGatewayModel(BaseModel): + """A model available through the LLM Gateway, as returned by `GET /v1/models`.""" + + id: str + name: str + description: str + default_parameters: LLMGatewayDefaultParameters + supported_parameters: List[str] = Field(default_factory=list) + top_provider: LLMGatewayTopProvider + context_length: int + pricing: LLMGatewayPricing + creator: str + retirement_date: int + "Unix timestamp of the model's retirement date, or 0 if not retiring" + + available_regions: List[str] = Field(default_factory=list) + providers: List[str] = Field(default_factory=list) + default_provider: str + + +class LLMGatewayModelList(BaseModel): + """The response of `GET /v1/models`.""" + + data: List[LLMGatewayModel] = Field(default_factory=list) + + +class LLMGatewayCacheCreation(BaseModel): + ephemeral_5m_input_tokens: int + ephemeral_1h_input_tokens: int + + +class LLMGatewayPromptTokensDetails(BaseModel): + cached_tokens: int + audio_tokens: int + cache_creation: Optional[LLMGatewayCacheCreation] = None + cache_write_tokens: Optional[int] = None + + +class LLMGatewayCompletionTokensDetails(BaseModel): + reasoning_tokens: int + audio_tokens: int + accepted_prediction_tokens: int + rejected_prediction_tokens: int + + +class LLMGatewayUsage(BaseModel): + """ + `input_tokens`/`output_tokens` are only populated for non-streaming + responses — the gateway forwards the upstream provider's own streaming + usage chunk unmodified, which carries `prompt_tokens`/`completion_tokens` + only. + """ + + input_tokens: Optional[int] = None + prompt_tokens: int + output_tokens: Optional[int] = None + completion_tokens: int + total_tokens: int + prompt_tokens_details: LLMGatewayPromptTokensDetails + completion_tokens_details: LLMGatewayCompletionTokensDetails + + +class LLMGatewayFunction(BaseModel): + name: str + arguments: Any = None + + +class LLMGatewayToolCall(BaseModel): + id: str + type: str + function: LLMGatewayFunction + + +class LLMGatewayResponseMessage(BaseModel): + role: str + content: Optional[str] = None + thinking: Optional[str] = None + tool_calls: Optional[List[LLMGatewayToolCall]] = None + tool_call_id: Optional[str] = None + name: Optional[str] = None + + +class LLMGatewayChoice(BaseModel): + index: int + finish_reason: Optional[str] = None + message: LLMGatewayResponseMessage + + +class LLMGatewayChatCompletion(BaseModel): + """The response of `POST /v1/chat/completions` (non-streaming).""" + + request_id: str + choices: List[LLMGatewayChoice] = Field(default_factory=list) + usage: LLMGatewayUsage + request: Any = None + "Loosely-typed echo of the request, as sent back by the server" + + http_status_code: Optional[int] = None + response_time: Optional[int] = None + "Nanoseconds, as sent by the server — not converted to seconds" + + llm_status_code: Optional[int] = None + + +class LLMGatewayChunkDelta(BaseModel): + content: Optional[str] = None + + +class LLMGatewayChunkChoice(BaseModel): + index: int + delta: LLMGatewayChunkDelta + finish_reason: Optional[str] = None + + +class LLMGatewayCompletionChunk(BaseModel): + """ + A streamed chat completion chunk. + + Verified against OpenAI-routed models only — Claude/Gemini/Bedrock + streaming framing was not verified and may not parse into this shape. + """ + + id: str + object: str + created: int + model: str + service_tier: Optional[str] = None + system_fingerprint: Any = None + choices: List[LLMGatewayChunkChoice] = Field(default_factory=list) + usage: Optional[LLMGatewayUsage] = None + obfuscation: Optional[str] = None + + +class LLMGatewayUnderstandingResponse(BaseModel): + """The response of `POST /v1/understanding`.""" + + speech_understanding: Dict[str, Any] = Field(default_factory=dict) + request_id: str + utterances: Optional[List[Dict[str, Any]]] = None + translated_texts: Optional[Dict[str, str]] = None diff --git a/assemblyai/llm_gateway/v1/params.py b/assemblyai/llm_gateway/v1/params.py new file mode 100644 index 00000000..0929c845 --- /dev/null +++ b/assemblyai/llm_gateway/v1/params.py @@ -0,0 +1,20 @@ +from typing import Any, Dict, List, Optional, Union + +from typing_extensions import Required, TypedDict + + +class LLMGatewayMessageParam(TypedDict, total=False): + """A chat message. Pass as a plain dict: `{"role": "user", "content": "..."}`.""" + + role: Required[str] + '"user", "assistant", "system", or "tool" (a tool call result, paired with `tool_call_id`)' + + content: Optional[Union[str, List[Dict[str, Any]]]] + tool_calls: List[Dict[str, Any]] + tool_call_id: str + name: str + thinking: str + cache_control: Dict[str, Any] + + +LLMGatewayMessageParamUnion = LLMGatewayMessageParam diff --git a/assemblyai/llm_gateway/v1/stream.py b/assemblyai/llm_gateway/v1/stream.py new file mode 100644 index 00000000..d9aa819e --- /dev/null +++ b/assemblyai/llm_gateway/v1/stream.py @@ -0,0 +1,90 @@ +from typing import AsyncIterator, Iterator, Optional + +from typing_extensions import Self + +from . import models + + +class LLMGatewayStreamAccumulator: + """The message assembled from a `LLMGatewayStream`'s chunks, as seen so far.""" + + def __init__(self) -> None: + self.id: Optional[str] = None + self.model: Optional[str] = None + self.content: str = "" + self.finish_reason: Optional[str] = None + self.usage: Optional[models.LLMGatewayUsage] = None + + def _absorb(self, chunk: models.LLMGatewayCompletionChunk) -> None: + if self.id is None: + self.id = chunk.id + if self.model is None: + self.model = chunk.model + for choice in chunk.choices: + if choice.delta.content: + self.content += choice.delta.content + if choice.finish_reason is not None: + self.finish_reason = choice.finish_reason + if chunk.usage is not None: + self.usage = chunk.usage + + +class LLMGatewayStream: + """ + Wraps a streamed chat completion's chunk iterator, accumulating a final + assembled message as it's iterated. + + Iterate it exactly like the raw chunk iterator it replaces: + `for chunk in stream: ...`. Once exhausted, call `get_final_message()` for + the assembled text/usage/finish_reason. + + Verified for OpenAI-routed models only, like raw chunk parsing. Tool calls + are not reconstructed here — streamed chunks don't carry `tool_calls` + today. + """ + + def __init__(self, chunks: Iterator[models.LLMGatewayCompletionChunk]) -> None: + self._chunks = chunks + self._accumulator = LLMGatewayStreamAccumulator() + + def __iter__(self) -> Self: + return self + + def __next__(self) -> models.LLMGatewayCompletionChunk: + chunk = next(self._chunks) + self._accumulator._absorb(chunk) + return chunk + + def get_final_message(self) -> LLMGatewayStreamAccumulator: + """ + Returns the message assembled from chunks seen so far. + + Call after exhausting the iterator for a complete result; calling it + earlier just returns whatever has streamed in up to that point. + """ + return self._accumulator + + +class AsyncLLMGatewayStream: + """The asyncio counterpart of `LLMGatewayStream`.""" + + def __init__(self, chunks: AsyncIterator[models.LLMGatewayCompletionChunk]) -> None: + self._chunks = chunks + self._accumulator = LLMGatewayStreamAccumulator() + + def __aiter__(self) -> Self: + return self + + async def __anext__(self) -> models.LLMGatewayCompletionChunk: + chunk = await self._chunks.__anext__() + self._accumulator._absorb(chunk) + return chunk + + def get_final_message(self) -> LLMGatewayStreamAccumulator: + """ + Returns the message assembled from chunks seen so far. + + Call after exhausting the iterator for a complete result; calling it + earlier just returns whatever has streamed in up to that point. + """ + return self._accumulator diff --git a/assemblyai/types.py b/assemblyai/types.py index b0af7309..2c7842ed 100644 --- a/assemblyai/types.py +++ b/assemblyai/types.py @@ -73,6 +73,30 @@ class RedactedAudioUnavailableError(AssemblyAIError): """ +class LLMGatewayError(AssemblyAIError): + """ + Error raised when an LLM Gateway request fails. + + The server uses two different error envelopes: business errors + (`{"message", "code", "request_id", "metadata": {"errors": [...]}}`) and + 401s raised by auth middleware before a request reaches a route handler + (`{"error", "status", "request_id"}`). Both resolve to this one class; + `request_id` and `errors` (validation messages, when present) are + populated from whichever shape the server sent. + """ + + def __init__( + self, + message: str, + status_code: Optional[int] = None, + request_id: Optional[str] = None, + errors: Optional[List[str]] = None, + ): + super().__init__(message, status_code) + self.request_id = request_id + self.errors = errors + + class SyncTranscriptError(AssemblyAIError): """ Error raised when a synchronous transcription request fails. @@ -119,6 +143,12 @@ class Settings(BaseSettings): sync_http_timeout: float = 60.0 "The HTTP timeout for synchronous transcription requests. Kept above the server's 30s deadline so the client doesn't race it." + llm_gateway_base_url: str = "https://llm-gateway.assemblyai.com" + "The base URL for the LLM Gateway API (used by `LLMGateway`/`AsyncLLMGateway`)" + + llm_gateway_http_timeout: float = 30.0 + "The HTTP timeout for LLM Gateway requests" + polling_interval: float = Field(default=3.0, gt=0.0) "The default polling interval for long-running requests (e.g. polling the `Transcript`'s status)" diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py new file mode 100644 index 00000000..8d4081ec --- /dev/null +++ b/tests/unit/test_llm_gateway.py @@ -0,0 +1,759 @@ +import httpx +import pytest +from pytest_httpx import HTTPXMock, IteratorStream + +import assemblyai as aai + +aai.settings.api_key = "test" + +MODELS_URL = f"{aai.settings.llm_gateway_base_url}/v1/models" +COMPLETIONS_URL = f"{aai.settings.llm_gateway_base_url}/v1/chat/completions" +UNDERSTANDING_URL = f"{aai.settings.llm_gateway_base_url}/v1/understanding" +UNDERSTANDING_VALIDATE_URL = ( + f"{aai.settings.llm_gateway_base_url}/v1/understanding/validate" +) + +_OK_RESPONSE = { + "data": [ + { + "id": "claude-sonnet-5", + "name": "Claude Sonnet 5", + "description": "Anthropic's Claude Sonnet 5", + "default_parameters": { + "temperature": 1.0, + "top_p": 0.9, + "frequency_penalty": 0, + }, + "supported_parameters": ["temperature", "top_p", "tools"], + "top_provider": { + "is_moderated": False, + "context_length": 200000, + "max_completion_tokens": 8192, + }, + "context_length": 200000, + "pricing": { + "us": {"completions": 15.0, "prompt": 3.0}, + "eu": {"completions": 16.5, "prompt": 3.3}, + "global": {"completions": 15.0, "prompt": 3.0}, + "regional_increase_percent": 10.0, + }, + "creator": "anthropic", + "retirement_date": 0, + "available_regions": ["us", "eu"], + "providers": ["anthropic"], + "default_provider": "anthropic", + }, + { + "id": "gpt-4o", + "name": "GPT-4o", + "description": "OpenAI's GPT-4o", + "default_parameters": { + "temperature": None, + "top_p": None, + "frequency_penalty": None, + }, + "supported_parameters": ["temperature"], + "top_provider": { + "is_moderated": True, + "context_length": 128000, + "max_completion_tokens": 4096, + }, + "context_length": 128000, + "pricing": { + "global": {"completions": 10.0, "prompt": 2.5}, + }, + "creator": "openai", + "retirement_date": 0, + "available_regions": ["global"], + "providers": ["openai"], + "default_provider": "openai", + }, + ] +} + + +def _mock_ok(httpx_mock: HTTPXMock) -> None: + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.OK, + json=_OK_RESPONSE, + ) + + +def test_list_models_parses_response(httpx_mock: HTTPXMock): + # Given a mocked models endpoint + _mock_ok(httpx_mock) + + # When listing models + result = aai.LLMGateway().models.list() + + # Then the response is parsed into an LLMGatewayModelList + assert isinstance(result, aai.LLMGatewayModelList) + assert len(result.data) == 2 + first = result.data[0] + assert first.id == "claude-sonnet-5" + assert first.default_parameters.temperature == 1.0 + assert first.pricing.us.completions == 15.0 + assert first.pricing.eu.prompt == 3.3 + assert first.pricing.global_.completions == 15.0 + assert first.pricing.regional_increase_percent == 10.0 + + +def test_list_models_handles_minimal_payload(httpx_mock: HTTPXMock): + # Given a mocked endpoint with all optional fields omitted + _mock_ok(httpx_mock) + + # When listing models + result = aai.LLMGateway().models.list() + + # Then omitted optional fields default to None + second = result.data[1] + assert second.pricing.us is None + assert second.pricing.eu is None + assert second.pricing.regional_increase_percent is None + assert second.default_parameters.temperature is None + + +def test_models_requires_explicit_list_call(httpx_mock: HTTPXMock): + # Given a mocked models endpoint + _mock_ok(httpx_mock) + + # When calling .list() explicitly (gateway.models itself is not iterable) + result = aai.LLMGateway().models.list() + + # Then it returns the response envelope, matching openai-python's client.models.list() + assert isinstance(result, aai.LLMGatewayModelList) + ids = [m.id for m in result.data] + assert ids == ["claude-sonnet-5", "gpt-4o"] + assert not hasattr(aai.LLMGateway().models, "__iter__") + + +def test_gateway_uses_llm_gateway_base_url(httpx_mock: HTTPXMock): + # Given a mocked models endpoint + _mock_ok(httpx_mock) + + # When listing models + aai.LLMGateway().models.list() + + # Then the request targets llm_gateway_base_url, not the generic base_url + request = httpx_mock.get_requests()[0] + assert str(request.url) == MODELS_URL + assert aai.settings.llm_gateway_base_url != aai.settings.base_url + + +def test_list_models_raises_llm_gateway_error_on_500(httpx_mock: HTTPXMock): + # Given a mocked 500 response with the real business-error envelope + # ({"message", "code", "request_id", "metadata": {"errors": [...]}}) + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.INTERNAL_SERVER_ERROR, + json={ + "message": "something went wrong", + "code": 500, + "request_id": "req-abc-123", + }, + ) + + # When listing models + with pytest.raises(aai.LLMGatewayError) as exc_info: + aai.LLMGateway().models.list() + + # Then the error carries the server's status code, message, and request id + error = exc_info.value + assert error.status_code == httpx.codes.INTERNAL_SERVER_ERROR + assert error.request_id == "req-abc-123" + assert error.errors is None + assert "something went wrong" in str(error) + + +def test_error_carries_metadata_errors_list(httpx_mock: HTTPXMock): + # Given a mocked 400 response with metadata.errors populated + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.BAD_REQUEST, + json={ + "message": "invalid request body", + "code": 400, + "request_id": "req-abc-124", + "metadata": {"errors": ["target_languages is empty"]}, + }, + ) + + # When listing models + with pytest.raises(aai.LLMGatewayError) as exc_info: + aai.LLMGateway().models.list() + + # Then .errors carries the validation messages + assert exc_info.value.errors == ["target_languages is empty"] + + +def test_error_handles_legacy_auth_error_envelope(httpx_mock: HTTPXMock): + # Given a mocked 401 response using the legacy auth-middleware envelope + # ({"error", "status", "request_id"} — raised before any route handler runs) + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.UNAUTHORIZED, + json={ + "error": "Authentication error, API token missing/invalid", + "status": "error", + "request_id": "req-abc-125", + }, + ) + + # When listing models + with pytest.raises(aai.LLMGatewayError) as exc_info: + aai.LLMGateway().models.list() + + # Then the legacy shape still resolves to a normal LLMGatewayError + error = exc_info.value + assert error.status_code == httpx.codes.UNAUTHORIZED + assert error.request_id == "req-abc-125" + assert "Authentication error" in str(error) + + +def test_list_models_raises_on_malformed_error_body(httpx_mock: HTTPXMock): + # Given a mocked 500 response with a non-JSON body + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.INTERNAL_SERVER_ERROR, + content=b"upstream timeout", + ) + + # When listing models + with pytest.raises(aai.LLMGatewayError) as exc_info: + aai.LLMGateway().models.list() + + # Then a fallback message is still raised + error = exc_info.value + assert error.status_code == httpx.codes.INTERNAL_SERVER_ERROR + assert "upstream timeout" in str(error) + + +_COMPLETION_RESPONSE = { + "request_id": "req-completion-1", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "Hello there."}, + } + ], + "usage": { + "input_tokens": 10, + "prompt_tokens": 10, + "output_tokens": 5, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + "http_status_code": 200, + "response_time": 123456789, + "llm_status_code": 200, +} + + +def test_chat_completions_create_parses_response(httpx_mock: HTTPXMock): + # Given a mocked non-streaming completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion + result = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then the response is parsed into an LLMGatewayChatCompletion + assert isinstance(result, aai.LLMGatewayChatCompletion) + assert result.request_id == "req-completion-1" + assert result.choices[0].message.content == "Hello there." + assert result.usage.total_tokens == 15 + assert result.response_time == 123456789 + + +def test_chat_completions_create_sends_extension_kwargs(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion with AssemblyAI extension kwargs + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + fallbacks=[ + {"model": "gpt-4o", "messages": [{"role": "user", "content": "Hi"}]} + ], + zero_data_retention=True, + transcript_id="transcript-123", + ) + + # Then the extension fields land in the posted JSON body untouched + body = httpx_mock.get_requests()[0].read() + import json as _json + + payload = _json.loads(body) + assert payload["zero_data_retention"] is True + assert payload["transcript_id"] == "transcript-123" + assert payload["fallbacks"][0]["model"] == "gpt-4o" + assert payload["stream"] is False + + +def test_chat_completions_stream_yields_chunks(httpx_mock: HTTPXMock): + # Given a mocked streaming completions endpoint (OpenAI-shaped SSE) + chunk1 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "service_tier": None, + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], + "usage": None, + "obfuscation": "", + } + chunk2 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "service_tier": None, + "system_fingerprint": None, + "choices": [ + {"index": 0, "delta": {"content": " there."}, "finish_reason": "stop"} + ], + "usage": None, + "obfuscation": "", + } + import json as _json + + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [ + f"data: {_json.dumps(chunk1)}\n\n".encode(), + f"data: {_json.dumps(chunk2)}\n\n".encode(), + b"data: [DONE]\n\n", + ] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When creating a streaming completion + chunks = list( + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + ) + + # Then it yields parsed chunks and stops at [DONE] + assert len(chunks) == 2 + assert all(isinstance(c, aai.LLMGatewayCompletionChunk) for c in chunks) + assert chunks[0].choices[0].delta.content == "Hello" + assert chunks[1].choices[0].finish_reason == "stop" + + +def test_chat_completions_stream_raises_on_setup_error(httpx_mock: HTTPXMock): + # Given a mocked 500 response before any streaming bytes are sent + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.INTERNAL_SERVER_ERROR, + json={"message": "something went wrong", "code": 500, "request_id": "req-x"}, + ) + + # When creating a streaming completion and consuming the iterator + with pytest.raises(aai.LLMGatewayError): + list( + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + ) + + +_UNDERSTANDING_RESPONSE = { + "speech_understanding": { + "request": {"speaker_identification": {"speaker_type": "name"}}, + "response": { + "speaker_identification": { + "status": "success", + "mapping": {"A": "John Doe"}, + } + }, + }, + "utterances": [{"speaker": "John Doe", "text": "Hello"}], + "request_id": "req-understanding-1", +} + + +def test_understanding_create_parses_response(httpx_mock: HTTPXMock): + # Given a mocked understanding endpoint + httpx_mock.add_response( + url=UNDERSTANDING_URL, + method="POST", + status_code=httpx.codes.OK, + json=_UNDERSTANDING_RESPONSE, + ) + + # When running speaker identification against an existing transcript + result = aai.LLMGateway().understanding.create( + transcript_id="transcript-123", + request={"speaker_identification": {"speaker_type": "name"}}, + ) + + # Then the response is parsed into an LLMGatewayUnderstandingResponse + assert isinstance(result, aai.LLMGatewayUnderstandingResponse) + assert result.request_id == "req-understanding-1" + assert result.utterances[0]["speaker"] == "John Doe" + assert result.speech_understanding["response"]["speaker_identification"][ + "mapping" + ] == {"A": "John Doe"} + + +def test_understanding_create_sends_transcript_id_and_request(httpx_mock: HTTPXMock): + # Given a mocked understanding endpoint + httpx_mock.add_response( + url=UNDERSTANDING_URL, + method="POST", + status_code=httpx.codes.OK, + json=_UNDERSTANDING_RESPONSE, + ) + + # When calling create() + aai.LLMGateway().understanding.create( + transcript_id="transcript-123", + request={"speaker_identification": {"speaker_type": "name", "effort": "low"}}, + ) + + # Then the posted body carries transcript_id and the nested feature request + import json as _json + + payload = _json.loads(httpx_mock.get_requests()[0].read()) + assert payload["transcript_id"] == "transcript-123" + assert ( + payload["speech_understanding"]["request"]["speaker_identification"]["effort"] + == "low" + ) + + +def test_understanding_validate_returns_none_on_success(httpx_mock: HTTPXMock): + # Given a mocked validate endpoint returning an empty 200 + httpx_mock.add_response( + url=UNDERSTANDING_VALIDATE_URL, + method="POST", + status_code=httpx.codes.OK, + ) + + # When validating a request + result = aai.LLMGateway().understanding.validate( + request={"speaker_identification": {"speaker_type": "role"}} + ) + + # Then it returns None + assert result is None + + +def test_understanding_validate_raises_with_errors_list(httpx_mock: HTTPXMock): + # Given a mocked validate endpoint returning a 400 with validation messages + httpx_mock.add_response( + url=UNDERSTANDING_VALIDATE_URL, + method="POST", + status_code=httpx.codes.BAD_REQUEST, + json={ + "message": "invalid request body", + "code": 400, + "request_id": "req-validate-1", + "metadata": {"errors": ["target_languages is empty"]}, + }, + ) + + # When validating an invalid request + with pytest.raises(aai.LLMGatewayError) as exc_info: + aai.LLMGateway().understanding.validate(request={"translation": {}}) + + # Then .errors carries the validation messages + assert exc_info.value.errors == ["target_languages is empty"] + + +def test_create_fallbacks_sent_as_given(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion with the API-expected fallbacks shape + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + fallbacks=[ + {"model": "gpt-4o", "messages": [{"role": "user", "content": "Hi"}]} + ], + ) + + # Then it is forwarded into the request body untouched + import json as _json + + payload = _json.loads(httpx_mock.get_requests()[0].read()) + assert payload["fallbacks"] == [ + {"model": "gpt-4o", "messages": [{"role": "user", "content": "Hi"}]} + ] + + +def test_create_without_fallbacks_omits_field(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion without fallbacks + aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then the field is omitted entirely, not sent as null + import json as _json + + payload = _json.loads(httpx_mock.get_requests()[0].read()) + assert "fallbacks" not in payload + + +_TOOL_CALL_RESPONSE = { + "request_id": "req-tool-1", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + } + ], + "usage": _COMPLETION_RESPONSE["usage"], +} + + +def test_run_tools_completes_after_tool_call(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint that first requests a tool call, then answers + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + calls = [] + + def get_weather(city): + calls.append(city) + return {"forecast": "sunny"} + + # When running the tool loop + result = aai.LLMGateway().chat.completions.run_tools( + model="claude-sonnet-5", + messages=messages, + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={"get_weather": get_weather}, + ) + + # Then the tool was invoked with parsed arguments and the final completion is returned + assert calls == ["Paris"] + assert isinstance(result, aai.LLMGatewayChatCompletion) + assert result.request_id == "req-completion-1" + + +def test_run_tools_sends_chat_message_shaped_items(httpx_mock: HTTPXMock): + # Given the same two-round mock + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + + # When running the tool loop + aai.LLMGateway().chat.completions.run_tools( + model="claude-sonnet-5", + messages=messages, + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={"get_weather": lambda city: {"forecast": "sunny"}}, + ) + + # Then the follow-up request replays the tool call and its output as plain chat messages + import json as _json + + second_request_body = _json.loads(httpx_mock.get_requests()[1].read()) + replayed = second_request_body["messages"][1:] + assert replayed[0]["role"] == "assistant" + assert replayed[0]["tool_calls"][0]["id"] == "call_1" + assert replayed[0]["tool_calls"][0]["function"]["name"] == "get_weather" + assert replayed[1]["role"] == "tool" + assert replayed[1]["tool_call_id"] == "call_1" + assert _json.loads(replayed[1]["content"]) == {"forecast": "sunny"} + + # And the caller's messages list now holds the complete transcript + assert messages[-1] == {"role": "assistant", "content": "Hello there."} + + +def test_run_tools_raises_on_unknown_tool(httpx_mock: HTTPXMock): + # Given a mocked response requesting an unregistered tool + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + + # When running the tool loop with no matching function + with pytest.raises(ValueError, match="get_weather"): + aai.LLMGateway().chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={}, + ) + + +def test_run_tools_raises_after_max_rounds(httpx_mock: HTTPXMock): + # Given a mock that always requests the same tool call + for _ in range(3): + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + + # When running the tool loop with a max_rounds it can't finish within + with pytest.raises(RuntimeError, match="max_rounds"): + aai.LLMGateway().chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={"get_weather": lambda city: "sunny"}, + max_rounds=3, + ) + + +def test_run_tools_rejects_streaming(): + # When running the tool loop with stream=True + with pytest.raises(ValueError, match="stream"): + aai.LLMGateway().chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[], + functions={}, + stream=True, + ) + + +def test_chat_completions_stream_get_final_message(httpx_mock: HTTPXMock): + # Given a mocked streaming completions endpoint + chunk1 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], + } + chunk2 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "choices": [ + {"index": 0, "delta": {"content": " there."}, "finish_reason": "stop"} + ], + "usage": _COMPLETION_RESPONSE["usage"], + } + import json as _json + + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [ + f"data: {_json.dumps(chunk1)}\n\n".encode(), + f"data: {_json.dumps(chunk2)}\n\n".encode(), + b"data: [DONE]\n\n", + ] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When consuming the stream and asking for the final assembled message + stream = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + chunks = list(stream) + final = stream.get_final_message() + + # Then it still yields the raw chunks (unchanged behavior)... + assert len(chunks) == 2 + assert all(isinstance(c, aai.LLMGatewayCompletionChunk) for c in chunks) + + # ...and the accumulator has the assembled text, usage, and finish_reason + assert final.content == "Hello there." + assert final.finish_reason == "stop" + assert final.usage.total_tokens == 15 + assert final.id == "chatcmpl-1" + assert final.model == "claude-sonnet-5" diff --git a/tests/unit/test_llm_gateway_async.py b/tests/unit/test_llm_gateway_async.py new file mode 100644 index 00000000..939681a1 --- /dev/null +++ b/tests/unit/test_llm_gateway_async.py @@ -0,0 +1,480 @@ +import json + +import httpx +import pytest +from pytest_httpx import HTTPXMock, IteratorStream + +import assemblyai as aai + +pytestmark = pytest.mark.asyncio + +aai.settings.api_key = "test" + +MODELS_URL = f"{aai.settings.llm_gateway_base_url}/v1/models" +COMPLETIONS_URL = f"{aai.settings.llm_gateway_base_url}/v1/chat/completions" +UNDERSTANDING_URL = f"{aai.settings.llm_gateway_base_url}/v1/understanding" +UNDERSTANDING_VALIDATE_URL = ( + f"{aai.settings.llm_gateway_base_url}/v1/understanding/validate" +) + +_OK_RESPONSE = { + "data": [ + { + "id": "claude-sonnet-5", + "name": "Claude Sonnet 5", + "description": "Anthropic's Claude Sonnet 5", + "default_parameters": { + "temperature": 1.0, + "top_p": 0.9, + "frequency_penalty": 0, + }, + "supported_parameters": ["temperature", "top_p", "tools"], + "top_provider": { + "is_moderated": False, + "context_length": 200000, + "max_completion_tokens": 8192, + }, + "context_length": 200000, + "pricing": { + "global": {"completions": 15.0, "prompt": 3.0}, + }, + "creator": "anthropic", + "retirement_date": 0, + "available_regions": ["us", "eu"], + "providers": ["anthropic"], + "default_provider": "anthropic", + }, + ] +} + + +def _mock_ok(httpx_mock: HTTPXMock) -> None: + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.OK, + json=_OK_RESPONSE, + ) + + +async def test_list_models_parses_response(httpx_mock: HTTPXMock): + # Given a mocked models endpoint + _mock_ok(httpx_mock) + + # When listing models + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.models.list() + + # Then the response is parsed into an LLMGatewayModelList + assert isinstance(result, aai.LLMGatewayModelList) + assert result.data[0].id == "claude-sonnet-5" + assert result.data[0].pricing.global_.completions == 15.0 + + +async def test_models_requires_explicit_list_call(httpx_mock: HTTPXMock): + # Given a mocked models endpoint + _mock_ok(httpx_mock) + + # When calling .list() explicitly (gateway.models itself is not iterable) + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.models.list() + + # Then it returns the response envelope, matching openai-python's client.models.list() + assert isinstance(result, aai.LLMGatewayModelList) + assert [m.id for m in result.data] == ["claude-sonnet-5"] + assert not hasattr(gateway.models, "__aiter__") + + +async def test_list_models_raises_llm_gateway_error_on_500(httpx_mock: HTTPXMock): + # Given a mocked 500 response with the real business-error envelope + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.INTERNAL_SERVER_ERROR, + json={ + "message": "something went wrong", + "code": 500, + "request_id": "req-abc-123", + }, + ) + + # When listing models + async with aai.AsyncLLMGateway() as gateway: + with pytest.raises(aai.LLMGatewayError) as exc_info: + await gateway.models.list() + + # Then the error carries the server's status code, message, and request id + error = exc_info.value + assert error.status_code == httpx.codes.INTERNAL_SERVER_ERROR + assert error.request_id == "req-abc-123" + assert "something went wrong" in str(error) + + +async def test_error_handles_legacy_auth_error_envelope(httpx_mock: HTTPXMock): + # Given a mocked 401 response using the legacy auth-middleware envelope + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.UNAUTHORIZED, + json={ + "error": "Authentication error, API token missing/invalid", + "status": "error", + "request_id": "req-abc-125", + }, + ) + + # When listing models + async with aai.AsyncLLMGateway() as gateway: + with pytest.raises(aai.LLMGatewayError) as exc_info: + await gateway.models.list() + + # Then the legacy shape still resolves to a normal LLMGatewayError + error = exc_info.value + assert error.status_code == httpx.codes.UNAUTHORIZED + assert error.request_id == "req-abc-125" + + +_COMPLETION_RESPONSE = { + "request_id": "req-completion-1", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "Hello there."}, + } + ], + "usage": { + "input_tokens": 10, + "prompt_tokens": 10, + "output_tokens": 5, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, +} + + +async def test_chat_completions_create_parses_response(httpx_mock: HTTPXMock): + # Given a mocked non-streaming completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then the response is parsed into an LLMGatewayChatCompletion + assert isinstance(result, aai.LLMGatewayChatCompletion) + assert result.choices[0].message.content == "Hello there." + + +async def test_chat_completions_stream_yields_chunks(httpx_mock: HTTPXMock): + # Given a mocked streaming completions endpoint (OpenAI-shaped SSE) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "service_tier": None, + "system_fingerprint": None, + "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], + "usage": None, + "obfuscation": "", + } + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [f"data: {json.dumps(chunk)}\n\n".encode(), b"data: [DONE]\n\n"] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When creating a streaming completion + async with aai.AsyncLLMGateway() as gateway: + stream = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + chunks = [c async for c in stream] + + # Then it yields parsed chunks + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "Hello" + + +_UNDERSTANDING_RESPONSE = { + "speech_understanding": { + "request": {"speaker_identification": {"speaker_type": "name"}}, + "response": { + "speaker_identification": { + "status": "success", + "mapping": {"A": "John Doe"}, + } + }, + }, + "utterances": [{"speaker": "John Doe", "text": "Hello"}], + "request_id": "req-understanding-1", +} + + +async def test_understanding_create_parses_response(httpx_mock: HTTPXMock): + # Given a mocked understanding endpoint + httpx_mock.add_response( + url=UNDERSTANDING_URL, + method="POST", + status_code=httpx.codes.OK, + json=_UNDERSTANDING_RESPONSE, + ) + + # When running speaker identification against an existing transcript + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.understanding.create( + transcript_id="transcript-123", + request={"speaker_identification": {"speaker_type": "name"}}, + ) + + # Then the response is parsed into an LLMGatewayUnderstandingResponse + assert isinstance(result, aai.LLMGatewayUnderstandingResponse) + assert result.utterances[0]["speaker"] == "John Doe" + + +async def test_understanding_validate_returns_none_on_success(httpx_mock: HTTPXMock): + # Given a mocked validate endpoint returning an empty 200 + httpx_mock.add_response( + url=UNDERSTANDING_VALIDATE_URL, + method="POST", + status_code=httpx.codes.OK, + ) + + # When validating a request + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.understanding.validate( + request={"speaker_identification": {"speaker_type": "role"}} + ) + + # Then it returns None + assert result is None + + +async def test_create_fallbacks_sent_as_given(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + # When creating a completion with the API-expected fallbacks shape + async with aai.AsyncLLMGateway() as gateway: + await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + fallbacks=[ + {"model": "gpt-4o", "messages": [{"role": "user", "content": "Hi"}]} + ], + ) + + # Then it is forwarded into the request body untouched + payload = json.loads(httpx_mock.get_requests()[0].read()) + assert payload["fallbacks"] == [ + {"model": "gpt-4o", "messages": [{"role": "user", "content": "Hi"}]} + ] + + +_TOOL_CALL_RESPONSE = { + "request_id": "req-tool-1", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + } + ], + "usage": _COMPLETION_RESPONSE["usage"], +} + + +async def test_run_tools_completes_after_tool_call(httpx_mock: HTTPXMock): + # Given a mocked completions endpoint that first requests a tool call, then answers + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_COMPLETION_RESPONSE, + ) + + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + calls = [] + + def get_weather(city): + calls.append(city) + return {"forecast": "sunny"} + + # When running the tool loop + async with aai.AsyncLLMGateway() as gateway: + result = await gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=messages, + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={"get_weather": get_weather}, + ) + + # Then the tool was invoked with parsed arguments and the final completion is returned + assert calls == ["Paris"] + assert isinstance(result, aai.LLMGatewayChatCompletion) + assert result.request_id == "req-completion-1" + assert messages[-1] == {"role": "assistant", "content": "Hello there."} + + +async def test_run_tools_raises_on_unknown_tool(httpx_mock: HTTPXMock): + # Given a mocked response requesting an unregistered tool + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + + # When running the tool loop with no matching function + async with aai.AsyncLLMGateway() as gateway: + with pytest.raises(ValueError, match="get_weather"): + await gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={}, + ) + + +async def test_run_tools_raises_after_max_rounds(httpx_mock: HTTPXMock): + # Given a mock that always requests the same tool call + for _ in range(2): + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=_TOOL_CALL_RESPONSE, + ) + + # When running the tool loop with a max_rounds it can't finish within + async with aai.AsyncLLMGateway() as gateway: + with pytest.raises(RuntimeError, match="max_rounds"): + await gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[{"type": "function", "function": {"name": "get_weather"}}], + functions={"get_weather": lambda city: "sunny"}, + max_rounds=2, + ) + + +async def test_run_tools_rejects_streaming(): + # When running the tool loop with stream=True + async with aai.AsyncLLMGateway() as gateway: + with pytest.raises(ValueError, match="stream"): + await gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + tools=[], + functions={}, + stream=True, + ) + + +async def test_chat_completions_stream_get_final_message(httpx_mock: HTTPXMock): + # Given a mocked streaming completions endpoint + chunk1 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], + } + chunk2 = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-sonnet-5", + "choices": [ + {"index": 0, "delta": {"content": " there."}, "finish_reason": "stop"} + ], + "usage": _COMPLETION_RESPONSE["usage"], + } + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [ + f"data: {json.dumps(chunk1)}\n\n".encode(), + f"data: {json.dumps(chunk2)}\n\n".encode(), + b"data: [DONE]\n\n", + ] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When consuming the stream and asking for the final assembled message + async with aai.AsyncLLMGateway() as gateway: + stream = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + chunks = [c async for c in stream] + final = stream.get_final_message() + + # Then it still yields the raw chunks, and the accumulator has the assembled result + assert len(chunks) == 2 + assert final.content == "Hello there." + assert final.finish_reason == "stop" + assert final.usage.total_tokens == 15 + + +async def test_aclose_closes_owned_client_only(): + # Given a gateway with its own client, and one sharing a caller-owned client + async with aai.AsyncLLMGateway() as owned: + pass + assert owned.client.http_client.is_closed + + async with aai.AsyncClient(settings=aai.settings) as shared_client: + async with aai.AsyncLLMGateway(client=shared_client) as gateway: + pass + # Then a caller-supplied client is left open for the caller to close + assert not shared_client.http_client.is_closed From f95a112b365e8a407af73826c69410875afa690a Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Tue, 8 Sep 2026 12:54:47 -0600 Subject: [PATCH 2/8] fix(llm_gateway): narrow chat completion return types Co-Authored-By: Claude Opus 5 (1M context) --- assemblyai/llm_gateway/v1/async_client.py | 42 +++++++++++++++++++---- assemblyai/llm_gateway/v1/client.py | 42 +++++++++++++++++++---- 2 files changed, 72 insertions(+), 12 deletions(-) diff --git a/assemblyai/llm_gateway/v1/async_client.py b/assemblyai/llm_gateway/v1/async_client.py index aca155ca..e4a06ffa 100644 --- a/assemblyai/llm_gateway/v1/async_client.py +++ b/assemblyai/llm_gateway/v1/async_client.py @@ -2,13 +2,13 @@ import json from types import TracebackType -from typing import Any, Callable, Dict, List, Optional, Type, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, overload from typing_extensions import Self from ... import async_client as _async_client from . import async_api, models -from .params import LLMGatewayMessageParamUnion +from .params import LLMGatewayMessageParam from .stream import AsyncLLMGatewayStream @@ -39,11 +39,41 @@ class AsyncCompletionsResource: def __init__(self, gateway: "AsyncLLMGateway") -> None: self._gateway = gateway + @overload async def create( self, *, model: str, - messages: List[LLMGatewayMessageParamUnion], + messages: List[LLMGatewayMessageParam], + stream: Literal[False] = False, + **kwargs: Any, + ) -> models.LLMGatewayChatCompletion: ... + + @overload + async def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], + stream: Literal[True], + **kwargs: Any, + ) -> AsyncLLMGatewayStream: ... + + @overload + async def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], + stream: bool, + **kwargs: Any, + ) -> Union[models.LLMGatewayChatCompletion, AsyncLLMGatewayStream]: ... + + async def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], stream: bool = False, **kwargs: Any, ) -> Union[models.LLMGatewayChatCompletion, AsyncLLMGatewayStream]: @@ -98,12 +128,12 @@ async def run_tools( self, *, model: str, - messages: List[LLMGatewayMessageParamUnion], + messages: List[LLMGatewayMessageParam], tools: List[Dict[str, Any]], functions: Dict[str, Callable[..., Any]], max_rounds: int = 10, **kwargs: Any, - ) -> models.LLMGatewayChatCompletion | AsyncLLMGatewayStream: + ) -> models.LLMGatewayChatCompletion: """ Runs the tool-calling loop to completion: calls the model, invokes any requested tools from `functions`, feeds their results back, and repeats @@ -130,7 +160,7 @@ async def run_tools( raise ValueError("run_tools() does not support stream=True") for _ in range(max_rounds): - completion = await self.create( + completion: models.LLMGatewayChatCompletion = await self.create( model=model, messages=messages, tools=tools, **kwargs ) message = completion.choices[0].message diff --git a/assemblyai/llm_gateway/v1/client.py b/assemblyai/llm_gateway/v1/client.py index 4817e0cc..e6c9e9b7 100644 --- a/assemblyai/llm_gateway/v1/client.py +++ b/assemblyai/llm_gateway/v1/client.py @@ -1,9 +1,9 @@ import json -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Union, overload from ... import client as _client from . import api, models -from .params import LLMGatewayMessageParamUnion +from .params import LLMGatewayMessageParam from .stream import LLMGatewayStream @@ -34,11 +34,41 @@ class CompletionsResource: def __init__(self, gateway: "LLMGateway") -> None: self._gateway = gateway + @overload def create( self, *, model: str, - messages: List[LLMGatewayMessageParamUnion], + messages: List[LLMGatewayMessageParam], + stream: Literal[False] = False, + **kwargs: Any, + ) -> models.LLMGatewayChatCompletion: ... + + @overload + def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], + stream: Literal[True], + **kwargs: Any, + ) -> LLMGatewayStream: ... + + @overload + def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], + stream: bool, + **kwargs: Any, + ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: ... + + def create( + self, + *, + model: str, + messages: List[LLMGatewayMessageParam], stream: bool = False, **kwargs: Any, ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: @@ -92,12 +122,12 @@ def run_tools( self, *, model: str, - messages: List[LLMGatewayMessageParamUnion], + messages: List[LLMGatewayMessageParam], tools: List[Dict[str, Any]], functions: Dict[str, Callable[..., Any]], max_rounds: int = 10, **kwargs: Any, - ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: + ) -> models.LLMGatewayChatCompletion: """ Runs the tool-calling loop to completion: calls the model, invokes any requested tools from `functions`, feeds their results back, and repeats @@ -124,7 +154,7 @@ def run_tools( raise ValueError("run_tools() does not support stream=True") for _ in range(max_rounds): - completion = self.create( + completion: models.LLMGatewayChatCompletion = self.create( model=model, messages=messages, tools=tools, **kwargs ) message = completion.choices[0].message From 250329af531e14889d1416844b141bae38422cda Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Tue, 8 Sep 2026 12:54:55 -0600 Subject: [PATCH 3/8] fix(llm_gateway): harden response model parsing Co-Authored-By: Claude Opus 5 (1M context) --- assemblyai/__init__.py | 4 + assemblyai/llm_gateway/__init__.py | 4 + assemblyai/llm_gateway/v1/__init__.py | 4 + assemblyai/llm_gateway/v1/models.py | 151 +++++++++++++++--- assemblyai/llm_gateway/v1/stream.py | 8 +- tests/unit/test_llm_gateway.py | 210 ++++++++++++++++++++++++++ 6 files changed, 357 insertions(+), 24 deletions(-) diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index 0125ac39..cf3eed25 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -10,6 +10,8 @@ LLMGatewayChoice, LLMGatewayChunkChoice, LLMGatewayChunkDelta, + LLMGatewayChunkFunction, + LLMGatewayChunkToolCall, LLMGatewayCompletionChunk, LLMGatewayCompletionTokensDetails, LLMGatewayDefaultParameters, @@ -156,6 +158,8 @@ def __getattr__(name: str): "LLMGatewayChoice", "LLMGatewayChunkChoice", "LLMGatewayChunkDelta", + "LLMGatewayChunkFunction", + "LLMGatewayChunkToolCall", "LLMGatewayCompletionChunk", "LLMGatewayCompletionTokensDetails", "LLMGatewayDefaultParameters", diff --git a/assemblyai/llm_gateway/__init__.py b/assemblyai/llm_gateway/__init__.py index f8cb8fad..22df2f61 100644 --- a/assemblyai/llm_gateway/__init__.py +++ b/assemblyai/llm_gateway/__init__.py @@ -13,6 +13,8 @@ LLMGatewayChoice, LLMGatewayChunkChoice, LLMGatewayChunkDelta, + LLMGatewayChunkFunction, + LLMGatewayChunkToolCall, LLMGatewayCompletionChunk, LLMGatewayCompletionTokensDetails, LLMGatewayDefaultParameters, @@ -43,6 +45,8 @@ "LLMGatewayChoice", "LLMGatewayChunkChoice", "LLMGatewayChunkDelta", + "LLMGatewayChunkFunction", + "LLMGatewayChunkToolCall", "LLMGatewayCompletionChunk", "LLMGatewayCompletionTokensDetails", "LLMGatewayDefaultParameters", diff --git a/assemblyai/llm_gateway/v1/__init__.py b/assemblyai/llm_gateway/v1/__init__.py index 4c3a22b8..3f5cab81 100644 --- a/assemblyai/llm_gateway/v1/__init__.py +++ b/assemblyai/llm_gateway/v1/__init__.py @@ -7,6 +7,8 @@ LLMGatewayChoice, LLMGatewayChunkChoice, LLMGatewayChunkDelta, + LLMGatewayChunkFunction, + LLMGatewayChunkToolCall, LLMGatewayCompletionChunk, LLMGatewayCompletionTokensDetails, LLMGatewayDefaultParameters, @@ -34,6 +36,8 @@ "LLMGatewayChoice", "LLMGatewayChunkChoice", "LLMGatewayChunkDelta", + "LLMGatewayChunkFunction", + "LLMGatewayChunkToolCall", "LLMGatewayCompletionChunk", "LLMGatewayCompletionTokensDetails", "LLMGatewayDefaultParameters", diff --git a/assemblyai/llm_gateway/v1/models.py b/assemblyai/llm_gateway/v1/models.py index 5a5c640b..c13382c2 100644 --- a/assemblyai/llm_gateway/v1/models.py +++ b/assemblyai/llm_gateway/v1/models.py @@ -2,8 +2,29 @@ from ...types import BaseModel, ConfigDict, Field, pydantic_v2 +if pydantic_v2: + from ...types import field_validator +else: + from ...types import validator -class LLMGatewayPricingData(BaseModel): + +class _LLMGatewayResponseModel(BaseModel): + """ + Base for the LLM Gateway response models. + + Keeps unknown server fields reachable instead of dropping them, so a gateway + that adds a field doesn't need an SDK release before callers can read it. + """ + + if pydantic_v2: + model_config = ConfigDict(extra="allow") + else: + + class Config: + extra = "allow" + + +class LLMGatewayPricingData(_LLMGatewayResponseModel): completions: float prompt: float input_cache_read: Optional[float] = None @@ -11,7 +32,7 @@ class LLMGatewayPricingData(BaseModel): input_cache_write_1h: Optional[float] = None -class LLMGatewayPricing(BaseModel): +class LLMGatewayPricing(_LLMGatewayResponseModel): us: Optional[LLMGatewayPricingData] = None eu: Optional[LLMGatewayPricingData] = None global_: LLMGatewayPricingData = Field(alias="global") @@ -25,19 +46,19 @@ class Config: allow_population_by_field_name = True -class LLMGatewayDefaultParameters(BaseModel): +class LLMGatewayDefaultParameters(_LLMGatewayResponseModel): temperature: Optional[float] = None top_p: Optional[float] = None frequency_penalty: Optional[int] = None -class LLMGatewayTopProvider(BaseModel): +class LLMGatewayTopProvider(_LLMGatewayResponseModel): is_moderated: bool context_length: int max_completion_tokens: int -class LLMGatewayModel(BaseModel): +class LLMGatewayModel(_LLMGatewayResponseModel): """A model available through the LLM Gateway, as returned by `GET /v1/models`.""" id: str @@ -56,33 +77,61 @@ class LLMGatewayModel(BaseModel): providers: List[str] = Field(default_factory=list) default_provider: str + # The gateway is Go: a nil slice marshals as `null`, not `[]`, and an + # explicit null bypasses the field default. + if pydantic_v2: + + @field_validator( + "supported_parameters", "available_regions", "providers", mode="before" + ) + def set_collection_default(cls, v): + return [] if v is None else v -class LLMGatewayModelList(BaseModel): + else: + + @validator("supported_parameters", "available_regions", "providers", pre=True) + def set_collection_default(cls, v): + return [] if v is None else v + + +class LLMGatewayModelList(_LLMGatewayResponseModel): """The response of `GET /v1/models`.""" data: List[LLMGatewayModel] = Field(default_factory=list) + if pydantic_v2: + + @field_validator("data", mode="before") + def set_collection_default(cls, v): + return [] if v is None else v -class LLMGatewayCacheCreation(BaseModel): + else: + + @validator("data", pre=True) + def set_collection_default(cls, v): + return [] if v is None else v + + +class LLMGatewayCacheCreation(_LLMGatewayResponseModel): ephemeral_5m_input_tokens: int ephemeral_1h_input_tokens: int -class LLMGatewayPromptTokensDetails(BaseModel): +class LLMGatewayPromptTokensDetails(_LLMGatewayResponseModel): cached_tokens: int audio_tokens: int cache_creation: Optional[LLMGatewayCacheCreation] = None cache_write_tokens: Optional[int] = None -class LLMGatewayCompletionTokensDetails(BaseModel): +class LLMGatewayCompletionTokensDetails(_LLMGatewayResponseModel): reasoning_tokens: int audio_tokens: int accepted_prediction_tokens: int rejected_prediction_tokens: int -class LLMGatewayUsage(BaseModel): +class LLMGatewayUsage(_LLMGatewayResponseModel): """ `input_tokens`/`output_tokens` are only populated for non-streaming responses — the gateway forwards the upstream provider's own streaming @@ -95,22 +144,22 @@ class LLMGatewayUsage(BaseModel): output_tokens: Optional[int] = None completion_tokens: int total_tokens: int - prompt_tokens_details: LLMGatewayPromptTokensDetails - completion_tokens_details: LLMGatewayCompletionTokensDetails + prompt_tokens_details: Optional[LLMGatewayPromptTokensDetails] = None + completion_tokens_details: Optional[LLMGatewayCompletionTokensDetails] = None -class LLMGatewayFunction(BaseModel): +class LLMGatewayFunction(_LLMGatewayResponseModel): name: str arguments: Any = None -class LLMGatewayToolCall(BaseModel): +class LLMGatewayToolCall(_LLMGatewayResponseModel): id: str type: str function: LLMGatewayFunction -class LLMGatewayResponseMessage(BaseModel): +class LLMGatewayResponseMessage(_LLMGatewayResponseModel): role: str content: Optional[str] = None thinking: Optional[str] = None @@ -119,13 +168,13 @@ class LLMGatewayResponseMessage(BaseModel): name: Optional[str] = None -class LLMGatewayChoice(BaseModel): +class LLMGatewayChoice(_LLMGatewayResponseModel): index: int finish_reason: Optional[str] = None message: LLMGatewayResponseMessage -class LLMGatewayChatCompletion(BaseModel): +class LLMGatewayChatCompletion(_LLMGatewayResponseModel): """The response of `POST /v1/chat/completions` (non-streaming).""" request_id: str @@ -140,18 +189,54 @@ class LLMGatewayChatCompletion(BaseModel): llm_status_code: Optional[int] = None + if pydantic_v2: + + @field_validator("choices", mode="before") + def set_collection_default(cls, v): + return [] if v is None else v -class LLMGatewayChunkDelta(BaseModel): + else: + + @validator("choices", pre=True) + def set_collection_default(cls, v): + return [] if v is None else v + + +class LLMGatewayChunkFunction(_LLMGatewayResponseModel): + name: Optional[str] = None + arguments: Optional[str] = None + """A fragment of the JSON arguments, not the whole value — concatenate the + fragments carrying the same `index` across chunks.""" + + +class LLMGatewayChunkToolCall(_LLMGatewayResponseModel): + """ + A tool call as it arrives mid-stream. + + Everything but `index` is optional: providers send the `id`/`type`/`name` + once and then stream `function.arguments` in fragments, so a single chunk + carries only part of the call. + """ + + index: int + id: Optional[str] = None + type: Optional[str] = None + function: Optional[LLMGatewayChunkFunction] = None + + +class LLMGatewayChunkDelta(_LLMGatewayResponseModel): content: Optional[str] = None + role: Optional[str] = None + tool_calls: Optional[List[LLMGatewayChunkToolCall]] = None -class LLMGatewayChunkChoice(BaseModel): +class LLMGatewayChunkChoice(_LLMGatewayResponseModel): index: int delta: LLMGatewayChunkDelta finish_reason: Optional[str] = None -class LLMGatewayCompletionChunk(BaseModel): +class LLMGatewayCompletionChunk(_LLMGatewayResponseModel): """ A streamed chat completion chunk. @@ -169,11 +254,35 @@ class LLMGatewayCompletionChunk(BaseModel): usage: Optional[LLMGatewayUsage] = None obfuscation: Optional[str] = None + if pydantic_v2: + + @field_validator("choices", mode="before") + def set_collection_default(cls, v): + return [] if v is None else v + + else: + + @validator("choices", pre=True) + def set_collection_default(cls, v): + return [] if v is None else v + -class LLMGatewayUnderstandingResponse(BaseModel): +class LLMGatewayUnderstandingResponse(_LLMGatewayResponseModel): """The response of `POST /v1/understanding`.""" speech_understanding: Dict[str, Any] = Field(default_factory=dict) request_id: str utterances: Optional[List[Dict[str, Any]]] = None translated_texts: Optional[Dict[str, str]] = None + + if pydantic_v2: + + @field_validator("speech_understanding", mode="before") + def set_collection_default(cls, v): + return {} if v is None else v + + else: + + @validator("speech_understanding", pre=True) + def set_collection_default(cls, v): + return {} if v is None else v diff --git a/assemblyai/llm_gateway/v1/stream.py b/assemblyai/llm_gateway/v1/stream.py index d9aa819e..6db07dc4 100644 --- a/assemblyai/llm_gateway/v1/stream.py +++ b/assemblyai/llm_gateway/v1/stream.py @@ -38,9 +38,11 @@ class LLMGatewayStream: `for chunk in stream: ...`. Once exhausted, call `get_final_message()` for the assembled text/usage/finish_reason. - Verified for OpenAI-routed models only, like raw chunk parsing. Tool calls - are not reconstructed here — streamed chunks don't carry `tool_calls` - today. + Verified for OpenAI-routed models only, like raw chunk parsing. Chunks can + carry tool calls (`chunk.choices[0].delta.tool_calls`), but they arrive as + fragments and are not reassembled here — `get_final_message()` covers text, + usage and `finish_reason` only. For assembled tool calls, use `run_tools()` + or a non-streaming `create()`. """ def __init__(self, chunks: Iterator[models.LLMGatewayCompletionChunk]) -> None: diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py index 8d4081ec..8d445631 100644 --- a/tests/unit/test_llm_gateway.py +++ b/tests/unit/test_llm_gateway.py @@ -757,3 +757,213 @@ def test_chat_completions_stream_get_final_message(httpx_mock: HTTPXMock): assert final.usage.total_tokens == 15 assert final.id == "chatcmpl-1" assert final.model == "claude-sonnet-5" + + +# The gateway is written in Go, where a nil slice/map marshals as `null` rather +# than `[]`/`{}`. An explicit null bypasses a pydantic field default, so each +# collection on the response path has to tolerate it. + + +def test_list_models_tolerates_null_data(httpx_mock: HTTPXMock): + # Given a models response whose `data` slice came back null + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.OK, + json={"data": None}, + ) + + # When listing models + result = aai.LLMGateway().models.list() + + # Then it parses to an empty list rather than raising + assert result.data == [] + + +def test_list_models_tolerates_null_model_collections(httpx_mock: HTTPXMock): + # Given a model entry whose list fields all came back null + httpx_mock.add_response( + url=MODELS_URL, + method="GET", + status_code=httpx.codes.OK, + json={ + "data": [ + { + "id": "claude-sonnet-5", + "name": "Claude Sonnet 5", + "description": "", + "default_parameters": {}, + "supported_parameters": None, + "top_provider": { + "is_moderated": False, + "context_length": 200000, + "max_completion_tokens": 8192, + }, + "context_length": 200000, + "pricing": {"global": {"completions": 15.0, "prompt": 3.0}}, + "creator": "anthropic", + "retirement_date": 0, + "available_regions": None, + "providers": None, + "default_provider": "bedrock", + } + ] + }, + ) + + # When listing models + model = aai.LLMGateway().models.list().data[0] + + # Then every null collection defaults to empty + assert model.supported_parameters == [] + assert model.available_regions == [] + assert model.providers == [] + + +def test_chat_completions_tolerates_null_choices(httpx_mock: HTTPXMock): + # Given a completion whose `choices` slice came back null + response = dict(_COMPLETION_RESPONSE, choices=None) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=response, + ) + + # When creating a completion + result = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then it parses to an empty list rather than raising + assert result.choices == [] + + +def test_chat_completions_tolerates_usage_without_token_details( + httpx_mock: HTTPXMock, +): + # Given a completion whose usage omits the OpenAI-shaped detail objects + response = dict( + _COMPLETION_RESPONSE, + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + ) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=response, + ) + + # When creating a completion + result = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then the totals still parse and the detail objects are None + assert result.usage.total_tokens == 15 + assert result.usage.prompt_tokens_details is None + assert result.usage.completion_tokens_details is None + + +def test_understanding_tolerates_null_speech_understanding(httpx_mock: HTTPXMock): + # Given an understanding response whose `speech_understanding` map came back null + httpx_mock.add_response( + url=UNDERSTANDING_URL, + method="POST", + status_code=httpx.codes.OK, + json={"speech_understanding": None, "request_id": "req-1"}, + ) + + # When running understanding + result = aai.LLMGateway().understanding.create( + transcript_id="transcript-1", + request={"translation": {"target_languages": ["es"]}}, + ) + + # Then it parses to an empty dict rather than raising + assert result.speech_understanding == {} + + +def test_chat_completions_preserves_unknown_server_fields(httpx_mock: HTTPXMock): + # Given a completion carrying a field this SDK version doesn't model + response = dict(_COMPLETION_RESPONSE, brand_new_field={"nested": 1}) + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + json=response, + ) + + # When creating a completion + result = aai.LLMGateway().chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Hi"}], + ) + + # Then the unknown field is still reachable, not dropped + assert result.brand_new_field == {"nested": 1} + + +def test_stream_chunk_exposes_tool_call_delta(httpx_mock: HTTPXMock): + # Given a Claude-shaped streaming chunk carrying a fragmented tool call + # (see the gateway's pkg/models/claude/stream.go Delta/OpenAIToolCalls) + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "claude-haiku-4-5-20251001", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"ci', + }, + } + ], + }, + "finish_reason": None, + } + ], + } + import json as _json + + httpx_mock.add_response( + url=COMPLETIONS_URL, + method="POST", + status_code=httpx.codes.OK, + stream=IteratorStream( + [f"data: {_json.dumps(chunk)}\n\n".encode(), b"data: [DONE]\n\n"] + ), + headers={"content-type": "text/event-stream"}, + ) + + # When consuming the stream + chunks = list( + aai.LLMGateway().chat.completions.create( + model="claude-haiku-4-5-20251001", + messages=[{"role": "user", "content": "Weather in Paris?"}], + stream=True, + ) + ) + + # Then the delta's role and tool-call fragment are typed and reachable + delta = chunks[0].choices[0].delta + assert delta.role == "assistant" + assert delta.content is None + tool_call = delta.tool_calls[0] + assert isinstance(tool_call, aai.LLMGatewayChunkToolCall) + assert tool_call.index == 0 + assert tool_call.id == "call_1" + assert tool_call.function.name == "get_weather" + # arguments stream in fragments; the SDK doesn't reassemble them + assert tool_call.function.arguments == '{"ci' From 29d2d79cb3759c202359dd99601e29ae9cb2f1ef Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Tue, 8 Sep 2026 14:22:23 -0600 Subject: [PATCH 4/8] refactor(llm_gateway): share client logic and trim the public surface Move the logic duplicated between client.py and async_client.py into a new _base.py, matching sync/v1, prerecorded/v2 and streaming/v3: the resource base holding the gateway and its base_url/timeout pair, the three request body builders, and the tool-calling loop's core. Trim the exported names from 29 to 18, keeping the clients, the message param, the stream types, the response envelopes and the layer below them a caller annotates; the nested leaf models stay reachable through assemblyai.llm_gateway.v1.models. Reduce llm_gateway/__init__.py to a docstring, matching the empty streaming and prerecorded package inits, and drop the vestigial LLMGatewayMessageParamUnion alias. Remove a committed pprint of every understanding response, and reformat the one line in api.py that failed ruff format --check. Co-Authored-By: Claude Opus 5 (1M context) --- assemblyai/__init__.py | 22 ---- assemblyai/llm_gateway/__init__.py | 70 +--------- assemblyai/llm_gateway/v1/__init__.py | 32 ++--- assemblyai/llm_gateway/v1/_base.py | 153 ++++++++++++++++++++++ assemblyai/llm_gateway/v1/api.py | 9 +- assemblyai/llm_gateway/v1/async_client.py | 116 +++------------- assemblyai/llm_gateway/v1/client.py | 116 +++------------- assemblyai/llm_gateway/v1/params.py | 3 - tests/unit/test_llm_gateway.py | 3 +- 9 files changed, 209 insertions(+), 315 deletions(-) create mode 100644 assemblyai/llm_gateway/v1/_base.py diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index cf3eed25..9c4a40cb 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -5,29 +5,18 @@ AsyncLLMGateway, AsyncLLMGatewayStream, LLMGateway, - LLMGatewayCacheCreation, LLMGatewayChatCompletion, LLMGatewayChoice, LLMGatewayChunkChoice, LLMGatewayChunkDelta, - LLMGatewayChunkFunction, - LLMGatewayChunkToolCall, LLMGatewayCompletionChunk, - LLMGatewayCompletionTokensDetails, - LLMGatewayDefaultParameters, - LLMGatewayFunction, LLMGatewayMessageParam, - LLMGatewayMessageParamUnion, LLMGatewayModel, LLMGatewayModelList, - LLMGatewayPricing, - LLMGatewayPricingData, - LLMGatewayPromptTokensDetails, LLMGatewayResponseMessage, LLMGatewayStream, LLMGatewayStreamAccumulator, LLMGatewayToolCall, - LLMGatewayTopProvider, LLMGatewayUnderstandingResponse, LLMGatewayUsage, ) @@ -153,30 +142,19 @@ def __getattr__(name: str): "ListTranscriptParameters", "ListTranscriptResponse", "LLMGateway", - "LLMGatewayCacheCreation", "LLMGatewayChatCompletion", "LLMGatewayChoice", "LLMGatewayChunkChoice", "LLMGatewayChunkDelta", - "LLMGatewayChunkFunction", - "LLMGatewayChunkToolCall", "LLMGatewayCompletionChunk", - "LLMGatewayCompletionTokensDetails", - "LLMGatewayDefaultParameters", "LLMGatewayError", - "LLMGatewayFunction", "LLMGatewayMessageParam", - "LLMGatewayMessageParamUnion", "LLMGatewayModel", "LLMGatewayModelList", - "LLMGatewayPricing", - "LLMGatewayPricingData", - "LLMGatewayPromptTokensDetails", "LLMGatewayResponseMessage", "LLMGatewayStream", "LLMGatewayStreamAccumulator", "LLMGatewayToolCall", - "LLMGatewayTopProvider", "LLMGatewayUnderstandingResponse", "LLMGatewayUsage", "PageDetails", diff --git a/assemblyai/llm_gateway/__init__.py b/assemblyai/llm_gateway/__init__.py index 22df2f61..967bf27a 100644 --- a/assemblyai/llm_gateway/__init__.py +++ b/assemblyai/llm_gateway/__init__.py @@ -1,69 +1 @@ -"""Re-exports of the LLM Gateway client surface. - -The canonical location is ``assemblyai.llm_gateway.v1``, mirroring -``assemblyai/sync/v1``. -""" - -from .v1 import ( # noqa: F401 - AsyncLLMGateway, - AsyncLLMGatewayStream, - LLMGateway, - LLMGatewayCacheCreation, - LLMGatewayChatCompletion, - LLMGatewayChoice, - LLMGatewayChunkChoice, - LLMGatewayChunkDelta, - LLMGatewayChunkFunction, - LLMGatewayChunkToolCall, - LLMGatewayCompletionChunk, - LLMGatewayCompletionTokensDetails, - LLMGatewayDefaultParameters, - LLMGatewayError, - LLMGatewayFunction, - LLMGatewayMessageParam, - LLMGatewayMessageParamUnion, - LLMGatewayModel, - LLMGatewayModelList, - LLMGatewayPricing, - LLMGatewayPricingData, - LLMGatewayPromptTokensDetails, - LLMGatewayResponseMessage, - LLMGatewayStream, - LLMGatewayStreamAccumulator, - LLMGatewayToolCall, - LLMGatewayTopProvider, - LLMGatewayUnderstandingResponse, - LLMGatewayUsage, -) - -__all__ = [ - "AsyncLLMGateway", - "AsyncLLMGatewayStream", - "LLMGateway", - "LLMGatewayCacheCreation", - "LLMGatewayChatCompletion", - "LLMGatewayChoice", - "LLMGatewayChunkChoice", - "LLMGatewayChunkDelta", - "LLMGatewayChunkFunction", - "LLMGatewayChunkToolCall", - "LLMGatewayCompletionChunk", - "LLMGatewayCompletionTokensDetails", - "LLMGatewayDefaultParameters", - "LLMGatewayError", - "LLMGatewayFunction", - "LLMGatewayMessageParam", - "LLMGatewayMessageParamUnion", - "LLMGatewayModel", - "LLMGatewayModelList", - "LLMGatewayPricing", - "LLMGatewayPricingData", - "LLMGatewayPromptTokensDetails", - "LLMGatewayResponseMessage", - "LLMGatewayStream", - "LLMGatewayStreamAccumulator", - "LLMGatewayToolCall", - "LLMGatewayTopProvider", - "LLMGatewayUnderstandingResponse", - "LLMGatewayUsage", -] +"""LLM Gateway support. The client surface lives in ``assemblyai.llm_gateway.v1``.""" diff --git a/assemblyai/llm_gateway/v1/__init__.py b/assemblyai/llm_gateway/v1/__init__.py index 3f5cab81..504cfe5a 100644 --- a/assemblyai/llm_gateway/v1/__init__.py +++ b/assemblyai/llm_gateway/v1/__init__.py @@ -1,60 +1,48 @@ +"""LLM Gateway client surface. + +Exports the clients, the request param type, and the response types a caller +receives or annotates. The nested leaf models — token-detail breakdowns, +pricing, provider metadata, streamed tool-call fragments — are reached by +attribute access off these types and stay in `models` rather than the public +surface. +""" + from ...types import LLMGatewayError from .async_client import AsyncLLMGateway from .client import LLMGateway from .models import ( - LLMGatewayCacheCreation, LLMGatewayChatCompletion, LLMGatewayChoice, LLMGatewayChunkChoice, LLMGatewayChunkDelta, - LLMGatewayChunkFunction, - LLMGatewayChunkToolCall, LLMGatewayCompletionChunk, - LLMGatewayCompletionTokensDetails, - LLMGatewayDefaultParameters, - LLMGatewayFunction, LLMGatewayModel, LLMGatewayModelList, - LLMGatewayPricing, - LLMGatewayPricingData, - LLMGatewayPromptTokensDetails, LLMGatewayResponseMessage, LLMGatewayToolCall, - LLMGatewayTopProvider, LLMGatewayUnderstandingResponse, LLMGatewayUsage, ) -from .params import LLMGatewayMessageParam, LLMGatewayMessageParamUnion +from .params import LLMGatewayMessageParam from .stream import AsyncLLMGatewayStream, LLMGatewayStream, LLMGatewayStreamAccumulator __all__ = [ "AsyncLLMGateway", "AsyncLLMGatewayStream", "LLMGateway", - "LLMGatewayCacheCreation", "LLMGatewayChatCompletion", "LLMGatewayChoice", "LLMGatewayChunkChoice", "LLMGatewayChunkDelta", - "LLMGatewayChunkFunction", - "LLMGatewayChunkToolCall", "LLMGatewayCompletionChunk", - "LLMGatewayCompletionTokensDetails", - "LLMGatewayDefaultParameters", "LLMGatewayError", - "LLMGatewayFunction", "LLMGatewayMessageParam", - "LLMGatewayMessageParamUnion", "LLMGatewayModel", "LLMGatewayModelList", - "LLMGatewayPricing", - "LLMGatewayPricingData", - "LLMGatewayPromptTokensDetails", "LLMGatewayResponseMessage", "LLMGatewayStream", "LLMGatewayStreamAccumulator", "LLMGatewayToolCall", - "LLMGatewayTopProvider", "LLMGatewayUnderstandingResponse", "LLMGatewayUsage", ] diff --git a/assemblyai/llm_gateway/v1/_base.py b/assemblyai/llm_gateway/v1/_base.py new file mode 100644 index 00000000..9cbe0031 --- /dev/null +++ b/assemblyai/llm_gateway/v1/_base.py @@ -0,0 +1,153 @@ +"""Logic shared by `client.py` and `async_client.py`. + +Holds the request body shapes and the tool-calling loop's core, so a +wire-contract or tool-loop change lands in one place. Everything here is +I/O-free — the sync and async clients supply the HTTP calls themselves. +""" + +import json +from typing import Any, Callable, Dict, List, Optional + +from . import models +from .params import LLMGatewayMessageParam + + +class _BaseResource: + """ + Shared base for the gateway's resource classes. + + Subclasses hold a gateway (sync or async) and implement the I/O. + """ + + def __init__(self, gateway: Any) -> None: + self._gateway = gateway + + @property + def _http_options(self) -> Dict[str, Any]: + """The `base_url`/`timeout` keyword arguments every `api` call takes.""" + settings = self._gateway.client.settings + + return { + "base_url": settings.llm_gateway_base_url, + "timeout": settings.llm_gateway_http_timeout, + } + + +def completion_body( + *, + model: str, + messages: List[LLMGatewayMessageParam], + stream: bool, + extra: Dict[str, Any], +) -> Dict[str, Any]: + """Builds the request body of `POST /v1/chat/completions`.""" + return { + "model": model, + "messages": messages, + "stream": stream, + **extra, + } + + +def understanding_body( + *, + transcript_id: str, + request: Dict[str, Any], + extra: Dict[str, Any], +) -> Dict[str, Any]: + """Builds the request body of `POST /v1/understanding`.""" + return { + "transcript_id": transcript_id, + "speech_understanding": {"request": request}, + **extra, + } + + +def validate_body( + *, + request: Dict[str, Any], + transcript_id: Optional[str], + extra: Dict[str, Any], +) -> Dict[str, Any]: + """Builds the request body of `POST /v1/understanding/validate`.""" + body: Dict[str, Any] = {"speech_understanding": {"request": request}, **extra} + if transcript_id is not None: + body["transcript_id"] = transcript_id + + return body + + +def reject_stream(kwargs: Dict[str, Any]) -> None: + """Raises `ValueError` if `run_tools()` was passed `stream=True`.""" + if kwargs.get("stream"): + raise ValueError("run_tools() does not support stream=True") + + +def apply_tool_round( + completion: models.LLMGatewayChatCompletion, + messages: List[LLMGatewayMessageParam], + functions: Dict[str, Callable[..., Any]], +) -> bool: + """ + Plays one round of the tool-calling loop into `messages`, in place. + + Appends the model's reply, and — when it requested tool calls — invokes each + one from `functions` and appends its result. + + Returns: True once the model has replied without requesting a tool call, + i.e. the loop is finished. + + Raises: + ValueError: if the model requests a tool with no matching entry in + `functions`. + """ + message = completion.choices[0].message + + if not message.tool_calls: + messages.append({"role": "assistant", "content": message.content}) + return True + + messages.append( + { + "role": "assistant", + "content": message.content, + "tool_calls": [ + { + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + if isinstance(tool_call.function.arguments, str) + else json.dumps(tool_call.function.arguments), + }, + } + for tool_call in message.tool_calls + ], + } + ) + + for tool_call in message.tool_calls: + function = functions.get(tool_call.function.name) + if function is None: + raise ValueError( + f"No function registered for tool call {tool_call.function.name!r}" + ) + + arguments = tool_call.function.arguments + if isinstance(arguments, str): + arguments = json.loads(arguments) + + result = function(**arguments) + output = result if isinstance(result, str) else json.dumps(result) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": output, + "name": tool_call.function.name, + } + ) + + return False diff --git a/assemblyai/llm_gateway/v1/api.py b/assemblyai/llm_gateway/v1/api.py index cd9c0fc0..db5f03aa 100644 --- a/assemblyai/llm_gateway/v1/api.py +++ b/assemblyai/llm_gateway/v1/api.py @@ -27,8 +27,6 @@ def _error_from_response(response: httpx.Response) -> types.LLMGatewayError: message: Optional[str] = None errors: Optional[List[str]] = None - print(response.text) - try: body = response.json() if isinstance(body, dict): @@ -36,7 +34,9 @@ def _error_from_response(response: httpx.Response) -> types.LLMGatewayError: if "message" in body: message = body.get("message") metadata = body.get("metadata") - if isinstance(metadata, dict) and isinstance(metadata.get("errors"), list): + if isinstance(metadata, dict) and isinstance( + metadata.get("errors"), list + ): errors = metadata["errors"] elif "error" in body: message = body.get("error") @@ -160,9 +160,6 @@ def create_understanding( timeout=timeout, ) - from pprint import pprint - pprint (response.json()) - if response.status_code != httpx.codes.OK: raise _error_from_response(response) diff --git a/assemblyai/llm_gateway/v1/async_client.py b/assemblyai/llm_gateway/v1/async_client.py index e4a06ffa..3786deee 100644 --- a/assemblyai/llm_gateway/v1/async_client.py +++ b/assemblyai/llm_gateway/v1/async_client.py @@ -1,44 +1,34 @@ from __future__ import annotations -import json from types import TracebackType from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, overload from typing_extensions import Self from ... import async_client as _async_client -from . import async_api, models +from . import _base, async_api, models from .params import LLMGatewayMessageParam from .stream import AsyncLLMGatewayStream -class AsyncModelsResource: +class AsyncModelsResource(_base._BaseResource): """The `models` resource of `AsyncLLMGateway`.""" - def __init__(self, gateway: "AsyncLLMGateway") -> None: - self._gateway = gateway - async def list(self) -> models.LLMGatewayModelList: """ Lists the models available through the LLM Gateway. Raises: `LLMGatewayError` if the request fails. """ - settings = self._gateway.client.settings - return await async_api.list_models( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, + **self._http_options, ) -class AsyncCompletionsResource: +class AsyncCompletionsResource(_base._BaseResource): """The `chat.completions` resource of `AsyncLLMGateway`.""" - def __init__(self, gateway: "AsyncLLMGateway") -> None: - self._gateway = gateway - @overload async def create( self, @@ -99,29 +89,23 @@ async def create( it fails before any bytes are received — a failure mid-stream just truncates the iterator). """ - settings = self._gateway.client.settings - body: Dict[str, Any] = { - "model": model, - "messages": messages, - "stream": stream, - **kwargs, - } + body = _base.completion_body( + model=model, messages=messages, stream=stream, extra=kwargs + ) if stream: return AsyncLLMGatewayStream( async_api.stream_completion( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, body=body, + **self._http_options, ) ) return await async_api.create_completion( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, body=body, + **self._http_options, ) async def run_tools( @@ -156,62 +140,15 @@ async def run_tools( RuntimeError: if the loop doesn't finish within `max_rounds`. LLMGatewayError: if any underlying `create()` call fails. """ - if kwargs.get("stream"): - raise ValueError("run_tools() does not support stream=True") + _base.reject_stream(kwargs) for _ in range(max_rounds): completion: models.LLMGatewayChatCompletion = await self.create( model=model, messages=messages, tools=tools, **kwargs ) - message = completion.choices[0].message - - if not message.tool_calls: - messages.append({"role": "assistant", "content": message.content}) + if _base.apply_tool_round(completion, messages, functions): return completion - messages.append( - { - "role": "assistant", - "content": message.content, - "tool_calls": [ - { - "id": tool_call.id, - "type": tool_call.type, - "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments - if isinstance(tool_call.function.arguments, str) - else json.dumps(tool_call.function.arguments), - }, - } - for tool_call in message.tool_calls - ], - } - ) - - for tool_call in message.tool_calls: - function = functions.get(tool_call.function.name) - if function is None: - raise ValueError( - f"No function registered for tool call {tool_call.function.name!r}" - ) - - arguments = tool_call.function.arguments - if isinstance(arguments, str): - arguments = json.loads(arguments) - - result = function(**arguments) - output = result if isinstance(result, str) else json.dumps(result) - - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "content": output, - "name": tool_call.function.name, - } - ) - raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") @@ -222,12 +159,9 @@ def __init__(self, gateway: "AsyncLLMGateway") -> None: self.completions = AsyncCompletionsResource(gateway) -class AsyncUnderstandingResource: +class AsyncUnderstandingResource(_base._BaseResource): """The `understanding` resource of `AsyncLLMGateway`.""" - def __init__(self, gateway: "AsyncLLMGateway") -> None: - self._gateway = gateway - async def create( self, *, @@ -244,18 +178,12 @@ async def create( Raises: `LLMGatewayError` if the request fails. """ - settings = self._gateway.client.settings - body: Dict[str, Any] = { - "transcript_id": transcript_id, - "speech_understanding": {"request": request}, - **kwargs, - } - return await async_api.create_understanding( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, - body=body, + body=_base.understanding_body( + transcript_id=transcript_id, request=request, extra=kwargs + ), + **self._http_options, ) async def validate( @@ -275,16 +203,12 @@ async def validate( Raises: `LLMGatewayError` (with `.errors` populated with the validation messages) if the request is invalid. """ - settings = self._gateway.client.settings - body: Dict[str, Any] = {"speech_understanding": {"request": request}, **kwargs} - if transcript_id is not None: - body["transcript_id"] = transcript_id - await async_api.validate_understanding( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, - body=body, + body=_base.validate_body( + request=request, transcript_id=transcript_id, extra=kwargs + ), + **self._http_options, ) diff --git a/assemblyai/llm_gateway/v1/client.py b/assemblyai/llm_gateway/v1/client.py index e6c9e9b7..c267f4fc 100644 --- a/assemblyai/llm_gateway/v1/client.py +++ b/assemblyai/llm_gateway/v1/client.py @@ -1,39 +1,29 @@ -import json from typing import Any, Callable, Dict, List, Literal, Optional, Union, overload from ... import client as _client -from . import api, models +from . import _base, api, models from .params import LLMGatewayMessageParam from .stream import LLMGatewayStream -class ModelsResource: +class ModelsResource(_base._BaseResource): """The `models` resource of `LLMGateway`.""" - def __init__(self, gateway: "LLMGateway") -> None: - self._gateway = gateway - def list(self) -> models.LLMGatewayModelList: """ Lists the models available through the LLM Gateway. Raises: `LLMGatewayError` if the request fails. """ - settings = self._gateway.client.settings - return api.list_models( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, + **self._http_options, ) -class CompletionsResource: +class CompletionsResource(_base._BaseResource): """The `chat.completions` resource of `LLMGateway`.""" - def __init__(self, gateway: "LLMGateway") -> None: - self._gateway = gateway - @overload def create( self, @@ -93,29 +83,23 @@ def create( it fails before any bytes are received — a failure mid-stream just truncates the iterator). """ - settings = self._gateway.client.settings - body: Dict[str, Any] = { - "model": model, - "messages": messages, - "stream": stream, - **kwargs, - } + body = _base.completion_body( + model=model, messages=messages, stream=stream, extra=kwargs + ) if stream: return LLMGatewayStream( api.stream_completion( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, body=body, + **self._http_options, ) ) return api.create_completion( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, body=body, + **self._http_options, ) def run_tools( @@ -150,62 +134,15 @@ def run_tools( RuntimeError: if the loop doesn't finish within `max_rounds`. LLMGatewayError: if any underlying `create()` call fails. """ - if kwargs.get("stream"): - raise ValueError("run_tools() does not support stream=True") + _base.reject_stream(kwargs) for _ in range(max_rounds): completion: models.LLMGatewayChatCompletion = self.create( model=model, messages=messages, tools=tools, **kwargs ) - message = completion.choices[0].message - - if not message.tool_calls: - messages.append({"role": "assistant", "content": message.content}) + if _base.apply_tool_round(completion, messages, functions): return completion - messages.append( - { - "role": "assistant", - "content": message.content, - "tool_calls": [ - { - "id": tool_call.id, - "type": tool_call.type, - "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments - if isinstance(tool_call.function.arguments, str) - else json.dumps(tool_call.function.arguments), - }, - } - for tool_call in message.tool_calls - ], - } - ) - - for tool_call in message.tool_calls: - function = functions.get(tool_call.function.name) - if function is None: - raise ValueError( - f"No function registered for tool call {tool_call.function.name!r}" - ) - - arguments = tool_call.function.arguments - if isinstance(arguments, str): - arguments = json.loads(arguments) - - result = function(**arguments) - output = result if isinstance(result, str) else json.dumps(result) - - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "content": output, - "name": tool_call.function.name, - } - ) - raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") @@ -216,12 +153,9 @@ def __init__(self, gateway: "LLMGateway") -> None: self.completions = CompletionsResource(gateway) -class UnderstandingResource: +class UnderstandingResource(_base._BaseResource): """The `understanding` resource of `LLMGateway`.""" - def __init__(self, gateway: "LLMGateway") -> None: - self._gateway = gateway - def create( self, *, @@ -238,18 +172,12 @@ def create( Raises: `LLMGatewayError` if the request fails. """ - settings = self._gateway.client.settings - body: Dict[str, Any] = { - "transcript_id": transcript_id, - "speech_understanding": {"request": request}, - **kwargs, - } - return api.create_understanding( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, - body=body, + body=_base.understanding_body( + transcript_id=transcript_id, request=request, extra=kwargs + ), + **self._http_options, ) def validate( @@ -269,16 +197,12 @@ def validate( Raises: `LLMGatewayError` (with `.errors` populated with the validation messages) if the request is invalid. """ - settings = self._gateway.client.settings - body: Dict[str, Any] = {"speech_understanding": {"request": request}, **kwargs} - if transcript_id is not None: - body["transcript_id"] = transcript_id - api.validate_understanding( self._gateway.client.http_client, - base_url=settings.llm_gateway_base_url, - timeout=settings.llm_gateway_http_timeout, - body=body, + body=_base.validate_body( + request=request, transcript_id=transcript_id, extra=kwargs + ), + **self._http_options, ) diff --git a/assemblyai/llm_gateway/v1/params.py b/assemblyai/llm_gateway/v1/params.py index 0929c845..8f67d91a 100644 --- a/assemblyai/llm_gateway/v1/params.py +++ b/assemblyai/llm_gateway/v1/params.py @@ -15,6 +15,3 @@ class LLMGatewayMessageParam(TypedDict, total=False): name: str thinking: str cache_control: Dict[str, Any] - - -LLMGatewayMessageParamUnion = LLMGatewayMessageParam diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py index 8d445631..78247489 100644 --- a/tests/unit/test_llm_gateway.py +++ b/tests/unit/test_llm_gateway.py @@ -3,6 +3,7 @@ from pytest_httpx import HTTPXMock, IteratorStream import assemblyai as aai +from assemblyai.llm_gateway.v1 import models aai.settings.api_key = "test" @@ -961,7 +962,7 @@ def test_stream_chunk_exposes_tool_call_delta(httpx_mock: HTTPXMock): assert delta.role == "assistant" assert delta.content is None tool_call = delta.tool_calls[0] - assert isinstance(tool_call, aai.LLMGatewayChunkToolCall) + assert isinstance(tool_call, models.LLMGatewayChunkToolCall) assert tool_call.index == 0 assert tool_call.id == "call_1" assert tool_call.function.name == "get_weather" From d797ccfbda12fc55301ab42bd793d0bd20b3b5b5 Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Wed, 9 Sep 2026 09:07:42 -0600 Subject: [PATCH 5/8] feat(llm_gateway): documentation --- CLAUDE.md | 180 ++++++++++++++++++++++++++++++++- MIGRATION.md | 30 ++++-- README.md | 275 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 476 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 515c7869..fa02dcb2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,6 +1,6 @@ # AssemblyAI Python SDK -Speech-to-text and audio intelligence SDK. Supports pre-recorded transcription, real-time streaming, and audio analysis features. +Speech-to-text and audio intelligence SDK. Supports pre-recorded transcription, real-time streaming, audio analysis features, and LLM access through AssemblyAI's LLM Gateway. ## Quick start @@ -50,6 +50,11 @@ aai.settings.api_key = "your-key" - `aai.SyncTranscriptResponse` — Sync result: `.text`, `.words` (`SyncWord` with `confidence` always, `start`/`end` only when `timestamps=True`), `.confidence`, `.audio_duration_ms`, `.session_id`, `.request_time_ms` - `assemblyai.streaming.v3.RealTimeTranscriber` — Real-time streaming with event-based API (threaded) - `assemblyai.streaming.v3.AsyncRealTimeTranscriber` — Asyncio-native counterpart; same options/events +- `aai.LLMGateway` — LLM Gateway client. Resources: `models.list()`, `chat.completions.create()` (incl. `stream=True`), `chat.completions.run_tools()`, `understanding.create()`/`.validate()` +- `aai.AsyncLLMGateway` — Asyncio counterpart of `LLMGateway`. Same resources, every API call a coroutine; owns an HTTP pool (`async with` or `await aclose()`) +- `assemblyai.llm_gateway.v1` — Canonical module for both gateways, matching the gateway's `/v1` API. The top-level `aai.*` names re-export it +- `aai.LLMGatewayMessageParam` — `TypedDict` for a chat message; pass plain dicts (`{"role": "user", "content": "..."}`) +- `aai.LLMGatewayStream` / `aai.AsyncLLMGatewayStream` — What `create(stream=True)` returns: a chunk iterator plus `get_final_message()` ## Common patterns @@ -410,6 +415,176 @@ async with AsyncRealTimeTranscriber(RealTimeTranscriberOptions(token=token_from_ - `format_turns=True` enables punctuation/casing on confirmed end-of-turns. Toggle mid-session via `client.set_params(RealTimeSessionParameters(format_turns=True))`. - `AsyncRealTimeTranscriber` used as `async with` calls `disconnect(terminate=True)` on normal block exit and `disconnect(terminate=False)` on exception — no explicit `disconnect()` needed inside the block. +## LLM Gateway + +`aai.LLMGateway` is a client for AssemblyAI's LLM Gateway: model listing, OpenAI-shaped chat +completions routed to OpenAI/Claude/Gemini/Bedrock models, and speech understanding over an +existing transcript. It targets `llm-gateway.assemblyai.com` and rides on the same +`Client`/`AsyncClient` as the rest of the SDK. + +```python +import assemblyai as aai + +aai.settings.api_key = os.environ["ASSEMBLYAI_API_KEY"] + +gateway = aai.LLMGateway() + +for model in gateway.models.list().data: + print(model.id, model.context_length, model.pricing.global_.prompt) + +completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], +) +print(completion.choices[0].message.content) +``` + +**Messages are plain dicts** (`aai.LLMGatewayMessageParam`) — no message object to build. +`role` is `"user"`, `"assistant"`, `"system"`, or `"tool"` (a tool result, paired with +`tool_call_id`); `content` is a string or a list of content-part dicts. + +**Request params**: everything past `model`/`messages`/`stream` is a keyword argument forwarded +verbatim into the request body — `max_tokens`, `temperature`, `tools`, `tool_choice`, +`response_format`, plus the AssemblyAI extensions (`fallbacks`, `fallback_config`, +`zero_data_retention`, `transcript_id`, `reasoning`, `model_region`, `fail_fast`, `max_timeout`, +…). `fallbacks` is a list of dicts in the API's own shape: +`[{"model": "gpt-5-mini", "messages": [...]}]`. + +```python +completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + max_tokens=512, + fallbacks=[{"model": "gpt-5-mini", "messages": [{"role": "user", "content": "..."}]}], + zero_data_retention=True, + transcript_id="transcript_abc123", +) +``` + +**Streaming**: `stream=True` returns an `LLMGatewayStream` — iterate it for +`LLMGatewayCompletionChunk`s, then call `get_final_message()` for the accumulated +text/usage/`finish_reason`: +```python +stream = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + stream=True, +) +for chunk in stream: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") + +final = stream.get_final_message() +print(final.content, final.finish_reason, final.usage, final.model) +``` +The accumulator also carries `.id` and `.model`. + +Chunk parsing is verified for OpenAI-routed models only. Chunks can carry tool-call fragments +(`chunk.choices[0].delta.tool_calls`), but they are not reassembled — `get_final_message()` +covers text, usage, and `finish_reason`. For assembled tool calls use `run_tools()` or a +non-streaming `create()`. + +**Tool calling**: `run_tools()` drives the whole loop — call the model, invoke the requested +tools, feed the results back, repeat until the model answers without a tool call: +```python +def get_weather(city: str) -> dict: + return {"city": city, "temp_c": 21} + +messages = [{"role": "user", "content": "What's the weather in Denver?"}] + +completion = gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=messages, + tools=[{ + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }], + functions={"get_weather": get_weather}, + max_rounds=10, +) +print(completion.choices[0].message.content) +``` +`functions` maps tool names to Python callables, each invoked with the model's parsed arguments +as keyword arguments; a non-string return value is JSON-encoded on the way back. +`messages` is **mutated in place** with every round-trip, so it is a full transcript afterwards. +`stream=True` raises `ValueError`, an unregistered tool name raises `ValueError`, and blowing +past `max_rounds` raises `RuntimeError`. + +**Understanding**: speaker identification, translation, and custom formatting over an existing +transcript — `transcript_id` is required, there is no raw-text path. +```python +result = gateway.understanding.create( + transcript_id="transcript_abc123", + request={"speaker_identification": {"speaker_type": "name", "known_values": ["Ana", "Peter"]}}, +) +print(result.speech_understanding["response"], result.utterances) + +gateway.understanding.validate(request={"translation": {"target_languages": ["es"]}}) +``` +`request` is a plain dict — the server has no fixed schema for it. `validate()` returns `None` +when the request is valid and raises `aai.LLMGatewayError` with `.errors` populated otherwise; +it costs the same auth/balance as a real call, so it is not a free dry run. + +**Errors**: every failure raises `aai.LLMGatewayError` (an `aai.AssemblyAIError`) with +`.status_code`, `.request_id`, and `.errors` (validation messages, when the server sent any). +For a stream, only a failure before the first byte raises — a mid-stream failure truncates the +iterator. + +**Unknown server fields survive**: response models allow extras, so a field the gateway adds is +readable as an attribute without an SDK upgrade, and collections the server sends as `null` are +read as empty (`completion.choices` is always a list). + +**Settings**: `aai.settings.llm_gateway_base_url` (default +`https://llm-gateway.assemblyai.com`) and `aai.settings.llm_gateway_http_timeout` (default +`30.0`), both overridable via `ASSEMBLYAI_LLM_GATEWAY_BASE_URL` / +`ASSEMBLYAI_LLM_GATEWAY_HTTP_TIMEOUT`. + +## Asyncio LLM Gateway (`AsyncLLMGateway`) + +`aai.AsyncLLMGateway` mirrors `aai.LLMGateway` with coroutines — same resources, same request +params, same response models. Every API call is awaited; `stream=True` resolves to an +`AsyncLLMGatewayStream` you `async for` over. + +```python +import asyncio +import assemblyai as aai + +aai.settings.api_key = os.environ["ASSEMBLYAI_API_KEY"] + +async def main(): + async with aai.AsyncLLMGateway() as gateway: + completion = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], + ) + print(completion.choices[0].message.content) + + stream = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + stream=True, + ) + async for chunk in stream: + print(chunk.choices[0].delta.content or "", end="") + +asyncio.run(main()) +``` + +**Lifecycle**: the gateway owns an HTTP connection pool — use `async with`, or call +`await gateway.aclose()`. Pass `client=aai.AsyncClient(settings=aai.settings)` to share one +pool; a client you pass in stays yours to close. Sync `LLMGateway` needs no such handling: it +borrows the shared default `Client` and has no `close()`. + +**Per-gateway key**: both gateways take `api_key=...` instead of a client, and build their own +client for it. Given alongside `client`, the key wins and the passed client is left untouched. + ## Important gotchas - **Always check status**: `if transcript.status == aai.TranscriptStatus.error` — accessing `.text` on a failed transcript returns None, not an exception @@ -418,6 +593,8 @@ async with AsyncRealTimeTranscriber(RealTimeTranscriberOptions(token=token_from_ - **Streaming v3 lives in its own module**: `assemblyai.streaming.v3` (there is no other streaming API in this SDK). See the "Streaming (real-time)" section above. - **The SDK does not capture microphone audio**: bring your own capture (e.g. `pyaudio`, `sounddevice`) and stream the PCM chunks - **`transcribe_async()` returns a `concurrent.futures.Future`**, not an asyncio coroutine. In asyncio code use `aai.AsyncTranscriber` (see "Asyncio transcription" above) — or `aai.AsyncSyncTranscriber` for the sync API +- **LLM Gateway request params are `**kwargs`**: everything past `model`/`messages`/`stream` is forwarded verbatim, so there is no client-side check on the name — spell params carefully +- **`run_tools()` mutates the `messages` list you pass in**, appending every round-trip - **Timestamps are in milliseconds** throughout the SDK - **Minimum Python**: 3.8+ @@ -429,4 +606,5 @@ async with AsyncRealTimeTranscriber(RealTimeTranscriberOptions(token=token_from_ - [Full documentation](https://www.assemblyai.com/docs) - [API reference](https://www.assemblyai.com/docs/api-reference) +- [LLM Gateway docs](https://www.assemblyai.com/docs/llm-gateway/quickstart) - [llms-full.txt](https://www.assemblyai.com/docs/llms-full.txt?lang=python) (Python-filtered docs for LLMs) diff --git a/MIGRATION.md b/MIGRATION.md index 92f199f7..b7013b8c 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -6,7 +6,7 @@ Two things in one guide: what genuinely breaks moving from 0.x to 1.0.0, and the | 0.x pattern | 1.0 best practice | | --- | --- | -| `aai.Lemur(...)` | Removed. Feed `transcript.text` to the LLM of your choice | +| `aai.Lemur(...)` | Removed. Feed `transcript.text` to `aai.LLMGateway`, or to the LLM of your choice | | `from assemblyai.extras import MicrophoneStream` | Removed. Capture PCM yourself (`pyaudio`, `sounddevice`) and pass it to `stream(...)` | | `pip install "assemblyai[extras]"` | `pip install -U assemblyai` | | `StreamingClient`, `StreamingClientOptions` | `RealTimeTranscriber`, `RealTimeTranscriberOptions` | @@ -23,25 +23,39 @@ Two things in one guide: what genuinely breaks moving from 0.x to 1.0.0, and the ### LeMUR support removed -The LeMUR API and every `aai.Lemur*` name are gone from the SDK. There is no drop-in replacement in this package. +The LeMUR API and every `aai.Lemur*` name are gone from the SDK. There is no drop-in replacement. -A transcript is still just text, so the migration is to send it to whichever LLM you already use: +A transcript is still just text, so the migration is to send it to an LLM yourself: ```python -from assemblyai import TranscriptStatus -from assemblyai.prerecorded.v2 import Transcriber +import assemblyai as aai -transcript = Transcriber(api_key="YOUR_API_KEY").transcribe( +transcript = aai.Transcriber(api_key="YOUR_API_KEY").transcribe( "https://example.org/audio.wav" ) -if transcript.status == TranscriptStatus.error: +if transcript.status == aai.TranscriptStatus.error: raise RuntimeError(transcript.error) prompt = f"Summarize this call transcript:\n\n{transcript.text}" -# ...hand `prompt` to your LLM client of choice. ``` +`aai.LLMGateway` covers that in-package, across OpenAI/Claude/Gemini/Bedrock models — see [LLM Gateway Examples](README.md#llm-gateway-examples): + +```python +gateway = aai.LLMGateway(api_key="YOUR_API_KEY") + +completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": prompt}], +) +print(completion.choices[0].message.content) +``` + +Speaker identification, translation, and custom formatting run off the transcript id instead, via `gateway.understanding.create(transcript_id=transcript.id, request={...})`. + +Or hand `prompt` to whichever LLM client you already use. + ### Audio-capture extras removed `assemblyai.extras` (including `MicrophoneStream`) and the `[extras]` install option no longer exist. `pip install "assemblyai[extras]"` fails; use `pip install -U assemblyai`. diff --git a/README.md b/README.md index 4cda0e60..07f8df69 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ See [Coding agent prompts](https://www.assemblyai.com/docs/coding-agent-prompts) - [**Core Examples**](#core-examples) - [**Asyncio Examples**](#asyncio-examples) - [**Speech Understanding Examples**](#speech-understanding-examples) + - [**LLM Gateway Examples**](#llm-gateway-examples) - [**Streaming Examples**](#streaming-examples) - [**Change the default settings**](#change-the-default-settings) - [Playground](#playground) @@ -1024,6 +1025,280 @@ for result in transcript.auto_highlights.results: --- +### **LLM Gateway Examples** + +`aai.LLMGateway` is a client for AssemblyAI's LLM Gateway — model listing, OpenAI-shaped chat completions routed to OpenAI/Claude/Gemini/Bedrock models, and speech understanding over an existing transcript. `aai.AsyncLLMGateway` is the asyncio counterpart, with the same resources and every API call a coroutine. + +[Read more about the LLM Gateway.](https://www.assemblyai.com/docs/llm-gateway/quickstart) + +
+ List the available models + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +for model in gateway.models.list().data: + print(model.id, model.context_length, model.pricing.global_.prompt) +``` + +`models.list()` returns the full response envelope (`LLMGatewayModelList`) with the models on `.data`. + +
+ +
+ Create a chat completion + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], + max_tokens=512, +) + +print(completion.choices[0].message.content) +print(completion.usage.total_tokens, completion.request_id) +``` + +Messages are plain dicts (`aai.LLMGatewayMessageParam`) — `role` is `"user"`, `"assistant"`, `"system"`, or `"tool"` (a tool result, paired with `tool_call_id`), and `content` is a string or a list of content-part dicts. + +Everything past `model`/`messages`/`stream` is a keyword argument forwarded verbatim into the request body: `max_tokens`, `temperature`, `tools`, `tool_choice`, `response_format`, plus the AssemblyAI extensions (`fallbacks`, `fallback_config`, `zero_data_retention`, `transcript_id`, `reasoning`, `model_region`, `fail_fast`, `max_timeout`, …). There is no client-side check on the name, so spell them carefully. + +```python +completion = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + fallbacks=[{"model": "gpt-5-mini", "messages": [{"role": "user", "content": "..."}]}], + zero_data_retention=True, + transcript_id="transcript_abc123", +) +``` + +[Read more about model fallbacks](https://www.assemblyai.com/docs/llm-gateway/fallback) and the [available models](https://www.assemblyai.com/docs/llm-gateway/available-models). + +
+ +
+ Stream a chat completion + +`stream=True` returns an `aai.LLMGatewayStream`: iterate it for chunks, then call `get_final_message()` for the accumulated text, usage, and finish reason. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +stream = gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + stream=True, +) + +for chunk in stream: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") + +final = stream.get_final_message() +print(final.content, final.finish_reason, final.usage, final.model) +``` + +The accumulator also carries `.id` and `.model`. + +Chunk parsing is verified for OpenAI-routed models only. Chunks can carry tool-call fragments on `chunk.choices[0].delta.tool_calls`, but the SDK does not reassemble them — `get_final_message()` covers text, usage, and `finish_reason`. For assembled tool calls use `run_tools()` or a non-streaming `create()`. + +
+ +
+ Run a tool-calling loop + +`run_tools()` drives the whole loop: it calls the model, invokes the requested tools, feeds their results back, and repeats until the model replies without a tool call. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +def get_weather(city: str) -> dict: + return {"city": city, "temp_c": 21} + +messages = [{"role": "user", "content": "What's the weather in Denver?"}] + +completion = gateway.chat.completions.run_tools( + model="claude-sonnet-5", + messages=messages, + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + functions={"get_weather": get_weather}, + max_rounds=10, +) + +print(completion.choices[0].message.content) +``` + +`functions` maps tool names to the Python callables that implement them; each is called with the model's parsed arguments as keyword arguments, and a non-string return value is JSON-encoded on the way back. The `messages` list is **mutated in place** with every round-trip, so it holds the full transcript once this returns. + +`stream=True` raises `ValueError`, as does a tool call with no matching entry in `functions`; exceeding `max_rounds` raises `RuntimeError`. + +[Read more about tool calling here.](https://www.assemblyai.com/docs/llm-gateway/tool-calling) + +
+ +
+ Run speech understanding on a transcript + +Speaker identification, translation, and custom formatting run against an existing transcript, so `transcript_id` is required. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +result = gateway.understanding.create( + transcript_id="transcript_abc123", + request={ + "speaker_identification": { + "speaker_type": "name", + "known_values": ["Michel Martin", "Peter DeCarlo"], + }, + }, +) + +print(result.speech_understanding["response"]) +print(result.utterances) +``` + +`request` is a plain dict — the server has no fixed schema for it. `validate()` checks a request without running it, returning `None` when it is valid: + +```python +try: + gateway.understanding.validate(request={"translation": {"target_languages": ["es"]}}) +except aai.LLMGatewayError as error: + print(error.errors) # list of validation messages +``` + +`validate()` needs the same authorization and balance as a real call — it is not a free dry run. + +
+ +
+ Use it from asyncio (`AsyncLLMGateway`) + +```python +import asyncio +import assemblyai as aai + +aai.settings.api_key = "" + +async def main(): + async with aai.AsyncLLMGateway() as gateway: + completion = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Summarize this call."}], + ) + print(completion.choices[0].message.content) + + stream = await gateway.chat.completions.create( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "..."}], + stream=True, + ) + async for chunk in stream: + print(chunk.choices[0].delta.content or "", end="") + +asyncio.run(main()) +``` + +The gateway owns an HTTP connection pool: use `async with`, or call `await gateway.aclose()`. Pass `client=aai.AsyncClient(settings=aai.settings)` to share one pool between gateways — a client you pass in stays yours to close. Sync `LLMGateway` needs none of this; it borrows the shared default `Client`. + +Both gateways also accept `api_key=...` in place of a client and build their own client for it. Given alongside `client`, the key takes precedence and the client you passed is left untouched. + +
+ +
+ Handle errors + +Every failure raises `aai.LLMGatewayError`, a subclass of `aai.AssemblyAIError`. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +gateway = aai.LLMGateway() + +try: + gateway.chat.completions.create(model="claude-sonnet-5", messages=[{"role": "user", "content": "hi"}]) +except aai.LLMGatewayError as error: + print(error.status_code, error.request_id, error.errors, str(error)) +``` + +- `.status_code` — the HTTP status +- `.request_id` — for support correlation +- `.errors` — validation messages when the server sent any (e.g. a failed `understanding.validate()`), `None` otherwise + +For a stream, only a failure before the first byte raises; a mid-stream failure truncates the iterator. + +
+ +
+ Change the gateway base URL and timeout + +```python +import assemblyai as aai + +# default: https://llm-gateway.assemblyai.com +aai.settings.llm_gateway_base_url = "https://llm-gateway.assemblyai.com" + +# The HTTP timeout in seconds for LLM Gateway requests, default is 30.0 +aai.settings.llm_gateway_http_timeout = 60.0 +``` + +Both are also settable via the `ASSEMBLYAI_LLM_GATEWAY_BASE_URL` and `ASSEMBLYAI_LLM_GATEWAY_HTTP_TIMEOUT` environment variables. + +
+ +
+ Read fields the SDK doesn't model + +Response models keep unrecognized fields instead of dropping them, so a field the gateway adds is readable without an SDK upgrade: + +```python +completion = gateway.chat.completions.create(model="claude-sonnet-5", messages=[...]) +completion.some_new_field # reachable as an attribute +``` + +Collections the server sends as `null` are read as empty, so `completion.choices` is always a list. + +
+ +--- + ### **Streaming Examples** Real-time speech-to-text via WebSocket against the `universal-3-5-pro` model. The SDK ships two clients with identical option/event/handler surfaces — `RealTimeTranscriber` (threaded) and `AsyncRealTimeTranscriber` (asyncio). Pick whichever fits your codebase. From dfb37701987cd4606929b13979712426982b9b05 Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Wed, 9 Sep 2026 10:31:41 -0600 Subject: [PATCH 6/8] Remove helper feature to reduce v1 complexity --- CLAUDE.md | 57 +----- README.md | 57 +----- assemblyai/__init__.py | 6 - assemblyai/llm_gateway/v1/__init__.py | 4 - assemblyai/llm_gateway/v1/_base.py | 86 +-------- assemblyai/llm_gateway/v1/async_client.py | 86 +++------ assemblyai/llm_gateway/v1/client.py | 74 ++------ assemblyai/llm_gateway/v1/stream.py | 92 ---------- tests/unit/test_llm_gateway.py | 211 ---------------------- tests/unit/test_llm_gateway_async.py | 169 ----------------- 10 files changed, 58 insertions(+), 784 deletions(-) delete mode 100644 assemblyai/llm_gateway/v1/stream.py diff --git a/CLAUDE.md b/CLAUDE.md index fa02dcb2..808b0b2e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,11 +50,10 @@ aai.settings.api_key = "your-key" - `aai.SyncTranscriptResponse` — Sync result: `.text`, `.words` (`SyncWord` with `confidence` always, `start`/`end` only when `timestamps=True`), `.confidence`, `.audio_duration_ms`, `.session_id`, `.request_time_ms` - `assemblyai.streaming.v3.RealTimeTranscriber` — Real-time streaming with event-based API (threaded) - `assemblyai.streaming.v3.AsyncRealTimeTranscriber` — Asyncio-native counterpart; same options/events -- `aai.LLMGateway` — LLM Gateway client. Resources: `models.list()`, `chat.completions.create()` (incl. `stream=True`), `chat.completions.run_tools()`, `understanding.create()`/`.validate()` +- `aai.LLMGateway` — LLM Gateway client. Resources: `models.list()`, `chat.completions.create()` (incl. `stream=True`), `understanding.create()`/`.validate()` - `aai.AsyncLLMGateway` — Asyncio counterpart of `LLMGateway`. Same resources, every API call a coroutine; owns an HTTP pool (`async with` or `await aclose()`) - `assemblyai.llm_gateway.v1` — Canonical module for both gateways, matching the gateway's `/v1` API. The top-level `aai.*` names re-export it - `aai.LLMGatewayMessageParam` — `TypedDict` for a chat message; pass plain dicts (`{"role": "user", "content": "..."}`) -- `aai.LLMGatewayStream` / `aai.AsyncLLMGatewayStream` — What `create(stream=True)` returns: a chunk iterator plus `get_final_message()` ## Common patterns @@ -461,9 +460,7 @@ completion = gateway.chat.completions.create( ) ``` -**Streaming**: `stream=True` returns an `LLMGatewayStream` — iterate it for -`LLMGatewayCompletionChunk`s, then call `get_final_message()` for the accumulated -text/usage/`finish_reason`: +**Streaming**: `stream=True` returns an iterator of `LLMGatewayCompletionChunk`: ```python stream = gateway.chat.completions.create( model="claude-sonnet-5", @@ -473,49 +470,16 @@ stream = gateway.chat.completions.create( for chunk in stream: if chunk.choices[0].delta.content: print(chunk.choices[0].delta.content, end="") - -final = stream.get_final_message() -print(final.content, final.finish_reason, final.usage, final.model) ``` -The accumulator also carries `.id` and `.model`. - Chunk parsing is verified for OpenAI-routed models only. Chunks can carry tool-call fragments -(`chunk.choices[0].delta.tool_calls`), but they are not reassembled — `get_final_message()` -covers text, usage, and `finish_reason`. For assembled tool calls use `run_tools()` or a -non-streaming `create()`. - -**Tool calling**: `run_tools()` drives the whole loop — call the model, invoke the requested -tools, feed the results back, repeat until the model answers without a tool call: -```python -def get_weather(city: str) -> dict: - return {"city": city, "temp_c": 21} - -messages = [{"role": "user", "content": "What's the weather in Denver?"}] +(`chunk.choices[0].delta.tool_calls`), but this SDK does not reassemble them or accumulate the +streamed text for you — collect chunks yourself, or use a non-streaming `create()`. -completion = gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=messages, - tools=[{ - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - }, - }], - functions={"get_weather": get_weather}, - max_rounds=10, -) -print(completion.choices[0].message.content) -``` -`functions` maps tool names to Python callables, each invoked with the model's parsed arguments -as keyword arguments; a non-string return value is JSON-encoded on the way back. -`messages` is **mutated in place** with every round-trip, so it is a full transcript afterwards. -`stream=True` raises `ValueError`, an unregistered tool name raises `ValueError`, and blowing -past `max_rounds` raises `RuntimeError`. +There is no built-in tool-calling loop: `tools`/`tool_choice` are forwarded like any other +request param (see "Request params" above), and replaying a tool result is a plain +`{"role": "tool", "tool_call_id": ..., "content": ...}` message per +https://www.assemblyai.com/docs/llm-gateway/tool-calling — drive the loop yourself around +`create()`. **Understanding**: speaker identification, translation, and custom formatting over an existing transcript — `transcript_id` is required, there is no raw-text path. @@ -550,7 +514,7 @@ read as empty (`completion.choices` is always a list). `aai.AsyncLLMGateway` mirrors `aai.LLMGateway` with coroutines — same resources, same request params, same response models. Every API call is awaited; `stream=True` resolves to an -`AsyncLLMGatewayStream` you `async for` over. +async iterator of `LLMGatewayCompletionChunk` you `async for` over. ```python import asyncio @@ -594,7 +558,6 @@ client for it. Given alongside `client`, the key wins and the passed client is l - **The SDK does not capture microphone audio**: bring your own capture (e.g. `pyaudio`, `sounddevice`) and stream the PCM chunks - **`transcribe_async()` returns a `concurrent.futures.Future`**, not an asyncio coroutine. In asyncio code use `aai.AsyncTranscriber` (see "Asyncio transcription" above) — or `aai.AsyncSyncTranscriber` for the sync API - **LLM Gateway request params are `**kwargs`**: everything past `model`/`messages`/`stream` is forwarded verbatim, so there is no client-side check on the name — spell params carefully -- **`run_tools()` mutates the `messages` list you pass in**, appending every round-trip - **Timestamps are in milliseconds** throughout the SDK - **Minimum Python**: 3.8+ diff --git a/README.md b/README.md index e401c698..bf3ac808 100644 --- a/README.md +++ b/README.md @@ -1090,7 +1090,7 @@ completion = gateway.chat.completions.create(
Stream a chat completion -`stream=True` returns an `aai.LLMGatewayStream`: iterate it for chunks, then call `get_final_message()` for the accumulated text, usage, and finish reason. +`stream=True` returns an iterator of `LLMGatewayCompletionChunk`. ```python import assemblyai as aai @@ -1108,62 +1108,11 @@ stream = gateway.chat.completions.create( for chunk in stream: if chunk.choices[0].delta.content: print(chunk.choices[0].delta.content, end="") - -final = stream.get_final_message() -print(final.content, final.finish_reason, final.usage, final.model) ``` -The accumulator also carries `.id` and `.model`. - -Chunk parsing is verified for OpenAI-routed models only. Chunks can carry tool-call fragments on `chunk.choices[0].delta.tool_calls`, but the SDK does not reassemble them — `get_final_message()` covers text, usage, and `finish_reason`. For assembled tool calls use `run_tools()` or a non-streaming `create()`. - -
- -
- Run a tool-calling loop - -`run_tools()` drives the whole loop: it calls the model, invokes the requested tools, feeds their results back, and repeats until the model replies without a tool call. - -```python -import assemblyai as aai - -aai.settings.api_key = "" - -gateway = aai.LLMGateway() - -def get_weather(city: str) -> dict: - return {"city": city, "temp_c": 21} - -messages = [{"role": "user", "content": "What's the weather in Denver?"}] - -completion = gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=messages, - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - }, - } - ], - functions={"get_weather": get_weather}, - max_rounds=10, -) - -print(completion.choices[0].message.content) -``` - -`functions` maps tool names to the Python callables that implement them; each is called with the model's parsed arguments as keyword arguments, and a non-string return value is JSON-encoded on the way back. The `messages` list is **mutated in place** with every round-trip, so it holds the full transcript once this returns. - -`stream=True` raises `ValueError`, as does a tool call with no matching entry in `functions`; exceeding `max_rounds` raises `RuntimeError`. +Chunk parsing is verified for OpenAI-routed models only. Chunks can carry tool-call fragments on `chunk.choices[0].delta.tool_calls`, but this SDK does not reassemble them or accumulate the streamed text for you — collect chunks yourself, or use a non-streaming `create()`. -[Read more about tool calling here.](https://www.assemblyai.com/docs/llm-gateway/tool-calling) +There is no built-in tool-calling loop: `tools`/`tool_choice` are forwarded like any other request param, and a tool result is a plain `{"role": "tool", "tool_call_id": ..., "content": ...}` message — drive the loop yourself around `create()`. [Read more about tool calling here.](https://www.assemblyai.com/docs/llm-gateway/tool-calling)
diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index 9c4a40cb..5596bd56 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -3,7 +3,6 @@ from .client import Client from .llm_gateway.v1 import ( AsyncLLMGateway, - AsyncLLMGatewayStream, LLMGateway, LLMGatewayChatCompletion, LLMGatewayChoice, @@ -14,8 +13,6 @@ LLMGatewayModel, LLMGatewayModelList, LLMGatewayResponseMessage, - LLMGatewayStream, - LLMGatewayStreamAccumulator, LLMGatewayToolCall, LLMGatewayUnderstandingResponse, LLMGatewayUsage, @@ -117,7 +114,6 @@ def __getattr__(name: str): "AssemblyAIError", "AsyncClient", "AsyncLLMGateway", - "AsyncLLMGatewayStream", "AsyncSyncTranscriber", "AsyncTranscriber", "AsyncTranscript", @@ -152,8 +148,6 @@ def __getattr__(name: str): "LLMGatewayModel", "LLMGatewayModelList", "LLMGatewayResponseMessage", - "LLMGatewayStream", - "LLMGatewayStreamAccumulator", "LLMGatewayToolCall", "LLMGatewayUnderstandingResponse", "LLMGatewayUsage", diff --git a/assemblyai/llm_gateway/v1/__init__.py b/assemblyai/llm_gateway/v1/__init__.py index 504cfe5a..ef27485e 100644 --- a/assemblyai/llm_gateway/v1/__init__.py +++ b/assemblyai/llm_gateway/v1/__init__.py @@ -24,11 +24,9 @@ LLMGatewayUsage, ) from .params import LLMGatewayMessageParam -from .stream import AsyncLLMGatewayStream, LLMGatewayStream, LLMGatewayStreamAccumulator __all__ = [ "AsyncLLMGateway", - "AsyncLLMGatewayStream", "LLMGateway", "LLMGatewayChatCompletion", "LLMGatewayChoice", @@ -40,8 +38,6 @@ "LLMGatewayModel", "LLMGatewayModelList", "LLMGatewayResponseMessage", - "LLMGatewayStream", - "LLMGatewayStreamAccumulator", "LLMGatewayToolCall", "LLMGatewayUnderstandingResponse", "LLMGatewayUsage", diff --git a/assemblyai/llm_gateway/v1/_base.py b/assemblyai/llm_gateway/v1/_base.py index 9cbe0031..8775b431 100644 --- a/assemblyai/llm_gateway/v1/_base.py +++ b/assemblyai/llm_gateway/v1/_base.py @@ -1,14 +1,12 @@ """Logic shared by `client.py` and `async_client.py`. -Holds the request body shapes and the tool-calling loop's core, so a -wire-contract or tool-loop change lands in one place. Everything here is -I/O-free — the sync and async clients supply the HTTP calls themselves. +Holds the request body shapes, so a wire-contract change lands in one place. +Everything here is I/O-free — the sync and async clients supply the HTTP +calls themselves. """ -import json -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Dict, List, Optional -from . import models from .params import LLMGatewayMessageParam @@ -75,79 +73,3 @@ def validate_body( body["transcript_id"] = transcript_id return body - - -def reject_stream(kwargs: Dict[str, Any]) -> None: - """Raises `ValueError` if `run_tools()` was passed `stream=True`.""" - if kwargs.get("stream"): - raise ValueError("run_tools() does not support stream=True") - - -def apply_tool_round( - completion: models.LLMGatewayChatCompletion, - messages: List[LLMGatewayMessageParam], - functions: Dict[str, Callable[..., Any]], -) -> bool: - """ - Plays one round of the tool-calling loop into `messages`, in place. - - Appends the model's reply, and — when it requested tool calls — invokes each - one from `functions` and appends its result. - - Returns: True once the model has replied without requesting a tool call, - i.e. the loop is finished. - - Raises: - ValueError: if the model requests a tool with no matching entry in - `functions`. - """ - message = completion.choices[0].message - - if not message.tool_calls: - messages.append({"role": "assistant", "content": message.content}) - return True - - messages.append( - { - "role": "assistant", - "content": message.content, - "tool_calls": [ - { - "id": tool_call.id, - "type": tool_call.type, - "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments - if isinstance(tool_call.function.arguments, str) - else json.dumps(tool_call.function.arguments), - }, - } - for tool_call in message.tool_calls - ], - } - ) - - for tool_call in message.tool_calls: - function = functions.get(tool_call.function.name) - if function is None: - raise ValueError( - f"No function registered for tool call {tool_call.function.name!r}" - ) - - arguments = tool_call.function.arguments - if isinstance(arguments, str): - arguments = json.loads(arguments) - - result = function(**arguments) - output = result if isinstance(result, str) else json.dumps(result) - - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "content": output, - "name": tool_call.function.name, - } - ) - - return False diff --git a/assemblyai/llm_gateway/v1/async_client.py b/assemblyai/llm_gateway/v1/async_client.py index 3786deee..064397f7 100644 --- a/assemblyai/llm_gateway/v1/async_client.py +++ b/assemblyai/llm_gateway/v1/async_client.py @@ -1,14 +1,23 @@ from __future__ import annotations from types import TracebackType -from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, overload +from typing import ( + Any, + AsyncIterator, + Dict, + List, + Literal, + Optional, + Type, + Union, + overload, +) from typing_extensions import Self from ... import async_client as _async_client from . import _base, async_api, models from .params import LLMGatewayMessageParam -from .stream import AsyncLLMGatewayStream class AsyncModelsResource(_base._BaseResource): @@ -47,7 +56,7 @@ async def create( messages: List[LLMGatewayMessageParam], stream: Literal[True], **kwargs: Any, - ) -> AsyncLLMGatewayStream: ... + ) -> AsyncIterator[models.LLMGatewayCompletionChunk]: ... @overload async def create( @@ -57,7 +66,9 @@ async def create( messages: List[LLMGatewayMessageParam], stream: bool, **kwargs: Any, - ) -> Union[models.LLMGatewayChatCompletion, AsyncLLMGatewayStream]: ... + ) -> Union[ + models.LLMGatewayChatCompletion, AsyncIterator[models.LLMGatewayCompletionChunk] + ]: ... async def create( self, @@ -66,7 +77,9 @@ async def create( messages: List[LLMGatewayMessageParam], stream: bool = False, **kwargs: Any, - ) -> Union[models.LLMGatewayChatCompletion, AsyncLLMGatewayStream]: + ) -> Union[ + models.LLMGatewayChatCompletion, AsyncIterator[models.LLMGatewayCompletionChunk] + ]: """ Creates a chat completion. @@ -78,12 +91,10 @@ async def create( `fallbacks` is a list of dicts, each with the API's expected shape, e.g. `[{"model": "gpt-4o", "messages": [...]}]`. - With `stream=True`, resolves to an `AsyncLLMGatewayStream` — iterate it - for `LLMGatewayCompletionChunk`s exactly like the raw stream it - replaces (`async for chunk in await gateway.chat.completions.create(..., stream=True)`), - then call `.get_final_message()` for the assembled text/usage — - instead of a single `LLMGatewayChatCompletion`. Chunk parsing is - verified for OpenAI-routed models only. + With `stream=True`, resolves to an async iterator of + `LLMGatewayCompletionChunk` instead of a single + `LLMGatewayChatCompletion`. Chunk parsing is verified for + OpenAI-routed models only. Raises: `LLMGatewayError` if the request fails (for a stream, only if it fails before any bytes are received — a failure mid-stream @@ -94,12 +105,10 @@ async def create( ) if stream: - return AsyncLLMGatewayStream( - async_api.stream_completion( - self._gateway.client.http_client, - body=body, - **self._http_options, - ) + return async_api.stream_completion( + self._gateway.client.http_client, + body=body, + **self._http_options, ) return await async_api.create_completion( @@ -108,49 +117,6 @@ async def create( **self._http_options, ) - async def run_tools( - self, - *, - model: str, - messages: List[LLMGatewayMessageParam], - tools: List[Dict[str, Any]], - functions: Dict[str, Callable[..., Any]], - max_rounds: int = 10, - **kwargs: Any, - ) -> models.LLMGatewayChatCompletion: - """ - Runs the tool-calling loop to completion: calls the model, invokes any - requested tools from `functions`, feeds their results back, and repeats - until the model responds without requesting a tool call. - - `messages` is mutated in place with every tool round-trip and the final - assistant reply, so it's a complete transcript once this returns. Tool - round-trips are replayed as `{"role": "assistant", "tool_calls": [...]}` - followed by one `{"role": "tool", "tool_call_id": ..., "content": ...}` - per call, per https://www.assemblyai.com/docs/llm-gateway/tool-calling. - - `functions` maps tool names (as declared in `tools`) to the Python - callables that implement them; each is called with the model's parsed - arguments as keyword arguments, and its return value is sent back to - the model (JSON-encoded, unless it's already a string). - - Raises: - ValueError: if `stream=True` is passed (unsupported), or the model - requests a tool with no matching entry in `functions`. - RuntimeError: if the loop doesn't finish within `max_rounds`. - LLMGatewayError: if any underlying `create()` call fails. - """ - _base.reject_stream(kwargs) - - for _ in range(max_rounds): - completion: models.LLMGatewayChatCompletion = await self.create( - model=model, messages=messages, tools=tools, **kwargs - ) - if _base.apply_tool_round(completion, messages, functions): - return completion - - raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") - class AsyncChatResource: """The `chat` resource of `AsyncLLMGateway`.""" diff --git a/assemblyai/llm_gateway/v1/client.py b/assemblyai/llm_gateway/v1/client.py index c267f4fc..76224563 100644 --- a/assemblyai/llm_gateway/v1/client.py +++ b/assemblyai/llm_gateway/v1/client.py @@ -1,9 +1,8 @@ -from typing import Any, Callable, Dict, List, Literal, Optional, Union, overload +from typing import Any, Dict, Iterator, List, Literal, Optional, Union, overload from ... import client as _client from . import _base, api, models from .params import LLMGatewayMessageParam -from .stream import LLMGatewayStream class ModelsResource(_base._BaseResource): @@ -42,7 +41,7 @@ def create( messages: List[LLMGatewayMessageParam], stream: Literal[True], **kwargs: Any, - ) -> LLMGatewayStream: ... + ) -> Iterator[models.LLMGatewayCompletionChunk]: ... @overload def create( @@ -52,7 +51,9 @@ def create( messages: List[LLMGatewayMessageParam], stream: bool, **kwargs: Any, - ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: ... + ) -> Union[ + models.LLMGatewayChatCompletion, Iterator[models.LLMGatewayCompletionChunk] + ]: ... def create( self, @@ -61,7 +62,9 @@ def create( messages: List[LLMGatewayMessageParam], stream: bool = False, **kwargs: Any, - ) -> Union[models.LLMGatewayChatCompletion, LLMGatewayStream]: + ) -> Union[ + models.LLMGatewayChatCompletion, Iterator[models.LLMGatewayCompletionChunk] + ]: """ Creates a chat completion. @@ -73,11 +76,9 @@ def create( `fallbacks` is a list of dicts, each with the API's expected shape, e.g. `[{"model": "gpt-4o", "messages": [...]}]`. - With `stream=True`, returns a `LLMGatewayStream` — iterate it for - `LLMGatewayCompletionChunk`s exactly like the raw stream it replaces, - then call `.get_final_message()` for the assembled text/usage — instead - of a single `LLMGatewayChatCompletion`. Chunk parsing is verified for - OpenAI-routed models only. + With `stream=True`, returns an iterator of `LLMGatewayCompletionChunk` + instead of a single `LLMGatewayChatCompletion`. Chunk parsing is + verified for OpenAI-routed models only. Raises: `LLMGatewayError` if the request fails (for a stream, only if it fails before any bytes are received — a failure mid-stream @@ -88,12 +89,10 @@ def create( ) if stream: - return LLMGatewayStream( - api.stream_completion( - self._gateway.client.http_client, - body=body, - **self._http_options, - ) + return api.stream_completion( + self._gateway.client.http_client, + body=body, + **self._http_options, ) return api.create_completion( @@ -102,49 +101,6 @@ def create( **self._http_options, ) - def run_tools( - self, - *, - model: str, - messages: List[LLMGatewayMessageParam], - tools: List[Dict[str, Any]], - functions: Dict[str, Callable[..., Any]], - max_rounds: int = 10, - **kwargs: Any, - ) -> models.LLMGatewayChatCompletion: - """ - Runs the tool-calling loop to completion: calls the model, invokes any - requested tools from `functions`, feeds their results back, and repeats - until the model responds without requesting a tool call. - - `messages` is mutated in place with every tool round-trip and the final - assistant reply, so it's a complete transcript once this returns. Tool - round-trips are replayed as `{"role": "assistant", "tool_calls": [...]}` - followed by one `{"role": "tool", "tool_call_id": ..., "content": ...}` - per call, per https://www.assemblyai.com/docs/llm-gateway/tool-calling. - - `functions` maps tool names (as declared in `tools`) to the Python - callables that implement them; each is called with the model's parsed - arguments as keyword arguments, and its return value is sent back to - the model (JSON-encoded, unless it's already a string). - - Raises: - ValueError: if `stream=True` is passed (unsupported), or the model - requests a tool with no matching entry in `functions`. - RuntimeError: if the loop doesn't finish within `max_rounds`. - LLMGatewayError: if any underlying `create()` call fails. - """ - _base.reject_stream(kwargs) - - for _ in range(max_rounds): - completion: models.LLMGatewayChatCompletion = self.create( - model=model, messages=messages, tools=tools, **kwargs - ) - if _base.apply_tool_round(completion, messages, functions): - return completion - - raise RuntimeError(f"Tool loop did not finish within max_rounds={max_rounds}") - class ChatResource: """The `chat` resource of `LLMGateway`.""" diff --git a/assemblyai/llm_gateway/v1/stream.py b/assemblyai/llm_gateway/v1/stream.py deleted file mode 100644 index 6db07dc4..00000000 --- a/assemblyai/llm_gateway/v1/stream.py +++ /dev/null @@ -1,92 +0,0 @@ -from typing import AsyncIterator, Iterator, Optional - -from typing_extensions import Self - -from . import models - - -class LLMGatewayStreamAccumulator: - """The message assembled from a `LLMGatewayStream`'s chunks, as seen so far.""" - - def __init__(self) -> None: - self.id: Optional[str] = None - self.model: Optional[str] = None - self.content: str = "" - self.finish_reason: Optional[str] = None - self.usage: Optional[models.LLMGatewayUsage] = None - - def _absorb(self, chunk: models.LLMGatewayCompletionChunk) -> None: - if self.id is None: - self.id = chunk.id - if self.model is None: - self.model = chunk.model - for choice in chunk.choices: - if choice.delta.content: - self.content += choice.delta.content - if choice.finish_reason is not None: - self.finish_reason = choice.finish_reason - if chunk.usage is not None: - self.usage = chunk.usage - - -class LLMGatewayStream: - """ - Wraps a streamed chat completion's chunk iterator, accumulating a final - assembled message as it's iterated. - - Iterate it exactly like the raw chunk iterator it replaces: - `for chunk in stream: ...`. Once exhausted, call `get_final_message()` for - the assembled text/usage/finish_reason. - - Verified for OpenAI-routed models only, like raw chunk parsing. Chunks can - carry tool calls (`chunk.choices[0].delta.tool_calls`), but they arrive as - fragments and are not reassembled here — `get_final_message()` covers text, - usage and `finish_reason` only. For assembled tool calls, use `run_tools()` - or a non-streaming `create()`. - """ - - def __init__(self, chunks: Iterator[models.LLMGatewayCompletionChunk]) -> None: - self._chunks = chunks - self._accumulator = LLMGatewayStreamAccumulator() - - def __iter__(self) -> Self: - return self - - def __next__(self) -> models.LLMGatewayCompletionChunk: - chunk = next(self._chunks) - self._accumulator._absorb(chunk) - return chunk - - def get_final_message(self) -> LLMGatewayStreamAccumulator: - """ - Returns the message assembled from chunks seen so far. - - Call after exhausting the iterator for a complete result; calling it - earlier just returns whatever has streamed in up to that point. - """ - return self._accumulator - - -class AsyncLLMGatewayStream: - """The asyncio counterpart of `LLMGatewayStream`.""" - - def __init__(self, chunks: AsyncIterator[models.LLMGatewayCompletionChunk]) -> None: - self._chunks = chunks - self._accumulator = LLMGatewayStreamAccumulator() - - def __aiter__(self) -> Self: - return self - - async def __anext__(self) -> models.LLMGatewayCompletionChunk: - chunk = await self._chunks.__anext__() - self._accumulator._absorb(chunk) - return chunk - - def get_final_message(self) -> LLMGatewayStreamAccumulator: - """ - Returns the message assembled from chunks seen so far. - - Call after exhausting the iterator for a complete result; calling it - earlier just returns whatever has streamed in up to that point. - """ - return self._accumulator diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py index 78247489..31177da4 100644 --- a/tests/unit/test_llm_gateway.py +++ b/tests/unit/test_llm_gateway.py @@ -549,217 +549,6 @@ def test_create_without_fallbacks_omits_field(httpx_mock: HTTPXMock): assert "fallbacks" not in payload -_TOOL_CALL_RESPONSE = { - "request_id": "req-tool-1", - "choices": [ - { - "index": 0, - "finish_reason": "tool_calls", - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Paris"}', - }, - } - ], - }, - } - ], - "usage": _COMPLETION_RESPONSE["usage"], -} - - -def test_run_tools_completes_after_tool_call(httpx_mock: HTTPXMock): - # Given a mocked completions endpoint that first requests a tool call, then answers - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_COMPLETION_RESPONSE, - ) - - messages = [{"role": "user", "content": "What's the weather in Paris?"}] - calls = [] - - def get_weather(city): - calls.append(city) - return {"forecast": "sunny"} - - # When running the tool loop - result = aai.LLMGateway().chat.completions.run_tools( - model="claude-sonnet-5", - messages=messages, - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={"get_weather": get_weather}, - ) - - # Then the tool was invoked with parsed arguments and the final completion is returned - assert calls == ["Paris"] - assert isinstance(result, aai.LLMGatewayChatCompletion) - assert result.request_id == "req-completion-1" - - -def test_run_tools_sends_chat_message_shaped_items(httpx_mock: HTTPXMock): - # Given the same two-round mock - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_COMPLETION_RESPONSE, - ) - - messages = [{"role": "user", "content": "What's the weather in Paris?"}] - - # When running the tool loop - aai.LLMGateway().chat.completions.run_tools( - model="claude-sonnet-5", - messages=messages, - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={"get_weather": lambda city: {"forecast": "sunny"}}, - ) - - # Then the follow-up request replays the tool call and its output as plain chat messages - import json as _json - - second_request_body = _json.loads(httpx_mock.get_requests()[1].read()) - replayed = second_request_body["messages"][1:] - assert replayed[0]["role"] == "assistant" - assert replayed[0]["tool_calls"][0]["id"] == "call_1" - assert replayed[0]["tool_calls"][0]["function"]["name"] == "get_weather" - assert replayed[1]["role"] == "tool" - assert replayed[1]["tool_call_id"] == "call_1" - assert _json.loads(replayed[1]["content"]) == {"forecast": "sunny"} - - # And the caller's messages list now holds the complete transcript - assert messages[-1] == {"role": "assistant", "content": "Hello there."} - - -def test_run_tools_raises_on_unknown_tool(httpx_mock: HTTPXMock): - # Given a mocked response requesting an unregistered tool - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - - # When running the tool loop with no matching function - with pytest.raises(ValueError, match="get_weather"): - aai.LLMGateway().chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={}, - ) - - -def test_run_tools_raises_after_max_rounds(httpx_mock: HTTPXMock): - # Given a mock that always requests the same tool call - for _ in range(3): - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - - # When running the tool loop with a max_rounds it can't finish within - with pytest.raises(RuntimeError, match="max_rounds"): - aai.LLMGateway().chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={"get_weather": lambda city: "sunny"}, - max_rounds=3, - ) - - -def test_run_tools_rejects_streaming(): - # When running the tool loop with stream=True - with pytest.raises(ValueError, match="stream"): - aai.LLMGateway().chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[], - functions={}, - stream=True, - ) - - -def test_chat_completions_stream_get_final_message(httpx_mock: HTTPXMock): - # Given a mocked streaming completions endpoint - chunk1 = { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-sonnet-5", - "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], - } - chunk2 = { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-sonnet-5", - "choices": [ - {"index": 0, "delta": {"content": " there."}, "finish_reason": "stop"} - ], - "usage": _COMPLETION_RESPONSE["usage"], - } - import json as _json - - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - stream=IteratorStream( - [ - f"data: {_json.dumps(chunk1)}\n\n".encode(), - f"data: {_json.dumps(chunk2)}\n\n".encode(), - b"data: [DONE]\n\n", - ] - ), - headers={"content-type": "text/event-stream"}, - ) - - # When consuming the stream and asking for the final assembled message - stream = aai.LLMGateway().chat.completions.create( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - stream=True, - ) - chunks = list(stream) - final = stream.get_final_message() - - # Then it still yields the raw chunks (unchanged behavior)... - assert len(chunks) == 2 - assert all(isinstance(c, aai.LLMGatewayCompletionChunk) for c in chunks) - - # ...and the accumulator has the assembled text, usage, and finish_reason - assert final.content == "Hello there." - assert final.finish_reason == "stop" - assert final.usage.total_tokens == 15 - assert final.id == "chatcmpl-1" - assert final.model == "claude-sonnet-5" - - # The gateway is written in Go, where a nil slice/map marshals as `null` rather # than `[]`/`{}`. An explicit null bypasses a pydantic field default, so each # collection on the response path has to tolerate it. diff --git a/tests/unit/test_llm_gateway_async.py b/tests/unit/test_llm_gateway_async.py index 939681a1..fef77f3c 100644 --- a/tests/unit/test_llm_gateway_async.py +++ b/tests/unit/test_llm_gateway_async.py @@ -298,175 +298,6 @@ async def test_create_fallbacks_sent_as_given(httpx_mock: HTTPXMock): ] -_TOOL_CALL_RESPONSE = { - "request_id": "req-tool-1", - "choices": [ - { - "index": 0, - "finish_reason": "tool_calls", - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Paris"}', - }, - } - ], - }, - } - ], - "usage": _COMPLETION_RESPONSE["usage"], -} - - -async def test_run_tools_completes_after_tool_call(httpx_mock: HTTPXMock): - # Given a mocked completions endpoint that first requests a tool call, then answers - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_COMPLETION_RESPONSE, - ) - - messages = [{"role": "user", "content": "What's the weather in Paris?"}] - calls = [] - - def get_weather(city): - calls.append(city) - return {"forecast": "sunny"} - - # When running the tool loop - async with aai.AsyncLLMGateway() as gateway: - result = await gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=messages, - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={"get_weather": get_weather}, - ) - - # Then the tool was invoked with parsed arguments and the final completion is returned - assert calls == ["Paris"] - assert isinstance(result, aai.LLMGatewayChatCompletion) - assert result.request_id == "req-completion-1" - assert messages[-1] == {"role": "assistant", "content": "Hello there."} - - -async def test_run_tools_raises_on_unknown_tool(httpx_mock: HTTPXMock): - # Given a mocked response requesting an unregistered tool - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - - # When running the tool loop with no matching function - async with aai.AsyncLLMGateway() as gateway: - with pytest.raises(ValueError, match="get_weather"): - await gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={}, - ) - - -async def test_run_tools_raises_after_max_rounds(httpx_mock: HTTPXMock): - # Given a mock that always requests the same tool call - for _ in range(2): - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - json=_TOOL_CALL_RESPONSE, - ) - - # When running the tool loop with a max_rounds it can't finish within - async with aai.AsyncLLMGateway() as gateway: - with pytest.raises(RuntimeError, match="max_rounds"): - await gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[{"type": "function", "function": {"name": "get_weather"}}], - functions={"get_weather": lambda city: "sunny"}, - max_rounds=2, - ) - - -async def test_run_tools_rejects_streaming(): - # When running the tool loop with stream=True - async with aai.AsyncLLMGateway() as gateway: - with pytest.raises(ValueError, match="stream"): - await gateway.chat.completions.run_tools( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - tools=[], - functions={}, - stream=True, - ) - - -async def test_chat_completions_stream_get_final_message(httpx_mock: HTTPXMock): - # Given a mocked streaming completions endpoint - chunk1 = { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-sonnet-5", - "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}], - } - chunk2 = { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-sonnet-5", - "choices": [ - {"index": 0, "delta": {"content": " there."}, "finish_reason": "stop"} - ], - "usage": _COMPLETION_RESPONSE["usage"], - } - httpx_mock.add_response( - url=COMPLETIONS_URL, - method="POST", - status_code=httpx.codes.OK, - stream=IteratorStream( - [ - f"data: {json.dumps(chunk1)}\n\n".encode(), - f"data: {json.dumps(chunk2)}\n\n".encode(), - b"data: [DONE]\n\n", - ] - ), - headers={"content-type": "text/event-stream"}, - ) - - # When consuming the stream and asking for the final assembled message - async with aai.AsyncLLMGateway() as gateway: - stream = await gateway.chat.completions.create( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "Hi"}], - stream=True, - ) - chunks = [c async for c in stream] - final = stream.get_final_message() - - # Then it still yields the raw chunks, and the accumulator has the assembled result - assert len(chunks) == 2 - assert final.content == "Hello there." - assert final.finish_reason == "stop" - assert final.usage.total_tokens == 15 - - async def test_aclose_closes_owned_client_only(): # Given a gateway with its own client, and one sharing a caller-owned client async with aai.AsyncLLMGateway() as owned: From 3f35d6250527c539e2be8f2757e2ea58303756bb Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Wed, 9 Sep 2026 11:12:17 -0600 Subject: [PATCH 7/8] Version bump. Field annotations --- assemblyai/__version__.py | 2 +- assemblyai/llm_gateway/v1/async_client.py | 42 +++++++++++++++++++++++ assemblyai/llm_gateway/v1/client.py | 42 +++++++++++++++++++++++ assemblyai/llm_gateway/v1/models.py | 9 +++-- tests/unit/test_llm_gateway.py | 10 +++--- tests/unit/test_llm_gateway_async.py | 4 +-- 6 files changed, 99 insertions(+), 10 deletions(-) diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 67bc602a..3e8d9f94 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.3.0" +__version__ = "1.4.0" diff --git a/assemblyai/llm_gateway/v1/async_client.py b/assemblyai/llm_gateway/v1/async_client.py index 064397f7..92c41ac5 100644 --- a/assemblyai/llm_gateway/v1/async_client.py +++ b/assemblyai/llm_gateway/v1/async_client.py @@ -38,6 +38,9 @@ async def list(self) -> models.LLMGatewayModelList: class AsyncCompletionsResource(_base._BaseResource): """The `chat.completions` resource of `AsyncLLMGateway`.""" + # Named keyword annotations below are for IDE/mypy discoverability only. They add no + # runtime validation: unnamed fields still flow through **kwargs untouched, and the + # gateway remains the sole source of truth for what a request accepts. @overload async def create( self, @@ -45,6 +48,19 @@ async def create( model: str, messages: List[LLMGatewayMessageParam], stream: Literal[False] = False, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> models.LLMGatewayChatCompletion: ... @@ -55,6 +71,19 @@ async def create( model: str, messages: List[LLMGatewayMessageParam], stream: Literal[True], + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> AsyncIterator[models.LLMGatewayCompletionChunk]: ... @@ -65,6 +94,19 @@ async def create( model: str, messages: List[LLMGatewayMessageParam], stream: bool, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> Union[ models.LLMGatewayChatCompletion, AsyncIterator[models.LLMGatewayCompletionChunk] diff --git a/assemblyai/llm_gateway/v1/client.py b/assemblyai/llm_gateway/v1/client.py index 76224563..f5d8f5bf 100644 --- a/assemblyai/llm_gateway/v1/client.py +++ b/assemblyai/llm_gateway/v1/client.py @@ -23,6 +23,9 @@ def list(self) -> models.LLMGatewayModelList: class CompletionsResource(_base._BaseResource): """The `chat.completions` resource of `LLMGateway`.""" + # Named keyword annotations below are for IDE/mypy discoverability only. They add no + # runtime validation: unnamed fields still flow through **kwargs untouched, and the + # gateway remains the sole source of truth for what a request accepts. @overload def create( self, @@ -30,6 +33,19 @@ def create( model: str, messages: List[LLMGatewayMessageParam], stream: Literal[False] = False, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> models.LLMGatewayChatCompletion: ... @@ -40,6 +56,19 @@ def create( model: str, messages: List[LLMGatewayMessageParam], stream: Literal[True], + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> Iterator[models.LLMGatewayCompletionChunk]: ... @@ -50,6 +79,19 @@ def create( model: str, messages: List[LLMGatewayMessageParam], stream: bool, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, + response_format: Optional[Dict[str, Any]] = None, + fallbacks: Optional[List[Dict[str, Any]]] = None, + fallback_config: Optional[Dict[str, Any]] = None, + zero_data_retention: Optional[bool] = None, + transcript_id: Optional[str] = None, + reasoning: Optional[Dict[str, Any]] = None, + model_region: Optional[str] = None, + fail_fast: Optional[bool] = None, + max_timeout: Optional[float] = None, **kwargs: Any, ) -> Union[ models.LLMGatewayChatCompletion, Iterator[models.LLMGatewayCompletionChunk] diff --git a/assemblyai/llm_gateway/v1/models.py b/assemblyai/llm_gateway/v1/models.py index c13382c2..8f1a3e60 100644 --- a/assemblyai/llm_gateway/v1/models.py +++ b/assemblyai/llm_gateway/v1/models.py @@ -58,6 +58,11 @@ class LLMGatewayTopProvider(_LLMGatewayResponseModel): max_completion_tokens: int +class LLMGatewayProvider(_LLMGatewayResponseModel): + id: str + name: str + + class LLMGatewayModel(_LLMGatewayResponseModel): """A model available through the LLM Gateway, as returned by `GET /v1/models`.""" @@ -74,8 +79,8 @@ class LLMGatewayModel(_LLMGatewayResponseModel): "Unix timestamp of the model's retirement date, or 0 if not retiring" available_regions: List[str] = Field(default_factory=list) - providers: List[str] = Field(default_factory=list) - default_provider: str + providers: List[LLMGatewayProvider] = Field(default_factory=list) + default_provider: LLMGatewayProvider # The gateway is Go: a nil slice marshals as `null`, not `[]`, and an # explicit null bypasses the field default. diff --git a/tests/unit/test_llm_gateway.py b/tests/unit/test_llm_gateway.py index 31177da4..93961923 100644 --- a/tests/unit/test_llm_gateway.py +++ b/tests/unit/test_llm_gateway.py @@ -41,8 +41,8 @@ "creator": "anthropic", "retirement_date": 0, "available_regions": ["us", "eu"], - "providers": ["anthropic"], - "default_provider": "anthropic", + "providers": [{"id": "anthropic", "name": "Anthropic"}], + "default_provider": {"id": "anthropic", "name": "Anthropic"}, }, { "id": "gpt-4o", @@ -66,8 +66,8 @@ "creator": "openai", "retirement_date": 0, "available_regions": ["global"], - "providers": ["openai"], - "default_provider": "openai", + "providers": [{"id": "openai", "name": "Open AI"}], + "default_provider": {"id": "openai", "name": "Open AI"}, }, ] } @@ -595,7 +595,7 @@ def test_list_models_tolerates_null_model_collections(httpx_mock: HTTPXMock): "retirement_date": 0, "available_regions": None, "providers": None, - "default_provider": "bedrock", + "default_provider": {"id": "bedrock", "name": "AWS Bedrock"}, } ] }, diff --git a/tests/unit/test_llm_gateway_async.py b/tests/unit/test_llm_gateway_async.py index fef77f3c..3cabd09a 100644 --- a/tests/unit/test_llm_gateway_async.py +++ b/tests/unit/test_llm_gateway_async.py @@ -41,8 +41,8 @@ "creator": "anthropic", "retirement_date": 0, "available_regions": ["us", "eu"], - "providers": ["anthropic"], - "default_provider": "anthropic", + "providers": [{"id": "anthropic", "name": "Anthropic"}], + "default_provider": {"id": "anthropic", "name": "Anthropic"}, }, ] } From ec3b52d235a1d4311a8d987621d7c4e304715e2b Mon Sep 17 00:00:00 2001 From: Daniel Morrissey Date: Mon, 21 Sep 2026 09:19:21 -0600 Subject: [PATCH 8/8] docs(llm_gateway): use speakers key instead of deprecated known_values Addresses PR review feedback to stop showing the deprecated known_values field in speaker identification examples. --- CLAUDE.md | 2 +- README.md | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 2c2717d9..51592ca2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -524,7 +524,7 @@ transcript — `transcript_id` is required, there is no raw-text path. ```python result = gateway.understanding.create( transcript_id="transcript_abc123", - request={"speaker_identification": {"speaker_type": "name", "known_values": ["Ana", "Peter"]}}, + request={"speaker_identification": {"speaker_type": "name", "speakers": [{"name": "Ana"}, {"name": "Peter"}]}}, ) print(result.speech_understanding["response"], result.utterances) diff --git a/README.md b/README.md index d6b1f458..e243627c 100644 --- a/README.md +++ b/README.md @@ -1448,7 +1448,10 @@ result = gateway.understanding.create( request={ "speaker_identification": { "speaker_type": "name", - "known_values": ["Michel Martin", "Peter DeCarlo"], + "speakers": [ + {"name": "Michel Martin", "description": "Hosts the program and interviews the guests"}, + {"name": "Peter DeCarlo", "description": "Answers questions from the interview"}, + ], }, }, )