mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
87d91d78b7
Strategies, pre-call checks, cooldown and fallback helpers, and the auto, complexity and quality sub-routers now take a RoutingContext instead of the Router, and read the router only through registry(), settings(), health() and the context's forwarding methods. Router-free helpers move from router_utils to litellm_core_utils, outside callers of router state switch to entry points, and the surface allowlist shrinks to the endpoint methods plus the entry points
207 lines
8.9 KiB
Python
207 lines
8.9 KiB
Python
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from typing import Final, Literal
|
|
|
|
import pytest
|
|
from pydantic import TypeAdapter
|
|
|
|
import litellm
|
|
from litellm import _v2
|
|
from litellm._v2.cache import NativeBackend
|
|
from litellm.caching.caching_handler import (
|
|
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
|
|
)
|
|
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
|
|
from litellm.types.utils import ModelResponse
|
|
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
|
|
from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT
|
|
from tests.test_litellm_rust.cache.support import cache_key, invoke, payload
|
|
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
|
|
@pytest.mark.parametrize("backend", ("memory", "redis"))
|
|
async def test_v2_cache_skips_provider_and_reports_one_success_per_call(
|
|
recording_server: RecordingServer,
|
|
route: Literal["chat", "messages", "responses"],
|
|
backend: Literal["memory", "redis"],
|
|
redis_url: str,
|
|
) -> None:
|
|
recording_server.expected_requests = 2
|
|
litellm.cache = _v2.Cache.memory() if backend == "memory" else _v2.Cache.redis(redis_url, namespace="headers")
|
|
recorder: Final = RecordingLogger()
|
|
first: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
|
|
await recorder.wait_for_async("async_log_success_event")
|
|
second: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
|
|
assert payload(first) == payload(second)
|
|
assert cache_key(first) is None
|
|
key: Final = cache_key(second)
|
|
assert isinstance(key, str)
|
|
assert key == get_hidden_params_dict(second)["cache_key"]
|
|
assert len(recording_server.requests) == 1
|
|
await drain_logging()
|
|
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
|
assert len(successes) == 2
|
|
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
|
|
assert cached_log["cache_hit"] is True
|
|
assert cached_log["response_cost"] == 0
|
|
await litellm.cache.delete_cache_keys([key])
|
|
refreshed: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
|
|
assert cache_key(refreshed) is None
|
|
assert len(recording_server.requests) == 2
|
|
assert len(await recorder.wait_for_async("async_log_success_event", count=3)) == 3
|
|
await litellm.cache.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_global_cache_leaves_legacy_only_calls_usable() -> None:
|
|
litellm.cache = _v2.Cache.memory()
|
|
response: Final = await litellm.aembedding(
|
|
model="openai/cache-test-embedding",
|
|
input=["hello"],
|
|
api_key="test-key",
|
|
mock_response=[0.25, 0.75],
|
|
)
|
|
assert response.model_dump(include={"data"}) == {
|
|
"data": [{"embedding": [0.25, 0.75], "index": 0, "object": "embedding"}]
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("route", ("chat", "responses"))
|
|
def test_v2_cache_works_through_python_inference(
|
|
recording_server: RecordingServer, route: Literal["chat", "messages", "responses"], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_RUST", "0")
|
|
litellm.cache = _v2.Cache.memory()
|
|
common: Final = {"api_key": "test-key", "api_base": recording_server.base_url}
|
|
if route == "responses":
|
|
recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
|
|
parameters: Final = {"model": RESPONSES_MODEL, "input": "hello", **common}
|
|
first: Final = litellm.responses(**parameters)
|
|
second: Final = litellm.responses(**parameters)
|
|
assert payload(first) == payload(second)
|
|
else:
|
|
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
|
|
arguments: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common}
|
|
initial: Final = litellm.completion(**arguments)
|
|
cached: Final = litellm.completion(**arguments)
|
|
assert isinstance(initial, ModelResponse) and isinstance(cached, ModelResponse)
|
|
assert (
|
|
initial.choices[0].message.content
|
|
== cached.choices[0].message.content
|
|
== MESSAGES_RESPONSE["content"][0]["text"]
|
|
)
|
|
assert len(recording_server.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_facade_and_backend_share_storage_and_management() -> None:
|
|
cache: Final = _v2.Cache.memory()
|
|
await cache.async_add_cache({"answer": 7}, cache_key="shared")
|
|
assert cache.get_cache(cache_key="shared") == {"answer": 7}
|
|
assert await cache.ping() is True
|
|
await cache.delete_cache_keys(["shared"])
|
|
assert await cache.async_get_cache(cache_key="shared") is None
|
|
cache.add_cache({"answer": 8}, cache_key="flush")
|
|
backend: Final = cache.cache
|
|
assert isinstance(backend, NativeBackend)
|
|
backend.flush_cache()
|
|
assert cache.get_cache(cache_key="flush") is None
|
|
await cache.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("control", ("s-maxage", "s-max-age"))
|
|
async def test_v2_native_cache_accepts_existing_freshness_aliases(
|
|
recording_server: RecordingServer, control: str
|
|
) -> None:
|
|
litellm.cache = _v2.Cache.memory()
|
|
first: Final = await invoke("responses", recording_server, {})
|
|
second: Final = await invoke("responses", recording_server, {"cache": {control: 600}})
|
|
assert payload(first) == payload(second)
|
|
assert len(recording_server.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_cache_does_not_force_native_responses_streaming(
|
|
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_RUST", "0")
|
|
litellm.cache = _v2.Cache.memory()
|
|
recording_server.default_response = ResponseSpec(
|
|
body=None,
|
|
events=(
|
|
("response.created", {"type": "response.created", "sequence_number": 0, "response": RESPONSES_RESPONSE}),
|
|
(
|
|
"response.completed",
|
|
{"type": "response.completed", "sequence_number": 1, "response": RESPONSES_RESPONSE},
|
|
),
|
|
),
|
|
)
|
|
response: Final = await litellm.aresponses(
|
|
model=RESPONSES_MODEL,
|
|
input="hello",
|
|
stream=True,
|
|
caching=False,
|
|
api_key="test-key",
|
|
api_base=recording_server.base_url,
|
|
)
|
|
assert isinstance(response, AsyncIterator)
|
|
chunks: Final = [chunk async for chunk in response]
|
|
assert chunks[-1].type == "response.completed"
|
|
assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_cache_works_through_python_messages(
|
|
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_RUST", "0")
|
|
litellm.cache = _v2.Cache.memory()
|
|
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
|
|
parameters: Final = {
|
|
"model": MESSAGES_MODEL,
|
|
"messages": list(MESSAGES),
|
|
"max_tokens": 32,
|
|
"api_key": "test-key",
|
|
"api_base": recording_server.base_url,
|
|
}
|
|
first: Final = await litellm.anthropic_messages(**parameters)
|
|
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
|
second: Final = await litellm.anthropic_messages(**parameters)
|
|
assert payload(first) == payload(second)
|
|
assert len(recording_server.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
async def test_v2_redis_flush_only_removes_its_namespace(
|
|
redis_url: str, recording_server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
own: Final = _v2.Cache.redis(redis_url, namespace="flush-own")
|
|
other: Final = _v2.Cache.redis(redis_url, namespace="flush-other")
|
|
litellm.cache = own
|
|
recording_server.expected_requests = 2
|
|
await own.async_add_cache({"answer": "own"}, cache_key="shared")
|
|
await other.async_add_cache({"answer": "other"}, cache_key="shared")
|
|
await invoke("responses", recording_server, {})
|
|
hit: Final = await invoke("responses", recording_server, {})
|
|
assert isinstance(cache_key(hit), str)
|
|
assert await own.async_get_cache(cache_key="shared") == {"answer": "own"}
|
|
backend: Final = own.cache
|
|
assert isinstance(backend, NativeBackend)
|
|
if asynchronous:
|
|
await backend.async_flush_cache()
|
|
else:
|
|
backend.flush_cache()
|
|
assert await own.async_get_cache(cache_key="shared") is None
|
|
assert await other.async_get_cache(cache_key="shared") == {"answer": "other"}
|
|
refreshed: Final = await invoke("responses", recording_server, {})
|
|
assert cache_key(refreshed) is None
|
|
assert len(recording_server.requests) == 2
|
|
await own.disconnect()
|
|
await other.disconnect()
|