scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix failing tests on PR-4760

Browse source
This commit is contained in:
skucherlapati 2024-07-17 19:32:17 -07:00
parent 5297d33426
commit ce474ff008
1 changed files with 14 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_completion_cost.py
View file

@ -707,12 +707,23 @@ def test_vertex_ai_completion_cost():
def test_vertex_ai_medlm_completion_cost():
model="medlm-medium"
"""Test for medlm completion cost."""
with pytest.raises(Exception) as e:
model="vertex_ai/medlm-medium"
messages = [{"role": "user", "content": "Test MedLM completion cost."}]
predictive_cost = completion_cost(model=model, messages=messages, custom_llm_provider="vertex_ai")
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model="vertex_ai/medlm-medium"
messages = [{"role": "user", "content": "Test MedLM completion cost."}]
predictive_cost = completion_cost(model=model, messages=messages)
predictive_cost = completion_cost(model=model, messages=messages, custom_llm_provider="vertex_ai")
assert predictive_cost > 0
model="medlm-large"
model="vertex_ai/medlm-large"
messages = [{"role": "user", "content": "Test MedLM completion cost."}]
predictive_cost = completion_cost(model=model, messages=messages)
assert predictive_cost > 0