mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 11:43:54 +00:00
101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
import sys, os
|
|
import traceback
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import os, io
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import pytest
|
|
import litellm
|
|
from litellm import embedding, completion, completion_cost, Timeout
|
|
from litellm import RateLimitError
|
|
|
|
litellm.num_retries = 3
|
|
|
|
|
|
# FYI - cohere_chat looks quite unstable, even when testing locally
|
|
def test_chat_completion_cohere():
|
|
try:
|
|
litellm.set_verbose = True
|
|
messages = [
|
|
{"role": "system", "content": "You're a good bot"},
|
|
{
|
|
"role": "user",
|
|
"content": "Hey",
|
|
},
|
|
]
|
|
response = completion(
|
|
model="cohere_chat/command-r",
|
|
messages=messages,
|
|
max_tokens=10,
|
|
)
|
|
print(response)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_chat_completion_cohere_stream():
|
|
try:
|
|
litellm.set_verbose = False
|
|
messages = [
|
|
{"role": "system", "content": "You're a good bot"},
|
|
{
|
|
"role": "user",
|
|
"content": "Hey",
|
|
},
|
|
]
|
|
response = completion(
|
|
model="cohere_chat/command-r",
|
|
messages=messages,
|
|
max_tokens=10,
|
|
stream=True,
|
|
)
|
|
print(response)
|
|
for chunk in response:
|
|
print(chunk)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_chat_completion_cohere_tool_calling():
|
|
try:
|
|
litellm.set_verbose = True
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": "What is the weather like in Boston?",
|
|
},
|
|
]
|
|
response = completion(
|
|
model="cohere_chat/command-r",
|
|
messages=messages,
|
|
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"],
|
|
},
|
|
},
|
|
}
|
|
],
|
|
)
|
|
print(response)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|