Skip to content

Commit 07b1b50

Browse files
author
dori
committed
feat: try to fix build
1 parent 1fd2a73 commit 07b1b50

File tree

1 file changed

+22
-7
lines changed

1 file changed

+22
-7
lines changed

src/mcp_as_a_judge/db/dynamic_token_limits.py

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)