scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

test: fix local tests

Browse source
This commit is contained in:
Krrish Dholakia 2024-05-06 07:14:33 -07:00
parent 5f119f2abb
commit d83f0b02da
1 changed files with 55 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/tests/test_completion.py
View file

@ -1441,13 +1441,18 @@ def test_completion_ollama_hosted():
# test_completion_ollama_hosted()
@pytest.mark.parametrize(("model"), [
@pytest.mark.skip(reason="Local test")
@pytest.mark.parametrize(
("model"),
[
"ollama/llama2",
"ollama_chat/llama2",
]
],
)
def test_completion_ollama_function_call(model):
messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
messages = [
{"role": "user", "content": "What's the weather like in San Francisco?"}
]
tools = [
{
"type": "function",
@ -1473,19 +1478,27 @@ def test_completion_ollama_function_call(model):
response = litellm.completion(model=model, messages=messages, tools=tools)
print(response)
assert response.choices[0].message.tool_calls
assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather"
assert (
response.choices[0].message.tool_calls[0].function.name
== "get_current_weather"
)
assert response.choices[0].finish_reason == "tool_calls"
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize(("model"), [
@pytest.mark.skip(reason="Local test")
@pytest.mark.parametrize(
("model"),
[
"ollama/llama2",
"ollama_chat/llama2",
]
],
)
def test_completion_ollama_function_call_stream(model):
messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
messages = [
{"role": "user", "content": "What's the weather like in San Francisco?"}
]
tools = [
{
"type": "function",
@ -1508,24 +1521,33 @@ def test_completion_ollama_function_call_stream(model):
]
try:
litellm.set_verbose = True
response = litellm.completion(model=model, messages=messages, tools=tools, stream=True)
response = litellm.completion(
model=model, messages=messages, tools=tools, stream=True
)
print(response)
first_chunk = next(response)
assert first_chunk.choices[0].delta.tool_calls
assert first_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_weather"
assert (
first_chunk.choices[0].delta.tool_calls[0].function.name
== "get_current_weather"
)
assert first_chunk.choices[0].finish_reason == "tool_calls"
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize(("model"), [
@pytest.mark.parametrize(
("model"),
[
"ollama/llama2",
"ollama_chat/llama2",
]
],
)
@pytest.mark.asyncio
async def test_acompletion_ollama_function_call(model):
messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
messages = [
{"role": "user", "content": "What's the weather like in San Francisco?"}
]
tools = [
{
"type": "function",
@ -1548,23 +1570,32 @@ async def test_acompletion_ollama_function_call(model):
]
try:
litellm.set_verbose = True
response = await litellm.acompletion(model=model, messages=messages, tools=tools)
response = await litellm.acompletion(
model=model, messages=messages, tools=tools
)
print(response)
assert response.choices[0].message.tool_calls
assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather"
assert (
response.choices[0].message.tool_calls[0].function.name
== "get_current_weather"
)
assert response.choices[0].finish_reason == "tool_calls"
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize(("model"), [
@pytest.mark.parametrize(
("model"),
[
"ollama/llama2",
"ollama_chat/llama2",
]
],
)
@pytest.mark.asyncio
async def test_acompletion_ollama_function_call_stream(model):
messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
messages = [
{"role": "user", "content": "What's the weather like in San Francisco?"}
]
tools = [
{
"type": "function",
@ -1587,11 +1618,16 @@ async def test_acompletion_ollama_function_call_stream(model):
]
try:
litellm.set_verbose = True
response = await litellm.acompletion(model=model, messages=messages, tools=tools, stream=True)
response = await litellm.acompletion(
model=model, messages=messages, tools=tools, stream=True
)
print(response)
first_chunk = await anext(response)
assert first_chunk.choices[0].delta.tool_calls
assert first_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_weather"
assert (
first_chunk.choices[0].delta.tool_calls[0].function.name
== "get_current_weather"
)
assert first_chunk.choices[0].finish_reason == "tool_calls"
except Exception as e:
pytest.fail(f"Error occurred: {e}")