Updated message trim to handle message better and support tool calling.

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
 

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}"
                 )