test_cohere_generate_api.py

importos
importsys
importtraceback

fromdotenvimportload_dotenv

load_dotenv()
importio
importos

sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
importjson

importpytest

importlitellm
fromlitellmimportcompletion
fromlitellm.llms.cohere.completion.transformationimportCohereTextConfig


deftest_cohere_generate_api_completion():
try:
fromlitellm.llms.custom_httpx.http_handlerimportHTTPHandler
fromunittest.mockimportpatch, MagicMock

client=HTTPHandler()
litellm.set_verbose=True
messages= [
 {"role": "system", "content": "You're a good bot"},
 {
"role": "user",
"content": "Hey",
 },
 ]

withpatch.object(client, "post") asmock_client:
try:
completion(
model="cohere/command",
messages=messages,
max_tokens=10,
client=client,
 )
exceptExceptionase:
print(e)
mock_client.assert_called_once()
print("mock_client.call_args.kwargs", mock_client.call_args.kwargs)

assert (
mock_client.call_args.kwargs["url"]
=="https://api.cohere.ai/v1/generate"
 )
json_data=json.loads(mock_client.call_args.kwargs["data"])
assertjson_data["model"] =="command"
assertjson_data["prompt"] =="You're a good bot Hey"
assertjson_data["max_tokens"] ==10
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


@pytest.mark.asyncio
asyncdeftest_cohere_generate_api_stream():
try:
litellm.set_verbose=True
messages= [
 {"role": "system", "content": "You're a good bot"},
 {
"role": "user",
"content": "Hey",
 },
 ]
response=awaitlitellm.acompletion(
model="cohere/command",
messages=messages,
max_tokens=10,
stream=True,
 )
print("async cohere stream response", response)
asyncforchunkinresponse:
print(chunk)
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


deftest_completion_cohere_stream_bad_key():
try:
api_key="bad-key"
messages= [
 {"role": "system", "content": "You are a helpful assistant."},
 {
"role": "user",
"content": "how does a court case get to the Supreme Court?",
 },
 ]
completion(
model="command",
messages=messages,
stream=True,
max_tokens=50,
api_key=api_key,
 )

exceptlitellm.AuthenticationErrorase:
pass
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


deftest_cohere_transform_request():
try:
config=CohereTextConfig()
messages= [
 {"role": "system", "content": "You're a helpful bot"},
 {"role": "user", "content": "Hello"},
 ]
optional_params= {"max_tokens": 10, "temperature": 0.7}
headers= {}

transformed_request=config.transform_request(
model="command",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers=headers,
 )

print("transformed_request", json.dumps(transformed_request, indent=4))

asserttransformed_request["model"] =="command"
asserttransformed_request["prompt"] =="You're a helpful bot Hello"
asserttransformed_request["max_tokens"] ==10
asserttransformed_request["temperature"] ==0.7
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


deftest_cohere_transform_request_with_tools():
try:
config=CohereTextConfig()
messages= [{"role": "user", "content": "What's the weather?"}]
tools= [
 {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
 },
 },
 }
 ]
optional_params= {"tools": tools}

transformed_request=config.transform_request(
model="command",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={},
 )

print("transformed_request", json.dumps(transformed_request, indent=4))
assert"tools"intransformed_request
asserttransformed_request["tools"] == {"tools": tools}
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


deftest_cohere_map_openai_params():
try:
config=CohereTextConfig()
openai_params= {
"temperature": 0.7,
"max_tokens": 100,
"n": 2,
"top_p": 0.9,
"frequency_penalty": 0.5,
"presence_penalty": 0.5,
"stop": ["END"],
"stream": True,
 }

mapped_params=config.map_openai_params(
non_default_params=openai_params,
optional_params={},
model="command",
drop_params=False,
 )

assertmapped_params["temperature"] ==0.7
assertmapped_params["max_tokens"] ==100
assertmapped_params["num_generations"] ==2
assertmapped_params["p"] ==0.9
assertmapped_params["frequency_penalty"] ==0.5
assertmapped_params["presence_penalty"] ==0.5
assertmapped_params["stop_sequences"] == ["END"]
assertmapped_params["stream"] ==True
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")