mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
1052 lines
44 KiB
Python
1052 lines
44 KiB
Python
import asyncio
|
|
import json
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from queue import SimpleQueue
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import yaml
|
|
from hypothesis import Phase, example, given, settings
|
|
from hypothesis import strategies as st
|
|
from integration._support.client import Gateway, eventually, object_value, string_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
from openai import OpenAI
|
|
from pydantic import JsonValue, TypeAdapter
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
|
|
|
|
|
def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
|
|
value: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
|
|
}
|
|
return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n"
|
|
|
|
|
|
def text_stream(identity: str) -> tuple[bytes, ...]:
|
|
usage: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
|
|
}
|
|
return (
|
|
frame(identity, {"role": "assistant", "content": "Hello "}),
|
|
frame(identity, {"content": "雪 café"}),
|
|
frame(identity, {}, finish="stop"),
|
|
b"data: " + json.dumps(usage).encode() + b"\n\n",
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage")
|
|
def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage() -> None:
|
|
import litellm
|
|
|
|
body: Final = b"".join(text_stream("stream-partition-control"))
|
|
|
|
@settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
|
|
@example(cuts=tuple(range(1, len(body))))
|
|
@example(cuts=())
|
|
@given(cuts=st.lists(st.integers(min_value=1, max_value=len(body) - 1), max_size=35, unique=True).map(tuple))
|
|
def check(cuts: tuple[int, ...]) -> None:
|
|
boundaries: Final = (0, *sorted(cuts), len(body))
|
|
pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:]))
|
|
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire:
|
|
stream: Final = litellm.completion(
|
|
model="openai/gpt-4o-mini",
|
|
api_base=wire.url + "/v1",
|
|
api_key="synthetic-stream-key",
|
|
messages=[{"role": "user", "content": "partition control"}],
|
|
stream=True,
|
|
stream_options={"include_usage": True},
|
|
timeout=5,
|
|
num_retries=0,
|
|
)
|
|
try:
|
|
chunks: Final = tuple(stream)
|
|
finally:
|
|
asyncio.run(stream.aclose())
|
|
assert (
|
|
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
|
|
)
|
|
assert {chunk.id for chunk in chunks} == {"stream-partition-control"}
|
|
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
|
|
"stop"
|
|
]
|
|
usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
|
|
assert len(usages) == 1
|
|
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
|
|
assert len(wire.drain()) == 1
|
|
|
|
check()
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.tools.fragmented_calls_keep_independent_arguments")
|
|
def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None:
|
|
import litellm
|
|
|
|
identity: Final = "stream-tools-control"
|
|
deltas: Final = (
|
|
{
|
|
"role": "assistant",
|
|
"tool_calls": [
|
|
{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}},
|
|
{"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}},
|
|
],
|
|
},
|
|
{
|
|
"tool_calls": [
|
|
{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}},
|
|
{"index": 0, "function": {"arguments": '{"x":1,'}},
|
|
]
|
|
},
|
|
{
|
|
"tool_calls": [
|
|
{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}},
|
|
{"index": 1, "function": {"arguments": '"y":4}'}},
|
|
]
|
|
},
|
|
)
|
|
frames: Final = (
|
|
*tuple(frame(identity, delta) for delta in deltas),
|
|
frame(identity, {}, finish="tool_calls"),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire:
|
|
stream: Final = litellm.completion(
|
|
model="openai/gpt-4o-mini",
|
|
api_base=wire.url + "/v1",
|
|
api_key="synthetic-stream-key",
|
|
messages=[{"role": "user", "content": "tool control"}],
|
|
stream=True,
|
|
timeout=5,
|
|
num_retries=0,
|
|
)
|
|
try:
|
|
chunks: Final = tuple(stream)
|
|
finally:
|
|
asyncio.run(stream.aclose())
|
|
events: Final = tuple(
|
|
(choice.index, tool)
|
|
for chunk in chunks
|
|
for choice in chunk.choices
|
|
for tool in (choice.delta.tool_calls or ())
|
|
)
|
|
for index, name, call_id, arguments in (
|
|
(0, "add", "call-add", {"x": 1, "y": 2}),
|
|
(1, "multiply", "call-multiply", {"x": 3, "y": 4}),
|
|
):
|
|
selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index))
|
|
assert "".join(tool.id or "" for tool in selected) == call_id
|
|
assert "".join(tool.function.name or "" for tool in selected) == name
|
|
assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments
|
|
assert {tool.index for _, tool in events} == {0, 1}
|
|
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
|
|
"tool_calls"
|
|
]
|
|
assert len(wire.drain()) == 1
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.usage.client_visibility_preserves_persisted_accounting")
|
|
def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
for include in (None, False, True):
|
|
identity: Final = "stream-usage-" + uuid.uuid4().hex
|
|
with wire_server(
|
|
lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))
|
|
) as wire:
|
|
model: Final = scenario.model(
|
|
api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
|
|
)
|
|
eventually(
|
|
lambda: tuple(
|
|
object_value(entry).get("model_name") for entry in gateway.get("/model/info")["data"]
|
|
),
|
|
lambda model_names, model=model: model in model_names,
|
|
seconds=60,
|
|
)
|
|
with OpenAI(
|
|
api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0
|
|
) as client:
|
|
stream: Final = client.chat.completions.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": identity}],
|
|
stream=True,
|
|
**({} if include is None else {"stream_options": {"include_usage": include}}),
|
|
)
|
|
with stream:
|
|
chunks: Final = tuple(stream)
|
|
assert (
|
|
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
|
|
== "Hello 雪 café"
|
|
)
|
|
assert {chunk.id for chunk in chunks} == {identity}
|
|
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
|
|
assert len(usages) == (1 if include else 0)
|
|
if include:
|
|
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True
|
|
rows: Final = eventually(
|
|
lambda identity=identity: read_rows(
|
|
'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(identity,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4
|
|
assert float(rows[0]["spend"]) == pytest.approx(0.019)
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream")
|
|
def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None:
|
|
identity: Final = "messages-empty-choices-" + uuid.uuid4().hex
|
|
metadata: Final = (
|
|
b"data: "
|
|
+ json.dumps(
|
|
{
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [],
|
|
"prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
|
|
},
|
|
ensure_ascii=False,
|
|
).encode()
|
|
+ b"\n\n"
|
|
)
|
|
frames: Final = (metadata, *text_stream(identity))
|
|
with (
|
|
wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
|
|
gateway.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model="azure/gpt-4o-mini", api_base=wire.url + "/v1")
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/messages",
|
|
json={
|
|
"model": model,
|
|
"max_tokens": 64,
|
|
"stream": True,
|
|
"messages": [{"role": "user", "content": identity}],
|
|
},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read().decode()
|
|
events: Final = tuple(
|
|
json.loads(line.removeprefix("data: ")) for line in response.iter_lines() if line.startswith("data: ")
|
|
)
|
|
assert tuple(event["type"] for event in events) == (
|
|
"message_start",
|
|
"content_block_start",
|
|
"content_block_delta",
|
|
"content_block_delta",
|
|
"content_block_stop",
|
|
"message_delta",
|
|
"message_stop",
|
|
), f"observed events: {events!r}"
|
|
assert (
|
|
"".join(event["delta"]["text"] for event in events if event["type"] == "content_block_delta") == "Hello 雪 café"
|
|
)
|
|
message_delta: Final = next(event for event in events if event["type"] == "message_delta")
|
|
assert message_delta["usage"] == {"input_tokens": 11, "output_tokens": 4}
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
outbound: Final = json.loads(requests[0].body)
|
|
assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
|
|
f"observed outbound body: {outbound!r}"
|
|
)
|
|
|
|
|
|
def reasoning_first_stream(identity: str) -> tuple[bytes, ...]:
|
|
usage: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 11, "completion_tokens": 6, "total_tokens": 17},
|
|
}
|
|
return (
|
|
frame(identity, {"role": "assistant", "content": None, "reasoning_content": "Let me "}),
|
|
frame(identity, {"content": None, "reasoning_content": "think."}),
|
|
frame(identity, {"content": "Hello "}),
|
|
frame(identity, {"content": "there"}),
|
|
frame(identity, {}, finish="stop"),
|
|
b"data: " + json.dumps(usage).encode() + b"\n\n",
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
|
|
@pytest.mark.covers("streaming.messages_bridge.reasoning_content_only_chunks_open_a_thinking_block_first")
|
|
def test_messages_stream_opens_thinking_block_at_index_zero_for_reasoning_content_only_chunks(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
identity: Final = "messages-reasoning-first-" + uuid.uuid4().hex
|
|
with (
|
|
wire_server(
|
|
lambda request: Reply(content_type="text/event-stream", chunks=reasoning_first_stream(identity))
|
|
) as wire,
|
|
gateway.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model="hosted_vllm/reasoning-model", api_base=wire.url + "/v1")
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/messages",
|
|
json={
|
|
"model": model,
|
|
"max_tokens": 64,
|
|
"stream": True,
|
|
"messages": [{"role": "user", "content": identity}],
|
|
},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
text: Final = response.read().decode()
|
|
assert response.status_code == 200, text
|
|
assert response.headers["content-type"].startswith("text/event-stream"), text
|
|
events: Final = tuple(json.loads(line) for line in sse_data_lines(text))
|
|
blocks: Final = tuple(
|
|
(event["index"], event.get("content_block") or event["delta"])
|
|
for event in events
|
|
if event["type"] in ("content_block_start", "content_block_delta")
|
|
)
|
|
assert blocks == (
|
|
(0, {"type": "thinking", "thinking": "", "signature": ""}),
|
|
(0, {"type": "thinking_delta", "thinking": "Let me "}),
|
|
(0, {"type": "thinking_delta", "thinking": "think."}),
|
|
(1, {"type": "text", "text": ""}),
|
|
(1, {"type": "text_delta", "text": "Hello "}),
|
|
(1, {"type": "text_delta", "text": "there"}),
|
|
), text
|
|
assert tuple(event["type"] for event in events) == (
|
|
"message_start",
|
|
"content_block_start",
|
|
"content_block_delta",
|
|
"content_block_delta",
|
|
"content_block_stop",
|
|
"content_block_start",
|
|
"content_block_delta",
|
|
"content_block_delta",
|
|
"content_block_stop",
|
|
"message_delta",
|
|
"message_stop",
|
|
), text
|
|
message_delta: Final = next(event for event in events if event["type"] == "message_delta")
|
|
assert message_delta["usage"] == {"input_tokens": 11, "output_tokens": 6}, text
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
outbound: Final = json.loads(requests[0].body)
|
|
assert outbound["stream"] is True and outbound["messages"] == [{"role": "user", "content": identity}], outbound
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.responses_bridge.empty_choices_chunks_complete_stream")
|
|
def test_responses_stream_completes_through_empty_choices_metadata_and_usage_chunks(gateway: Gateway) -> None:
|
|
identity: Final = "responses-empty-choices-" + uuid.uuid4().hex
|
|
metadata: Final = (
|
|
b"data: "
|
|
+ json.dumps(
|
|
{
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [],
|
|
"prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
|
|
},
|
|
ensure_ascii=False,
|
|
).encode()
|
|
+ b"\n\n"
|
|
)
|
|
frames: Final = (metadata, *text_stream(identity))
|
|
with (
|
|
wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
|
|
gateway.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1")
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/responses",
|
|
json={"model": model, "input": identity, "stream": True},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read().decode()
|
|
events: Final = tuple(
|
|
json.loads(line.removeprefix("data: "))
|
|
for line in response.iter_lines()
|
|
if line.startswith("data: ") and line != "data: [DONE]"
|
|
)
|
|
assert (
|
|
"".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == "Hello 雪 café"
|
|
), f"observed events: {events!r}"
|
|
assert tuple(event["type"] for event in events if event["type"] != "response.output_text.delta") == (
|
|
"response.created",
|
|
"response.in_progress",
|
|
"response.output_item.added",
|
|
"response.content_part.added",
|
|
"response.output_text.done",
|
|
"response.content_part.done",
|
|
"response.output_item.done",
|
|
"response.completed",
|
|
), f"observed events: {events!r}"
|
|
assert events[-1]["type"] == "response.completed"
|
|
assert events[-1]["response"]["usage"] == {
|
|
"input_tokens": 11,
|
|
"output_tokens": 4,
|
|
"output_tokens_details": {"reasoning_tokens": 0, "text_tokens": 4},
|
|
"total_tokens": 15,
|
|
}
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
outbound: Final = json.loads(requests[0].body)
|
|
assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
|
|
f"observed outbound body: {outbound!r}"
|
|
)
|
|
|
|
|
|
def provider_cost_object_stream(identity: str, total_cost: float) -> tuple[bytes, ...]:
|
|
cost: Final = {
|
|
"input_tokens_cost": 0.0001,
|
|
"output_tokens_cost": 0.0002,
|
|
"request_cost": 0.012,
|
|
"total_cost": total_cost,
|
|
}
|
|
usage: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "sonar",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15, "cost": cost},
|
|
}
|
|
return (
|
|
frame(identity, {"role": "assistant", "content": "Hello "}),
|
|
frame(identity, {"content": "from search"}),
|
|
frame(identity, {}, finish="stop"),
|
|
b"data: " + json.dumps(usage).encode() + b"\n\n",
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
|
|
def sse_data_lines(text: str) -> tuple[str, ...]:
|
|
return tuple(line.removeprefix("data: ") for line in text.splitlines() if line.startswith("data: "))
|
|
|
|
|
|
def _native_anthropic_message(model: str, identity: str, text: str) -> dict[str, object]:
|
|
return {
|
|
"id": identity,
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": model,
|
|
"content": [{"type": "text", "text": text}],
|
|
"stop_reason": "end_turn",
|
|
"stop_sequence": None,
|
|
"usage": {"input_tokens": 11, "output_tokens": 4},
|
|
}
|
|
|
|
|
|
def _native_responses_response(model: str, identity: str, text: str) -> dict[str, object]:
|
|
return {
|
|
"id": identity,
|
|
"object": "response",
|
|
"created_at": 1,
|
|
"status": "completed",
|
|
"model": model,
|
|
"output": [
|
|
{
|
|
"id": "msg_" + identity,
|
|
"type": "message",
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
|
}
|
|
],
|
|
"usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15},
|
|
}
|
|
|
|
|
|
def _native_sse_event(event: dict[str, object]) -> bytes:
|
|
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
|
|
|
|
|
|
def _native_anthropic_stream(model: str, identity: str, text: str) -> tuple[bytes, ...]:
|
|
return (
|
|
_native_sse_event(
|
|
{
|
|
"type": "message_start",
|
|
"message": {
|
|
**_native_anthropic_message(model, identity, ""),
|
|
"content": [],
|
|
"stop_reason": None,
|
|
"usage": {"input_tokens": 11, "output_tokens": 0},
|
|
},
|
|
}
|
|
),
|
|
_native_sse_event({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
|
_native_sse_event(
|
|
{
|
|
"type": "content_block_delta",
|
|
"index": 0,
|
|
"delta": {"type": "text_delta", "text": text},
|
|
}
|
|
),
|
|
_native_sse_event({"type": "content_block_stop", "index": 0}),
|
|
_native_sse_event(
|
|
{
|
|
"type": "message_delta",
|
|
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
|
"usage": {"output_tokens": 4},
|
|
}
|
|
),
|
|
_native_sse_event({"type": "message_stop"}),
|
|
)
|
|
|
|
|
|
def _native_responses_stream(model: str, identity: str, text: str) -> tuple[bytes, ...]:
|
|
response: Final = _native_responses_response(model, identity, text)
|
|
return (
|
|
_native_sse_event(
|
|
{
|
|
"type": "response.created",
|
|
"sequence_number": 0,
|
|
"response": {**response, "status": "in_progress", "output": []},
|
|
}
|
|
),
|
|
_native_sse_event(
|
|
{
|
|
"type": "response.output_text.delta",
|
|
"sequence_number": 1,
|
|
"item_id": "msg_" + identity,
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": text,
|
|
}
|
|
),
|
|
_native_sse_event({"type": "response.completed", "sequence_number": 2, "response": response}),
|
|
)
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.usage.provider_cost_object_completes_stream_and_bills_total_cost")
|
|
def test_perplexity_stream_with_cost_breakdown_object_completes_and_bills_total_cost(gateway: Gateway) -> None:
|
|
identity: Final = "stream-cost-object-" + uuid.uuid4().hex
|
|
total_cost: Final = 0.0123
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
wire_server(
|
|
lambda request: Reply(
|
|
content_type="text/event-stream", chunks=provider_cost_object_stream(identity, total_cost)
|
|
)
|
|
) as wire,
|
|
):
|
|
model: Final = scenario.model(model="perplexity/sonar", api_base=wire.url + "/v1")
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": identity}],
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
text: Final = response.read().decode()
|
|
assert response.status_code == 200, text
|
|
lines: Final = sse_data_lines(text)
|
|
assert lines[-1] == "[DONE]", text
|
|
events: Final = tuple(json.loads(line) for line in lines[:-1])
|
|
assert [event for event in events if "error" in event] == [], text
|
|
assert (
|
|
"".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"])
|
|
== "Hello from search"
|
|
), text
|
|
assert [
|
|
choice.get("finish_reason")
|
|
for event in events
|
|
for choice in event["choices"]
|
|
if choice.get("finish_reason")
|
|
] == ["stop"], text
|
|
usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
|
|
assert len(usages) == 1, text
|
|
assert (usages[0]["prompt_tokens"], usages[0]["completion_tokens"], usages[0]["total_tokens"]) == (11, 4, 15), (
|
|
text
|
|
)
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
outbound: Final = json.loads(requests[0].body)
|
|
assert outbound["model"] == "sonar" and outbound["stream"] is True, outbound
|
|
assert outbound["messages"] == [{"role": "user", "content": identity}], outbound
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(identity,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (11, 4)
|
|
assert float(rows[0]["spend"]) == pytest.approx(total_cost)
|
|
|
|
|
|
@pytest.mark.covers(
|
|
"other.streaming.fallback.empty_leading_chunk_then_disconnect_streams_fallback_with_usage_and_spend"
|
|
)
|
|
def test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bills_the_fallback(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
identity: Final = "stream-empty-fallback-" + uuid.uuid4().hex
|
|
empty_first: Final = (
|
|
b"data: "
|
|
+ json.dumps(
|
|
{
|
|
"id": identity + "-primary",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 11, "completion_tokens": 0, "total_tokens": 11},
|
|
}
|
|
).encode()
|
|
+ b"\n\n"
|
|
)
|
|
with (
|
|
wire_server(
|
|
lambda request: Reply(
|
|
content_type="text/event-stream",
|
|
chunks=(empty_first, b":" + b"x" * 4_000_000 + b"\n\n", empty_first),
|
|
abort_after=2,
|
|
)
|
|
) as primary,
|
|
wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as fallback,
|
|
):
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
config["model_list"] = [
|
|
{
|
|
"model_name": name,
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o-mini",
|
|
"api_key": "synthetic-fallback-key",
|
|
"api_base": server.url + "/v1",
|
|
"input_cost_per_token": 0.001,
|
|
"output_cost_per_token": 0.002,
|
|
},
|
|
}
|
|
for name, server in (("primary", primary), ("fallback", fallback))
|
|
]
|
|
config["router_settings"] = {
|
|
"num_retries": 0,
|
|
"disable_cooldowns": True,
|
|
"fallbacks": [{"primary": ["fallback"]}],
|
|
}
|
|
path: Final = tmp_path / "fallbacks.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate:
|
|
body: Final = {
|
|
"model": "primary",
|
|
"messages": [{"role": "user", "content": identity}],
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
}
|
|
with candidate.client.stream(
|
|
"POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {candidate.key}"}
|
|
) as response:
|
|
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data:"))
|
|
assert response.status_code == 200, lines
|
|
assert lines[-1] == "data: [DONE]", lines
|
|
events: Final = tuple(json.loads(line.removeprefix("data:")) for line in lines[:-1])
|
|
assert all("error" not in event for event in events), lines
|
|
assert (
|
|
"".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"])
|
|
== "Hello 雪 café"
|
|
), lines
|
|
usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
|
|
assert (usages[-1]["prompt_tokens"], usages[-1]["completion_tokens"]) == (11, 4), lines
|
|
assert tuple(
|
|
json.loads(request.body)["messages"]
|
|
for request in primary.drain()
|
|
if request.target.endswith("/chat/completions")
|
|
) == (body["messages"],)
|
|
assert tuple(
|
|
json.loads(request.body)["messages"]
|
|
for request in fallback.drain()
|
|
if request.target.endswith("/chat/completions")
|
|
) == (body["messages"],)
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT spend, prompt_tokens, completion_tokens, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(identity,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["status"]) == (11, 4, "success"), (
|
|
rows
|
|
)
|
|
assert float(rows[0]["spend"]) == pytest.approx(0.019), rows
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("route", "failure_mode"),
|
|
(
|
|
("messages", "non-streaming"),
|
|
("messages", "before-first-byte"),
|
|
("messages", "mid-stream"),
|
|
("messages", "request-body-fallbacks"),
|
|
("responses", "non-streaming"),
|
|
("responses", "before-first-byte"),
|
|
("responses", "mid-stream"),
|
|
("responses", "request-body-fallbacks"),
|
|
),
|
|
)
|
|
def test_native_messages_and_responses_routes_fall_back_and_bill_the_fallback(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
route: str,
|
|
failure_mode: str,
|
|
) -> None:
|
|
identity: Final = f"native-fallback-{route}-{failure_mode}-{uuid.uuid4().hex}"
|
|
is_messages: Final = route == "messages"
|
|
path: Final = "/v1/messages" if is_messages else "/v1/responses"
|
|
primary_model: Final = "claude-sonnet-4-5-20250929" if is_messages else "gpt-5.6"
|
|
fallback_model: Final = "claude-opus-4-5-20251101" if is_messages else "gpt-5.6"
|
|
stream: Final = failure_mode in ("before-first-byte", "mid-stream")
|
|
client_body: Final = (
|
|
{
|
|
"model": "primary",
|
|
"max_tokens": 12,
|
|
"messages": [{"role": "user", "content": identity}],
|
|
"stream": stream,
|
|
**({"fallbacks": ["fallback"]} if failure_mode == "request-body-fallbacks" else {}),
|
|
}
|
|
if is_messages
|
|
else {
|
|
"model": "primary",
|
|
"input": identity,
|
|
"stream": stream,
|
|
**({"fallbacks": ["fallback"]} if failure_mode == "request-body-fallbacks" else {}),
|
|
}
|
|
)
|
|
primary_body: Final = (
|
|
{
|
|
"model": primary_model,
|
|
"max_tokens": 12,
|
|
"messages": [{"role": "user", "content": identity}],
|
|
"stream": stream,
|
|
}
|
|
if is_messages
|
|
else {"model": primary_model, "input": identity, "stream": stream}
|
|
)
|
|
fallback_body: Final = {**primary_body, "model": fallback_model}
|
|
call_order: Final[SimpleQueue[str]] = SimpleQueue()
|
|
received: Final[SimpleQueue[tuple[str, Request]]] = SimpleQueue()
|
|
fallback_text: Final = f"fallback response {identity}"
|
|
primary_text: Final = f"primary response {identity}"
|
|
|
|
def response_for(model: str, text: str) -> Reply:
|
|
if is_messages:
|
|
if stream:
|
|
return Reply(content_type="text/event-stream", chunks=_native_anthropic_stream(model, identity, text))
|
|
return Reply(body=json.dumps(_native_anthropic_message(model, identity, text)).encode())
|
|
if stream:
|
|
return Reply(content_type="text/event-stream", chunks=_native_responses_stream(model, identity, text))
|
|
return Reply(body=json.dumps(_native_responses_response(model, identity, text)).encode())
|
|
|
|
def upstream(name: str, model: str) -> Callable[[Request], Reply]:
|
|
def handle(request: Request) -> Reply:
|
|
if request.method == "GET":
|
|
return Reply(body=b'{"object":"list","data":[]}')
|
|
assert request.method == "POST" and request.target == path, (request.method, request.target)
|
|
expected_body: Final = primary_body if name == "primary" else fallback_body
|
|
assert json.loads(request.body) == expected_body, request.body.decode()
|
|
if is_messages:
|
|
assert request.headers.get("x-api-key") == "synthetic-native-fallback-key", request.headers
|
|
assert request.headers.get("anthropic-version") == "2023-06-01", request.headers
|
|
assert "authorization" not in request.headers, request.headers
|
|
else:
|
|
assert request.headers.get("authorization") == "Bearer synthetic-native-fallback-key", request.headers
|
|
call_order.put(name)
|
|
received.put((name, request))
|
|
if name == "primary" and failure_mode == "mid-stream":
|
|
primary_events: Final = (
|
|
_native_anthropic_stream(primary_model, identity + "-primary", primary_text)
|
|
if is_messages
|
|
else _native_responses_stream(primary_model, identity + "-primary", primary_text)
|
|
)
|
|
first_event: Final = primary_events[0]
|
|
second_event: Final = (
|
|
primary_events[1]
|
|
if is_messages
|
|
else _native_sse_event(
|
|
{
|
|
"type": "response.in_progress",
|
|
"sequence_number": 1,
|
|
"response": {
|
|
**_native_responses_response(primary_model, identity + "-primary", primary_text),
|
|
"status": "in_progress",
|
|
"output": [],
|
|
},
|
|
}
|
|
)
|
|
)
|
|
return Reply(
|
|
content_type="text/event-stream",
|
|
chunks=(first_event, second_event, first_event),
|
|
abort_after=2,
|
|
pause_between_chunks=0.05,
|
|
)
|
|
if name == "primary" and failure_mode in (
|
|
"non-streaming",
|
|
"before-first-byte",
|
|
"request-body-fallbacks",
|
|
):
|
|
return Reply(
|
|
status=500,
|
|
body=b'{"error":{"message":"synthetic primary failure","type":"server_error","code":"internal_error"}}',
|
|
)
|
|
return response_for(model, fallback_text if name == "fallback" else primary_text)
|
|
|
|
return handle
|
|
|
|
with (
|
|
wire_server(upstream("primary", primary_model)) as primary,
|
|
wire_server(upstream("fallback", fallback_model)) as fallback,
|
|
):
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
config["model_list"] = [
|
|
{
|
|
"model_name": model_group,
|
|
"litellm_params": {
|
|
"model": f"{'anthropic' if is_messages else 'openai'}/{model}",
|
|
"api_key": "synthetic-native-fallback-key",
|
|
"api_base": server.url if is_messages else server.url + "/v1",
|
|
"input_cost_per_token": 0.001,
|
|
"output_cost_per_token": 0.002,
|
|
},
|
|
}
|
|
for model_group, model, server in (
|
|
("primary", primary_model, primary),
|
|
("fallback", fallback_model, fallback),
|
|
)
|
|
]
|
|
config["router_settings"] = {
|
|
"num_retries": 0,
|
|
"disable_cooldowns": True,
|
|
"fallbacks": [] if failure_mode == "request-body-fallbacks" else [{"primary": ["fallback"]}],
|
|
}
|
|
proxy_config: Final = tmp_path / f"native-{route}-{failure_mode}.yaml"
|
|
proxy_config.write_text(yaml.safe_dump(config))
|
|
with owned_proxy(gateway, tmp_path, {}, config=proxy_config) as candidate:
|
|
request_headers: Final = {
|
|
"x-litellm-call-id": identity,
|
|
**(
|
|
{"x-api-key": candidate.key, "anthropic-version": "2023-06-01"}
|
|
if is_messages
|
|
else {"Authorization": f"Bearer {candidate.key}"}
|
|
),
|
|
}
|
|
with candidate.client.stream("POST", path, json=client_body, headers=request_headers) as response:
|
|
response_text: Final = response.read().decode()
|
|
assert response.status_code == 200, response_text
|
|
assert response.headers.get("x-litellm-attempted-fallbacks") == "1", response_text
|
|
assert tuple(call_order.get_nowait() for _ in range(call_order.qsize())) == ("primary", "fallback")
|
|
assert tuple(
|
|
(name, request.method, request.target, json.loads(request.body))
|
|
for name, request in (received.get_nowait() for _ in range(received.qsize()))
|
|
) == (("primary", "POST", path, primary_body), ("fallback", "POST", path, fallback_body))
|
|
if stream:
|
|
event_lines: Final = tuple(line for line in sse_data_lines(response_text) if line != "[DONE]")
|
|
events: Final = tuple(JSON_OBJECT.validate_json(line) for line in event_lines)
|
|
event_types: Final = tuple(event["type"] for event in events)
|
|
expected_types: Final = (
|
|
(
|
|
"message_start",
|
|
"content_block_start",
|
|
"content_block_delta",
|
|
"content_block_stop",
|
|
"message_delta",
|
|
"message_stop",
|
|
)
|
|
if is_messages
|
|
else ("response.created", "response.output_text.delta", "response.completed")
|
|
)
|
|
assert event_types.count("message_start" if is_messages else "response.created") == 1, response_text
|
|
assert event_types.count("message_stop" if is_messages else "response.completed") == 1, response_text
|
|
assert event_types[-1] == ("message_stop" if is_messages else "response.completed"), response_text
|
|
assert event_types == expected_types, response_text
|
|
output_text: Final = (
|
|
"".join(
|
|
string_value(object_value(event["delta"])["text"])
|
|
for event in events
|
|
if event["type"] == "content_block_delta"
|
|
)
|
|
if is_messages
|
|
else "".join(
|
|
string_value(event["delta"])
|
|
for event in events
|
|
if event["type"] == "response.output_text.delta"
|
|
)
|
|
)
|
|
assert output_text == fallback_text, response_text
|
|
else:
|
|
response_body: Final = JSON_OBJECT.validate_json(response.content)
|
|
expected_response: Final = (
|
|
_native_anthropic_message(fallback_model, identity, fallback_text)
|
|
if is_messages
|
|
else {
|
|
"id": response_body["id"],
|
|
"created_at": 1,
|
|
"error": None,
|
|
"incomplete_details": None,
|
|
"instructions": None,
|
|
"metadata": None,
|
|
"model": fallback_model,
|
|
"object": "response",
|
|
"output": [
|
|
{
|
|
"id": "msg_" + identity,
|
|
"content": [
|
|
{
|
|
"annotations": [],
|
|
"text": fallback_text,
|
|
"type": "output_text",
|
|
"logprobs": None,
|
|
}
|
|
],
|
|
"role": "assistant",
|
|
"status": "completed",
|
|
"type": "message",
|
|
"phase": None,
|
|
}
|
|
],
|
|
"parallel_tool_calls": None,
|
|
"temperature": None,
|
|
"tool_choice": None,
|
|
"tools": None,
|
|
"top_p": None,
|
|
"max_output_tokens": None,
|
|
"previous_response_id": None,
|
|
"reasoning": None,
|
|
"status": "completed",
|
|
"text": None,
|
|
"truncation": None,
|
|
"usage": {
|
|
"input_tokens": 11,
|
|
"input_tokens_details": None,
|
|
"output_tokens": 4,
|
|
"output_tokens_details": None,
|
|
"total_tokens": 15,
|
|
"cost": None,
|
|
},
|
|
"user": None,
|
|
"store": None,
|
|
}
|
|
)
|
|
assert response_body == expected_response, response_text
|
|
if not is_messages:
|
|
assert response_body["id"].startswith("resp_"), response_text
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
"SELECT model, model_group, spend, prompt_tokens, completion_tokens, status "
|
|
'FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s',
|
|
(identity, identity),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert rows[0]["model_group"] == "fallback", rows
|
|
assert rows[0]["status"] == "success", rows
|
|
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (11, 4), rows
|
|
assert float(rows[0]["spend"]) == pytest.approx(0.019), rows
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
|
|
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
|
|
import litellm
|
|
|
|
for truncated in (True, False):
|
|
with wire_server(
|
|
lambda request, truncated=truncated: Reply(
|
|
content_type="text/event-stream",
|
|
chunks=text_stream("stream-truncated"),
|
|
abort_after=1 if truncated else None,
|
|
)
|
|
) as wire:
|
|
stream: Final = litellm.completion(
|
|
model="openai/gpt-4o-mini",
|
|
api_base=wire.url + "/v1",
|
|
api_key="synthetic-stream-key",
|
|
messages=[{"role": "user", "content": "truncation control"}],
|
|
stream=True,
|
|
timeout=5,
|
|
num_retries=0,
|
|
)
|
|
try:
|
|
if truncated:
|
|
with pytest.raises(
|
|
litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read"
|
|
) as failure:
|
|
tuple(stream)
|
|
assert isinstance(failure.value.original_exception, litellm.APIConnectionError)
|
|
assert failure.value.generated_content == "Hello "
|
|
assert failure.value.is_pre_first_chunk is False
|
|
else:
|
|
chunks: Final = tuple(stream)
|
|
assert (
|
|
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
|
|
== "Hello 雪 café"
|
|
)
|
|
assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices)
|
|
finally:
|
|
asyncio.run(stream.aclose())
|
|
assert len(wire.drain()) == 1
|
|
|
|
|
|
@pytest.mark.covers("other.streaming.cancellation.closes_actual_provider_connection")
|
|
def test_client_cancellation_releases_the_actual_provider_connection() -> None:
|
|
import litellm
|
|
|
|
gate: Final = threading.Event()
|
|
frames: Final = (
|
|
frame("stream-cancel", {"role": "assistant", "content": "first"}),
|
|
b":" + b"x" * 4_000_000 + b"\n\n",
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
with wire_server(
|
|
lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)
|
|
) as wire:
|
|
stream: Final = litellm.completion(
|
|
model="openai/gpt-4o-mini",
|
|
api_base=wire.url + "/v1",
|
|
api_key="synthetic-stream-key",
|
|
messages=[{"role": "user", "content": "cancellation control"}],
|
|
stream=True,
|
|
timeout=5,
|
|
num_retries=0,
|
|
)
|
|
try:
|
|
first: Final = next(stream)
|
|
assert first.choices[0].delta.content == "first"
|
|
finally:
|
|
try:
|
|
asyncio.run(stream.aclose())
|
|
finally:
|
|
gate.set()
|
|
assert wire.disconnected.get(timeout=5) == "/v1/chat/completions"
|
|
assert len(wire.drain()) == 1
|