diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index 43cb236ad4..618287e195 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -565,48 +565,22 @@ def test_together_ai_qwen_completion_cost(): assert response == "together-ai-41.1b-80b" -@pytest.mark.parametrize("above_128k", [False, True]) @pytest.mark.parametrize("provider", ["gemini"]) -def test_gemini_completion_cost(above_128k, provider): +def test_gemini_completion_cost(provider): """ Check if cost correctly calculated for gemini models based on context window """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - if provider == "gemini": - model_name = "gemini-1.5-flash-latest" - else: - model_name = "gemini-1.5-flash-preview-0514" - if above_128k: - prompt_tokens = 128001.0 - output_tokens = 228001.0 - else: - prompt_tokens = 128.0 - output_tokens = 228.0 + model_name = "gemini-2.0-flash" + prompt_tokens = 128.0 + output_tokens = 228.0 ## GET MODEL FROM LITELLM.MODEL_INFO model_info = litellm.get_model_info(model=model_name, custom_llm_provider=provider) ## EXPECTED COST - if above_128k: - assert ( - model_info["input_cost_per_token_above_128k_tokens"] is not None - ), "model info for model={} does not have pricing for > 128k tokens\nmodel_info={}".format( - model_name, model_info - ) - assert ( - model_info["output_cost_per_token_above_128k_tokens"] is not None - ), "model info for model={} does not have pricing for > 128k tokens\nmodel_info={}".format( - model_name, model_info - ) - input_cost = ( - prompt_tokens * model_info["input_cost_per_token_above_128k_tokens"] - ) - output_cost = ( - output_tokens * model_info["output_cost_per_token_above_128k_tokens"] - ) - else: - input_cost = prompt_tokens * model_info["input_cost_per_token"] - output_cost = output_tokens * model_info["output_cost_per_token"] + input_cost = prompt_tokens * model_info["input_cost_per_token"] + output_cost = output_tokens * model_info["output_cost_per_token"] ## CALCULATED COST calculated_input_cost, calculated_output_cost = cost_per_token( @@ -630,21 +604,20 @@ def test_vertex_ai_completion_cost(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - text = "The quick brown fox jumps over the lazy dog." - characters = _count_characters(text=text) + prompt_tokens = 100 - model_info = litellm.get_model_info(model="gemini-1.5-flash") + model_info = litellm.get_model_info(model="gemini-2.0-flash") print("\nExpected model info:\n{}\n\n".format(model_info)) - expected_input_cost = characters * model_info["input_cost_per_character"] + expected_input_cost = prompt_tokens * model_info["input_cost_per_token"] ## CALCULATED COST calculated_input_cost, calculated_output_cost = cost_per_token( - model="gemini-1.5-flash", + model="gemini-2.0-flash", custom_llm_provider="vertex_ai", - prompt_characters=characters, - completion_characters=0, + prompt_tokens=prompt_tokens, + completion_tokens=0, ) assert round(expected_input_cost, 6) == round(calculated_input_cost, 6) @@ -2289,14 +2262,14 @@ def test_completion_cost_params(): """ litellm.set_verbose = True resp1_prompt_cost, resp1_completion_cost = cost_per_token( - model="gemini-1.5-pro-002", + model="gemini-2.0-flash", prompt_tokens=1000, completion_tokens=1000, custom_llm_provider="vertex_ai_beta", ) resp2_prompt_cost, resp2_completion_cost = cost_per_token( - model="gemini-1.5-pro-002", prompt_tokens=1000, completion_tokens=1000 + model="gemini-2.0-flash", prompt_tokens=1000, completion_tokens=1000 ) assert resp2_prompt_cost > 0 @@ -2305,7 +2278,7 @@ def test_completion_cost_params(): assert resp1_completion_cost == resp2_completion_cost resp3_prompt_cost, resp3_completion_cost = cost_per_token( - model="vertex_ai/gemini-1.5-pro-002", prompt_tokens=1000, completion_tokens=1000 + model="vertex_ai/gemini-2.0-flash", prompt_tokens=1000, completion_tokens=1000 ) assert resp3_prompt_cost > 0 @@ -2320,24 +2293,22 @@ def test_completion_cost_params_2(): """ litellm.set_verbose = True - prompt_characters = 1000 - completion_characters = 1000 + prompt_tokens = 1000 + completion_tokens = 1000 resp1_prompt_cost, resp1_completion_cost = cost_per_token( - model="gemini-1.5-pro-002", - prompt_characters=prompt_characters, - completion_characters=completion_characters, - prompt_tokens=1000, - completion_tokens=1000, + model="gemini-2.0-flash", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, ) print(resp1_prompt_cost, resp1_completion_cost) - model_info = litellm.get_model_info("gemini-1.5-pro-002") - input_cost_per_character = model_info["input_cost_per_character"] - output_cost_per_character = model_info["output_cost_per_character"] + model_info = litellm.get_model_info("gemini-2.0-flash") + input_cost_per_token = model_info["input_cost_per_token"] + output_cost_per_token = model_info["output_cost_per_token"] - assert resp1_prompt_cost == input_cost_per_character * prompt_characters - assert resp1_completion_cost == output_cost_per_character * completion_characters + assert resp1_prompt_cost == input_cost_per_token * prompt_tokens + assert resp1_completion_cost == output_cost_per_token * completion_tokens def test_completion_cost_params_gemini_3(): @@ -2371,7 +2342,7 @@ def test_completion_cost_params_gemini_3(): ) ], created=1728529259, - model="gemini-1.5-flash", + model="gemini-2.0-flash", object="chat.completion", system_fingerprint=None, usage=usage, @@ -2395,7 +2366,7 @@ def test_completion_cost_params_gemini_3(): pc, cc = cost_per_character( **{ - "model": "gemini-1.5-flash", + "model": "gemini-2.0-flash", "custom_llm_provider": "vertex_ai", "prompt_characters": None, "completion_characters": 3, @@ -2403,11 +2374,13 @@ def test_completion_cost_params_gemini_3(): } ) - model_info = litellm.get_model_info("gemini-1.5-flash") + model_info = litellm.get_model_info("gemini-2.0-flash") + # gemini-2.0-flash has no per-character pricing, so cost_per_character + # falls back to per-token pricing using usage.prompt_tokens / usage.completion_tokens assert round(pc, 10) == round(3771 * model_info["input_cost_per_token"], 10) assert round(cc, 10) == round( - 3 * model_info["output_cost_per_character"], + 2 * model_info["output_cost_per_token"], 10, ) @@ -2461,16 +2434,16 @@ async def test_test_completion_cost_gpt4o_audio_output_from_model(stream): ) ], created=1729282652, - model="gpt-4o-audio-preview-2024-10-01", + model="gpt-4o-audio-preview", object="chat.completion", system_fingerprint="fp_4eafc16e9d", usage=usage_object, service_tier=None, ) - cost = completion_cost(completion, model="gpt-4o-audio-preview-2024-10-01") + cost = completion_cost(completion, model="gpt-4o-audio-preview") - model_info = litellm.get_model_info("gpt-4o-audio-preview-2024-10-01") + model_info = litellm.get_model_info("gpt-4o-audio-preview") print(f"model_info: {model_info}") ## input cost diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index d46a087eb7..37c38b074b 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -55,7 +55,7 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): def test_get_model_info_shows_correct_supports_vision(): - info = litellm.get_model_info("gemini/gemini-1.5-flash") + info = litellm.get_model_info("gemini/gemini-2.0-flash") print("info", info) assert info["supports_vision"] is True @@ -83,9 +83,9 @@ def test_get_model_info_finetuned_models(): def test_get_model_info_gemini_pro(): - info = litellm.get_model_info("gemini-1.5-pro-002") + info = litellm.get_model_info("gemini-2.0-flash") print("info", info) - assert info["key"] == "gemini-1.5-pro-002" + assert info["key"] == "gemini-2.0-flash" def test_get_model_info_ollama_chat():