mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
test(integration): cover chat_completions endpoint contracts (#44688)
* test(integration): cover chat_completions endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin per-call prompts and both n=2 chat choices Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover mixed path-model auth, string guardrail rejection and typed text-completion forms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop the prompt-spellings cache workaround and cover Azure SDK text completions on the deployments alias --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
11e91f464a
commit
d3cc77def4
5 changed files with 1198 additions and 3 deletions
|
|
@ -11,6 +11,7 @@ import pytest
|
|||
import websockets
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
|
|
@ -382,3 +383,115 @@ def test_max_budget_allows_then_denies_and_records_spend(gateway: Gateway) -> No
|
|||
second: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert second.status_code == 422, second.text
|
||||
assert "budget_exceeded" in second.text, second.text
|
||||
|
||||
|
||||
_PATH_MODEL_BACKEND: Final = "gpt-5.4-mini"
|
||||
_PATH_MODEL_ROUTES: Final = (
|
||||
("/openai/deployments/{model}/chat/completions", "chat"),
|
||||
("/engines/{model}/chat/completions", "chat"),
|
||||
("/openai/deployments/{model}/completions", "text"),
|
||||
("/engines/{model}/completions", "text"),
|
||||
)
|
||||
|
||||
|
||||
def _path_model_reply(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/chat/completions"), request
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-" + _marker(),
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _PATH_MODEL_BACKEND,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "path"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _path_model_body(kind: str, prompt: str) -> dict[str, JsonValue]:
|
||||
return {"messages": [{"role": "user", "content": prompt}]} if kind == "chat" else {"prompt": prompt}
|
||||
|
||||
|
||||
def _upstream_bodies(wire: Wire) -> list[dict[str, JsonValue]]:
|
||||
return [json.loads(request.body) for request in wire.drain()]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("route", "kind"), _PATH_MODEL_ROUTES)
|
||||
def test_path_only_model_alias_refuses_an_unallowed_deployment(gateway: Gateway, route: str, kind: str) -> None:
|
||||
with wire_server(_path_model_reply) as allowed_wire, wire_server(_path_model_reply) as other_wire:
|
||||
with gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=allowed_wire.url)
|
||||
other: Final = scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=other_wire.url)
|
||||
key: Final = scenario.key(models=[allowed])
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
route.format(model=other),
|
||||
_path_model_body(kind, _marker()),
|
||||
key=key,
|
||||
params={"api-version": "2024-10-21"},
|
||||
)
|
||||
assert (_upstream_bodies(other_wire), _upstream_bodies(allowed_wire)) == ([], []), response.text
|
||||
assert response.status_code == 403, response.text
|
||||
assert f"The requested model '{other}' is not available for this API key" in response.text, response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("route", "kind"), _PATH_MODEL_ROUTES)
|
||||
def test_mixed_body_and_path_model_is_refused_or_served_by_the_body_model(
|
||||
gateway: Gateway, route: str, kind: str
|
||||
) -> None:
|
||||
pytest.skip(
|
||||
"BUG: a key scoped to model A that sends model A in the body and an unallowed model B in the "
|
||||
"/openai/deployments or /engines path is served by deployment B"
|
||||
)
|
||||
with wire_server(_path_model_reply) as allowed_wire, wire_server(_path_model_reply) as other_wire:
|
||||
with gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=allowed_wire.url)
|
||||
other: Final = scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=other_wire.url)
|
||||
key: Final = scenario.key(models=[allowed])
|
||||
prompt: Final = _marker()
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
route.format(model=other),
|
||||
{"model": allowed, **_path_model_body(kind, prompt)},
|
||||
key=key,
|
||||
params={"api-version": "2024-10-21"},
|
||||
)
|
||||
other_bodies: Final = _upstream_bodies(other_wire)
|
||||
allowed_bodies: Final = _upstream_bodies(allowed_wire)
|
||||
outcome: Final = (response.status_code, other_bodies, allowed_bodies)
|
||||
refused: Final = outcome[0] in (401, 403) and outcome[1:] == ([], [])
|
||||
served_by_body_model: Final = outcome == (
|
||||
200,
|
||||
[],
|
||||
[{"model": _PATH_MODEL_BACKEND, "messages": [{"role": "user", "content": prompt}]}],
|
||||
)
|
||||
assert refused or served_by_body_model, (outcome, response.text)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("route", "kind"), _PATH_MODEL_ROUTES)
|
||||
def test_path_only_model_alias_reaches_the_allowed_deployment(gateway: Gateway, route: str, kind: str) -> None:
|
||||
with wire_server(_path_model_reply) as allowed_wire, wire_server(_path_model_reply) as other_wire:
|
||||
with gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=allowed_wire.url)
|
||||
scenario.model(model=f"openai/{_PATH_MODEL_BACKEND}", api_base=other_wire.url)
|
||||
key: Final = scenario.key(models=[allowed])
|
||||
prompt: Final = _marker()
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
route.format(model=allowed),
|
||||
_path_model_body(kind, prompt),
|
||||
key=key,
|
||||
params={"api-version": "2024-10-21"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = response.json()
|
||||
assert [
|
||||
choice["message"]["content"] if kind == "chat" else choice["text"] for choice in payload["choices"]
|
||||
] == ["path"], response.text
|
||||
assert payload["object"] == ("chat.completion" if kind == "chat" else "text_completion"), response.text
|
||||
assert _upstream_bodies(other_wire) == [], response.text
|
||||
assert _upstream_bodies(allowed_wire) == [
|
||||
{"model": _PATH_MODEL_BACKEND, "messages": [{"role": "user", "content": prompt}]}
|
||||
], response.text
|
||||
|
|
|
|||
|
|
@ -1,14 +1,14 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from importlib.metadata import version
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
|
||||
|
||||
|
||||
@pytest.mark.covers("other.compatibility.openai.retained_client_parses_tools_and_usage")
|
||||
|
|
@ -107,3 +107,373 @@ async def test_retained_openai_clients_parse_real_proxy_tool_and_usage_responses
|
|||
assert response.choices[0].message.content == "Synthetic answer: 8"
|
||||
assert not response.choices[0].message.tool_calls
|
||||
assert len(wire.drain()) == 4
|
||||
|
||||
|
||||
_AZURE_API_VERSION: Final = "2024-10-21"
|
||||
_CHAT_BACKEND: Final = "gpt-5.4-mini"
|
||||
_TEXT_BACKEND: Final = "gpt-3.5-turbo-instruct"
|
||||
_USAGE: Final = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
|
||||
|
||||
|
||||
def _sse(events: tuple[dict[str, object], ...]) -> Reply:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"data: {json.dumps(event)}\n\n".encode() for event in events) + (b"data: [DONE]\n\n",),
|
||||
)
|
||||
|
||||
|
||||
def _data_lines(text: str) -> tuple[str, ...]:
|
||||
return tuple(line for line in text.splitlines() if line.startswith("data: "))
|
||||
|
||||
|
||||
def _chat_provider(identity: str, prompts: set[str]) -> Callable[[Request], Reply]:
|
||||
def provider(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/chat/completions"), request
|
||||
assert request.headers["authorization"] == "Bearer synthetic-azure-consumer-key"
|
||||
body: Final = json.loads(request.body)
|
||||
assert body["messages"][0]["content"] in prompts, body
|
||||
messages: Final = [{"role": "user", "content": body["messages"][0]["content"]}]
|
||||
if body.get("stream"):
|
||||
assert body == {
|
||||
"model": _CHAT_BACKEND,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
return _sse(
|
||||
(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _CHAT_BACKEND,
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "Hel"}, "finish_reason": None}
|
||||
],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _CHAT_BACKEND,
|
||||
"choices": [{"index": 0, "delta": {"content": "lo"}, "finish_reason": "stop"}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _CHAT_BACKEND,
|
||||
"choices": [],
|
||||
"usage": _USAGE,
|
||||
},
|
||||
)
|
||||
)
|
||||
assert body == {"model": _CHAT_BACKEND, "messages": messages}
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _CHAT_BACKEND,
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": _USAGE,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
return provider
|
||||
|
||||
|
||||
async def test_azure_openai_clients_reach_openai_deployment_through_deployments_alias(gateway: Gateway) -> None:
|
||||
identity: Final = f"chatcmpl-azure-{uuid.uuid4().hex}"
|
||||
prompts: Final = {f"{kind} {identity}" for kind in ("sync", "sync-stream", "async", "async-stream", "raw-stream")}
|
||||
with wire_server(_chat_provider(identity, prompts)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_CHAT_BACKEND}", api_base=wire.url, api_key="synthetic-azure-consumer-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
endpoint: Final = str(gateway.client.base_url).rstrip("/")
|
||||
sent: Final[list[str]] = []
|
||||
|
||||
def record(request: httpx.Request) -> None:
|
||||
assert request.headers["api-key"] == key
|
||||
sent.append(f"{request.url.path}?{request.url.query.decode()}")
|
||||
|
||||
async def record_async(request: httpx.Request) -> None:
|
||||
record(request)
|
||||
|
||||
with AzureOpenAI(
|
||||
azure_endpoint=endpoint,
|
||||
api_key=key,
|
||||
api_version=_AZURE_API_VERSION,
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=10, trust_env=False, event_hooks={"request": [record]}),
|
||||
) as sync:
|
||||
completion: Final = sync.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": f"sync {identity}"}]
|
||||
)
|
||||
sync_chunks: Final = list(
|
||||
sync.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": f"sync-stream {identity}"}],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
)
|
||||
async with AsyncAzureOpenAI(
|
||||
azure_endpoint=endpoint,
|
||||
api_key=key,
|
||||
api_version=_AZURE_API_VERSION,
|
||||
max_retries=0,
|
||||
http_client=httpx.AsyncClient(timeout=10, trust_env=False, event_hooks={"request": [record_async]}),
|
||||
) as asynchronous:
|
||||
async_completion: Final = await asynchronous.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": f"async {identity}"}]
|
||||
)
|
||||
async_stream: Final = await asynchronous.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": f"async-stream {identity}"}],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
async_chunks: Final = [chunk async for chunk in async_stream]
|
||||
raw: Final = gateway.client.post(
|
||||
f"/openai/deployments/{model}/chat/completions",
|
||||
params={"api-version": _AZURE_API_VERSION},
|
||||
headers={"api-key": key},
|
||||
json={
|
||||
"messages": [{"role": "user", "content": f"raw-stream {identity}"}],
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
},
|
||||
)
|
||||
assert sent == [f"/openai/deployments/{model}/chat/completions?api-version={_AZURE_API_VERSION}"] * 4
|
||||
for response in (completion, async_completion):
|
||||
assert (response.id, response.object, response.model) == (identity, "chat.completion", model), response
|
||||
assert [(choice.message.content, choice.finish_reason) for choice in response.choices] == [
|
||||
("Hello", "stop")
|
||||
]
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (
|
||||
3,
|
||||
2,
|
||||
5,
|
||||
), response.usage
|
||||
for chunks in (sync_chunks, async_chunks):
|
||||
assert {(chunk.id, chunk.object, chunk.model) for chunk in chunks} == {
|
||||
(identity, "chat.completion.chunk", model)
|
||||
}, chunks
|
||||
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello"
|
||||
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
|
||||
"stop"
|
||||
], chunks
|
||||
assert [chunk.usage is None for chunk in chunks] == [True] * (len(chunks) - 1) + [False], chunks
|
||||
assert (
|
||||
chunks[-1].usage.prompt_tokens,
|
||||
chunks[-1].usage.completion_tokens,
|
||||
chunks[-1].usage.total_tokens,
|
||||
) == (3, 2, 5), chunks[-1]
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert raw.headers["content-type"].startswith("text/event-stream"), raw.text
|
||||
lines: Final = _data_lines(raw.text)
|
||||
assert lines[-1] == "data: [DONE]", raw.text
|
||||
assert (
|
||||
"".join(
|
||||
choice["delta"].get("content") or ""
|
||||
for line in lines[:-1]
|
||||
for choice in json.loads(line.removeprefix("data: "))["choices"]
|
||||
)
|
||||
== "Hello"
|
||||
), raw.text
|
||||
assert json.loads(lines[-2].removeprefix("data: "))["usage"]["total_tokens"] == 5, raw.text
|
||||
assert [
|
||||
(request.method, request.target, json.loads(request.body)["messages"][0]["content"])
|
||||
for request in wire.drain()
|
||||
] == [
|
||||
("POST", "/chat/completions", f"{kind} {identity}")
|
||||
for kind in ("sync", "sync-stream", "async", "async-stream", "raw-stream")
|
||||
]
|
||||
|
||||
|
||||
def test_azure_openai_text_completion_reaches_openai_text_deployment_through_deployments_alias(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
identity: Final = f"cmpl-azure-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"azure text completion {uuid.uuid4().hex}"
|
||||
|
||||
def provider(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == "Bearer synthetic-text-consumer-key"
|
||||
assert json.loads(request.body) == {"model": _TEXT_BACKEND, "prompt": prompt}
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [{"index": 0, "text": "Azure text", "finish_reason": "stop", "logprobs": None}],
|
||||
"usage": _USAGE,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(provider) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key="synthetic-text-consumer-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
endpoint: Final = str(gateway.client.base_url).rstrip("/")
|
||||
with AzureOpenAI(
|
||||
api_key=key,
|
||||
azure_endpoint=endpoint,
|
||||
api_version=_AZURE_API_VERSION,
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=10, trust_env=False),
|
||||
) as client:
|
||||
completion: Final = client.completions.create(model=model, prompt=prompt)
|
||||
assert completion.id == identity, completion
|
||||
assert completion.object == "text_completion", completion
|
||||
assert [(choice.index, choice.text) for choice in completion.choices] == [(0, "Azure text")], completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
|
||||
|
||||
def _text_provider(identity: str, prompts: set[str]) -> Callable[[Request], Reply]:
|
||||
def provider(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == "Bearer synthetic-text-consumer-key"
|
||||
body: Final = json.loads(request.body)
|
||||
assert body["prompt"] in prompts, body
|
||||
if body.get("stream"):
|
||||
assert body == {
|
||||
"model": _TEXT_BACKEND,
|
||||
"prompt": body["prompt"],
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
return _sse(
|
||||
(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [{"index": 0, "text": "Hel", "finish_reason": None, "logprobs": None}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [{"index": 0, "text": "lo", "finish_reason": "stop", "logprobs": None}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [],
|
||||
"usage": _USAGE,
|
||||
},
|
||||
)
|
||||
)
|
||||
assert body == {"model": _TEXT_BACKEND, "prompt": body["prompt"], "n": 2}
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [
|
||||
{"index": 0, "text": "first", "finish_reason": "stop", "logprobs": None},
|
||||
{"index": 1, "text": "second", "finish_reason": "length", "logprobs": None},
|
||||
],
|
||||
"usage": _USAGE,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
return provider
|
||||
|
||||
|
||||
async def test_openai_clients_parse_text_completion_streams_and_choices(gateway: Gateway) -> None:
|
||||
identity: Final = f"cmpl-consumer-{uuid.uuid4().hex}"
|
||||
prompts: Final = {f"{kind} {identity}" for kind in ("sync-stream", "async-stream", "raw-stream", "sync-n")}
|
||||
with wire_server(_text_provider(identity, prompts)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key="synthetic-text-consumer-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
base_url: Final = str(gateway.client.base_url).rstrip("/") + "/v1"
|
||||
streamed: Final = {"stream": True, "stream_options": {"include_usage": True}}
|
||||
with OpenAI(
|
||||
api_key=key, base_url=base_url, max_retries=0, http_client=httpx.Client(timeout=10, trust_env=False)
|
||||
) as sync:
|
||||
sync_chunks: Final = list(
|
||||
sync.completions.create(model=model, prompt=f"sync-stream {identity}", **streamed)
|
||||
)
|
||||
choices: Final = sync.completions.create(model=model, prompt=f"sync-n {identity}", n=2)
|
||||
async with AsyncOpenAI(
|
||||
api_key=key, base_url=base_url, max_retries=0, http_client=httpx.AsyncClient(timeout=10, trust_env=False)
|
||||
) as asynchronous:
|
||||
async_stream: Final = await asynchronous.completions.create(
|
||||
model=model, prompt=f"async-stream {identity}", **streamed
|
||||
)
|
||||
async_chunks: Final = [chunk async for chunk in async_stream]
|
||||
raw: Final = gateway.request(
|
||||
"POST", "/v1/completions", {"model": model, "prompt": f"raw-stream {identity}", **streamed}, key=key
|
||||
)
|
||||
for chunks in (sync_chunks, async_chunks):
|
||||
assert {(chunk.id, chunk.object, chunk.model) for chunk in chunks} == {
|
||||
(identity, "text_completion", model)
|
||||
}, chunks
|
||||
assert "".join(choice.text or "" for chunk in chunks for choice in chunk.choices) == "Hello", chunks
|
||||
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
|
||||
"stop"
|
||||
], chunks
|
||||
assert [chunk.usage is None for chunk in chunks] == [True] * (len(chunks) - 1) + [False], chunks
|
||||
assert (
|
||||
chunks[-1].usage.prompt_tokens,
|
||||
chunks[-1].usage.completion_tokens,
|
||||
chunks[-1].usage.total_tokens,
|
||||
) == (3, 2, 5), chunks[-1]
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert raw.headers["content-type"].startswith("text/event-stream"), raw.text
|
||||
lines: Final = _data_lines(raw.text)
|
||||
assert lines[-1] == "data: [DONE]", raw.text
|
||||
frames: Final = [json.loads(line.removeprefix("data: ")) for line in lines[:-1]]
|
||||
assert {frame["object"] for frame in frames} == {"text_completion"}, raw.text
|
||||
assert "".join(choice.get("text") or "" for frame in frames for choice in frame["choices"]) == "Hello"
|
||||
assert (choices.id, choices.object, choices.model) == (identity, "text_completion", model), choices
|
||||
assert [(choice.index, choice.text, choice.finish_reason) for choice in choices.choices] == [
|
||||
(0, "first", "stop"),
|
||||
(1, "second", "length"),
|
||||
], choices
|
||||
assert [(request.method, request.target, json.loads(request.body)["prompt"]) for request in wire.drain()] == [
|
||||
("POST", "/completions", f"{kind} {identity}")
|
||||
for kind in ("sync-stream", "sync-n", "async-stream", "raw-stream")
|
||||
]
|
||||
|
||||
|
||||
def test_text_completion_usage_chunk_has_empty_choices(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: the /v1/completions include_usage chunk carries choices [{index: 0}] instead of []")
|
||||
identity: Final = f"cmpl-usage-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"raw-stream {identity}"
|
||||
with wire_server(_text_provider(identity, {prompt})) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key="synthetic-text-consumer-key"
|
||||
)
|
||||
raw: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
{"model": model, "prompt": prompt, "stream": True, "stream_options": {"include_usage": True}},
|
||||
)
|
||||
assert raw.status_code == 200, raw.text
|
||||
lines: Final = _data_lines(raw.text)
|
||||
assert lines[-1] == "data: [DONE]", raw.text
|
||||
usage_frame: Final = json.loads(lines[-2].removeprefix("data: "))
|
||||
assert (usage_frame["choices"], usage_frame["usage"]["total_tokens"]) == ([], 5), raw.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from integration._support.mcp import mcp_peer, register_mcp, tool_names
|
|||
from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
from pydantic import JsonValue
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.rewrite_reaches_correct_anthropic_positions")
|
||||
|
|
@ -1658,6 +1658,203 @@ def test_chat_pre_call_denial_returns_content_filter(gateway: Gateway, tmp_path:
|
|||
) == (0, 0, 0), body["usage"]
|
||||
|
||||
|
||||
def _refusing_provider(request: Request) -> Reply:
|
||||
raise AssertionError(f"denied text completion reached the provider: {request!r}")
|
||||
|
||||
|
||||
_STRING_REJECTION: Final = "__DENIAL__"
|
||||
_STRING_REJECTER_MODULE: Final = """from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
|
||||
class StringRejecter(CustomGuardrail):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
return "__DENIAL__"
|
||||
"""
|
||||
|
||||
|
||||
class _TextCompletionChoice(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class _TextCompletionResponse(BaseModel):
|
||||
object: str
|
||||
choices: list[_TextCompletionChoice]
|
||||
|
||||
|
||||
class _TextCompletionStreamChoice(BaseModel):
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _TextCompletionStreamFrame(BaseModel):
|
||||
object: str
|
||||
choices: list[_TextCompletionStreamChoice]
|
||||
|
||||
|
||||
def _string_rejection_config(tmp_path: Path, identity: str) -> Path:
|
||||
module: Final = "string_rejecter_" + uuid.uuid4().hex
|
||||
(tmp_path / f"{module}.py").write_text(_STRING_REJECTER_MODULE)
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": identity,
|
||||
"litellm_params": {"guardrail": f"{module}.StringRejecter", "mode": "pre_call"},
|
||||
}
|
||||
]
|
||||
path: Final = tmp_path / "string-rejection.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def test_text_completion_pre_call_denial_returns_text_completion_without_upstream(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
with (
|
||||
wire_server(_refusing_provider) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=config) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="openai/gpt-3.5-turbo-instruct", api_base=wire.url)
|
||||
key: Final = scenario.key(models=[model])
|
||||
with OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(candidate.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
) as client:
|
||||
completion: Final = client.completions.create(
|
||||
model=model, prompt="say hi", extra_body={"guardrails": [identity]}
|
||||
)
|
||||
assert completion.object == "text_completion", completion
|
||||
assert [choice.text for choice in completion.choices] == [_RESPONSES_DENIAL], completion
|
||||
assert wire.drain() == (), completion
|
||||
|
||||
|
||||
def test_text_completion_pre_call_denial_streams_text_completion_without_upstream(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
pytest.skip("BUG: a streamed /v1/completions pre_call denial sends one frame with no text and no [DONE]")
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
with (
|
||||
wire_server(_refusing_provider) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=config) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="openai/gpt-3.5-turbo-instruct", api_base=wire.url)
|
||||
key: Final = scenario.key(models=[model])
|
||||
raw: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
{"model": model, "prompt": "say hi", "stream": True, "guardrails": [identity]},
|
||||
key=key,
|
||||
)
|
||||
with OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(candidate.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
) as client:
|
||||
chunks: Final = list(
|
||||
client.completions.create(
|
||||
model=model, prompt="say hi", stream=True, extra_body={"guardrails": [identity]}
|
||||
)
|
||||
)
|
||||
assert {chunk.object for chunk in chunks} == {"text_completion"}, chunks
|
||||
assert "".join(choice.text or "" for chunk in chunks for choice in chunk.choices) == _RESPONSES_DENIAL, chunks
|
||||
assert raw.headers["content-type"].startswith("text/event-stream"), raw.text
|
||||
lines: Final = tuple(line for line in raw.text.split("\n") if line.startswith("data: "))
|
||||
assert lines[-1] == "data: [DONE]", raw.text
|
||||
frames: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1])
|
||||
assert {frame["object"] for frame in frames} == {"text_completion"}, raw.text
|
||||
assert "".join(choice.get("text") or "" for frame in frames for choice in frame["choices"]) == _RESPONSES_DENIAL
|
||||
assert wire.drain() == (), raw.text
|
||||
|
||||
|
||||
def test_text_completion_pre_call_string_rejection_returns_text_completion_without_upstream(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
pytest.skip("BUG: /v1/completions pre_call string rejection returns HTTP 400 instead of a text_completion response")
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _string_rejection_config(tmp_path, identity)
|
||||
with (
|
||||
wire_server(_refusing_provider) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=config) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="openai/gpt-3.5-turbo-instruct", api_base=wire.url)
|
||||
key: Final = scenario.key(models=[model])
|
||||
raw: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
{"model": model, "prompt": "say hi", "guardrails": [identity]},
|
||||
key=key,
|
||||
)
|
||||
assert raw.status_code == 200, raw.text
|
||||
payload: Final = _TextCompletionResponse.model_validate_json(raw.content)
|
||||
assert payload.object == "text_completion", raw.text
|
||||
assert [choice.text for choice in payload.choices] == [_STRING_REJECTION], raw.text
|
||||
with OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(candidate.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
) as client:
|
||||
completion: Final = client.completions.create(
|
||||
model=model, prompt="say hi", extra_body={"guardrails": [identity]}
|
||||
)
|
||||
assert completion.object == "text_completion", completion
|
||||
assert [choice.text for choice in completion.choices] == [_STRING_REJECTION], completion
|
||||
assert wire.drain() == (), raw.text
|
||||
|
||||
|
||||
def test_text_completion_pre_call_string_rejection_streams_text_completion_without_upstream(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
pytest.skip(
|
||||
"BUG: /v1/completions pre_call string rejection returns HTTP 400 JSON instead of a text/event-stream denial"
|
||||
)
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _string_rejection_config(tmp_path, identity)
|
||||
with (
|
||||
wire_server(_refusing_provider) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, config=config) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="openai/gpt-3.5-turbo-instruct", api_base=wire.url)
|
||||
key: Final = scenario.key(models=[model])
|
||||
raw: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
{"model": model, "prompt": "say hi", "stream": True, "guardrails": [identity]},
|
||||
key=key,
|
||||
)
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert raw.headers["content-type"].startswith("text/event-stream"), raw.text
|
||||
lines: Final = tuple(line for line in raw.text.split("\n") if line.startswith("data: "))
|
||||
assert lines[-1] == "data: [DONE]", raw.text
|
||||
frames: Final = tuple(
|
||||
_TextCompletionStreamFrame.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]
|
||||
)
|
||||
assert tuple(frame.object for frame in frames) == ("text_completion",) * len(frames), raw.text
|
||||
assert "".join(choice.text or "" for frame in frames for choice in frame.choices) == _STRING_REJECTION, raw.text
|
||||
with OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(candidate.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
) as client:
|
||||
chunks: Final = tuple(
|
||||
client.completions.create(
|
||||
model=model, prompt="say hi", stream=True, extra_body={"guardrails": [identity]}
|
||||
)
|
||||
)
|
||||
assert tuple(chunk.object for chunk in chunks) == ("text_completion",) * len(chunks), chunks
|
||||
assert "".join(choice.text or "" for chunk in chunks for choice in chunk.choices) == _STRING_REJECTION, chunks
|
||||
assert wire.drain() == (), raw.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.messages_pre_call_denial_returns_message")
|
||||
def test_messages_pre_call_denial_returns_message(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
|
|
|
|||
|
|
@ -1,13 +1,17 @@
|
|||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from hashlib import sha256
|
||||
from itertools import chain, count
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_BACKEND: Final = "gpt-5.4-mini"
|
||||
|
|
@ -255,6 +259,198 @@ def test_azure_gpt_6_bridged_stream_returns_text_and_tool_call_on_one_choice(gat
|
|||
]
|
||||
|
||||
|
||||
_SPEND_ROW: Final = 'SELECT request_tags, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
|
||||
_TAGGED_METADATA: Final = json.dumps({"tags": ["t1"]})
|
||||
|
||||
|
||||
def _spend_row(request_id: str) -> tuple[list[JsonValue], dict[str, JsonValue]]:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (request_id,)), lambda values: len(values) == 1, seconds=70)
|
||||
tags: Final = rows[0]["request_tags"]
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
assert isinstance(tags, list) and isinstance(metadata, dict), rows
|
||||
return tags, metadata
|
||||
|
||||
|
||||
def _openai(gateway: Gateway, key: str) -> OpenAI:
|
||||
return OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
)
|
||||
|
||||
|
||||
def _tagged_chat_wire(identity: str, expected: dict[str, JsonValue]) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/chat/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert request.headers["content-type"] == "application/json"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == expected
|
||||
return Reply(body=_completion(identity, "tagged"))
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def test_openai_chat_json_string_metadata_reaches_upstream_and_keeps_key_identity(gateway: Gateway) -> None:
|
||||
identity: Final = f"openai-string-metadata-{uuid.uuid4().hex}"
|
||||
expected: Final = {"model": _BACKEND, "messages": [{"role": "user", "content": identity}]}
|
||||
|
||||
with wire_server(_tagged_chat_wire(identity, expected)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
team: Final = scenario.team()
|
||||
user_id: Final = f"integration-{uuid.uuid4().hex}"
|
||||
key: Final = scenario.key(models=[model], team_id=team, user_id=user_id)
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": identity}],
|
||||
extra_body={"metadata": _TAGGED_METADATA},
|
||||
)
|
||||
assert (completion.id, completion.model, completion.choices[0].message.content) == (identity, model, "tagged")
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
||||
_, metadata = _spend_row(identity)
|
||||
assert (
|
||||
metadata["user_api_key"],
|
||||
metadata["user_api_key_user_id"],
|
||||
metadata["user_api_key_team_id"],
|
||||
) == (sha256(key.encode()).hexdigest(), user_id, team), metadata
|
||||
|
||||
|
||||
def test_openai_chat_json_string_metadata_tags_the_spend_row(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: /v1/chat/completions replaces JSON-string metadata with {} so the client tag never reaches SpendLogs"
|
||||
)
|
||||
identity: Final = f"openai-string-metadata-tags-{uuid.uuid4().hex}"
|
||||
expected: Final = {"model": _BACKEND, "messages": [{"role": "user", "content": identity}]}
|
||||
|
||||
with wire_server(_tagged_chat_wire(identity, expected)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": identity}], "metadata": _TAGGED_METADATA},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["id"] == identity, response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
||||
tags, _ = _spend_row(identity)
|
||||
assert "t1" in tags, tags
|
||||
|
||||
|
||||
def test_openai_chat_multipart_form_fields_reach_upstream_as_json_types(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: multipart /v1/chat/completions returns 400 because messages stays a JSON string")
|
||||
identity: Final = f"openai-form-chat-{uuid.uuid4().hex}"
|
||||
messages: Final = [{"role": "user", "content": identity}]
|
||||
expected: Final = {"model": _BACKEND, "messages": messages, "temperature": 0.2}
|
||||
|
||||
with wire_server(_tagged_chat_wire(identity, expected)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
fields: Final = {
|
||||
"model": model,
|
||||
"messages": json.dumps(messages),
|
||||
"temperature": "0.2",
|
||||
"stream": "false",
|
||||
"metadata": _TAGGED_METADATA,
|
||||
}
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/chat/completions",
|
||||
files=[(name, (None, value)) for name, value in fields.items()],
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.request.headers["content-type"].startswith("multipart/form-data"), response.request.headers
|
||||
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
assert (payload["id"], payload["object"]) == (identity, "chat.completion"), response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
||||
tags, _ = _spend_row(identity)
|
||||
assert "t1" in tags, tags
|
||||
|
||||
|
||||
def test_openai_chat_sdk_parameter_set_reaches_upstream_whole(gateway: Gateway) -> None:
|
||||
identity: Final = f"openai-full-params-{uuid.uuid4().hex}"
|
||||
tools: Final = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Add two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
response_format: Final = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "sum",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"total": {"type": "integer"}},
|
||||
"required": ["total"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
messages: Final = [{"role": "user", "content": f"Add 3 and 5. {identity}"}]
|
||||
parameters: Final = {
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"tool_choice": {"type": "function", "function": {"name": "add"}},
|
||||
"parallel_tool_calls": False,
|
||||
"response_format": response_format,
|
||||
"max_completion_tokens": 64,
|
||||
"seed": 7,
|
||||
"n": 2,
|
||||
"logprobs": True,
|
||||
"top_logprobs": 2,
|
||||
"user": "integration-end-user",
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/chat/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {"model": _BACKEND, **parameters}
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _BACKEND,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": '{"total": 8}'},
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"message": {"role": "assistant", "content": '{"total": 9}'},
|
||||
"finish_reason": "length",
|
||||
},
|
||||
],
|
||||
"usage": {"prompt_tokens": 19, "completion_tokens": 14, "total_tokens": 33},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.chat.completions.create(model=model, **parameters)
|
||||
assert (completion.id, completion.model) == (identity, model), completion
|
||||
assert [(choice.index, choice.finish_reason, choice.message.content) for choice in completion.choices] == [
|
||||
(0, "stop", '{"total": 8}'),
|
||||
(1, "length", '{"total": 9}'),
|
||||
], completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
||||
|
||||
|
||||
_RESPONSES_TARGET: Final = "/openai/responses?api-version=2025-04-01-preview"
|
||||
|
||||
|
||||
|
|
|
|||
319
tests/integration/providers/test_openai_completion_wire.py
Normal file
319
tests/integration/providers/test_openai_completion_wire.py
Normal file
|
|
@ -0,0 +1,319 @@
|
|||
import json
|
||||
import uuid
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
_TEXT_BACKEND: Final = "gpt-3.5-turbo-instruct"
|
||||
_CHAT_BACKEND: Final = "gpt-5.4-mini"
|
||||
_API_KEY: Final = "synthetic-openai-key"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_SPEND_ROW: Final = 'SELECT request_tags, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
|
||||
|
||||
|
||||
class _TextChoice(BaseModel):
|
||||
index: int
|
||||
text: str
|
||||
finish_reason: str
|
||||
|
||||
|
||||
class _TextCompletion(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
model: str
|
||||
choices: list[_TextChoice]
|
||||
|
||||
|
||||
def _text_completion(identity: str, texts: tuple[str, ...]) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": _TEXT_BACKEND,
|
||||
"choices": [
|
||||
{"index": index, "text": text, "finish_reason": "stop", "logprobs": None}
|
||||
for index, text in enumerate(texts)
|
||||
],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def _chat_completion(identity: str, content: str) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _CHAT_BACKEND,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def _openai(gateway: Gateway, key: str) -> OpenAI:
|
||||
return OpenAI(
|
||||
api_key=key,
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
)
|
||||
|
||||
|
||||
def _form_fields(fields: dict[str, str]) -> list[tuple[str, tuple[None, str]]]:
|
||||
return [(name, (None, value)) for name, value in fields.items()]
|
||||
|
||||
|
||||
def _assert_text_completion_request(request: Request, prompt: str) -> None:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert request.headers["content-type"] == "application/json"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {"model": _TEXT_BACKEND, "prompt": prompt}
|
||||
|
||||
|
||||
def _assert_spend_row_carries_client_tag_and_key_identity(
|
||||
request_id: str, key: str, user_id: str, team_id: str
|
||||
) -> None:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (request_id,)), lambda values: len(values) == 1, seconds=70)
|
||||
tags: Final = rows[0]["request_tags"]
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
assert isinstance(tags, list) and isinstance(metadata, dict), rows
|
||||
assert "t1" in tags, rows
|
||||
assert (
|
||||
metadata["user_api_key"],
|
||||
metadata["user_api_key_user_id"],
|
||||
metadata["user_api_key_team_id"],
|
||||
) == (sha256(key.encode()).hexdigest(), user_id, team_id), metadata
|
||||
|
||||
|
||||
def test_text_completion_json_string_metadata_reaches_upstream_and_tags_spend(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: /v1/completions with metadata sent as a JSON string returns 500 before dispatch")
|
||||
identity: Final = f"cmpl-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"tagged prompt {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
_assert_text_completion_request(request, prompt)
|
||||
return Reply(body=_text_completion(identity, ("tagged",)))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
team: Final = scenario.team()
|
||||
user_id: Final = f"integration-{uuid.uuid4().hex}"
|
||||
key: Final = scenario.key(models=[model], team_id=team, user_id=user_id)
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.completions.create(
|
||||
model=model, prompt=prompt, extra_body={"metadata": json.dumps({"tags": ["t1"]})}
|
||||
)
|
||||
assert [(choice.index, choice.text) for choice in completion.choices] == [(0, "tagged")], completion
|
||||
assert completion.id == identity, completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
_assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team)
|
||||
|
||||
|
||||
def test_text_completion_form_metadata_reaches_upstream_and_tags_spend(gateway: Gateway) -> None:
|
||||
identity: Final = f"cmpl-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"tagged prompt {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
_assert_text_completion_request(request, prompt)
|
||||
return Reply(body=_text_completion(identity, ("tagged",)))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
team: Final = scenario.team()
|
||||
user_id: Final = f"integration-{uuid.uuid4().hex}"
|
||||
key: Final = scenario.key(models=[model], team_id=team, user_id=user_id)
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/completions",
|
||||
files=_form_fields({"model": model, "prompt": prompt, "metadata": json.dumps({"tags": ["t1"]})}),
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.request.headers["content-type"].startswith("multipart/form-data"), response.request.headers
|
||||
payload: Final = _TextCompletion.model_validate_json(response.content)
|
||||
assert (payload.id, payload.object, payload.model) == (identity, "text_completion", model), response.text
|
||||
assert [(choice.index, choice.text) for choice in payload.choices] == [(0, "tagged")], response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
_assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team)
|
||||
|
||||
|
||||
def test_text_completion_typed_form_fields_reach_upstream_as_json_types(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: multipart /v1/completions forwards temperature, max_tokens, n and echo upstream as strings, "
|
||||
"so echo='false' is truthy"
|
||||
)
|
||||
identity: Final = f"cmpl-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"typed form prompt {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert request.headers["content-type"] == "application/json"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {
|
||||
"model": _TEXT_BACKEND,
|
||||
"prompt": prompt,
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 5,
|
||||
"n": 1,
|
||||
"echo": False,
|
||||
}, request.body
|
||||
return Reply(body=_text_completion(identity, ("tagged",)))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
team: Final = scenario.team()
|
||||
user_id: Final = f"integration-{uuid.uuid4().hex}"
|
||||
key: Final = scenario.key(models=[model], team_id=team, user_id=user_id)
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/completions",
|
||||
files=_form_fields(
|
||||
{
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"temperature": "0.2",
|
||||
"max_tokens": "5",
|
||||
"n": "1",
|
||||
"echo": "false",
|
||||
"metadata": json.dumps({"tags": ["t1"]}),
|
||||
}
|
||||
),
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.request.headers["content-type"].startswith("multipart/form-data"), response.request.headers
|
||||
payload: Final = _TextCompletion.model_validate_json(response.content)
|
||||
assert (payload.id, payload.object, payload.model) == (identity, "text_completion", model), response.text
|
||||
assert [(choice.index, choice.text) for choice in payload.choices] == [(0, "tagged")], response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
_assert_spend_row_carries_client_tag_and_key_identity(identity, key, user_id, team)
|
||||
|
||||
|
||||
def test_text_completion_sdk_parameter_set_reaches_upstream_whole(gateway: Gateway) -> None:
|
||||
identity: Final = f"cmpl-full-params-{uuid.uuid4().hex}"
|
||||
prompt: Final = f"full parameter prompt {identity}"
|
||||
expected: Final = {
|
||||
"model": _TEXT_BACKEND,
|
||||
"prompt": prompt,
|
||||
"suffix": " tail",
|
||||
"echo": True,
|
||||
"logprobs": 2,
|
||||
"best_of": 2,
|
||||
"n": 1,
|
||||
"logit_bias": {"50256": -100},
|
||||
"seed": 7,
|
||||
"user": "integration-user",
|
||||
"stop": ["\n", "END"],
|
||||
"presence_penalty": 0.5,
|
||||
"frequency_penalty": 0.25,
|
||||
"temperature": 0.3,
|
||||
"top_p": 0.9,
|
||||
"max_tokens": 12,
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert request.headers["content-type"] == "application/json"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == expected, request.body
|
||||
return Reply(body=_text_completion(identity, ("one",)))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.completions.create(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
suffix=" tail",
|
||||
echo=True,
|
||||
logprobs=2,
|
||||
best_of=2,
|
||||
n=1,
|
||||
logit_bias={"50256": -100},
|
||||
seed=7,
|
||||
user="integration-user",
|
||||
stop=["\n", "END"],
|
||||
presence_penalty=0.5,
|
||||
frequency_penalty=0.25,
|
||||
temperature=0.3,
|
||||
top_p=0.9,
|
||||
max_tokens=12,
|
||||
)
|
||||
assert (completion.id, completion.object) == (identity, "text_completion"), completion
|
||||
assert [(choice.index, choice.text) for choice in completion.choices] == [(0, "one")], completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompt_shape",
|
||||
(
|
||||
pytest.param("string-list", id="string-list"),
|
||||
pytest.param("token-ids", id="token-ids"),
|
||||
pytest.param("token-id-batch", id="token-id-batch"),
|
||||
),
|
||||
)
|
||||
def test_text_completion_prompt_spellings_reach_upstream_unchanged(gateway: Gateway, prompt_shape: str) -> None:
|
||||
identity: Final = f"cmpl-{uuid.uuid4().hex}"
|
||||
prompt_identity: Final = uuid.uuid4()
|
||||
token_ids: Final = tuple(
|
||||
int.from_bytes(prompt_identity.bytes[offset : offset + 4], "big") % 50_000 + 50_000 * (offset // 4)
|
||||
for offset in range(0, 16, 4)
|
||||
)
|
||||
prompt: Final[JsonValue] = (
|
||||
[f"first {prompt_identity.hex}", f"second {prompt_identity.hex}"]
|
||||
if prompt_shape == "string-list"
|
||||
else list(token_ids[:3])
|
||||
if prompt_shape == "token-ids"
|
||||
else [list(token_ids[:2]), list(token_ids[2:])]
|
||||
)
|
||||
texts: Final = ("first", "second") if isinstance(prompt, list) and isinstance(prompt[0], (str, list)) else ("one",)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {"model": _TEXT_BACKEND, "prompt": prompt}
|
||||
return Reply(body=_text_completion(identity, texts))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_TEXT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.completions.create(model=model, prompt=prompt)
|
||||
assert completion.id == identity, completion
|
||||
assert completion.object == "text_completion", completion
|
||||
assert [(choice.index, choice.text) for choice in completion.choices] == list(enumerate(texts)), completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
|
||||
|
||||
|
||||
def test_text_completion_string_list_prompt_becomes_user_messages_for_chat_only_deployment(gateway: Gateway) -> None:
|
||||
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
prompts: Final = [f"first {identity}", f"second {identity}"]
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/chat/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {
|
||||
"model": _CHAT_BACKEND,
|
||||
"messages": [{"role": "user", "content": prompts[0]}, {"role": "user", "content": prompts[1]}],
|
||||
}
|
||||
return Reply(body=_chat_completion(identity, "answered"))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"openai/{_CHAT_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
with _openai(gateway, key) as client:
|
||||
completion: Final = client.completions.create(model=model, prompt=prompts)
|
||||
assert (completion.id, completion.object) == (identity, "text_completion"), completion
|
||||
assert [(choice.index, choice.text, choice.finish_reason) for choice in completion.choices] == [
|
||||
(0, "answered", "stop")
|
||||
], completion
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
||||
Loading…
Add table
Reference in a new issue