phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test: refactor gemini test to use mock, prevent ratelimit error

Browse source
This commit is contained in:
Krrish Dholakia 2024-10-30 11:09:50 -07:00
parent 03c56804a0
commit aa4c2827e8
2 changed files with 52 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

52
tests/llm_translation/test_vertex.py
View file

@ -178,3 +178,55 @@ def test_vertex_function_translation(tool, expect_parameters):
assert (
"parameters" not in optional_params["tools"][0]["function_declarations"][0]
)
def test_function_calling_with_gemini():
from litellm.llms.custom_httpx.http_handler import HTTPHandler
litellm.set_verbose = True
client = HTTPHandler()
with patch.object(client, "post", new=MagicMock()) as mock_post:
try:
litellm.completion(
model="gemini/gemini-1.5-pro-002",
messages=[
{
"content": [
{
"type": "text",
"text": "You are a helpful assistant that can interact with a computer to solve tasks.\n<IMPORTANT>\n* If user provides a path, you should NOT assume it's relative to the current working directory. Instead, you should explore the file system to find the file before working on it.\n</IMPORTANT>\n",
}
],
"role": "system",
},
{
"content": [{"type": "text", "text": "Hey, how's it going?"}],
"role": "user",
},
],
tools=[
{
"type": "function",
"function": {
"name": "finish",
"description": "Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task.",
},
},
],
client=client,
)
except Exception as e:
print(e)
mock_post.assert_called_once()
print(mock_post.call_args.kwargs)
assert mock_post.call_args.kwargs["json"]["tools"] == [
{
"function_declarations": [
{
"name": "finish",
"description": "Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task.",
}
]
}
]

31
tests/local_testing/test_function_calling.py
View file

@ -612,34 +612,3 @@ def test_passing_tool_result_as_list():
print(resp)
assert resp.usage.prompt_tokens_details.cached_tokens > 0
def test_function_calling_with_gemini():
litellm.set_verbose = True
resp = litellm.completion(
model="gemini/gemini-1.5-pro-002",
messages=[
{
"content": [
{
"type": "text",
"text": "You are a helpful assistant that can interact with a computer to solve tasks.\n<IMPORTANT>\n* If user provides a path, you should NOT assume it's relative to the current working directory. Instead, you should explore the file system to find the file before working on it.\n</IMPORTANT>\n",
}
],
"role": "system",
},
{
"content": [{"type": "text", "text": "Hey, how's it going?"}],
"role": "user",
},
],
tools=[
{
"type": "function",
"function": {
"name": "finish",
"description": "Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task.",
},
},
],
)