diff --git a/CHANGELOG.md b/CHANGELOG.md index f080629..3d541da 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ See [keep a changelog] for information about writing changes to this log. ### Changed -* ... +* Updated trim message litellm guardrail to support tool call and minor bug fixes. ### Added diff --git a/applications/litellm/templates/message-trimming-config.yaml b/applications/litellm/templates/message-trimming-config.yaml index f1ef517..aed912b 100644 --- a/applications/litellm/templates/message-trimming-config.yaml +++ b/applications/litellm/templates/message-trimming-config.yaml @@ -8,19 +8,20 @@ data: message_overflow.py: |- from typing import Literal, Optional, Union from litellm.utils import trim_messages, get_max_tokens, token_counter - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - get_completion_messages, - ) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache import logging import yaml - from pathlib import Path logger = logging.getLogger(__name__) + # Why this guardrail exists: models reject requests whose conversation + # exceeds their context window. This guardrail trims history (and repairs + # tool-call pairings) so requests stay within the window. For the reasons + # behind it and the principles guiding how we trim, see: + # https://os2ai.github.io/documentation/technical/guardrails.html class MessageTrimmingGuardrail(CustomGuardrail): def __init__( self, @@ -39,8 +40,52 @@ data: self.trim_ratio = default_config.get("trim_ratio", 0.75) self.max_output_tokens = default_config.get("max_output_tokens", 2000) self.safety_buffer = default_config.get("safety_buffer", 500) + self.min_output_tokens = default_config.get("min_output_tokens", 256) self.debug = default_config.get("debug", False) + # Context-window resolution: per-model map wins, then litellm's + # built-in get_max_tokens, then this global default. 8192 is a + # conservative floor every modern chat model supports; only used + # when both the per-model map and litellm's lookup miss. + self.default_max_context_tokens = default_config.get( + "default_max_context_tokens", 8192 + ) + self.max_context_tokens_by_model = self._as_dict( + default_config.get("max_context_tokens_by_model"), + "max_context_tokens_by_model", + ) + + # Trailing role:tool handling. Default is to PRESERVE trailing tool + # messages — the agent-loop shape `User -> Asst{tool_calls} -> Tool` + # is the normal way to ask the model to reason from tool results. + # Only enable for upstream chat templates that explicitly reject + # tool-terminal conversations (e.g. strict HF Mistral v0.3 template). + self.pop_trailing_tool_messages = bool( + default_config.get("pop_trailing_tool_messages", False) + ) + self.pop_trailing_tool_messages_by_model = self._as_dict( + default_config.get("pop_trailing_tool_messages_by_model"), + "pop_trailing_tool_messages_by_model", + ) + + @staticmethod + def _as_dict(value, name: str) -> dict: + """Return `value` if it's a dict, else {} (warning on misconfig). + + Catches both an explicit YAML null and a truthy non-dict so the + per-model lookups stay safe regardless of config shape. + """ + if value is None: + return {} + if not isinstance(value, dict): + logger.warning( + "%s must be a mapping, got %s; ignoring it.", + name, + type(value).__name__, + ) + return {} + return value + def _load_config(self): with open("config.yaml", "r") as file: config = yaml.safe_load(file) @@ -64,6 +109,46 @@ data: if self.debug: print(f"[GUARDRAIL] {message}") + def _resolve_max_context_tokens(self, model: Optional[str]) -> int: + """Resolve the model's context-window size. + + Order: per-model override map -> litellm.get_max_tokens -> global default. + litellm's `model_prices_and_context_window.json` doesn't cover every + proxied model name (vLLM, Bedrock variants, custom deployments...), + so falling through to a single hardcoded number is a footgun on a + fleet with mixed 8k/32k/128k models. + """ + if model and model in self.max_context_tokens_by_model: + return int(self.max_context_tokens_by_model[model]) + try: + resolved = get_max_tokens(model) + if resolved: + return int(resolved) + # Succeeded but returned a falsy value (0/None) — log it so the + # silent fall-through to the default is visible. + logger.warning( + "get_max_tokens(%r) returned a falsy value (%r); falling " + "back to default_max_context_tokens=%d", + model, + resolved, + self.default_max_context_tokens, + ) + except Exception as e: + logger.warning( + "get_max_tokens(%r) failed (%s); falling back to " + "default_max_context_tokens=%d", + model, + e, + self.default_max_context_tokens, + ) + return int(self.default_max_context_tokens) + + def _resolve_pop_trailing_tools(self, model: Optional[str]) -> bool: + """Per-model override > global default. See `pop_trailing_tool_messages`.""" + if model and model in self.pop_trailing_tool_messages_by_model: + return bool(self.pop_trailing_tool_messages_by_model[model]) + return self.pop_trailing_tool_messages + def _calculate_safe_completion_tokens( self, max_context_tokens: int, @@ -84,7 +169,8 @@ data: max_context_tokens - int(current_input_tokens) - self.safety_buffer ) safe_completion_tokens = min( - requested_completion, max(256, int(available_for_completion * 0.75)) + requested_completion, + max(self.min_output_tokens, int(available_for_completion * self.trim_ratio)), ) return safe_completion_tokens @@ -120,6 +206,88 @@ data: ) data["max_tokens"] = safe_completion_tokens + def _repair_tool_call_pairings(self, messages: list) -> list: + """Strip orphan tool messages and orphan tool_calls from assistant messages. + + After `trim_messages` truncates history, tool-call/tool-response pairings + can end up broken. Mistral/vLLM rejects such conversations. We keep only + tool_calls that have a matching later `role: tool` response, and keep only + `role: tool` messages whose `tool_call_id` was advertised by a surviving + assistant `tool_calls` entry. + """ + satisfied_ids = { + m.get("tool_call_id") + for m in messages + if m.get("role") == "tool" and m.get("tool_call_id") + } + + result = [] + advertised_ids = set() + for msg in messages: + match msg.get("role"): + case "tool": + if msg.get("tool_call_id") in advertised_ids: + result.append(msg) + else: + self._log_debug( + f"Dropping orphan tool message (tool_call_id={msg.get('tool_call_id')})" + ) + + case "assistant": + tcs = msg.get("tool_calls") or [] + if not tcs: + result.append(msg) + continue + + kept_tcs = [tc for tc in tcs if tc.get("id") in satisfied_ids] + content = msg.get("content") + content_empty = not (content or "").strip() + if not kept_tcs and content_empty: + self._log_debug( + "Dropping assistant message with no content and all tool_calls orphaned" + ) + continue + + new_msg = msg.copy() + if kept_tcs: + new_msg["tool_calls"] = kept_tcs + advertised_ids.update(tc["id"] for tc in kept_tcs) + else: + new_msg.pop("tool_calls", None) + self._log_debug( + "Stripped all orphan tool_calls from assistant message" + ) + result.append(new_msg) + + case _: + result.append(msg) + return result + + def _sanitize_messages(self, messages: list, model: Optional[str] = None) -> list: + """Repair tool-call pairings and fix the terminus. Safe to call on + every request. + + Order matters: the second repair has to run *before* we decide + whether to append a user-continue, because popping a trailing tool + may expose an empty assistant whose only `tool_calls` were just + orphaned — repair will drop that assistant entirely, and we don't + want to append `"Please continue"` on top of a now-defunct terminus. + """ + if not messages: + return messages + pop = self._resolve_pop_trailing_tools(model) + result = self._repair_tool_call_pairings(messages) + if pop: + while result and result[-1].get("role") == "tool": + self._log_debug("Popping trailing tool message from terminus") + result.pop() + # Pop may have orphaned tool_calls on the now-terminal assistant. + result = self._repair_tool_call_pairings(result) + if result and result[-1].get("role") == "assistant": + self._log_debug("Appending user continue message after assistant terminus") + result.append({"role": "user", "content": "Please continue"}) + return result + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -140,11 +308,8 @@ data: if "messages" in data and data["messages"]: model = data.get("model") - # Get model's context window size - try: - max_context_tokens = get_max_tokens(model) - except: - max_context_tokens = 8192 # Default fallback + # Get model's context window size (per-model map -> litellm -> default). + max_context_tokens = self._resolve_max_context_tokens(model) self._log_debug(f"Model: {model}") self._log_debug(f"Max context tokens: {max_context_tokens}") @@ -179,7 +344,7 @@ data: self._log_debug(f"Safe completion tokens: {safe_completion_tokens}") self._log_debug( - f"Calculation: min({requested_completion}, max(512, ({max_context_tokens} - {int(current_tokens)} - {self.safety_buffer}) * 0.90))" + f"Calculation: min({requested_completion}, max({self.min_output_tokens}, ({max_context_tokens} - {int(current_tokens)} - {self.safety_buffer}) * {self.trim_ratio}))" ) # Update completion tokens in the request @@ -204,55 +369,40 @@ data: self._log_debug( f"Input tokens ({current_tokens}) exceed limit ({max_input_tokens}), trimming messages..." ) - # Trim messages to fit + # NOTE: this 0.90 stacks on the 0.85 above, so we trim to + # ~76.5% of the raw input headroom. The extra 10% absorbs + # tokens litellm/the chat template add after trimming. data["messages"] = trim_messages( data["messages"], model=model, - max_tokens=int( - max_input_tokens * 0.90 - ), # Trim to 90% of already conservative max + max_tokens=int(max_input_tokens * 0.90), trim_ratio=self.trim_ratio, ) - - # Ensure "ensure_alternating_roles" is fixed for message after trim - data["messages"] = get_completion_messages( - messages=data["messages"], - assistant_continue_message={"role": "assistant", "content": ""}, - user_continue_message={ - "role": "user", - "content": "Please continue", - }, - ensure_alternating_roles=True, - ) - # Recount after trimming - try: - new_token_count = token_counter( - model=model, messages=data["messages"] - ) - self._log_debug(f"After trimming, input tokens: {new_token_count}") - - # RECALCULATE safe completion tokens based on actual trimmed input - safe_completion_tokens = self._calculate_safe_completion_tokens( - max_context_tokens, new_token_count, requested_completion - ) - self._log_debug( - f"Recalculated safe completion tokens after trim: {safe_completion_tokens}" - ) - - # Update the data with the recalculated values - self._update_completion_tokens( - data, safe_completion_tokens, has_max_tokens, has_max_completion - ) - - # Update for final logging - current_tokens = new_token_count - except Exception as e: - self._log_debug(f"Failed to recount tokens after trim: {e}") else: self._log_debug( f"No trimming needed, but current={current_tokens}, max={max_input_tokens}" ) + # Always sanitize: repair tool-call pairings and (optionally, per + # `pop_trailing_tool_messages` config) coerce the terminus into a + # user message so strict chat templates accept the request. + data["messages"] = self._sanitize_messages(data["messages"], model=model) + + # Recount once after sanitize — messages may have grown (user continue + # appended) or shrunk (orphan tool / empty assistant dropped). + try: + final_tokens = token_counter(model=model, messages=data["messages"]) + self._log_debug(f"After sanitize, input tokens: {final_tokens}") + safe_completion_tokens = self._calculate_safe_completion_tokens( + max_context_tokens, final_tokens, requested_completion + ) + self._update_completion_tokens( + data, safe_completion_tokens, has_max_tokens, has_max_completion + ) + current_tokens = final_tokens + except Exception as e: + self._log_debug(f"Failed to recount tokens after sanitize: {e}") + self._log_debug( f"Expected total: input ~{int(current_tokens)} + completion {safe_completion_tokens} + buffer {self.safety_buffer} = ~{int(current_tokens) + safe_completion_tokens + self.safety_buffer}/{max_context_tokens}" )