mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
258 lines
11 KiB
Python
258 lines
11 KiB
Python
import json
|
|
import uuid
|
|
from collections.abc import Mapping
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
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 pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
|
|
|
_BACKENDS: Final = ("mimo-v2.6-pro", "mimo-v2.6-flash")
|
|
_API_KEY: Final = "synthetic-xiaomi-key"
|
|
_ARITHMETIC_PROMPT: Final = "What is 17 + 26? Answer with just the number."
|
|
_WEATHER_PROMPT: Final = "What is the weather in Paris? Use the tool."
|
|
_COUNTING_PROMPT: Final = "Count from 1 to 5, one number per line."
|
|
_WEATHER_TOOL: Final[JsonValue] = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Get the current weather for a city",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"city": {"type": "string"}},
|
|
"required": ["city"],
|
|
},
|
|
},
|
|
}
|
|
_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json"
|
|
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
|
_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]])
|
|
|
|
|
|
class _Delta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
content: str | None = None
|
|
reasoning_content: str | None = None
|
|
|
|
|
|
class _Choice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
delta: _Delta
|
|
finish_reason: str | None = None
|
|
|
|
|
|
class _Chunk(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
id: str
|
|
choices: tuple[_Choice, ...]
|
|
|
|
|
|
def _catalog_cost(backend: str, field: str) -> float:
|
|
cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())
|
|
cost_value: Final = cost_map[f"xiaomi_mimo/{backend}"][field]
|
|
assert isinstance(cost_value, (int, float))
|
|
return float(cost_value)
|
|
|
|
|
|
def _approx(value: float) -> object:
|
|
return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs
|
|
|
|
|
|
def _completion(identity: str, backend: str, message: Mapping[str, object], finish: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"id": identity,
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": backend,
|
|
"choices": [{"index": 0, "message": message, "finish_reason": finish}],
|
|
"usage": {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64},
|
|
}
|
|
).encode()
|
|
|
|
|
|
def _frame(identity: str, backend: str, delta: Mapping[str, object], finish: str | None = None) -> bytes:
|
|
value: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": backend,
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
|
|
}
|
|
return b"data: " + json.dumps(value).encode() + b"\n\n"
|
|
|
|
|
|
def _assert_provider_request(request: Request, backend: str, prompt: str) -> dict[str, JsonValue]:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
|
assert request.headers["content-type"] == "application/json"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == backend
|
|
assert body["messages"] == [{"role": "user", "content": prompt}]
|
|
return body
|
|
|
|
|
|
@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing")
|
|
@pytest.mark.parametrize("backend", _BACKENDS)
|
|
def test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price(gateway: Gateway, backend: str) -> None:
|
|
identity: Final = f"xiaomi-cost-{uuid.uuid4().hex}"
|
|
|
|
def respond(request: Request) -> Reply:
|
|
body: Final = _assert_provider_request(request, backend, _ARITHMETIC_PROMPT)
|
|
assert body["max_tokens"] == 256
|
|
assert "max_completion_tokens" not in body
|
|
return Reply(
|
|
body=_completion(
|
|
identity,
|
|
backend,
|
|
{"role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43."},
|
|
"stop",
|
|
)
|
|
)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _ARITHMETIC_PROMPT}],
|
|
"max_completion_tokens": 256,
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert payload["id"] == identity
|
|
assert payload["choices"] == [
|
|
{
|
|
"finish_reason": "stop",
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "43",
|
|
"reasoning_content": "17 plus 26 is 43.",
|
|
"provider_specific_fields": {"refusal": None},
|
|
},
|
|
"provider_specific_fields": {},
|
|
}
|
|
]
|
|
assert payload["usage"] == {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64}
|
|
expected_cost: Final = 23 * _catalog_cost(backend, "input_cost_per_token") + 41 * _catalog_cost(
|
|
backend, "output_cost_per_token"
|
|
)
|
|
assert float(response.headers["x-litellm-response-cost"]) == _approx(expected_cost)
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
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"]) == (23, 41)
|
|
spend: Final = rows[0]["spend"]
|
|
assert isinstance(spend, (int, float, str))
|
|
assert float(spend) == _approx(expected_cost)
|
|
|
|
|
|
@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_and_answer_stream_as_deltas")
|
|
def test_xiaomi_mimo_stream_delivers_reasoning_then_answer_deltas(gateway: Gateway) -> None:
|
|
backend: Final = _BACKENDS[0]
|
|
identity: Final = f"xiaomi-stream-{uuid.uuid4().hex}"
|
|
frames: Final = (
|
|
_frame(identity, backend, {"role": "assistant", "reasoning_content": "Count "}),
|
|
_frame(identity, backend, {"reasoning_content": "up by one."}),
|
|
_frame(identity, backend, {"content": "1\n2\n"}),
|
|
_frame(identity, backend, {"content": "3\n4\n5"}),
|
|
_frame(identity, backend, {}, finish="stop"),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
def respond(request: Request) -> Reply:
|
|
body: Final = _assert_provider_request(request, backend, _COUNTING_PROMPT)
|
|
assert body["stream"] is True
|
|
return Reply(content_type="text/event-stream", chunks=frames)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY)
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
json={"model": model, "messages": [{"role": "user", "content": _COUNTING_PROMPT}], "stream": True},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read()
|
|
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
|
|
assert lines[-1] == "data: [DONE]"
|
|
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
|
|
assert {chunk.id for chunk in chunks} == {identity}
|
|
choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices)
|
|
assert "".join(choice.delta.reasoning_content or "" for choice in choices) == "Count up by one."
|
|
assert "".join(choice.delta.content or "" for choice in choices) == "1\n2\n3\n4\n5"
|
|
assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",)
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
@pytest.mark.covers("other.provider_wire.xiaomi_mimo.tool_call_survives_translation")
|
|
def test_xiaomi_mimo_tool_call_is_forwarded_and_returned(gateway: Gateway) -> None:
|
|
backend: Final = _BACKENDS[1]
|
|
identity: Final = f"xiaomi-tool-{uuid.uuid4().hex}"
|
|
tool_call: Final = {
|
|
"id": "call_paris",
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": json.dumps({"city": "Paris"})},
|
|
}
|
|
|
|
def respond(request: Request) -> Reply:
|
|
body: Final = _assert_provider_request(request, backend, _WEATHER_PROMPT)
|
|
assert body["tools"] == [_WEATHER_TOOL]
|
|
assert body["tool_choice"] == "auto"
|
|
return Reply(
|
|
body=_completion(
|
|
identity,
|
|
backend,
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"reasoning_content": "Need the tool.",
|
|
"tool_calls": [tool_call],
|
|
},
|
|
"tool_calls",
|
|
)
|
|
)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _WEATHER_PROMPT}],
|
|
"tools": [_WEATHER_TOOL],
|
|
"tool_choice": "auto",
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert payload["choices"] == [
|
|
{
|
|
"finish_reason": "tool_calls",
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"reasoning_content": "Need the tool.",
|
|
"tool_calls": [tool_call],
|
|
"provider_specific_fields": {"refusal": None},
|
|
},
|
|
"provider_specific_fields": {},
|
|
}
|
|
]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|