mirror of https://github.com/microsoft/autogen.git
70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
#!/usr/bin/env python3 -m pytest
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
try:
|
|
from autogen.oai.cohere import CohereClient, calculate_cohere_cost
|
|
|
|
skip = False
|
|
except ImportError:
|
|
CohereClient = object
|
|
skip = True
|
|
|
|
|
|
reason = "Cohere dependency not installed!"
|
|
|
|
|
|
@pytest.fixture()
|
|
def cohere_client():
|
|
return CohereClient(api_key="dummy_api_key")
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason=reason)
|
|
def test_initialization_missing_api_key():
|
|
os.environ.pop("COHERE_API_KEY", None)
|
|
with pytest.raises(
|
|
AssertionError,
|
|
match="Please include the api_key in your config list entry for Cohere or set the COHERE_API_KEY env variable.",
|
|
):
|
|
CohereClient()
|
|
|
|
CohereClient(api_key="dummy_api_key")
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason=reason)
|
|
def test_intialization(cohere_client):
|
|
assert cohere_client.api_key == "dummy_api_key", "`api_key` should be correctly set in the config"
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason=reason)
|
|
def test_calculate_cohere_cost():
|
|
assert (
|
|
calculate_cohere_cost(0, 0, model="command-r") == 0.0
|
|
), "Cost should be 0 for 0 input_tokens and 0 output_tokens"
|
|
assert calculate_cohere_cost(100, 200, model="command-r-plus") == 0.0033
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason=reason)
|
|
def test_load_config(cohere_client):
|
|
params = {
|
|
"model": "command-r-plus",
|
|
"stream": False,
|
|
"temperature": 1,
|
|
"p": 0.8,
|
|
"max_tokens": 100,
|
|
}
|
|
expected_params = {
|
|
"model": "command-r-plus",
|
|
"temperature": 1,
|
|
"p": 0.8,
|
|
"seed": None,
|
|
"max_tokens": 100,
|
|
"frequency_penalty": 0,
|
|
"presence_penalty": 0,
|
|
"k": 0,
|
|
}
|
|
result = cohere_client.parse_params(params)
|
|
assert result == expected_params, "Config should be correctly loaded"
|