litellm-mirror/litellm/tests/test_register_model.py
2023-10-19 18:26:36 -07:00

27 lines
No EOL
779 B
Python

#### What this tests ####
# This tests calling batch_completions by running 100 messages together
import sys, os
import traceback
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
def test_update_model_cost():
try:
litellm.register_model({
"gpt-4": {
"max_tokens": 8192,
"input_cost_per_token": 0.00002,
"output_cost_per_token": 0.00006,
"litellm_provider": "openai",
"mode": "chat"
},
})
assert litellm.model_cost["gpt-4"]["input_cost_per_token"] == 0.00002
except Exception as e:
pytest.fail(f"An error occurred: {e}")
test_update_model_cost()