mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
314 lines
13 KiB
Python
314 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
from collections.abc import Callable, Mapping
|
|
from pathlib import Path
|
|
from typing import Final, cast
|
|
|
|
import pytest
|
|
import yaml
|
|
from integration._support.anthropic_sse import message_json
|
|
from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
|
|
from integration._support.openai_wire import chat_reply, responses_reply
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.provider import SharedProvider
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
from pydantic import JsonValue, TypeAdapter
|
|
from redis import Redis
|
|
from redis.exceptions import ResponseError
|
|
|
|
_Turns = Callable[[str], tuple[list[JsonValue], list[JsonValue]]]
|
|
_INDEXES: Final = TypeAdapter(tuple[bytes, ...])
|
|
_COUNT: Final = TypeAdapter(tuple[int])
|
|
|
|
|
|
def _redis_command(redis: Redis[bytes], *arguments: str | int) -> object:
|
|
execute: Final = cast(
|
|
Callable[..., object], redis.execute_command
|
|
) # cast-ok: Redis stubs omit command return types
|
|
return execute(*arguments)
|
|
|
|
|
|
def _search_indexes(redis: Redis[bytes]) -> tuple[bytes, ...]:
|
|
return _INDEXES.validate_python(_redis_command(redis, "FT._LIST"))
|
|
|
|
|
|
def _index_size(redis: Redis[bytes], index: str) -> int:
|
|
return _COUNT.validate_python(_redis_command(redis, "FT.SEARCH", index, "*", "LIMIT", 0, 0))[0]
|
|
|
|
|
|
def _claude_code_turns(task: str) -> tuple[list[JsonValue], list[JsonValue]]:
|
|
"""Turns 1 and 3 of a Claude Code session on /v1/messages: 1 and 5 messages"""
|
|
first: Final[list[JsonValue]] = [{"role": "user", "content": task}]
|
|
third: Final[list[JsonValue]] = [
|
|
*first,
|
|
{
|
|
"role": "assistant",
|
|
"content": [{"type": "tool_use", "id": "toolu_1", "name": "Bash", "input": {"command": "ls"}}],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "calc.py test_calc.py"}],
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": [{"type": "tool_use", "id": "toolu_2", "name": "Read", "input": {"file_path": "calc.py"}}],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "toolu_2", "content": "def add(a, b): return a - b"}],
|
|
},
|
|
]
|
|
return first, third
|
|
|
|
|
|
def _agent_turns(task: str) -> tuple[list[JsonValue], list[JsonValue]]:
|
|
"""Turns 1 and 3 of an OpenAI tool loop on /v1/chat/completions: 2 and 6 messages"""
|
|
|
|
def call(call_id: str, path: str) -> list[JsonValue]:
|
|
function: Final[JsonValue] = {"name": "write_file", "arguments": json.dumps({"path": path})}
|
|
return [
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [{"id": call_id, "type": "function", "function": function}],
|
|
},
|
|
{"role": "tool", "tool_call_id": call_id, "content": f"wrote {path}"},
|
|
]
|
|
|
|
first: Final[list[JsonValue]] = [
|
|
{"role": "system", "content": "You are a coding agent"},
|
|
{"role": "user", "content": task},
|
|
]
|
|
return first, [*first, *call("call_1", "a.yaml"), *call("call_2", "b.yaml")]
|
|
|
|
|
|
def _responses_turns(task: str) -> tuple[list[JsonValue], list[JsonValue]]:
|
|
"""Turns 1 and 3 of an agent on /v1/responses: 1 and 5 input items"""
|
|
|
|
def call(call_id: str, path: str) -> list[JsonValue]:
|
|
return [
|
|
{
|
|
"type": "function_call",
|
|
"call_id": call_id,
|
|
"name": "write_file",
|
|
"arguments": json.dumps({"path": path}),
|
|
},
|
|
{"type": "function_call_output", "call_id": call_id, "output": f"wrote {path}"},
|
|
]
|
|
|
|
first: Final[list[JsonValue]] = [{"role": "user", "content": task}]
|
|
return first, [*first, *call("call_1", "a.yaml"), *call("call_2", "b.yaml")]
|
|
|
|
|
|
def _body(path: str, model: str, conversation: list[JsonValue]) -> dict[str, JsonValue]:
|
|
if path == "/v1/responses":
|
|
return {"model": model, "input": conversation}
|
|
return {"model": model, "max_tokens": 16, "messages": conversation}
|
|
|
|
|
|
def _reply(path: str, text: str) -> Reply:
|
|
identity: Final = f"id_{uuid.uuid4().hex}"
|
|
if path == "/v1/messages":
|
|
return Reply(body=message_json(identity, "claude-sonnet-5-5", text))
|
|
if path == "/v1/responses":
|
|
return responses_reply(identity, "gpt-5.6-sol", text, stream=False)
|
|
return chat_reply(identity, "gpt-5.4", text, stream=False)
|
|
|
|
|
|
def _answer(path: str, payload: dict[str, JsonValue]) -> str:
|
|
if path == "/v1/messages":
|
|
return string_value(_first(payload["content"])["text"])
|
|
if path == "/v1/responses":
|
|
return string_value(_first(_first(payload["output"])["content"])["text"])
|
|
return string_value(object_value(_first(payload["choices"])["message"])["content"])
|
|
|
|
|
|
def _first(value: JsonValue) -> dict[str, JsonValue]:
|
|
assert isinstance(value, list), value
|
|
return object_value(value[0])
|
|
|
|
|
|
def _cached_responses(redis: Redis[bytes]) -> frozenset[bytes]:
|
|
digests: Final = tuple(key for key in redis.scan_iter() if len(key) == 64)
|
|
return frozenset(key for key in digests if b'"response"' in (redis.get(key) or b""))
|
|
|
|
|
|
_SCENARIOS: Final = (
|
|
pytest.param("/v1/messages", "anthropic/claude-sonnet-5-5", _claude_code_turns, id="messages"),
|
|
pytest.param("/v1/chat/completions", "openai/gpt-5.4", _agent_turns, id="chat-completions"),
|
|
pytest.param("/v1/responses", "openai/responses/gpt-5.6-sol", _responses_turns, id="responses"),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(("path", "model", "turns"), _SCENARIOS)
|
|
def test_cache_serves_a_turn_under_max_messages_and_skips_one_past_it(
|
|
gateway: Gateway, provider: SharedProvider, path: str, model: str, turns: _Turns
|
|
) -> None:
|
|
under_cap, past_cap = turns(f"update the config {uuid.uuid4().hex}")
|
|
provider.expect(_reply(path, "first answer"), _reply(path, "second answer"), _reply(path, "third answer"))
|
|
|
|
redis_client: Final[Redis[bytes]] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
|
with redis_client as redis:
|
|
cached_before: Final = _cached_responses(redis)
|
|
gateway.post(path, _body(path, model, under_cap))
|
|
eventually(lambda: _cached_responses(redis) - cached_before, lambda written: len(written) == 1)
|
|
repeated: Final = gateway.post(path, _body(path, model, under_cap))
|
|
past_cap_twice: Final = (
|
|
gateway.post(path, _body(path, model, past_cap)),
|
|
gateway.post(path, _body(path, model, past_cap)),
|
|
)
|
|
|
|
assert _answer(path, repeated) == "first answer", "a repeated turn under max_messages was not served from the cache"
|
|
assert tuple(_answer(path, answer) for answer in past_cap_twice) == ("second answer", "third answer"), (
|
|
"a turn past max_messages was served from the cache"
|
|
)
|
|
assert len(provider.received()) == 3
|
|
|
|
|
|
def _tool_reply(path: str, assistant: Mapping[str, JsonValue]) -> Reply:
|
|
base: Final = JSON_OBJECT.validate_json(_reply(path, "unused").body)
|
|
if path == "/v1/messages":
|
|
return Reply(body=json.dumps({**base, "content": assistant["content"], "stop_reason": "tool_use"}).encode())
|
|
if path == "/v1/responses":
|
|
return Reply(body=json.dumps({**base, "output": [{**assistant, "id": "fc_cached"}]}).encode())
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
**base,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": assistant,
|
|
"finish_reason": "tool_calls",
|
|
}
|
|
],
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
|
|
def _stable_payload(path: str, payload: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
|
if path == "/v1/responses":
|
|
assert string_value(payload["id"]).startswith("resp_")
|
|
return {key: value for key, value in payload.items() if key != "id"}
|
|
return payload
|
|
|
|
|
|
@pytest.mark.parametrize(("path", "model", "turns"), _SCENARIOS)
|
|
def test_semantic_cache_embeds_tool_results_and_does_not_replay_the_previous_turn(
|
|
gateway: Gateway,
|
|
provider: SharedProvider,
|
|
tmp_path: Path,
|
|
path: str,
|
|
model: str,
|
|
turns: _Turns,
|
|
) -> None:
|
|
redis_client: Final[Redis[bytes]] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
|
with redis_client as redis:
|
|
try:
|
|
_search_indexes(redis)
|
|
except ResponseError:
|
|
pytest.skip("Redis Search is required for semantic caching")
|
|
|
|
task: Final = f"update the config {uuid.uuid4().hex}"
|
|
first_turn, past_cap = turns(task)
|
|
equivalent_first, equivalent_past_cap = turns(task + " please")
|
|
tool_turn: Final = past_cap[:-2]
|
|
tool_result: Final = "calc.py test_calc.py" if path == "/v1/messages" else "wrote a.yaml"
|
|
index: Final = f"semantic-tool-turn-{uuid.uuid4().hex}"
|
|
|
|
def embed(request: Request) -> Reply:
|
|
assert (request.method, request.target) == ("POST", "/v1/embeddings")
|
|
body: Final = JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == "text-embedding-3-small"
|
|
prompt: Final = string_value(body["input"])
|
|
vector: Final = ([0.0, 1.0] if tool_result in prompt else [1.0, 0.0]) + [0.0] * 1534
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"object": "list",
|
|
"model": "text-embedding-3-small",
|
|
"data": [{"object": "embedding", "index": 0, "embedding": vector}],
|
|
"usage": {"prompt_tokens": 1, "total_tokens": 1},
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
with wire_server(embed) as embeddings:
|
|
config: Final = JSON_OBJECT.validate_python(
|
|
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
)
|
|
models: Final = config["model_list"]
|
|
assert isinstance(models, list)
|
|
settings: Final = object_value(config["litellm_settings"])
|
|
config_path: Final = tmp_path / "semantic-cache.yaml"
|
|
config_path.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
**config,
|
|
"model_list": [
|
|
*models,
|
|
{
|
|
"model_name": "semantic-embedding",
|
|
"litellm_params": {
|
|
"model": "openai/text-embedding-3-small",
|
|
"api_base": embeddings.url + "/v1",
|
|
"api_key": "scripted-embedding-key",
|
|
},
|
|
},
|
|
],
|
|
"litellm_settings": {
|
|
**settings,
|
|
"cache_params": {
|
|
"type": "redis-semantic",
|
|
"redis_url": f"redis://{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}",
|
|
"similarity_threshold": 0.8,
|
|
"redis_semantic_cache_embedding_model": "semantic-embedding",
|
|
"redis_semantic_cache_index_name": index,
|
|
},
|
|
},
|
|
}
|
|
)
|
|
)
|
|
first_reply: Final = _tool_reply(path, object_value(tool_turn[len(first_turn)]))
|
|
provider.expect(
|
|
first_reply,
|
|
_reply(path, "tool result used"),
|
|
_reply(path, "second answer"),
|
|
_reply(path, "third answer"),
|
|
)
|
|
try:
|
|
with owned_proxy(gateway, tmp_path, {}, config=config_path) as proxy:
|
|
initial: Final = proxy.post(path, _body(path, model, first_turn))
|
|
eventually(lambda: _index_size(redis, index), lambda count: count == 1)
|
|
cached: Final = proxy.post(path, _body(path, model, equivalent_first))
|
|
assert _stable_payload(path, cached) == _stable_payload(path, initial), (
|
|
"semantic cache lost the provider-returned tool call"
|
|
)
|
|
assert len(provider.received()) == 1
|
|
|
|
follow_up: Final = proxy.post(path, _body(path, model, tool_turn))
|
|
assert _answer(path, follow_up) == "tool result used", "semantic cache replayed the prior tool call"
|
|
eventually(lambda: _index_size(redis, index), lambda count: count == 2)
|
|
assert (
|
|
_answer(path, proxy.post(path, _body(path, model, equivalent_past_cap[:-2])))
|
|
== "tool result used"
|
|
)
|
|
assert len(provider.received()) == 1
|
|
|
|
past_cap_twice: Final = tuple(proxy.post(path, _body(path, model, past_cap)) for _ in range(2))
|
|
assert tuple(_answer(path, answer) for answer in past_cap_twice) == (
|
|
"second answer",
|
|
"third answer",
|
|
)
|
|
assert len(provider.received()) == 2
|
|
assert any(
|
|
tool_result in string_value(JSON_OBJECT.validate_json(request.body)["input"])
|
|
for request in embeddings.drain()
|
|
)
|
|
finally:
|
|
if index.encode() in _search_indexes(redis):
|
|
_redis_command(redis, "FT.DROPINDEX", index, "DD")
|