Skip to content

Commit 55d2e84

Browse files
emerybergerclaude
andcommitted
Fix KeyError for custom/enterprise models in trim_messages (#74)
Use litellm.get_model_info() with exception handling to gracefully handle models not in litellm's database (e.g., ollama_chat/gpt-oss). When model info is unavailable, skip trimming and let the model API handle context limits. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 06be363 commit 55d2e84

1 file changed

Lines changed: 10 additions & 1 deletion

File tree

src/chatdbg/util/trim.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,16 @@ def trim_messages(
7979

8080
messages = copy.deepcopy(messages)
8181

82-
max_tokens_for_model = litellm.model_cost[model]["max_input_tokens"]
82+
try:
83+
model_info = litellm.get_model_info(model)
84+
max_tokens_for_model = model_info.get("max_input_tokens")
85+
if max_tokens_for_model is None:
86+
# Model info exists but doesn't have max_input_tokens
87+
return messages
88+
except Exception:
89+
# Model not in litellm's database (e.g., custom/enterprise models).
90+
# Skip trimming and let the model API handle context limits.
91+
return messages
8392
max_tokens = int(max_tokens_for_model * trim_ratio)
8493

8594
if litellm.token_counter(model, messages=messages) < max_tokens:

0 commit comments

Comments
 (0)