diff --git a/tests/llm_translation/test_cohere_generate_api.py b/tests/llm_translation/test_cohere_generate_api.py deleted file mode 100644 index 837e077cdee..00000000000 --- a/tests/llm_translation/test_cohere_generate_api.py +++ /dev/null @@ -1,203 +0,0 @@ -import os -import sys -import traceback - -from dotenv import load_dotenv - -load_dotenv() -import io -import os - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path -import json - -import pytest - -import litellm -from litellm import completion -from litellm.llms.cohere.completion.transformation import CohereTextConfig - - -def test_cohere_generate_api_completion(): - try: - from litellm.llms.custom_httpx.http_handler import HTTPHandler - from unittest.mock import patch, MagicMock - - client = HTTPHandler() - litellm.set_verbose = True - messages = [ - {"role": "system", "content": "You're a good bot"}, - { - "role": "user", - "content": "Hey", - }, - ] - - with patch.object(client, "post") as mock_client: - try: - completion( - model="cohere/command", - messages=messages, - max_tokens=10, - client=client, - ) - except Exception as e: - 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"]) - assert json_data["model"] == "command" - assert json_data["prompt"] == "You're a good bot Hey" - assert json_data["max_tokens"] == 10 - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -@pytest.mark.asyncio -async def test_cohere_generate_api_stream(): - try: - litellm.set_verbose = True - messages = [ - {"role": "system", "content": "You're a good bot"}, - { - "role": "user", - "content": "Hey", - }, - ] - response = await litellm.acompletion( - model="cohere/command", - messages=messages, - max_tokens=10, - stream=True, - ) - print("async cohere stream response", response) - async for chunk in response: - print(chunk) - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -def test_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, - ) - - except litellm.AuthenticationError as e: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -def test_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)) - - assert transformed_request["model"] == "command" - assert transformed_request["prompt"] == "You're a helpful bot Hello" - assert transformed_request["max_tokens"] == 10 - assert transformed_request["temperature"] == 0.7 - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -def test_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" in transformed_request - assert transformed_request["tools"] == {"tools": tools} - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -def test_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, - ) - - assert mapped_params["temperature"] == 0.7 - assert mapped_params["max_tokens"] == 100 - assert mapped_params["num_generations"] == 2 - assert mapped_params["p"] == 0.9 - assert mapped_params["frequency_penalty"] == 0.5 - assert mapped_params["presence_penalty"] == 0.5 - assert mapped_params["stop_sequences"] == ["END"] - assert mapped_params["stream"] == True - except Exception as e: - pytest.fail(f"Error occurred: {e}")