mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
test_anthropic_cache_control_hook_specific_index
This commit is contained in:
parent
fd13464108
commit
00b36ac60b
1 changed files with 40 additions and 0 deletions
|
@ -149,3 +149,43 @@ async def test_anthropic_cache_control_hook_user_message():
|
|||
assert request_body["messages"][1]["content"][1]["cachePoint"] == {
|
||||
"type": "default"
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_control_hook_specific_index():
|
||||
"""Test injecting cache control at a specific message index"""
|
||||
anthropic_cache_control_hook = AnthropicCacheControlHook()
|
||||
litellm.callbacks = [anthropic_cache_control_hook]
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"output": {"message": {"role": "assistant", "content": "Test response"}},
|
||||
"stopReason": "stop_sequence",
|
||||
"usage": {"inputTokens": 100, "outputTokens": 200, "totalTokens": 300},
|
||||
}
|
||||
mock_response.status_code = 200
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
|
||||
messages=[
|
||||
{"role": "system", "content": "System message"},
|
||||
{"role": "user", "content": "User message"},
|
||||
{"role": "assistant", "content": "Assistant message"},
|
||||
],
|
||||
cache_control_injection_points=[
|
||||
{
|
||||
"location": "message",
|
||||
"index": 1, # Target the second message specifically
|
||||
}
|
||||
],
|
||||
client=client,
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
request_body = json.loads(mock_post.call_args.kwargs["data"])
|
||||
# Verify cache control was added only to the second message
|
||||
assert "cache_control" not in request_body["messages"][0]
|
||||
assert request_body["messages"][1]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in request_body["messages"][2]
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue