Spaces:
Configuration error
Configuration error
File size: 1,582 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import os
import sys
from unittest.mock import MagicMock
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.cohere.chat.transformation import CohereChatConfig
class TestCohereTransform:
def setup_method(self):
self.config = CohereChatConfig()
self.model = "command-r-plus-latest"
self.logging_obj = MagicMock()
def test_map_cohere_params(self):
"""Test that parameters are correctly mapped"""
test_params = {
"temperature": 0.7,
"max_tokens": 200,
"max_completion_tokens": 256,
}
result = self.config.map_openai_params(
non_default_params=test_params,
optional_params={},
model=self.model,
drop_params=False,
)
# The function should properly map max_completion_tokens to max_tokens and override max_tokens
assert result == {"temperature": 0.7, "max_tokens": 256}
def test_cohere_max_tokens_backward_compat(self):
"""Test that parameters are correctly mapped"""
test_params = {
"temperature": 0.7,
"max_tokens": 200,
}
result = self.config.map_openai_params(
non_default_params=test_params,
optional_params={},
model=self.model,
drop_params=False,
)
# The function should properly map max_tokens if max_completion_tokens is not provided
assert result == {"temperature": 0.7, "max_tokens": 200}
|