mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix: preserve explicit provider for async image, audio, and realtime routes (#36346)
* fix(images): preserve explicit provider for async generation * test(images): prove a bare model name reaches its explicit provider on the async path * fix(images): validate async provider hints instead of reading untyped kwargs Read custom_llm_provider and api_base in aimage_generation through a TypeAdapter over a small TypedDict, so the provider lookup gets typed values and the basedpyright gate stays within its merge-base ceiling * fix: answer bad provider hints with a 400 and keep the explicit provider for async speech and transcription * fix: keep the explicit provider for realtime and cover the audio bad-hint 400 * fix: keep the explicit provider on the realtime HTTP routes and read base_url in async audio * fix: stop async image generation from rejecting a base_url it never reads * fix(audio): skip base_url validation when api_base is set on async speech and transcription * test(integration): cover a declared provider with a bare model on image, audio, and realtime routes --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Wang Yunzhuo
GitHub
parent
d5531234ce
commit
a2f792634e
12 changed files with 899 additions and 19 deletions
|
|
@ -5,6 +5,8 @@ from collections.abc import Coroutine
|
|||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
|
||||
|
|
@ -17,6 +19,9 @@ from litellm import client
|
|||
from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
|
||||
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
|
||||
from litellm.exceptions import LiteLLMUnknownProvider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
invalid_provider_routing_hints_message,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values
|
||||
|
|
@ -52,6 +57,7 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
from litellm.types.utils import (
|
||||
LITELLM_IMAGE_VARIATION_PROVIDERS,
|
||||
LlmProviders,
|
||||
ProviderRoutingHints,
|
||||
)
|
||||
from litellm.utils import (
|
||||
ImageResponse,
|
||||
|
|
@ -60,6 +66,8 @@ from litellm.utils import (
|
|||
get_optional_params_image_gen,
|
||||
)
|
||||
|
||||
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
|
||||
|
||||
# Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
|
||||
_ImageEditRequestUtils_cache: Optional["ImageEditRequestUtils"] = None
|
||||
|
||||
|
|
@ -92,6 +100,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
|
|||
model: Final = args[0] if len(args) > 0 else kwargs["model"]
|
||||
### PASS ARGS TO Image Generation ###
|
||||
kwargs["aimg_generation"] = True
|
||||
try:
|
||||
provider_hints: Final = _PROVIDER_ROUTING_HINTS.validate_python(kwargs)
|
||||
except ValidationError as error:
|
||||
raise litellm.BadRequestError(
|
||||
message=invalid_provider_routing_hints_message(error), model=model, llm_provider=""
|
||||
) from error
|
||||
custom_llm_provider = None
|
||||
try:
|
||||
# Use a partial function to pass your keyword arguments
|
||||
|
|
@ -101,7 +115,11 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
|
|||
ctx: Final = contextvars.copy_context()
|
||||
func_with_context: Final = partial(ctx.run, func)
|
||||
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=provider_hints.get("custom_llm_provider"),
|
||||
api_base=provider_hints.get("api_base"),
|
||||
)
|
||||
|
||||
# Await normally
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from typing import Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
NADIR_DEFAULT_API_BASE,
|
||||
|
|
@ -16,6 +18,12 @@ from litellm.secret_managers.main import get_secret, get_secret_str
|
|||
from ..types.router import GenericLiteLLMParams, LiteLLM_Params
|
||||
|
||||
|
||||
def invalid_provider_routing_hints_message(error: ValidationError) -> str:
|
||||
return "; ".join(
|
||||
f"Invalid {'.'.join(str(part) for part in detail['loc'])}: {detail['msg']}" for detail in error.errors()
|
||||
)
|
||||
|
||||
|
||||
def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
|
||||
"""
|
||||
Match a registered openai-compatible endpoint against a caller-supplied
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ if TYPE_CHECKING:
|
|||
import dotenv
|
||||
import httpx
|
||||
import openai
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never, overload
|
||||
|
||||
import litellm
|
||||
|
|
@ -87,6 +87,9 @@ from litellm.litellm_core_utils.get_litellm_params import (
|
|||
parse_control_options,
|
||||
with_control_options,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
invalid_provider_routing_hints_message,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import (
|
||||
ProviderSpecificHeaderUtils,
|
||||
)
|
||||
|
|
@ -137,8 +140,10 @@ from litellm.types.litellm_params import ControlOptions, RetryStrategy
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
BaseUrlRoutingHint,
|
||||
CustomPricingLiteLLMParams,
|
||||
ModelResponseStream,
|
||||
ProviderRoutingHints,
|
||||
RawRequestTypedDict,
|
||||
StreamingChoices,
|
||||
is_decisions_model_mode,
|
||||
|
|
@ -4119,6 +4124,8 @@ def _complete_sagemaker(ctx: CompletionDispatchContext) -> _CompletionDispatchRe
|
|||
|
||||
_ADDITIONAL_DROP_PARAMS_ADAPTER: Final = TypeAdapter(list[str])
|
||||
_OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
|
||||
_BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint)
|
||||
|
||||
|
||||
def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
|
||||
|
|
@ -7918,6 +7925,13 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
|
|||
### PASS ARGS TO Image Generation ###
|
||||
kwargs["atranscription"] = True
|
||||
file: Final = kwargs.get("file", None)
|
||||
try:
|
||||
provider_hints: Final = _PROVIDER_ROUTING_HINTS.validate_python(kwargs)
|
||||
api_base_hint: Final = provider_hints.get("api_base") or _BASE_URL_HINT.validate_python(kwargs).get("base_url")
|
||||
except ValidationError as error:
|
||||
raise litellm.BadRequestError(
|
||||
message=invalid_provider_routing_hints_message(error), model=model, llm_provider=""
|
||||
) from error
|
||||
custom_llm_provider = None
|
||||
try:
|
||||
# Use a partial function to pass your keyword arguments
|
||||
|
|
@ -7927,7 +7941,11 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
|
|||
ctx: Final = contextvars.copy_context()
|
||||
func_with_context: Final = partial(ctx.run, func)
|
||||
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=provider_hints.get("custom_llm_provider"),
|
||||
api_base=api_base_hint,
|
||||
)
|
||||
|
||||
# Await normally
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
|
|
@ -8246,7 +8264,14 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent:
|
|||
model: Final = args[0] if len(args) > 0 else kwargs["model"]
|
||||
### PASS ARGS TO Image Generation ###
|
||||
kwargs["aspeech"] = True
|
||||
custom_llm_provider = kwargs.get("custom_llm_provider", None)
|
||||
try:
|
||||
provider_hints: Final = _PROVIDER_ROUTING_HINTS.validate_python(kwargs)
|
||||
api_base_hint: Final = provider_hints.get("api_base") or _BASE_URL_HINT.validate_python(kwargs).get("base_url")
|
||||
except ValidationError as error:
|
||||
raise litellm.BadRequestError(
|
||||
message=invalid_provider_routing_hints_message(error), model=model, llm_provider=""
|
||||
) from error
|
||||
custom_llm_provider = provider_hints.get("custom_llm_provider")
|
||||
try:
|
||||
# Use a partial function to pass your keyword arguments
|
||||
func: Final = partial(speech, *args, **kwargs)
|
||||
|
|
@ -8255,7 +8280,11 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent:
|
|||
ctx: Final = contextvars.copy_context()
|
||||
func_with_context: Final = partial(ctx.run, func)
|
||||
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=provider_hints.get("custom_llm_provider"),
|
||||
api_base=api_base_hint,
|
||||
)
|
||||
|
||||
# Await normally
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
|
|
|
|||
|
|
@ -153,6 +153,7 @@ async def acreate_realtime_client_secret(
|
|||
dynamic_api_base,
|
||||
) = get_llm_provider(
|
||||
model=model_name,
|
||||
custom_llm_provider=litellm_params.custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
)
|
||||
|
|
@ -221,6 +222,7 @@ async def acreate_realtime_transcription_session(
|
|||
dynamic_api_base,
|
||||
) = get_llm_provider(
|
||||
model=model_name,
|
||||
custom_llm_provider=litellm_params.custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
)
|
||||
|
|
@ -282,6 +284,7 @@ async def arealtime_calls(
|
|||
dynamic_api_base,
|
||||
) = get_llm_provider(
|
||||
model=model_name,
|
||||
custom_llm_provider=litellm_params.custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
)
|
||||
|
|
@ -461,6 +464,7 @@ async def _arealtime(
|
|||
|
||||
model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=litellm_params.custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4740,3 +4740,12 @@ class GenericGuardrailAPIInputs(TypedDict, total=False):
|
|||
stream_holdback_chars: list[
|
||||
int
|
||||
] # trailing chars to withhold from streaming emission per text (word-boundary safety)
|
||||
|
||||
|
||||
class ProviderRoutingHints(TypedDict, total=False):
|
||||
custom_llm_provider: ReadOnly[str | None]
|
||||
api_base: ReadOnly[str | None]
|
||||
|
||||
|
||||
class BaseUrlRoutingHint(TypedDict, total=False):
|
||||
base_url: ReadOnly[str | None]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,534 @@
|
|||
"""A deployment that declares ``custom_llm_provider`` and names a bare model, on every route the declaration reaches.
|
||||
|
||||
Each row drives the client a user drives, then reads what the scripted provider received. A declared deployment must
|
||||
send the same request its ``openai/`` prefixed twin sends, and a value only the caller controls must not change where
|
||||
the deployment's credentials go.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import socket
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from email.parser import BytesParser
|
||||
from email.policy import HTTP
|
||||
from itertools import product
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import websockets
|
||||
from openai import OpenAI
|
||||
from pydantic import JsonValue
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
|
||||
from tests.integration._support.upstream import delete_scenario, register_scenario
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse
|
||||
|
||||
_KEY: Final = "synthetic-declared-provider-key"
|
||||
_FORMS: Final = ("declared", "prefixed", "blank")
|
||||
_IMAGE: Final = "aW1n"
|
||||
_AUDIO: Final = bytes(range(256)) * 8
|
||||
_WAV: Final = (
|
||||
b"RIFF\x30\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00"
|
||||
b"\x40\x1f\x00\x00\x40\x1f\x00\x00\x01\x00\x08\x00data\x08\x00\x00\x00"
|
||||
b"\x00\x01\x02\x03\x04\x05\x06\x07"
|
||||
)
|
||||
_OFFER: Final = (
|
||||
"v=0\r\no=- 4611731400430051336 2 IN IP4 127.0.0.1\r\ns=-\r\nt=0 0\r\n"
|
||||
"a=group:BUNDLE 0\r\nm=audio 9 UDP/TLS/RTP/SAVPF 111\r\nc=IN IP4 0.0.0.0\r\na=mid:0\r\n"
|
||||
"a=rtpmap:111 opus/48000/2\r\n"
|
||||
)
|
||||
_ANSWER: Final = b"v=0\r\no=- 1 2 IN IP4 203.0.113.7\r\ns=-\r\nt=0 0\r\nm=audio 3478 UDP/TLS/RTP/SAVPF 111\r\n"
|
||||
_CALLS_STATUS: Final = 201
|
||||
_TRANSCRIPTION_SESSION_FIELDS: Final[dict[str, JsonValue]] = {"input_audio_format": "pcm16"}
|
||||
|
||||
|
||||
def _bare() -> str:
|
||||
return f"bare-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def _provider_params(form: str, bare: str) -> dict[str, JsonValue]:
|
||||
match form:
|
||||
case "declared":
|
||||
return {"model": bare, "custom_llm_provider": "openai"}
|
||||
case "prefixed":
|
||||
return {"model": f"openai/{bare}"}
|
||||
case "blank":
|
||||
return {"model": f"openai/{bare}", "custom_llm_provider": ""}
|
||||
case _:
|
||||
raise AssertionError(form)
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, wire: Wire, form: str, bare: str) -> str:
|
||||
return scenario.model(**_provider_params(form, bare), api_base=f"{wire.url}/v1", api_key=_KEY)
|
||||
|
||||
|
||||
def _sdk(gateway: Gateway, key: str) -> OpenAI:
|
||||
return OpenAI(
|
||||
base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1",
|
||||
api_key=key,
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(trust_env=False, timeout=30),
|
||||
)
|
||||
|
||||
|
||||
def _json_body(request: Request) -> dict[str, JsonValue]:
|
||||
return JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def _form_fields(request: Request) -> dict[str, str]:
|
||||
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
|
||||
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
|
||||
assert parsed.is_multipart(), request.headers
|
||||
return {
|
||||
str(part.get_param("name", header="content-disposition")): bytes(part.get_payload(decode=True)).decode()
|
||||
for part in parsed.iter_parts()
|
||||
if part.get_filename() is None
|
||||
}
|
||||
|
||||
|
||||
def _image_reply(bare: str, prompt: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/v1/images/generations"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_KEY}", request.headers
|
||||
assert _json_body(request) == {"model": bare, "prompt": prompt}, request.body
|
||||
return Reply(body=json.dumps({"created": 1700000000, "data": [{"b64_json": _IMAGE}]}).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _speech_reply(bare: str, text: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/v1/audio/speech"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_KEY}", request.headers
|
||||
assert _json_body(request) == {"model": bare, "input": text, "voice": "alloy", "response_format": "mp3"}
|
||||
return Reply(body=_AUDIO, content_type="audio/mpeg")
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _transcription_reply(bare: str, prompt: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/v1/audio/transcriptions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_KEY}", request.headers
|
||||
fields: Final = _form_fields(request)
|
||||
assert (fields["model"], fields["prompt"]) == (bare, prompt), fields
|
||||
return Reply(body=json.dumps({"text": prompt}).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _targets(wire: Wire) -> list[tuple[str, str]]:
|
||||
return [(request.method, request.target) for request in wire.drain()]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_image_generation_reaches_the_provider_for_a_bare_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
prompt: Final = f"a lighthouse at dawn {bare}"
|
||||
with wire_server(_image_reply(bare, prompt)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, bare)
|
||||
with _sdk(gateway, scenario.key(models=[model])) as sdk:
|
||||
result: Final = sdk.images.generate(model=model, prompt=prompt)
|
||||
assert [image.b64_json for image in result.data or []] == [_IMAGE], result
|
||||
assert _targets(wire) == [("POST", "/v1/images/generations")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_speech_reaches_the_provider_for_a_bare_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
text: Final = f"hello {bare}"
|
||||
with wire_server(_speech_reply(bare, text)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, bare)
|
||||
with _sdk(gateway, scenario.key(models=[model])) as sdk:
|
||||
response: Final = sdk.audio.speech.with_raw_response.create(
|
||||
model=model, input=text, voice="alloy", response_format="mp3"
|
||||
)
|
||||
assert response.status_code == 200, response.read()
|
||||
assert response.content == _AUDIO
|
||||
assert _targets(wire) == [("POST", "/v1/audio/speech")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_transcription_reaches_the_provider_for_a_bare_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
prompt: Final = f"hello {bare}"
|
||||
with wire_server(_transcription_reply(bare, prompt)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, bare)
|
||||
with _sdk(gateway, scenario.key(models=[model])) as sdk:
|
||||
transcript: Final = sdk.audio.transcriptions.create(
|
||||
model=model, file=("a.wav", _WAV, "audio/wav"), prompt=prompt
|
||||
)
|
||||
assert transcript.text == prompt
|
||||
assert _targets(wire) == [("POST", "/v1/audio/transcriptions")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_test_connect_reaches_the_provider_for_a_bare_image_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/v1/images/generations"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_KEY}", request.headers
|
||||
assert _json_body(request)["model"] == bare, request.body
|
||||
return Reply(body=json.dumps({"created": 1700000000, "data": [{"b64_json": _IMAGE}]}).encode())
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
api_base: Final = f"{wire.url}/v1"
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/health/test_connection",
|
||||
{
|
||||
"litellm_params": {**_provider_params(form, bare), "api_base": api_base, "api_key": _KEY},
|
||||
"mode": "image_generation",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert object_value(response.json())["status"] == "success", response.text
|
||||
assert _targets(wire) == [("POST", "/v1/images/generations")]
|
||||
|
||||
|
||||
def _ws_url(gateway: Gateway, model: str) -> str:
|
||||
base: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://").replace("https://", "wss://")
|
||||
return f"{base}/v1/realtime?model={model}"
|
||||
|
||||
|
||||
async def _first_event(url: str, key: str) -> dict[str, JsonValue]:
|
||||
try:
|
||||
async with websockets.connect(url, additional_headers={"Authorization": f"Bearer {key}"}) as connection:
|
||||
return JSON_OBJECT.validate_json(await asyncio.wait_for(connection.recv(), 30))
|
||||
except ConnectionClosed as closed:
|
||||
return {"type": "closed", "code": 1006 if closed.rcvd is None else closed.rcvd.code}
|
||||
|
||||
|
||||
def _upgrades(upstream_url: str, scenario_id: str) -> list[tuple[JsonValue, JsonValue]]:
|
||||
with httpx.Client(base_url=upstream_url, timeout=5, trust_env=False) as upstream:
|
||||
observed: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
|
||||
assert isinstance(observed, list), observed
|
||||
return [
|
||||
(request["path"], object_value(request["body"]).get("query"))
|
||||
for request in map(object_value, observed)
|
||||
if request["method"] == "WEBSOCKET" and request["authorization"] == f"Bearer {scenario_id}"
|
||||
]
|
||||
|
||||
|
||||
def _realtime_scenario(scenario: Scenario) -> str:
|
||||
scenario_id: Final = f"declared-provider-{uuid.uuid4().hex[:12]}"
|
||||
handle: Final = register_scenario(scenario_id, RealtimeResponse(content_type="application/x-realtime", events=()))
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
return scenario_id
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_realtime_session_opens_at_the_provider_for_a_bare_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
with gateway.scenario() as scenario:
|
||||
scenario_id: Final = _realtime_scenario(scenario)
|
||||
model: Final = scenario.model(**_provider_params(form, bare), api_key=scenario_id)
|
||||
first: Final = asyncio.run(_first_event(_ws_url(gateway, model), scenario.key(models=[model])))
|
||||
assert first["type"] == "session.created", first
|
||||
assert _upgrades(gateway.upstream_url, scenario_id) == [("/v1/realtime", [["model", bare]])]
|
||||
|
||||
|
||||
def test_realtime_follows_the_declared_provider_over_a_contradicting_prefix(gateway: Gateway) -> None:
|
||||
bare: Final = _bare()
|
||||
with gateway.scenario() as scenario:
|
||||
scenario_id: Final = _realtime_scenario(scenario)
|
||||
model: Final = scenario.model(
|
||||
model=f"azure/{bare}", custom_llm_provider="openai", api_key=scenario_id, api_version="2025-04-01-preview"
|
||||
)
|
||||
first: Final = asyncio.run(_first_event(_ws_url(gateway, model), scenario.key(models=[model])))
|
||||
assert first["type"] == "session.created", first
|
||||
assert _upgrades(gateway.upstream_url, scenario_id) == [("/v1/realtime", [["model", f"azure/{bare}"]])]
|
||||
|
||||
|
||||
def test_image_generation_follows_the_declared_provider_over_a_contradicting_prefix(gateway: Gateway) -> None:
|
||||
bare: Final = _bare()
|
||||
with wire_server(_image_reply(f"azure/{bare}", "declared wins")) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"azure/{bare}",
|
||||
custom_llm_provider="openai",
|
||||
api_base=f"{wire.url}/v1",
|
||||
api_key=_KEY,
|
||||
api_version="2025-04-01-preview",
|
||||
)
|
||||
response: Final = gateway.request("POST", "/v1/images/generations", {"model": model, "prompt": "declared wins"})
|
||||
assert response.status_code == 200, response.text
|
||||
assert _targets(wire) == [("POST", "/v1/images/generations")]
|
||||
|
||||
|
||||
def _client_secret(raw: str) -> dict[str, JsonValue]:
|
||||
return {"value": raw, "expires_at": int(time.time()) + 600, "session": {"type": "realtime", "model": "scripted"}}
|
||||
|
||||
|
||||
def _scripted_realtime(raw: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.target.endswith("/v1/realtime/client_secrets"):
|
||||
return Reply(body=json.dumps(_client_secret(raw)).encode())
|
||||
if request.target.endswith("/v1/realtime/transcription_sessions"):
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "sess_declared",
|
||||
"object": "realtime.transcription_session",
|
||||
"client_secret": {"value": raw, "expires_at": int(time.time()) + 600},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
if request.target.endswith("/v1/realtime/calls"):
|
||||
return Reply(status=_CALLS_STATUS, body=_ANSWER, content_type="application/sdp")
|
||||
raise AssertionError(f"Unexpected upstream request {request.method} {request.target}")
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _redeem(gateway: Gateway, token: str) -> httpx.Response:
|
||||
return gateway.client.post(
|
||||
"/v1/realtime/calls",
|
||||
content=_OFFER.encode(),
|
||||
headers={"Content-Type": "application/sdp", "Authorization": f"Bearer {token}"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_browser_call_reaches_the_provider_for_a_bare_realtime_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
raw: Final = f"ek_raw_{uuid.uuid4().hex}"
|
||||
with wire_server(_scripted_realtime(raw)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, bare)
|
||||
key: Final = scenario.key(models=[model])
|
||||
minted: Final = gateway.request("POST", "/v1/realtime/client_secrets", {"model": model}, key=key)
|
||||
assert minted.status_code == 200, minted.text
|
||||
answered: Final = _redeem(gateway, string_value(object_value(minted.json())["value"]))
|
||||
assert (answered.status_code, answered.content) == (_CALLS_STATUS, _ANSWER), answered.text
|
||||
mint_request, calls_request = wire.drain()
|
||||
assert (mint_request.target, mint_request.headers["authorization"]) == (
|
||||
"/v1/realtime/client_secrets",
|
||||
f"Bearer {_KEY}",
|
||||
)
|
||||
assert _json_body(mint_request) == {"session": {"type": "realtime", "model": bare}}, mint_request.body
|
||||
assert (calls_request.target, calls_request.headers["authorization"]) == ("/v1/realtime/calls", f"Bearer {raw}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
def test_transcription_session_reaches_the_provider_for_a_bare_model(gateway: Gateway, form: str) -> None:
|
||||
bare: Final = _bare()
|
||||
raw: Final = f"ek_raw_{uuid.uuid4().hex}"
|
||||
with wire_server(_scripted_realtime(raw)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, bare)
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/realtime/transcription_sessions",
|
||||
{**_TRANSCRIPTION_SESSION_FIELDS, "input_audio_transcription": {"model": model}},
|
||||
key=scenario.key(models=[model]),
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
assert object_value(created.json())["id"] == "sess_declared", created.text
|
||||
(session_request,) = wire.drain()
|
||||
assert (session_request.target, session_request.headers["authorization"]) == (
|
||||
"/v1/realtime/transcription_sessions",
|
||||
f"Bearer {_KEY}",
|
||||
)
|
||||
assert _json_body(session_request) == {
|
||||
**_TRANSCRIPTION_SESSION_FIELDS,
|
||||
"input_audio_transcription": {"model": bare},
|
||||
}, session_request.body
|
||||
|
||||
|
||||
_IMAGES: Final = "/v1/images/generations"
|
||||
_SPEECH: Final = "/v1/audio/speech"
|
||||
_TRANSCRIPTIONS: Final = "/v1/audio/transcriptions"
|
||||
_CLIENT_SECRETS: Final = "/v1/realtime/client_secrets"
|
||||
_NON_STRING_HINTS: Final = MappingProxyType(
|
||||
{"int": "5", "list": '["openai"]', "twice": '"openai", "custom_llm_provider": 5'}
|
||||
)
|
||||
_SECRET_HINTS: Final = MappingProxyType({"int": "5", "list": '["openai"]', "azure": '"azure"'})
|
||||
|
||||
|
||||
def _unexpected(request: Request) -> Reply:
|
||||
raise AssertionError(f"Unexpected upstream request {request.method} {request.target}")
|
||||
|
||||
|
||||
def _fields(route: str, marker: str) -> dict[str, JsonValue]:
|
||||
match route:
|
||||
case "/v1/images/generations":
|
||||
return {"prompt": marker}
|
||||
case "/v1/audio/speech":
|
||||
return {"input": marker, "voice": "alloy", "response_format": "mp3"}
|
||||
case "/v1/realtime/client_secrets":
|
||||
return {"session": {"type": "realtime"}}
|
||||
case _:
|
||||
raise AssertionError(route)
|
||||
|
||||
|
||||
def _post_with_hint(gateway: Gateway, route: str, model: str, marker: str, hint: str) -> httpx.Response:
|
||||
body: Final = json.dumps({"model": model, **_fields(route, marker)})
|
||||
return gateway.client.post(
|
||||
route,
|
||||
content=f'{body[:-1]}, "custom_llm_provider": {hint}}}'.encode(),
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", (_IMAGES, _SPEECH))
|
||||
@pytest.mark.parametrize("hint", tuple(_NON_STRING_HINTS))
|
||||
def test_non_string_provider_hint_from_the_caller_is_refused_before_the_provider(
|
||||
gateway: Gateway, route: str, hint: str
|
||||
) -> None:
|
||||
with wire_server(_unexpected) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, "prefixed", _bare())
|
||||
response: Final = _post_with_hint(gateway, route, model, "refused", _NON_STRING_HINTS[hint])
|
||||
assert response.status_code == 400, response.text
|
||||
assert "custom_llm_provider" in string_value(object_value(object_value(response.json())["error"])["message"])
|
||||
assert _targets(wire) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", _FORMS)
|
||||
@pytest.mark.parametrize("hint", tuple(_SECRET_HINTS))
|
||||
def test_provider_hint_from_the_caller_does_not_move_a_realtime_secret(gateway: Gateway, form: str, hint: str) -> None:
|
||||
with wire_server(_scripted_realtime(f"ek_raw_{uuid.uuid4().hex}")) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, form, _bare())
|
||||
response: Final = _post_with_hint(gateway, _CLIENT_SECRETS, model, "", _SECRET_HINTS[hint])
|
||||
assert response.status_code == 200, response.text
|
||||
assert [(request.target, request.headers["authorization"]) for request in wire.drain()] == [
|
||||
(_CLIENT_SECRETS, f"Bearer {_KEY}")
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", (_IMAGES, _SPEECH))
|
||||
def test_empty_provider_hint_from_the_caller_keeps_the_deployment_route(gateway: Gateway, route: str) -> None:
|
||||
bare: Final = _bare()
|
||||
respond: Final = _image_reply(bare, "kept") if route == _IMAGES else _speech_reply(bare, "kept")
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, "prefixed", bare)
|
||||
response: Final = _post_with_hint(gateway, route, model, "kept", '""')
|
||||
assert response.status_code == 200, response.text
|
||||
assert _targets(wire) == [("POST", route)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", (_IMAGES, _SPEECH))
|
||||
def test_oversized_provider_hint_from_the_caller_never_reaches_the_provider(gateway: Gateway, route: str) -> None:
|
||||
with wire_server(_unexpected) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, "prefixed", _bare())
|
||||
_post_with_hint(gateway, route, model, "oversized", json.dumps("x" * 5120))
|
||||
assert _targets(wire) == []
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
|
||||
|
||||
def _send(client: httpx.AsyncClient, route: str, model: str, marker: str) -> Awaitable[httpx.Response]:
|
||||
if route == _TRANSCRIPTIONS:
|
||||
return client.post(route, data={"model": model, "prompt": marker}, files={"file": ("a.wav", _WAV, "audio/wav")})
|
||||
return client.post(route, json={"model": model, **_fields(route, marker)})
|
||||
|
||||
|
||||
async def _send_all(gateway: Gateway, calls: tuple[tuple[str, str, str], ...]) -> tuple[httpx.Response, ...]:
|
||||
async with httpx.AsyncClient(
|
||||
base_url=gateway.client.base_url,
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
timeout=120,
|
||||
trust_env=False,
|
||||
) as client:
|
||||
return tuple(await asyncio.gather(*(_send(client, route, model, marker) for route, model, marker in calls)))
|
||||
|
||||
|
||||
def _error_message(response: httpx.Response) -> str:
|
||||
return string_value(object_value(object_value(response.json())["error"])["message"]).split("\n\n")[0]
|
||||
|
||||
|
||||
def _scripted_outage(request: Request) -> Reply:
|
||||
return Reply(status=500, body=json.dumps({"error": {"message": "scripted outage", "code": "500"}}).encode())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "known_model"), ((_IMAGES, "dall-e-3"), (_SPEECH, "tts-1"), (_TRANSCRIPTIONS, "whisper-1"))
|
||||
)
|
||||
def test_declared_azure_deployment_fails_like_its_prefixed_twin(gateway: Gateway, route: str, known_model: str) -> None:
|
||||
with wire_server(_scripted_outage) as wire, gateway.scenario() as scenario:
|
||||
declared: Final = scenario.model(
|
||||
model=known_model,
|
||||
custom_llm_provider="azure",
|
||||
api_base=wire.url,
|
||||
api_key=_KEY,
|
||||
api_version="2025-04-01-preview",
|
||||
num_retries=2,
|
||||
)
|
||||
prefixed: Final = scenario.model(
|
||||
model=f"azure/{known_model}",
|
||||
api_base=wire.url,
|
||||
api_key=_KEY,
|
||||
api_version="2025-04-01-preview",
|
||||
num_retries=2,
|
||||
)
|
||||
(declared_response,) = asyncio.run(_send_all(gateway, ((route, declared, "outage"),)))
|
||||
declared_targets: Final = _targets(wire)
|
||||
(prefixed_response,) = asyncio.run(_send_all(gateway, ((route, prefixed, "outage"),)))
|
||||
assert declared_response.status_code == 500, declared_response.text
|
||||
assert declared_targets, declared_response.text
|
||||
assert (declared_response.status_code, _error_message(declared_response), declared_targets) == (
|
||||
prefixed_response.status_code,
|
||||
_error_message(prefixed_response),
|
||||
_targets(wire),
|
||||
)
|
||||
|
||||
|
||||
def _burst_reply(request: Request) -> Reply:
|
||||
match request.target:
|
||||
case "/v1/images/generations":
|
||||
return Reply(body=json.dumps({"created": 1700000000, "data": [{"b64_json": _IMAGE}]}).encode())
|
||||
case "/v1/audio/speech":
|
||||
return Reply(body=_AUDIO, content_type="audio/mpeg")
|
||||
case "/v1/audio/transcriptions":
|
||||
return Reply(body=json.dumps({"text": _form_fields(request)["prompt"]}).encode())
|
||||
case _:
|
||||
raise AssertionError(request.target)
|
||||
|
||||
|
||||
def _marker(request: Request) -> str:
|
||||
if request.target == _TRANSCRIPTIONS:
|
||||
return _form_fields(request)["prompt"]
|
||||
body: Final = _json_body(request)
|
||||
return string_value(body["prompt"] if request.target == _IMAGES else body["input"])
|
||||
|
||||
|
||||
def _closed_port() -> int:
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
return int(probe.getsockname()[1])
|
||||
|
||||
|
||||
def test_burst_across_declared_and_dead_deployments_reaches_each_live_provider_once(gateway: Gateway) -> None:
|
||||
dead_base: Final = f"http://127.0.0.1:{_closed_port()}/v1"
|
||||
with wire_server(_burst_reply) as wire, gateway.scenario() as scenario:
|
||||
live: Final = tuple(
|
||||
(route, _deployment(scenario, wire, form, _bare()))
|
||||
for route, form in product((_IMAGES, _SPEECH, _TRANSCRIPTIONS), _FORMS)
|
||||
)
|
||||
dead: Final = tuple(
|
||||
(route, scenario.model(model=_bare(), custom_llm_provider="openai", api_base=dead_base, api_key=_KEY))
|
||||
for route in (_IMAGES, _SPEECH)
|
||||
)
|
||||
live_calls: Final = tuple(
|
||||
(route, model, f"live-{uuid.uuid4().hex}") for (route, model), _ in product(live, range(5))
|
||||
)
|
||||
dead_calls: Final = tuple(
|
||||
(route, model, f"dead-{uuid.uuid4().hex}") for (route, model), _ in product(dead, range(5))
|
||||
)
|
||||
responses: Final = asyncio.run(_send_all(gateway, live_calls + dead_calls))
|
||||
live_responses: Final = responses[: len(live_calls)]
|
||||
dead_responses: Final = responses[len(live_calls) :]
|
||||
assert [response.status_code for response in live_responses] == [200] * len(live_calls), [
|
||||
response.text for response in live_responses if response.status_code != 200
|
||||
]
|
||||
assert sorted(_marker(request) for request in wire.drain()) == sorted(marker for _, _, marker in live_calls)
|
||||
assert all(response.status_code >= 500 and _error_message(response) for response in dead_responses), [
|
||||
(response.status_code, response.text) for response in dead_responses
|
||||
]
|
||||
assert gateway.client.get("/health/liveliness").status_code == 200
|
||||
assert gateway.chat(scenario.model(), text=f"burst control {uuid.uuid4().hex}")["choices"]
|
||||
|
|
@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.image import Image
|
||||
|
||||
import litellm
|
||||
|
|
@ -38,6 +39,67 @@ def test_image_generation_keeps_an_internal_prefixed_kwarg_out_of_the_provider_r
|
|||
assert sent["prompt"] == "a red circle"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aimage_generation_routes_a_bare_model_name_to_its_explicit_provider(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12348/v1"
|
||||
mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/images/generations.*").mock(
|
||||
return_value=httpx.Response(status_code=200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]})
|
||||
)
|
||||
|
||||
response: Final = await litellm.aimage_generation(
|
||||
model="Qwen-Image",
|
||||
prompt="a red circle",
|
||||
custom_llm_provider="openai",
|
||||
api_base=api_base,
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert mock_route.called
|
||||
sent: Final = json.loads(respx_mock.calls[0].request.content)
|
||||
assert sent["model"] == "Qwen-Image"
|
||||
assert response.data[0].b64_json == "aW1n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hint", ["custom_llm_provider", "api_base"])
|
||||
async def test_aimage_generation_answers_a_non_string_provider_hint_with_a_bad_request(
|
||||
hint: str, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as raised:
|
||||
await litellm.aimage_generation(
|
||||
model="Qwen-Image", prompt="a red circle", api_key="fake_openai_api_key", **{hint: 123}
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert f"Invalid {hint}" in raised.value.message
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aimage_generation_ignores_a_non_string_base_url_like_the_sync_path(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12349/v1"
|
||||
mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/images/generations.*").mock(
|
||||
return_value=httpx.Response(status_code=200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]})
|
||||
)
|
||||
|
||||
response: Final = await litellm.aimage_generation(
|
||||
model="dall-e-3",
|
||||
prompt="a red circle",
|
||||
api_base=api_base,
|
||||
base_url=httpx.URL(api_base),
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert mock_route.called
|
||||
assert response.data[0].b64_json == "aW1n"
|
||||
|
||||
|
||||
def test_image_edit_prices_a_vertex_deployment_at_its_configured_location(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -441,7 +441,7 @@ async def test_arealtime_transcription_intent_defaults_to_ga(monkeypatch):
|
|||
MagicMock(async_realtime=mock_async_realtime),
|
||||
)
|
||||
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None):
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None, custom_llm_provider=None):
|
||||
return (
|
||||
"gpt-realtime-whisper",
|
||||
"azure",
|
||||
|
|
@ -605,7 +605,7 @@ async def test_arealtime_resolves_azure_ad_token_when_no_api_key(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
realtime_main,
|
||||
"get_llm_provider",
|
||||
lambda model, api_base=None, api_key=None: (
|
||||
lambda model, api_base=None, api_key=None, custom_llm_provider=None: (
|
||||
"gpt-realtime-whisper",
|
||||
"azure",
|
||||
None,
|
||||
|
|
@ -645,7 +645,7 @@ async def test_arealtime_does_not_resolve_azure_ad_token_when_api_key_present(mo
|
|||
monkeypatch.setattr(
|
||||
realtime_main,
|
||||
"get_llm_provider",
|
||||
lambda model, api_base=None, api_key=None: (
|
||||
lambda model, api_base=None, api_key=None, custom_llm_provider=None: (
|
||||
"gpt-realtime-whisper",
|
||||
"azure",
|
||||
"test-key",
|
||||
|
|
@ -721,7 +721,7 @@ async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambi
|
|||
monkeypatch.setattr(
|
||||
realtime_main,
|
||||
"get_llm_provider",
|
||||
lambda model, api_base=None, api_key=None: (
|
||||
lambda model, api_base=None, api_key=None, custom_llm_provider=None: (
|
||||
"gpt-realtime-whisper",
|
||||
"azure",
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -470,7 +470,7 @@ async def test_realtime_query_params_use_normalized_model_name(monkeypatch):
|
|||
MagicMock(async_realtime=mock_async_realtime),
|
||||
)
|
||||
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None):
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None, custom_llm_provider=None):
|
||||
return ("gpt-4o-realtime-preview", "openai", None, None)
|
||||
|
||||
monkeypatch.setattr(realtime_main, "get_llm_provider", fake_get_llm_provider)
|
||||
|
|
@ -509,7 +509,7 @@ async def test_realtime_query_params_preserve_missing_model(monkeypatch):
|
|||
MagicMock(async_realtime=mock_async_realtime),
|
||||
)
|
||||
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None):
|
||||
def fake_get_llm_provider(model, api_base=None, api_key=None, custom_llm_provider=None):
|
||||
return ("gpt-realtime-whisper", "openai", None, None)
|
||||
|
||||
monkeypatch.setattr(realtime_main, "get_llm_provider", fake_get_llm_provider)
|
||||
|
|
|
|||
|
|
@ -214,7 +214,7 @@ def test_realtime_config_uses_xai_key_when_provider_does_not_resolve_key(monkeyp
|
|||
async def mock_async_realtime(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider=None):
|
||||
return model, "xai", None, api_base
|
||||
|
||||
monkeypatch.setattr(litellm, "xai_key", "xai_key_value")
|
||||
|
|
@ -244,7 +244,7 @@ def test_realtime_config_uses_generic_key_when_provider_does_not_resolve_key(
|
|||
async def mock_async_realtime(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider=None):
|
||||
return model, "xai", None, api_base
|
||||
|
||||
monkeypatch.setattr(litellm, "xai_key", None)
|
||||
|
|
|
|||
|
|
@ -1,16 +1,18 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from types import TracebackType
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.models.credentials import CredentialItem
|
||||
from litellm.realtime_api import main as realtime_main
|
||||
from litellm.realtime_api.main import _with_resolved_session_model
|
||||
|
|
@ -83,7 +85,7 @@ def _run_client_secret(session, model, monkeypatch):
|
|||
captured.update(kwargs)
|
||||
return object()
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
||||
return model, "openai", None, api_base
|
||||
|
||||
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
||||
|
|
@ -180,7 +182,7 @@ async def test_meta_realtime_dispatches_to_base_handler_with_meta_config(monkeyp
|
|||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
||||
return model.removeprefix("meta/"), "meta", None, api_base
|
||||
|
||||
async def mock_async_realtime(**kwargs):
|
||||
|
|
@ -210,7 +212,7 @@ async def test_arealtime_vertex_branch_resolves_credentials_under_a_bound(monkey
|
|||
async def hanging_token_refresh(**kwargs):
|
||||
await asyncio.sleep(30)
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
||||
return model, "vertex_ai", None, api_base
|
||||
|
||||
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
||||
|
|
@ -252,6 +254,7 @@ async def test_arealtime_live_session_dispatches_to_openai_live_handler(monkeypa
|
|||
model: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
||||
|
||||
|
|
@ -315,6 +318,7 @@ async def test_arealtime_live_session_rejects_matching_guardrail_before_dispatch
|
|||
model: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
||||
|
||||
|
|
@ -358,6 +362,7 @@ async def test_arealtime_live_session_rejects_streaming_only_pre_call_guardrail(
|
|||
model: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
||||
|
||||
|
|
@ -383,6 +388,7 @@ async def test_arealtime_live_session_rejects_non_openai_provider(monkeypatch: p
|
|||
model: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
return "claude-live", "anthropic", None, None
|
||||
|
||||
|
|
@ -407,6 +413,7 @@ async def test_arealtime_without_live_start_preserves_openai_realtime_dispatch(
|
|||
model: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
return "gpt-realtime", "openai", "resolved-key", "https://realtime-gateway.example/v1"
|
||||
|
||||
|
|
@ -652,6 +659,55 @@ async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai
|
|||
assert connect.url == "wss://my-project.services.ai.azure.com/openai/v1/realtime?model=gpt-realtime-mini"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arealtime_routes_a_bare_model_name_to_its_explicit_provider():
|
||||
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
||||
with patch("websockets.connect", connect):
|
||||
await realtime_main._arealtime.__wrapped__(
|
||||
model="my-realtime",
|
||||
websocket=MagicMock(close=AsyncMock()),
|
||||
custom_llm_provider="openai",
|
||||
api_base="https://realtime-gateway.example/v1",
|
||||
api_key="fake-key",
|
||||
query_params={"model": "my-realtime"},
|
||||
litellm_logging_obj=FakeLogging(),
|
||||
)
|
||||
assert connect.url == "wss://realtime-gateway.example/v1/realtime?model=my-realtime"
|
||||
|
||||
|
||||
class _GatewayThatRecordsTheUrl:
|
||||
url: str | None = None
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.url = str(request.url)
|
||||
return httpx.Response(200, json={"value": "ek_test"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("route", "route_kwargs", "path"),
|
||||
[
|
||||
(realtime_main.acreate_realtime_client_secret, {}, "client_secrets"),
|
||||
(realtime_main.acreate_realtime_transcription_session, {}, "transcription_sessions"),
|
||||
(realtime_main.arealtime_calls, {"openai_ephemeral_key": "ek_test", "sdp_body": b"v=0"}, "calls"),
|
||||
],
|
||||
)
|
||||
async def test_realtime_http_routes_send_a_bare_model_name_to_its_explicit_provider(
|
||||
route: Callable[..., Awaitable[object]], route_kwargs: Mapping[str, object], path: str
|
||||
) -> None:
|
||||
gateway: Final = _GatewayThatRecordsTheUrl()
|
||||
await route.__wrapped__(
|
||||
model="my-realtime",
|
||||
custom_llm_provider="openai",
|
||||
api_base="https://realtime-gateway.example/v1",
|
||||
api_key="fake-key",
|
||||
litellm_logging_obj=FakeLogging(),
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)),
|
||||
**route_kwargs,
|
||||
)
|
||||
assert gateway.url == f"https://realtime-gateway.example/v1/realtime/{path}"
|
||||
|
||||
|
||||
class _ClientWebSocketWithHeaders:
|
||||
def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None:
|
||||
self.scope: Final = {"headers": headers}
|
||||
|
|
@ -714,7 +770,7 @@ async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location:
|
|||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def mock_get_llm_provider(model, api_base, api_key):
|
||||
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
||||
return model.removeprefix("vertex_ai/"), "vertex_ai", None, api_base
|
||||
|
||||
async def mock_token_resolver(**kwargs):
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from openai import APITimeoutError
|
||||
from openai import APITimeoutError, AsyncOpenAI
|
||||
from openai.types.chat.chat_completion import ChatCompletion
|
||||
|
||||
import litellm
|
||||
|
|
@ -8572,3 +8572,163 @@ def test_model_alias_map_resolves_the_outbound_model(
|
|||
|
||||
assert response.model == "resolved-model"
|
||||
assert json.loads(route.calls[0].request.content)["model"] == "resolved-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aspeech_routes_a_bare_model_name_to_its_explicit_provider(respx_mock: respx.MockRouter) -> None:
|
||||
api_base: Final = "http://localhost:12351/v1"
|
||||
route: Final = respx_mock.post(f"{api_base}/audio/speech").mock(
|
||||
return_value=httpx.Response(status_code=200, content=b"ID3-speech")
|
||||
)
|
||||
|
||||
response: Final = await litellm.aspeech(
|
||||
model="kokoro",
|
||||
input="hello",
|
||||
voice="alloy",
|
||||
custom_llm_provider="openai",
|
||||
api_base=api_base,
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert json.loads(route.calls[0].request.content)["model"] == "kokoro"
|
||||
assert response.content == b"ID3-speech"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_atranscription_routes_a_bare_model_name_to_its_explicit_provider(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12352/v1"
|
||||
route: Final = respx_mock.post(f"{api_base}/audio/transcriptions").mock(
|
||||
return_value=httpx.Response(status_code=200, json={"text": "four score"})
|
||||
)
|
||||
|
||||
response: Final = await litellm.atranscription(
|
||||
model="my-whisper",
|
||||
file=("speech.wav", b"RIFF-audio", "audio/wav"),
|
||||
custom_llm_provider="openai",
|
||||
api_base=api_base,
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert b'name="model"\r\n\r\nmy-whisper' in route.calls[0].request.content
|
||||
assert response.text == "four score"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hint", ["custom_llm_provider", "api_base", "base_url"])
|
||||
async def test_aspeech_answers_a_non_string_provider_hint_with_a_bad_request(
|
||||
hint: str, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as raised:
|
||||
await litellm.aspeech(model="kokoro", input="hello", voice="alloy", api_key="fake_openai_api_key", **{hint: 123})
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert f"Invalid {hint}" in raised.value.message
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hint", ["custom_llm_provider", "api_base", "base_url"])
|
||||
async def test_atranscription_answers_a_non_string_provider_hint_with_a_bad_request(
|
||||
hint: str, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as raised:
|
||||
await litellm.atranscription(
|
||||
model="my-whisper",
|
||||
file=("speech.wav", b"RIFF-audio", "audio/wav"),
|
||||
api_key="fake_openai_api_key",
|
||||
**{hint: 123},
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert f"Invalid {hint}" in raised.value.message
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aspeech_infers_the_provider_from_base_url(respx_mock: respx.MockRouter) -> None:
|
||||
base_url: Final = "https://api.groq.com/openai/v1"
|
||||
route: Final = respx_mock.post(f"{base_url}/audio/speech").mock(
|
||||
return_value=httpx.Response(status_code=200, content=b"ID3-speech")
|
||||
)
|
||||
|
||||
response: Final = await litellm.aspeech(
|
||||
model="my-tts",
|
||||
input="hello",
|
||||
voice="alloy",
|
||||
base_url=base_url,
|
||||
api_key="fake_groq_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_groq_api_key", base_url=base_url, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.content == b"ID3-speech"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_atranscription_infers_the_provider_from_base_url(respx_mock: respx.MockRouter) -> None:
|
||||
base_url: Final = "https://api.groq.com/openai/v1"
|
||||
route: Final = respx_mock.post(f"{base_url}/audio/transcriptions").mock(
|
||||
return_value=httpx.Response(status_code=200, json={"text": "four score"})
|
||||
)
|
||||
|
||||
response: Final = await litellm.atranscription(
|
||||
model="my-whisper",
|
||||
file=("speech.wav", b"RIFF-audio", "audio/wav"),
|
||||
base_url=base_url,
|
||||
api_key="fake_groq_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_groq_api_key", base_url=base_url, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.text == "four score"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aspeech_ignores_a_non_string_base_url_when_api_base_is_set_like_the_sync_path(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12353/v1"
|
||||
route: Final = respx_mock.post(f"{api_base}/audio/speech").mock(
|
||||
return_value=httpx.Response(status_code=200, content=b"ID3-speech")
|
||||
)
|
||||
|
||||
response: Final = await litellm.aspeech(
|
||||
model="tts-1",
|
||||
input="hello",
|
||||
voice="alloy",
|
||||
api_base=api_base,
|
||||
base_url=httpx.URL(api_base),
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.content == b"ID3-speech"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_atranscription_ignores_a_non_string_base_url_when_api_base_is_set_like_the_sync_path(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12354/v1"
|
||||
route: Final = respx_mock.post(f"{api_base}/audio/transcriptions").mock(
|
||||
return_value=httpx.Response(status_code=200, json={"text": "four score"})
|
||||
)
|
||||
|
||||
response: Final = await litellm.atranscription(
|
||||
model="whisper-1",
|
||||
file=("speech.wav", b"RIFF-audio", "audio/wav"),
|
||||
api_base=api_base,
|
||||
base_url=httpx.URL(api_base),
|
||||
api_key="fake_openai_api_key",
|
||||
client=AsyncOpenAI(api_key="fake_openai_api_key", base_url=api_base, http_client=httpx.AsyncClient()),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.text == "four score"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue