scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

add test for caching tool calls

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-14 16:19:14 -07:00
parent fccc6dc928
commit e0ff4823d0
2 changed files with 160 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
docs/my-website/docs/providers/anthropic.md
View file

@ -236,6 +236,50 @@ Use Anthropic Prompt Caching
### Caching - Tools definitions
<Tabs>
<TabItem value="sdk" label="LiteLLM SDK">
```python
import litellm
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet-20240620",
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
"cache_control": {"type": "ephemeral"}
},
}
],
extra_headers={
"anthropic-version": "2023-06-01",
"anthropic-beta": "prompt-caching-2024-07-31",
},
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
</TabItem>
</Tabs>
### Caching - Continuing Multi-Turn Convo

117
litellm/tests/test_completion.py
View file

@ -14,7 +14,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
import os
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -3513,6 +3513,121 @@ async def test_anthropic_api_prompt_caching_basic():
)
@pytest.mark.asyncio
async def test_litellm_acompletion_httpx_call():
# Arrange: Set up the MagicMock for the httpx.AsyncClient
mock_response = AsyncMock()
def return_val():
return {
"id": "msg_01XFDUDYJgAACzvnptvVoYEL",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "Hello!"}],
"model": "claude-3-5-sonnet-20240620",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
}
mock_response.json = return_val
litellm.set_verbose = True
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=mock_response,
) as mock_post:
# Act: Call the litellm.acompletion function
response = await litellm.acompletion(
api_key="mock_api_key",
model="anthropic/claude-3-5-sonnet-20240620",
messages=[
{"role": "user", "content": "What's the weather like in Boston today?"}
],
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
"cache_control": {"type": "ephemeral"},
},
}
],
extra_headers={
"anthropic-version": "2023-06-01",
"anthropic-beta": "prompt-caching-2024-07-31",
},
)
# Print what was called on the mock
print("call args=", mock_post.call_args)
expected_url = "https://api.anthropic.com/v1/messages"
expected_headers = {
"accept": "application/json",
"content-type": "application/json",
"anthropic-version": "2023-06-01",
"anthropic-beta": "prompt-caching-2024-07-31",
"x-api-key": "mock_api_key",
}
expected_json = {
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": "What's the weather like in Boston today?",
}
],
}
],
"tools": [
{
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"cache_control": {"type": "ephemeral"},
"input_schema": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
}
],
"max_tokens": 4096,
"model": "claude-3-5-sonnet-20240620",
}
mock_post.assert_called_once_with(
expected_url, json=expected_json, headers=expected_headers, timeout=600.0
)
@pytest.mark.parametrize(
"model",
[