@@ -57,14 +57,29 @@ def get_model_limits(model_name: str | None = None) -> ModelLimits:
5757
5858 model_info = litellm .get_model_info (model_name )
5959
60+ # Extract values with proper fallbacks
61+ context_window = model_info .get ("max_tokens" )
62+ if context_window is not None :
63+ context_window = int (context_window )
64+ else :
65+ context_window = limits .context_window
66+
67+ max_input_tokens = model_info .get ("max_input_tokens" )
68+ if max_input_tokens is not None :
69+ max_input_tokens = int (max_input_tokens )
70+ else :
71+ max_input_tokens = limits .max_input_tokens
72+
73+ max_output_tokens = model_info .get ("max_output_tokens" )
74+ if max_output_tokens is not None :
75+ max_output_tokens = int (max_output_tokens )
76+ else :
77+ max_output_tokens = limits .max_output_tokens
78+
6079 limits = ModelLimits (
61- context_window = model_info .get ("max_tokens" , limits .context_window ),
62- max_input_tokens = model_info .get (
63- "max_input_tokens" , limits .max_input_tokens
64- ),
65- max_output_tokens = model_info .get (
66- "max_output_tokens" , limits .max_output_tokens
67- ),
80+ context_window = context_window ,
81+ max_input_tokens = max_input_tokens ,
82+ max_output_tokens = max_output_tokens ,
6883 model_name = model_name ,
6984 source = "litellm" ,
7085 )
0 commit comments