diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 093dffccac07..d90703d15440 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -247,6 +247,8 @@ def _merge_tools_after_guardrail( """ Merge remapped guardrailed tools with original tools that were not sent to the guardrail (e.g. web_search, web_search_preview), preserving order. + Tools a guardrail appended (``remapped`` longer than ``original_tools``) + have no original slot and are kept so an injected tool is not dropped. """ if not original_tools: return remapped @@ -262,6 +264,8 @@ def _merge_tools_after_guardrail( if j < len(remapped): result.append(remapped[j]) j += 1 + # Keep guardrail-appended tools that matched no original slot above. + result.extend(remapped[j:]) return result def _apply_guardrailed_tools_to_data( diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py new file mode 100644 index 000000000000..73d31f7aec02 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import ( + GuardrailEventHooks, + Mode, + SupportedGuardrailIntegrations, +) + +from .compresr import CompresrGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def _coerce_event_hook( + mode: str | list[str] | Mode, +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: + if isinstance(mode, Mode): + return mode + if isinstance(mode, list): + return [GuardrailEventHooks(item) for item in mode] + return GuardrailEventHooks(mode) + + +def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object: + if optional_params is not None: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> CompresrGuardrail: + import litellm + + optional_params = getattr(litellm_params, "optional_params", None) + + _callback = CompresrGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + model=litellm_params.model, + target_compression_ratio=_get_optional_value(litellm_params, optional_params, "target_compression_ratio"), + coarse=_get_optional_value(litellm_params, optional_params, "coarse"), + min_chars_to_compress=_get_optional_value(litellm_params, optional_params, "min_chars_to_compress"), + compress_tool_outputs=_get_optional_value(litellm_params, optional_params, "compress_tool_outputs"), + compress_system=_get_optional_value(litellm_params, optional_params, "compress_system"), + compress_history=_get_optional_value(litellm_params, optional_params, "compress_history"), + compress_last_user=_get_optional_value(litellm_params, optional_params, "compress_last_user"), + enable_retrieval=_get_optional_value(litellm_params, optional_params, "enable_retrieval"), + max_bytes_per_call=_get_optional_value(litellm_params, optional_params, "max_bytes_per_call"), + allow_bypass_header=_get_optional_value(litellm_params, optional_params, "allow_bypass_header"), + dynamic=_get_optional_value(litellm_params, optional_params, "dynamic"), + dynamic_min_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_min_ratio"), + dynamic_max_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_max_ratio"), + compression_params=_get_optional_value(litellm_params, optional_params, "compression_params"), + guardrail_name=guardrail["guardrail_name"], + event_hook=_coerce_event_hook(litellm_params.mode), + default_on=litellm_params.default_on or False, + unreachable_fallback=litellm_params.unreachable_fallback, + ) + litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped + _callback + ) + return _callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.COMPRESR.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.COMPRESR.value: CompresrGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py new file mode 100644 index 000000000000..3cd381b5e006 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -0,0 +1,1180 @@ +"""Compresr guardrail — query-aware, recoverable context compression. + +Compresses bulky message content (tool outputs by default) through the +Compresr API before the request reaches the LLM. Each compressed message +carries a hash marker; a ``compresr_retrieve`` tool is injected so the model +can fetch the original content back through the agentic loop when the +compressed version is not enough — making compression recoverable instead +of lossy. + +Unlike gateway-side compressors that operate on whole message lists, each +target is compressed *query-aware*: the query sent to Compresr is the intent +of the tool call that produced the message (``name + arguments``, resolved +via ``tool_call_id``), falling back to the last user message. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import ipaddress +import json +import time +from collections import Counter, OrderedDict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal +from urllib.parse import urlparse + +import httpx +from fastapi import HTTPException +from httpx import Response as HttpxResponse +from typing_extensions import TypeGuard + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.litellm_core_utils.prompt_templates.factory import ( + get_attribute_or_key, + get_tool_calls_from_response, + has_tool_with_name, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.base import ( + GuardrailConfigModel, + ) + +BYPASS_HEADER = "x-compresr-bypass" +COMPRESR_RETRIEVE_TOOL_NAME = "compresr_retrieve" +DEFAULT_API_BASE = "https://api.compresr.ai" +DEFAULT_COMPRESSION_MODEL = "latte_v2" +DEFAULT_TARGET_COMPRESSION_RATIO = 0.5 +DEFAULT_MIN_CHARS_TO_COMPRESS = 500 +_ORIGINALS_TTL_SECONDS = 15 * 60 +_MAX_TRACKED_CALLS = 256 +_DEFAULT_MAX_BYTES_PER_CALL = 10 * 1024 * 1024 +# Aggregate ceiling across all recovery-store entries. max_bytes_per_call only +# bounds a single call; this caps the whole store so many calls cannot exhaust it. +_MAX_TOTAL_STORE_BYTES = 256 * 1024 * 1024 +# Max compresr_retrieve calls expanded into a single follow-up (repeats deduped). +_MAX_RETRIEVALS_PER_LOOP = 8 +# The shared client's 600s read timeout is far too long for an on-request +# guardrail; bound the compress call so a stall hits the fail policy quickly. +_COMPRESS_TIMEOUT_SECONDS = 60.0 +_SOURCE_TAG = "integration:litellm" +# Request-content fields the compression_params passthrough must never +# override — they carry the actual message content/queries being compressed. +_RESERVED_COMPRESSION_PARAM_KEYS = frozenset({"context", "query", "inputs"}) +_BLOCKED_METADATA_HOSTS = frozenset( + { + "metadata.google.internal", + "metadata.goog", + } +) +_BLOCKED_METADATA_IPS = frozenset( + ipaddress.ip_address(ip) for ip in ("169.254.169.254", "fd00:ec2::254", "100.100.100.200") +) + + +def _parse_ip_literal(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None: + """Parse ``host`` as an IP literal, covering the alternate spellings the + socket layer accepts (decimal/hex single-integer IPv4, IPv4-mapped IPv6) + so a blocked address cannot be smuggled past a string comparison.""" + try: + addr = ipaddress.ip_address(host) + except ValueError: + try: + addr = ipaddress.ip_address(int(host, 0)) + except (TypeError, ValueError): + return None + if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: + return addr.ipv4_mapped + return addr + + +def _validate_api_base(url: str) -> str: + """Return ``url`` if it passes basic outbound-target checks, else raise. + + Best-effort defense in depth for a mis/maliciously-configured ``api_base``: + rejects non-http(s) schemes and cloud-metadata IPs/hosts (incl. alternate IP + encodings); private ranges are allowed for on-prem deployments. NOT a complete + SSRF control — no DNS resolution, and the shared client follows redirects and + re-resolves DNS (TOCTOU / rebinding); ``api_base`` is trusted operator config, + so this is an accepted limitation. + """ + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Compresr guardrail api_base must be http or https, got scheme={parsed.scheme!r}") + host = (parsed.hostname or "").lower() + if not host: + raise ValueError("Compresr guardrail api_base has no host") + ip_literal = _parse_ip_literal(host) + if host in _BLOCKED_METADATA_HOSTS or (ip_literal is not None and ip_literal in _BLOCKED_METADATA_IPS): + raise ValueError(f"Compresr guardrail api_base {host!r} is a blocked cloud-metadata host") + return url + + +def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, dict) + + +def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + +def _content_to_text(content: object) -> str: + """Collapse a message ``content`` (str or list-of-parts) to plain text. + + For the multimodal list shape, joins ``{type: "text", text: ...}`` parts + with blank-line separators; non-text parts are ignored. + """ + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + text = part.get("text") + if isinstance(text, str): + parts.append(text) + return "\n\n".join(parts) + return "" + + +def _replace_text_in_content(content: object, new_text: str) -> object: + """Write ``new_text`` back into a ``content`` value, preserving shape. + + ``str`` content is replaced directly. For list-of-parts content the first + text part carries ``new_text``, later text parts are dropped, and + non-text parts (images, audio, files) pass through untouched. + """ + if isinstance(content, str): + return new_text + if isinstance(content, list): + out: list[object] = [] + replaced = False + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + if not replaced: + out.append({**part, "text": new_text}) + replaced = True + continue + out.append(part) + if not replaced: + out.insert(0, {"type": "text", "text": new_text}) + return out + return new_text + + +def _render_tool_intent(fn: dict[str, object]) -> str: + name = str(fn.get("name") or "").strip() + args = fn.get("arguments") + if isinstance(args, dict): + try: + args_str = json.dumps(args, separators=(",", ":")) + except (TypeError, ValueError): + args_str = str(args) + else: + args_str = str(args).strip() if args is not None else "" + if name and args_str: + return f"{name}: {args_str}" + return name or args_str + + +def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallback: str) -> str: + """Query used to compress ``messages[target_idx]``. + + Tool/function outputs are compressed against the intent of the tool call + that produced them (found via ``tool_call_id`` on a prior assistant + message); everything else uses the last user message. + """ + msg = messages[target_idx] + if msg.get("role") not in ("tool", "function"): + return fallback + + tool_call_id = msg.get("tool_call_id") + fn_name = msg.get("name") + for j in range(target_idx - 1, -1, -1): + prev = messages[j] + if prev.get("role") != "assistant": + continue + tool_calls = prev.get("tool_calls") + if isinstance(tool_calls, list): + for tc in tool_calls: + if not isinstance(tc, dict): + continue + if tool_call_id and tc.get("id") == tool_call_id: + fn = tc.get("function") + intent = _render_tool_intent(fn if isinstance(fn, dict) else {}) + if intent: + return intent + # Legacy function_call fallback: require a name match, else an earlier + # function_call turn would attribute the wrong intent. + fc = prev.get("function_call") + if isinstance(fc, dict) and fn_name and fc.get("name") == fn_name: + intent = _render_tool_intent(fc) + if intent: + return intent + return fallback + + +def _safe_int(value: object) -> int: + """Parse a token-stat field defensively. + + A malformed-but-200 response must not raise here: ``_call_compress`` has + already returned successfully, so the fail_open/fail_closed decision is + behind us. A bare ``int()`` on a non-numeric field would surface as an + unhandled 500 even when ``fail_open`` is configured. + """ + try: + return int(value) if value is not None else 0 + except (TypeError, ValueError): + return 0 + + +def _safe_response_text(response: object, limit: int = 500) -> str: + """Read a response body for error logging without letting the read itself + raise. A corrupt ``Content-Encoding`` makes ``httpx``'s ``.text`` raise a + ``DecodingError``; if that happened while building a failure detail it would + turn an already-handled error into an unhandled 500.""" + try: + text = getattr(response, "text", "") + except httpx.DecodingError: + return "" + return (text or "")[:limit] + + +def _content_hash(text: str) -> str: + # surrogatepass so a lone surrogate in untrusted content (valid via a JSON + # \uXXXX escape) hashes instead of raising past the fail policy. + return hashlib.sha256(text.encode("utf-8", "surrogatepass")).hexdigest()[:24] + + +def _entry_bytes(originals: dict[str, str]) -> int: + """UTF-8 byte size of one recovery-store entry (surrogatepass, like _content_hash).""" + return sum(len(value.encode("utf-8", "surrogatepass")) for value in originals.values()) + + +def _display_hash(hash_value: str) -> str: + """Bound a model-supplied hash for logs/fallback text. A real marker hash is + 24 hex chars; a prompt-injected ``compresr_retrieve`` call could pass a huge + or control-character-laden string, so strip non-printables (no forged log + lines / ANSI escapes) and cap length before echoing into logs and the + conversation.""" + printable = "".join(ch for ch in hash_value if ch.isprintable()) + return printable if len(printable) <= 32 else f"{printable[:32]}…" + + +def _recovery_marker(hash_value: str) -> str: + return ( + f"\n\n[compresr hash={hash_value}: parts of this content were compressed " + f"away. If you need the full original, call the " + f"{COMPRESR_RETRIEVE_TOOL_NAME} tool with this hash.]" + ) + + +def _build_compresr_retrieve_tool() -> dict[str, object]: + return { + "type": "function", + "function": { + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "description": ( + "Retrieve the original, uncompressed content behind a Compresr " + "compression marker. Call this when a compression marker's hash " + "points at content you need in full." + ), + "parameters": { + "type": "object", + "properties": { + "hash": { + "type": "string", + "description": "The 24-character hex hash from the compression marker.", + }, + }, + "required": ["hash"], + }, + }, + } + + +def has_compresr_retrieve_tool(tools: object) -> bool: + return has_tool_with_name(tools, COMPRESR_RETRIEVE_TOOL_NAME) + + +def _extract_compresr_tool_calls(response: object) -> list[dict[str, object]]: + return [ + {"id": tc["id"], "type": "function", "name": tc["name"], "arguments": tc["arguments"]} + for tc in get_tool_calls_from_response(response) + if tc["name"] == COMPRESR_RETRIEVE_TOOL_NAME + ] + + +def _resolve_call_id(logging_obj: object) -> str | None: + """The call id from the framework logging object. + + This value ultimately derives from the client-settable ``x-litellm-call-id`` + header and is echoed back in responses, so it is NOT a trust boundary on its + own — ``_scoped_store_key`` prefixes it with the caller's virtual-key hash to + partition the recovery store per tenant. Request-body/kwargs call ids are + deliberately not consulted here. + """ + logging_call_id = getattr(logging_obj, "litellm_call_id", None) + if isinstance(logging_call_id, str) and logging_call_id: + return logging_call_id + return None + + +def _caller_scope(logging_obj: object) -> str: + """The caller's virtual-key hash, used to partition the recovery store. + + Trust is anchored on the ``UserAPIKeyAuth`` object the proxy sets + server-side (``metadata.user_api_key_auth``, litellm_pre_call_utils). Its + ``api_key`` is the hash of the authenticated key. Both metadata spellings + are scanned (``/v1/messages`` and ``/v1/responses`` carry it under + ``litellm_metadata``), but the bare ``user_api_key`` *string* is never + trusted on its own: a JSON request body can place one in the client-supplied + ``metadata`` field, which is only sanitized on the route's canonical + container. Returns "" when the proxy runs without per-key auth, in which case + all traffic is a single trust domain and the call id alone suffices. + """ + details = getattr(logging_obj, "model_call_details", None) + if not _is_str_object_dict(details): + return "" + litellm_params = details.get("litellm_params") + for container in (litellm_params, details): + if not _is_str_object_dict(container): + continue + for meta_key in ("metadata", "litellm_metadata"): + metadata = container.get(meta_key) + if not _is_str_object_dict(metadata): + continue + auth = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and isinstance(auth.api_key, str) and auth.api_key: + return auth.api_key + return "" + + +def _scoped_store_key(logging_obj: object) -> str | None: + """Key for the recovery store: caller identity plus framework call id. + + Keying on the call id alone is unsafe: it comes from the client-settable + ``x-litellm-call-id`` header and is echoed back in responses, so one caller + could read or evict another's originals by reusing the id. Prefixing the + unforgeable virtual-key hash binds each entry to the tenant that created it. + Returns None when there is no call id, which disables recovery for the call. + """ + call_id = _resolve_call_id(logging_obj) + if call_id is None: + return None + scope = _caller_scope(logging_obj) + return f"{scope}\x00{call_id}" if scope else call_id + + +def _is_responses_api_response(response: object) -> bool: + return isinstance(get_attribute_or_key(response, "output", None), list) + + +def _is_anthropic_messages_response(response: object) -> bool: + return isinstance(get_attribute_or_key(response, "content", None), list) + + +def _assistant_text_from_response(response: object) -> str | None: + """The assistant's natural-language text from a model response, across chat, + Anthropic, and Responses shapes. Preserved when the turn is rebuilt for the + retrieval follow-up so the model's reasoning is not lost.""" + choices = get_attribute_or_key(response, "choices", None) + if isinstance(choices, list) and choices: + message = get_attribute_or_key(choices[0], "message", None) + if message is not None: + text = _content_to_text(get_attribute_or_key(message, "content", None)) + if text: + return text + content = get_attribute_or_key(response, "content", None) + if isinstance(content, list): + parts = [ + text + for block in content + if get_attribute_or_key(block, "type", None) == "text" + for text in (get_attribute_or_key(block, "text", None),) + if isinstance(text, str) and text + ] + if parts: + return "".join(parts) + output = get_attribute_or_key(response, "output", None) + if isinstance(output, list): + parts = [] + for item in output: + if get_attribute_or_key(item, "type", None) != "message": + continue + item_content = get_attribute_or_key(item, "content", None) + if not isinstance(item_content, list): + continue + for chunk in item_content: + if get_attribute_or_key(chunk, "type", None) == "output_text": + text = get_attribute_or_key(chunk, "text", None) + if isinstance(text, str) and text: + parts.append(text) + if parts: + return "".join(parts) + return None + + +def _build_assistant_message_from_response( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> dict[str, object]: + """Rebuild the chat-completions assistant turn for the retrieval follow-up. + + Only the ``compresr_retrieve`` calls are echoed, each answered by a tool + result below. Other tool calls made in the same turn are omitted on purpose: + the follow-up re-runs the model with the recovered content so it re-plans + them. Echoing them would leave tool_calls with no matching tool result and + the provider would reject the request. + """ + return { + "role": "assistant", + "content": _assistant_text_from_response(response), + "tool_calls": [ + { + "id": tool_call.get("id"), + "type": "function", + "function": { + "name": tool_call.get("name"), + "arguments": json.dumps(tool_call.get("arguments", {})), + }, + } + for tool_call, _ in retrieved + ], + } + + +def _build_anthropic_followup_messages( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> list[dict[str, object]]: + """Anthropic requires the tool_use block echoed back in an assistant + message paired with a tool_result block keyed by the same tool_use_id. The + assistant text is preserved; non-retrieve tool calls are re-planned by the + follow-up (see _build_assistant_message_from_response).""" + assistant_content: list[dict[str, object]] = [] + text = _assistant_text_from_response(response) + if text: + assistant_content.append({"type": "text", "text": text}) + assistant_content.extend( + { + "type": "tool_use", + "id": tool_call.get("id"), + "name": tool_call.get("name"), + "input": tool_call.get("arguments", {}), + } + for tool_call, _ in retrieved + ) + assistant_message: dict[str, object] = {"role": "assistant", "content": assistant_content} + user_message: dict[str, object] = { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": tool_call.get("id"), "content": content} + for tool_call, content in retrieved + ], + } + return [assistant_message, user_message] + + +def _build_responses_followup_items( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> list[dict[str, object]]: + """The Responses API requires the model's function_call echoed back paired + with a function_call_output keyed by the same call_id. The assistant text is + preserved; non-retrieve tool calls are re-planned by the follow-up.""" + items: list[dict[str, object]] = [] + text = _assistant_text_from_response(response) + if text: + items.append({"role": "assistant", "content": text}) + for tool_call, content in retrieved: + call_id = tool_call.get("id") + items.append( + { + "type": "function_call", + "call_id": call_id, + "name": tool_call.get("name"), + "arguments": json.dumps(tool_call.get("arguments", {})), + } + ) + items.append({"type": "function_call_output", "call_id": call_id, "output": content}) + return items + + +@dataclass +class _CompressionResult: + """Outcome of applying compression results to a message list.""" + + compressed_messages: list[dict[str, object]] + originals: dict[str, str] = field(default_factory=dict) + # original text -> compressed text, plus the machinery the Responses `texts` + # mirror needs to replace only where it is unambiguous. + text_replacements: dict[str, str] = field(default_factory=dict) + replaced_text_counts: dict[str, int] = field(default_factory=dict) + ambiguous_texts: set[str] = field(default_factory=set) + messages_compressed: int = 0 + tokens_before: int = 0 + tokens_after: int = 0 + + +class CompresrGuardrail(CustomGuardrail): + def __init__( + self, + api_base: str | None = None, + api_key: str | None = None, + model: str | None = None, + target_compression_ratio: float | None = None, + coarse: bool | None = None, + min_chars_to_compress: int | None = None, + compress_tool_outputs: bool | None = None, + compress_system: bool | None = None, + compress_history: bool | None = None, + compress_last_user: bool | None = None, + enable_retrieval: bool | None = None, + guardrail_name: str | None = None, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, + default_on: bool = False, + unreachable_fallback: str | None = None, + max_bytes_per_call: int | None = None, + allow_bypass_header: bool | None = None, + dynamic: bool | None = None, + dynamic_min_ratio: float | None = None, + dynamic_max_ratio: float | None = None, + compression_params: dict[str, object] | None = None, + ): + raw_api_base = (api_base or get_secret_str("COMPRESR_API_BASE") or DEFAULT_API_BASE).rstrip("/") + self.compresr_api_base = _validate_api_base(raw_api_base) + self.compresr_api_key = api_key or get_secret_str("COMPRESR_API_KEY") + if not self.compresr_api_key: + raise ValueError( + "Compresr guardrail requires an API key. Set `api_key` in the " + "guardrail config or the COMPRESR_API_KEY env var." + ) + self.compression_model = model or DEFAULT_COMPRESSION_MODEL + self.target_compression_ratio = ( + DEFAULT_TARGET_COMPRESSION_RATIO if target_compression_ratio is None else target_compression_ratio + ) + self.coarse = True if coarse is None else coarse + self.min_chars_to_compress = ( + DEFAULT_MIN_CHARS_TO_COMPRESS if min_chars_to_compress is None else min_chars_to_compress + ) + self.compress_tool_outputs = True if compress_tool_outputs is None else compress_tool_outputs + self.compress_system = False if compress_system is None else compress_system + self.compress_history = False if compress_history is None else compress_history + self.compress_last_user = False if compress_last_user is None else compress_last_user + self.enable_retrieval = True if enable_retrieval is None else enable_retrieval + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( + "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" + ) + self.max_bytes_per_call = _DEFAULT_MAX_BYTES_PER_CALL if max_bytes_per_call is None else max_bytes_per_call + if self.max_bytes_per_call < 0: + raise ValueError("max_bytes_per_call must be >= 0 (0 disables the cap; positive values enforce it)") + self.allow_bypass_header = False if allow_bypass_header is None else allow_bypass_header + # Dynamic (adaptive) compression — latte_v2 only, on by default: the server + # picks the ratio per input instead of honoring target_compression_ratio. + self.dynamic = True if dynamic is None else dynamic + self.dynamic_min_ratio = dynamic_min_ratio + self.dynamic_max_ratio = dynamic_max_ratio + # Passthrough of extra compression params forwarded verbatim, so a new + # Compresr feature works without changing this guardrail. Named fields win; + # request-content fields are stripped. + reserved_keys = _RESERVED_COMPRESSION_PARAM_KEYS.intersection(compression_params or {}) + if reserved_keys: + verbose_proxy_logger.warning( + "Compresr: ignoring reserved compression_params keys %s", sorted(reserved_keys) + ) + self.compression_params: dict[str, object] = { + k: v for k, v in (compression_params or {}).items() if k not in _RESERVED_COMPRESSION_PARAM_KEYS + } + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + self._originals_by_call_id: OrderedDict[str, tuple[dict[str, str], float]] = OrderedDict() + # Running byte size of the store, kept in sync to enforce the global cap cheaply. + self._store_total_bytes = 0 + # One-shot guard so the "recovery skipped, no auth scope" warning fires once. + self._warned_no_scope_recovery = False + if self.enable_retrieval: + verbose_proxy_logger.warning( + "Compresr: enable_retrieval is on; the recovery store is per-process. " + "For multi-worker deployments, set enable_retrieval=false or run with --workers 1." + ) + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + def _should_bypass(self, request_data: dict) -> bool: + if not self.allow_bypass_header: + return False + psr = request_data.get("proxy_server_request") + if not _is_str_object_dict(psr): + return False + headers = psr.get("headers") + if not _is_str_object_dict(headers): + return False + return str(headers.get(BYPASS_HEADER)).lower() == "true" + + def _request_headers(self) -> dict[str, str]: + return { + "Content-Type": "application/json", + "X-API-Key": self.compresr_api_key or "", + } + + def _handle_compress_failure(self, error: str, log_detail: dict[str, object]) -> None: + """fail_open logs and returns (caller forwards uncompressed); + fail_closed raises. ``log_detail`` may include upstream response bodies + and is written only to server logs; the raised ``HTTPException`` carries + a generic message so a malicious ``api_base`` cannot exfiltrate response + bytes through the client-visible error.""" + if self.unreachable_fallback == "fail_open": + verbose_proxy_logger.warning( + "Compresr: %s; fail_open configured, forwarding request uncompressed. detail=%s", + error, + log_detail, + ) + return + verbose_proxy_logger.error("Compresr: %s. detail=%s", error, log_detail) + raise HTTPException(status_code=502, detail={"error": error}) + + def _evict_oldest(self) -> None: + """Drop the front (oldest) entry and decrement the running byte total.""" + _key, (evicted, _expiry) = self._originals_by_call_id.popitem(last=False) + self._store_total_bytes -= _entry_bytes(evicted) + + def _prune_originals(self) -> None: + # Insertion order == expiry order (shared TTL); prune from the front. + now = time.monotonic() + store = self._originals_by_call_id + while store and store[next(iter(store))][1] <= now: + self._evict_oldest() + while len(store) > _MAX_TRACKED_CALLS: + self._evict_oldest() + # Global byte budget; keep the most-recent entry so the current call's + # originals survive (a single call is already bounded by max_bytes_per_call). + while len(store) > 1 and self._store_total_bytes > _MAX_TOTAL_STORE_BYTES: + self._evict_oldest() + + def _existing_originals(self, store_key: str | None) -> dict[str, str]: + """Originals already stored under this key, so the per-call byte budget + can account for an earlier turn that reused the store key.""" + if store_key is None: + return {} + return self._originals_by_call_id.get(store_key, ({}, 0.0))[0] + + def _store_originals(self, store_key: str, originals: dict[str, str]) -> None: + existing, _ = self._originals_by_call_id.get(store_key, ({}, 0.0)) + merged = self._bound_call_bytes({**existing, **originals}) + # Keep the running total in sync: drop the overwritten entry, add the new one. + self._store_total_bytes += _entry_bytes(merged) - _entry_bytes(existing) + self._originals_by_call_id[store_key] = ( + merged, + time.monotonic() + _ORIGINALS_TTL_SECONDS, + ) + self._originals_by_call_id.move_to_end(store_key) + self._prune_originals() + + def _bound_call_bytes(self, merged: dict[str, str]) -> dict[str, str]: + """Drop oldest entries (dict insertion order) until the aggregate byte + size fits ``self.max_bytes_per_call``. Prevents one call with many + large tool outputs from growing proxy memory without bound.""" + if self.max_bytes_per_call <= 0: + return merged + total = _entry_bytes(merged) + if total <= self.max_bytes_per_call: + return merged + bounded = dict(merged) + for key in list(bounded.keys()): + if total <= self.max_bytes_per_call: + break + total -= len(bounded[key].encode("utf-8", "surrogatepass")) + del bounded[key] + verbose_proxy_logger.warning("Compresr: originals-store byte cap hit, evicted hash=%s", key) + return bounded + + def _retrieve_original(self, store_key: str | None, hash_value: str) -> str | None: + """Stored original for a marker hash, or None if not issued for this + request (unknown, expired, or from another caller's scope).""" + if store_key: + originals, expiry = self._originals_by_call_id.get(store_key, ({}, 0.0)) + if expiry > time.monotonic() and hash_value in originals: + return originals[hash_value] + verbose_proxy_logger.warning( + "Compresr retrieve: rejecting hash=%s (not issued for this request, or expired)", + _display_hash(hash_value), + ) + return None + + def _resolve_retrievals( + self, store_key: str | None, tool_calls: list[dict[str, object]] + ) -> tuple[list[tuple[dict[str, object], str]], bool]: + """Resolve compresr_retrieve calls to (call, result_text) pairs, deduping + repeated hashes and capping the count so the follow-up cannot be amplified. + The bool is True iff at least one call resolved to real stored content.""" + retrieved: list[tuple[dict[str, object], str]] = [] + seen: set[str] = set() + resolved_any = False + for idx, tc in enumerate(tool_calls): + arguments = tc.get("arguments", {}) + hash_value = str(arguments.get("hash", "")) if isinstance(arguments, dict) else "" + if idx >= _MAX_RETRIEVALS_PER_LOOP: + result = "[compresr: retrieval limit reached for this turn]" + elif hash_value in seen: + result = "[compresr: already retrieved above for this hash]" + else: + content = self._retrieve_original(store_key, hash_value) + if content is None: + result = f"[compresr: hash={_display_hash(hash_value)} not found, expired, or not issued for this request]" + else: + seen.add(hash_value) + resolved_any = True + result = content + verbose_proxy_logger.debug("Compresr retrieve: hash=%s -> %d chars", _display_hash(hash_value), len(result)) + retrieved.append((tc, result)) + return retrieved, resolved_any + + async def _call_compress( + self, + contexts: list[str], + queries: list[str], + ) -> list[dict[str, object]] | None: + """Compress ``contexts`` (query-aware). Returns one result dict per + context, or None when the service failed and fail_open applies.""" + common: dict[str, object] = { + # Passthrough first so the named fields below always win on collision. + **self.compression_params, + "compression_model_name": self.compression_model, + "target_compression_ratio": self.target_compression_ratio, + "coarse": self.coarse, + "dynamic": self.dynamic, + "source": _SOURCE_TAG, + } + # Only send the bounds the operator actually set; otherwise let the + # server apply its own floor/ceiling. + if self.dynamic_min_ratio is not None: + common["dynamic_min_ratio"] = self.dynamic_min_ratio + if self.dynamic_max_ratio is not None: + common["dynamic_max_ratio"] = self.dynamic_max_ratio + if len(contexts) == 1: + url = f"{self.compresr_api_base}/api/compress/question-specific/" + payload: dict[str, object] = { + "context": contexts[0], + "query": queries[0], + **common, + } + else: + url = f"{self.compresr_api_base}/api/compress/question-specific/batch" + payload = { + "inputs": [{"context": ctx, "query": q} for ctx, q in zip(contexts, queries)], + **common, + } + + try: + raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped + url=url, + json=payload, + headers=self._request_headers(), + timeout=_COMPRESS_TIMEOUT_SECONDS, + ) + except asyncio.CancelledError: + raise + except httpx.HTTPStatusError as e: + # The shared handler calls raise_for_status(), so a non-2xx reply arrives + # here as an error carrying the upstream body + our API key header; route + # it through the fail policy so none of that reaches the client. + resp = getattr(e, "response", None) + self._handle_compress_failure( + "Compresr compression service returned an error", + { + "status_code": getattr(resp, "status_code", None), + "body": _safe_response_text(resp), + }, + ) + return None + except (httpx.RequestError, litellm.Timeout) as e: + # Every request-side httpx failure is a RequestError; route the whole + # class through the fail policy so none escapes as a 500 under fail_open. + # (HTTPStatusError is handled above and is not a RequestError.) + self._handle_compress_failure( + "Compresr compression service request failed", + {"detail": str(e)}, + ) + return None + if raw_response is None or not 200 <= raw_response.status_code < 300: + self._handle_compress_failure( + "Compresr compression service returned an error", + { + "status_code": getattr(raw_response, "status_code", None), + "body": _safe_response_text(raw_response), + }, + ) + return None + + try: + body: object = raw_response.json() + except (ValueError, httpx.DecodingError, RecursionError): + # RecursionError: a deeply nested JSON body overflows the parser; + # route it through the fail policy rather than let it escape as a 500. + self._handle_compress_failure( + "Compresr compression service returned an unreadable response", + {"body": _safe_response_text(raw_response)}, + ) + return None + if not _is_str_object_dict(body) or not _is_str_object_dict(body.get("data")): + self._handle_compress_failure( + "Compresr compression service returned unexpected response shape", + {"body": _safe_response_text(raw_response)}, + ) + return None + data: dict[str, object] = body["data"] # pyright: ignore[reportAssignmentType] # dict-guarded above; subscript does not narrow + + if len(contexts) == 1: + return [data] + results = data.get("results") + if ( + not _is_object_list(results) + or len(results) != len(contexts) + or not all(_is_str_object_dict(r) for r in results) + ): + # Anything but a 1:1 dict-per-context mapping would misalign + # results with their target messages. + self._handle_compress_failure( + "Compresr batch response missing or mismatched 'results'", + {"expected": len(contexts), "got": len(results) if _is_object_list(results) else None}, + ) + return None + return results # pyright: ignore[reportReturnType] # every element dict-checked above; list[object] does not narrow + + def _select_targets(self, messages: list[dict[str, object]], query_idx: int | None) -> list[int]: + """Indices of messages whose text content should be compressed.""" + targets: list[int] = [] + for idx, msg in enumerate(messages): + if idx == query_idx and not self.compress_last_user: + continue + role = msg.get("role") + if role in ("tool", "function"): + if not self.compress_tool_outputs: + continue + elif role == "system": + if not self.compress_system: + continue + elif role == "user": + if idx != query_idx and not self.compress_history: + continue + else: + continue + if len(_content_to_text(msg.get("content"))) < self.min_chars_to_compress: + continue + targets.append(idx) + return targets + + @staticmethod + def _extract_fallback_query( + messages: list[dict[str, object]], + ) -> tuple[str, int | None]: + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") == "user": + return _content_to_text(messages[idx].get("content")), idx + return "", None + + def _apply_compression_results( + self, + messages: list[dict[str, object]], + targets: list[int], + contexts: list[str], + results: list[dict[str, object]], + recovery_enabled: bool, + existing_originals: dict[str, str] | None = None, + ) -> _CompressionResult: + """Write each compression result into a copy of ``messages``. + + A result is a real compression only when it is a non-empty string that + differs from the original; identical text is treated as a no-op so an + untouched request is not needlessly rewritten downstream. + """ + out = _CompressionResult(compressed_messages=list(messages)) + existing = existing_originals or {} + cap = self.max_bytes_per_call + # Seed with what is already stored under this store key: markers are + # attached only while the store (existing + this call's originals) stays + # within the cap, so _store_originals never has to evict a hash this call + # just shipped a marker for -- including on a later turn that reuses the + # store key. A hash already stored (or repeated here) costs no new bytes. + recovery_bytes = _entry_bytes(existing) + for target_idx, original_text, result in zip(targets, contexts, results): + compressed_text = result.get("compressed_context") + if not isinstance(compressed_text, str) or not compressed_text or compressed_text == original_text: + continue + out.messages_compressed += 1 + if recovery_enabled: + hash_value = _content_hash(original_text) + already_stored = hash_value in existing or hash_value in out.originals + new_bytes = 0 if already_stored else len(original_text.encode("utf-8", "surrogatepass")) + if cap <= 0 or recovery_bytes + new_bytes <= cap: + recovery_bytes += new_bytes + out.originals[hash_value] = original_text + compressed_text += _recovery_marker(hash_value) + previous = out.text_replacements.get(original_text) + if previous is not None and previous != compressed_text: + # Two targets with identical text but different query-specific + # compressions; a value-keyed replacement cannot tell them apart. + out.ambiguous_texts.add(original_text) + else: + out.text_replacements[original_text] = compressed_text + out.replaced_text_counts[original_text] = out.replaced_text_counts.get(original_text, 0) + 1 + original_msg = out.compressed_messages[target_idx] + out.compressed_messages[target_idx] = { + **original_msg, + "content": _replace_text_in_content(original_msg.get("content"), compressed_text), + } + out.tokens_before += _safe_int(result.get("original_tokens")) + out.tokens_after += _safe_int(result.get("compressed_tokens")) + return out + + @staticmethod + def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: + """Compressed content mirrored into the Responses `texts` channel. + + The chat/Anthropic handlers round-trip ``structured_messages``; the + Responses translation cannot rebuild its input from chat messages and + instead writes back through ``texts``. This matches by value, so a + replacement is applied only when it is unambiguous: one compression per + text, and every occurrence in ``texts`` accounted for by a compressed + target. Anything else is left uncompressed rather than risk a wrong or + out-of-policy replacement. Returns None when nothing safe applies. + """ + if not applied.text_replacements or not isinstance(input_texts, list): + return None + counts = Counter(text for text in input_texts if isinstance(text, str)) + safe = { + text: replacement + for text, replacement in applied.text_replacements.items() + if text not in applied.ambiguous_texts and counts.get(text) == applied.replaced_text_counts.get(text) + } + if not safe: + return None + return [safe.get(text, text) if isinstance(text, str) else text for text in input_texts] + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + if self._should_bypass(request_data): + verbose_proxy_logger.debug("Compresr: %s header set; skipping compression", BYPASS_HEADER) + return inputs + + structured_messages = inputs.get("structured_messages") + if not _is_object_list(structured_messages) or not structured_messages: + return inputs + messages = [m for m in structured_messages if _is_str_object_dict(m)] + if len(messages) != len(structured_messages): + return inputs + + fallback_query, query_idx = self._extract_fallback_query(messages) + targets: list[int] = [] + queries: list[str] = [] + for idx in self._select_targets(messages, query_idx): + query = _query_for_target(messages, idx, fallback_query) + # latte models require a non-empty query; leave targets we cannot + # derive one for uncompressed rather than erroring. + if not query.strip(): + continue + targets.append(idx) + queries.append(query) + if not targets: + verbose_proxy_logger.debug("Compresr: no messages eligible for compression") + return inputs + + contexts = [_content_to_text(messages[idx].get("content")) for idx in targets] + + start_time = time.monotonic() + results = await self._call_compress(contexts=contexts, queries=queries) + end_time = time.monotonic() + if results is None: # service failed, fail_open configured + return inputs + + # Recovery needs a per-tenant scope; without per-key auth the key would fall + # back to the client-settable call id (cross-tenant reads), so skip it. + store_key = _scoped_store_key(logging_obj) + scope = _caller_scope(logging_obj) + recovery_enabled = self.enable_retrieval and store_key is not None and bool(scope) + if self.enable_retrieval and not scope and not self._warned_no_scope_recovery: + # Surface the silent no-recovery case once (compressed, but no auth + # scope to inject the retrieve tool). + self._warned_no_scope_recovery = True + verbose_proxy_logger.warning( + "Compresr: enable_retrieval is on but this request has no per-key auth scope; " + "compressing without recovery (compresr_retrieve tool not injected). " + "Configure virtual-key auth to enable recovery." + ) + + existing_originals = self._existing_originals(store_key) + applied = self._apply_compression_results( + messages, targets, contexts, results, recovery_enabled, existing_originals + ) + if applied.messages_compressed == 0: + # Nothing replaced: return the original inputs object (handlers detect + # edits by identity; a fresh list forces write-back that strips Anthropic + # cache_control from thinking blocks). + verbose_proxy_logger.debug("Compresr: service returned no compressed content; request unchanged") + return inputs + + stats: dict[str, object] = { + "messages_compressed": applied.messages_compressed, + "tokens_before": applied.tokens_before, + "tokens_after": applied.tokens_after, + "tokens_saved": applied.tokens_before - applied.tokens_after, + "compression_model": self.compression_model, + } + verbose_proxy_logger.debug( + "Compresr: compressed %s message(s), %s -> %s tokens", + applied.messages_compressed, + applied.tokens_before, + applied.tokens_after, + ) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=stats, + request_data=request_data, + guardrail_status="success", + guardrail_provider="compresr", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + + compressed_inputs: dict[str, object] = {**inputs, "structured_messages": applied.compressed_messages} + mirrored_texts = self._mirror_texts_channel(inputs.get("texts"), applied) + if mirrored_texts is not None: + compressed_inputs["texts"] = mirrored_texts + + originals = applied.originals + if not recovery_enabled or not originals or store_key is None: + return compressed_inputs # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + + self._store_originals(store_key, originals) + + existing_tools = inputs.get("tools") + retrieve_tool = _build_compresr_retrieve_tool() + if isinstance(existing_tools, list) and not has_compresr_retrieve_tool(existing_tools): + merged_tools: list[object] = list(existing_tools) + [retrieve_tool] + elif existing_tools is None: + merged_tools = [retrieve_tool] + else: + merged_tools = list(existing_tools) if isinstance(existing_tools, list) else [retrieve_tool] + + compressed_inputs["tools"] = merged_tools + return compressed_inputs # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + if not has_compresr_retrieve_tool(tools): + return False, {} + tool_calls = _extract_compresr_tool_calls(response) + if not tool_calls: + return False, {} + return True, {"tool_calls": tool_calls} + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: dict, + logging_obj: Any, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + tool_calls: list[dict[str, object]] = tools.get("tool_calls", []) # pyright: ignore[reportAssignmentType] # gate hook builds this dict with list values only + + self._prune_originals() + store_key = _scoped_store_key(logging_obj) + retrieved, resolved_any = self._resolve_retrievals(store_key, tool_calls) + if not resolved_any: + # Nothing this guardrail stored resolved; skip the extra provider round-trip. + return AgenticLoopPlan(run_agentic_loop=False) + + if _is_responses_api_response(response): + follow_up_messages = list(messages) + _build_responses_followup_items(response, retrieved) + elif _is_anthropic_messages_response(response): + follow_up_messages = list(messages) + _build_anthropic_followup_messages(response, retrieved) + else: + assistant_message = _build_assistant_message_from_response(response, retrieved) + tool_results = [ + {"role": "tool", "tool_call_id": tc.get("id"), "content": content} for tc, content in retrieved + ] + follow_up_messages = list(messages) + [assistant_message] + tool_results + + anthropic_max = anthropic_messages_optional_request_params.get("max_tokens") + max_tokens: int | None = anthropic_max if anthropic_max is not None else kwargs.get("max_tokens") + optional_params_without_max_tokens = { + k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" + } + + full_model_name = model + if logging_obj is not None: + agentic_params = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) + candidate = agentic_params.get("model", model) + if isinstance(candidate, str) and candidate: + full_model_name = candidate + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs={ + k: v for k, v in kwargs.items() if not k.startswith("_compresr") and k != "litellm_logging_obj" + }, + ), + metadata={"tool_type": "compresr_retrieve"}, + ) + + @staticmethod + def get_config_model() -> type[GuardrailConfigModel[object]] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, + ) + + return CompresrGuardrailConfigModel diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7e080b13633..9da161be4dfb 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -56,6 +56,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( HeadroomGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, +) """ Pydantic object defining how to set guardrails on litellm proxy @@ -123,6 +126,7 @@ class SupportedGuardrailIntegrations(Enum): VIGIL_GUARD = "vigil_guard" REPELLOAI = "repelloai" HEADROOM = "headroom" + COMPRESR = "compresr" class Role(Enum): @@ -697,7 +701,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default="fail_closed", description=( "Behavior when a guardrail endpoint is unreachable due to network errors. " - "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', and 'headroom'. " + "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. " "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed." ), ) @@ -790,6 +794,7 @@ class LitellmParams( BedrockGuardrailConfigModel, LakeraV2GuardrailConfigModel, HeadroomGuardrailConfigModel, + CompresrGuardrailConfigModel, RepelloAIGuardrailConfigModel, LassoGuardrailConfigModel, PillarGuardrailConfigModel, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py new file mode 100644 index 000000000000..dad61f83b7d7 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py @@ -0,0 +1,135 @@ +from typing import Any, Dict, Literal + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class CompresrGuardrailOptionalParams(BaseModel): + """Optional tuning knobs for the Compresr guardrail.""" + + target_compression_ratio: float | None = Field( + default=None, + description=( + "Compression strength. 0-1 is the fraction of tokens to remove " + "(0.5 = remove ~50%, the default); a value >1 is an Nx reduction " + "factor (e.g. 4 = ~4x smaller)." + ), + ) + coarse: bool | None = Field( + default=None, + description=("Paragraph-level compression (default, faster) instead of token-level (finer-grained)."), + ) + min_chars_to_compress: int | None = Field( + default=None, + description=("Skip messages whose text is shorter than this many characters. Defaults to 500."), + ) + compress_tool_outputs: bool | None = Field( + default=None, + description=("Compress tool/function result messages (search hits, RAG chunks, API dumps). Defaults to True."), + ) + compress_system: bool | None = Field( + default=None, + description="Also compress system messages. Defaults to False.", + ) + compress_history: bool | None = Field( + default=None, + description="Also compress prior (non-last) user messages. Defaults to False.", + ) + compress_last_user: bool | None = Field( + default=None, + description=( + "Also compress the last user message. The query sent to Compresr " + "is always the original verbatim text. Defaults to False." + ), + ) + enable_retrieval: bool | None = Field( + default=None, + description=( + "Make compression recoverable: inject a `compresr_retrieve` tool " + "so the model can fetch the original content behind a compression " + "marker via the agentic loop. Defaults to True. Set to False (or " + "run the proxy with --workers 1) for multi-worker deployments: " + "the recovery store is per-process, so pre-call and retrieval hooks " + "on different workers cannot see each other's originals." + ), + ) + max_bytes_per_call: int | None = Field( + default=None, + description=( + "Cap on aggregate bytes of stored originals per litellm_call_id. " + "When a call exceeds this, oldest entries are evicted so the " + "in-process store cannot grow without bound. Defaults to 10 MiB." + ), + ) + allow_bypass_header: bool | None = Field( + default=None, + description=( + "Honor the `x-compresr-bypass: true` request header to skip " + "compression for a single call. Off by default because the " + "header is caller-settable; enable only on trusted deployments." + ), + ) + dynamic: bool | None = Field( + default=None, + description=( + "latte_v2 only. Let the server choose the compression amount per input " + "(Kneedle elbow) instead of using target_compression_ratio. Defaults to True." + ), + ) + dynamic_min_ratio: float | None = Field( + default=None, + description=( + "latte_v2 only. Floor on the adaptive ratio when `dynamic` is on. " + "Unset lets the server default apply (~1.5)." + ), + ) + dynamic_max_ratio: float | None = Field( + default=None, + description=( + "latte_v2 only. Ceiling on the adaptive ratio when `dynamic` is on. " + "Unset lets the server default apply (~10.0)." + ), + ) + compression_params: Dict[str, Any] | None = Field( + default=None, + description=( + "Passthrough of extra parameters forwarded verbatim in the Compresr " + "compress payload (e.g. `heuristic_chunking`, or any newer knob), so " + "a new Compresr feature works without a guardrail update. The named " + "fields above take precedence on collision." + ), + ) + + +class CompresrGuardrailConfigModel(GuardrailConfigModel[CompresrGuardrailOptionalParams]): + api_key: str | None = Field( + default=None, + description=("Compresr API key. Falls back to the COMPRESR_API_KEY env var."), + ) + api_base: str | None = Field( + default=None, + description=( + "Base URL of the Compresr API. Falls back to the COMPRESR_API_BASE " + "env var, then https://api.compresr.ai. Point at your internal " + "service URL for on-prem deployments." + ), + ) + model: str | None = Field( + default=None, + description=( + "Compresr compression model (not the LLM). Defaults to 'latte_v2', the query-aware compression model." + ), + ) + unreachable_fallback: Literal["fail_closed", "fail_open"] = Field( + default="fail_closed", + description=( + "Behavior when the Compresr compression service is unreachable or errors. " + "'fail_closed' raises an error (default). 'fail_open' logs a critical error and " + "forwards the request uncompressed instead of blocking it." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Compresr (context compression)" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 49cd1b71ef28..4c45eaac7b99 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1168,3 +1168,69 @@ def test_should_return_none_for_none_input(self): data = {"input": None} result = handler.get_structured_messages(data) assert result is None + + +class ToolAppendingGuardrail(CustomGuardrail): + """Guardrail that appends a new function tool, mimicking a guardrail that + injects a retrieval/recovery tool the model can later call.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tools = list(inputs.get("tools") or []) + tools.append( + { + "type": "function", + "function": { + "name": "injected_tool", + "description": "injected by guardrail", + "parameters": {"type": "object", "properties": {}}, + }, + } + ) + inputs["tools"] = tools + return inputs + + +class TestOpenAIResponsesHandlerToolInjection: + """A tool a guardrail injects must survive the write-back to Responses format.""" + + def test_merge_keeps_guardrail_appended_tool(self): + """_merge_tools_after_guardrail must not drop the extra appended tool.""" + handler = OpenAIResponsesHandler() + original = [{"type": "function", "name": "a"}] + remapped = [ + {"type": "function", "name": "a"}, + {"type": "function", "name": "b"}, + ] + merged = handler._merge_tools_after_guardrail(original, remapped) + assert [t["name"] for t in merged] == ["a", "b"] + + @pytest.mark.asyncio + async def test_injected_tool_survives_when_request_already_has_tools(self): + """Regression: the merge dropped the injected tool whenever the request + already carried tools, so the model never saw it.""" + handler = OpenAIResponsesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="test") + + data = { + "input": [{"role": "user", "content": "hi", "type": "message"}], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + } + ], + "model": "gpt-4", + } + + result = await handler.process_input_messages(data, guardrail) + + names = [t.get("name") for t in result["tools"]] + assert "get_weather" in names + assert "injected_tool" in names diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py new file mode 100644 index 000000000000..ad27d8291a73 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py @@ -0,0 +1,2002 @@ +""" +Unit tests for the Compresr guardrail. + +Tests cover: +- apply_guardrail compresses eligible messages query-aware (tool-call intent + resolved via tool_call_id, falling back to the last user message) +- target selection: tool outputs by default, system/history opt-in, min-chars + threshold, targets without a derivable query are left uncompressed +- multimodal content: text parts replaced, non-text parts preserved +- recovery: hash marker appended, compresr_retrieve tool injected, originals + stored per litellm_call_id, agentic loop returns the original content and + rejects hashes not issued for the current request +- x-compresr-bypass header, response-type passthrough +- fail_closed raises HTTPException; fail_open forwards uncompressed +""" + +import hashlib +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, create_autospec, patch + +import httpx +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import ( + COMPRESR_RETRIEVE_TOOL_NAME, + CompresrGuardrail, + _content_hash, + _scoped_store_key, + has_compresr_retrieve_tool, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +FAKE_API_BASE = "https://compresr.example.com" +FAKE_API_KEY = "cmp_test-key" + +TOOL_OUTPUT = "Result 1: EV range comparison. " * 40 # > 500 chars +USER_QUESTION = "Which 2026 EV has the longest range?" + +AGENT_MESSAGES = [ + {"role": "system", "content": "You are a research assistant."}, + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "2026 EV range"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT}, +] + + +def _make_guardrail(**kwargs) -> CompresrGuardrail: + defaults = dict( + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + guardrail_name="compresr", + default_on=True, + ) + defaults.update(kwargs) + return CompresrGuardrail(**defaults) + + +def _make_single_compress_response( + compressed_context: str = "compressed summary", + original_tokens: int = 1000, + compressed_tokens: int = 400, + status: int = 200, +) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "success": True, + "data": { + "compressed_context": compressed_context, + "original_tokens": original_tokens, + "compressed_tokens": compressed_tokens, + "actual_compression_ratio": 0.6, + "tokens_saved": original_tokens - compressed_tokens, + "duration_ms": 42, + }, + } + mock.text = "" + return mock + + +def _make_batch_compress_response(compressed_contexts: list, status: int = 200) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "success": True, + "data": { + "results": [ + { + "compressed_context": ctx, + "original_tokens": 1000, + "compressed_tokens": 400, + "actual_compression_ratio": 0.6, + "tokens_saved": 600, + "duration_ms": 42, + } + for ctx in compressed_contexts + ], + "count": len(compressed_contexts), + }, + } + mock.text = "" + return mock + + +def _make_openai_response_with_tool_call(tool_name: str, arguments: dict, tool_id: str = "call_abc123") -> MagicMock: + fn = MagicMock() + fn.name = tool_name + fn.arguments = json.dumps(arguments) + + tc = MagicMock() + tc.id = tool_id + tc.type = "function" + tc.function = fn + + message = MagicMock() + message.content = None + message.tool_calls = [tc] + + choice = MagicMock() + choice.message = message + + response = MagicMock() + response.choices = [choice] + # Plain chat-completion shape: no responses-API `output` list, no + # anthropic `content` list. + response.output = None + response.content = None + return response + + +def _make_openai_response_with_tool_calls(tool_calls: list, content: object = None) -> MagicMock: + """Chat-completion response carrying several tool calls in one turn + (parallel tool calling). ``tool_calls`` items are (name, arguments, id).""" + tcs = [] + for name, arguments, tool_id in tool_calls: + fn = MagicMock() + fn.name = name + fn.arguments = json.dumps(arguments) + tc = MagicMock() + tc.id = tool_id + tc.type = "function" + tc.function = fn + tcs.append(tc) + + message = MagicMock() + message.content = content + message.tool_calls = tcs + + choice = MagicMock() + choice.message = message + + response = MagicMock() + response.choices = [choice] + response.output = None + response.content = None + return response + + +def _apply_inputs(messages: list) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs(structured_messages=[dict(m) for m in messages]) + + +def _logging_obj(call_id: str) -> SimpleNamespace: + # Default fixture models a proxy with per-key auth enabled (the production + # shape). Recovery requires a caller scope; tests that need the no-auth + # path should build the object explicitly. + from litellm.proxy._types import UserAPIKeyAuth + + return SimpleNamespace( + litellm_call_id=call_id, + model_call_details={ + "litellm_params": {"metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="hash-default")}} + }, + ) + + +def _logging_obj_with_key(call_id: str, user_api_key: str, meta_key: str = "metadata") -> SimpleNamespace: + """Logging object carrying the server-set UserAPIKeyAuth object, the way the + proxy populates it for an authenticated request (the bare user_api_key + string alone is never trusted — a client could forge that).""" + from litellm.proxy._types import UserAPIKeyAuth + + return SimpleNamespace( + litellm_call_id=call_id, + model_call_details={"litellm_params": {meta_key: {"user_api_key_auth": UserAPIKeyAuth(api_key=user_api_key)}}}, + ) + + +def _retrieve_tool_call(hash_value: str, tool_id: str) -> dict: + return { + "id": tool_id, + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + + +def _retrieve_tool_stub() -> dict: + return { + "type": "function", + "function": {"name": COMPRESR_RETRIEVE_TOOL_NAME, "parameters": {}}, + } + + +@pytest.fixture +def guardrail() -> CompresrGuardrail: + return _make_guardrail() + + +# ── init ────────────────────────────────────────────────────────────── + + +def test_init_raises_without_api_key(monkeypatch): + monkeypatch.delenv("COMPRESR_API_KEY", raising=False) + with pytest.raises(ValueError, match="API key"): + CompresrGuardrail(guardrail_name="compresr") + + +def test_init_defaults(): + g = _make_guardrail() + assert g.compresr_api_base == FAKE_API_BASE + assert g.compression_model == "latte_v2" + assert g.target_compression_ratio == 0.5 + assert g.coarse is True + assert g.min_chars_to_compress == 500 + assert g.compress_tool_outputs is True + assert g.compress_system is False + assert g.compress_history is False + assert g.compress_last_user is False + assert g.enable_retrieval is True + assert g.unreachable_fallback == "fail_closed" + + +def test_init_coerces_unknown_unreachable_fallback_to_fail_closed(): + g = _make_guardrail(unreachable_fallback="banana") + assert g.unreachable_fallback == "fail_closed" + + +# ── compression core ───────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_apply_guardrail_compresses_tool_output_with_intent_query( + guardrail: CompresrGuardrail, +): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + _, call_kwargs = mock_post.call_args + assert call_kwargs["url"] == f"{FAKE_API_BASE}/api/compress/question-specific/" + assert call_kwargs["headers"]["X-API-Key"] == FAKE_API_KEY + payload = call_kwargs["json"] + assert payload["context"] == TOOL_OUTPUT + # Query is the tool call's intent, not the user question. + assert payload["query"] == 'web_search: {"query": "2026 EV range"}' + assert payload["compression_model_name"] == "latte_v2" + assert payload["target_compression_ratio"] == 0.5 + + out = result["structured_messages"] + assert out[3]["content"].startswith("compressed summary") + # Untouched messages pass through byte-identical. + assert out[0] == AGENT_MESSAGES[0] + assert out[1] == AGENT_MESSAGES[1] + assert out[2] == AGENT_MESSAGES[2] + + +@pytest.mark.asyncio +async def test_apply_guardrail_mirrors_compression_into_texts_channel( + guardrail: CompresrGuardrail, +): + """The /v1/responses translation writes compressed output back through the + `texts` channel, not structured_messages. Compression must be mirrored there + or that surface silently forwards the original content uncompressed.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_unknown", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[USER_QUESTION, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + texts = result["texts"] + # The compressed tool output replaces the original in the texts channel... + assert texts[1].startswith("compressed summary") + assert texts[1] != TOOL_OUTPUT + # ...while untouched text passes through byte-identical. + assert texts[0] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_apply_guardrail_returns_inputs_unchanged_when_nothing_compressed( + guardrail: CompresrGuardrail, +): + """A 200 response whose compressed_context is empty is a functional no-op. + The exact inputs object must come back: handlers detect guardrail edits by + identity, and a fresh structured_messages list would force a full write-back + of an untouched request (on Anthropic, reconversion strips cache_control + from thinking blocks).""" + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response(compressed_context="")) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + assert result is inputs + + +@pytest.mark.asyncio +async def test_texts_mirror_skips_duplicate_content_with_diverging_compressions( + guardrail: CompresrGuardrail, +): + """Two targets with identical text but different query-specific compressions: + the value-keyed texts mirror cannot tell the occurrences apart, so it must + leave them uncompressed rather than apply an arbitrary variant to both.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "search_docs", "arguments": '{"q": "a"}'}}, + {"id": "call_2", "type": "function", "function": {"name": "search_web", "arguments": '{"q": "b"}'}}, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT}, + {"role": "tool", "tool_call_id": "call_2", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[USER_QUESTION, TOOL_OUTPUT, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_batch_compress_response(["compressed for docs", "compressed for web"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + # Each message position still gets its own query-specific compression... + out = result["structured_messages"] + assert out[2]["content"].startswith("compressed for docs") + assert out[3]["content"].startswith("compressed for web") + # ...but the texts mirror leaves the ambiguous occurrences untouched. + assert result["texts"] == [USER_QUESTION, TOOL_OUTPUT, TOOL_OUTPUT] + + +@pytest.mark.asyncio +async def test_texts_mirror_skips_text_that_also_appears_outside_targets( + guardrail: CompresrGuardrail, +): + """compress_system is off, so a system message whose text happens to equal + a compressed tool output must not be rewritten through the texts mirror.""" + messages = [ + {"role": "system", "content": TOOL_OUTPUT}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[TOOL_OUTPUT, USER_QUESTION, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + out = result["structured_messages"] + assert out[0]["content"] == TOOL_OUTPUT # system message untouched + assert out[2]["content"].startswith("compressed summary") + # One compressed target cannot account for two occurrences in texts. + assert result["texts"] == [TOOL_OUTPUT, USER_QUESTION, TOOL_OUTPUT] + + +@pytest.mark.asyncio +async def test_tool_output_without_matching_call_uses_user_question( + guardrail: CompresrGuardrail, +): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_unknown", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["query"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_function_result_without_name_does_not_bind_unrelated_call( + guardrail: CompresrGuardrail, +): + """A legacy function-role result missing its name must not adopt the intent + of an arbitrary earlier assistant function_call; it falls back to the last + user message.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "function_call": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + {"role": "function", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["query"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_target_without_derivable_query_left_uncompressed( + guardrail: CompresrGuardrail, +): + # No user message and no tool-call intent anywhere -> nothing to compress. + messages = [{"role": "tool", "tool_call_id": "call_x", "content": TOOL_OUTPUT}] + mock_post = AsyncMock() + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + mock_post.assert_not_called() + assert result["structured_messages"][0]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_system_and_history_not_compressed_by_default( + guardrail: CompresrGuardrail, +): + long_system = "Rules. " * 200 + messages = [ + {"role": "system", "content": long_system}, + {"role": "user", "content": "Old question? " * 100}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + # Only one (single, non-batch) call: the tool output. + assert mock_post.call_count == 1 + assert mock_post.call_args.kwargs["json"]["context"] == TOOL_OUTPUT + out = result["structured_messages"] + assert out[0]["content"] == long_system + assert out[2]["content"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_opt_in_system_uses_batch_endpoint(): + guardrail = _make_guardrail(compress_system=True) + long_system = "Rules. " * 200 + messages = [ + {"role": "system", "content": long_system}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_batch_compress_response(["short system", "short tool"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"].endswith("/api/compress/question-specific/batch") + batch_inputs = call_kwargs["json"]["inputs"] + assert [i["context"] for i in batch_inputs] == [long_system, TOOL_OUTPUT] + out = result["structured_messages"] + assert out[0]["content"].startswith("short system") + assert out[2]["content"].startswith("short tool") + + +@pytest.mark.asyncio +async def test_short_messages_skipped(guardrail: CompresrGuardrail): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": "tiny result"}, + ] + mock_post = AsyncMock() + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + mock_post.assert_not_called() + assert result["structured_messages"][1]["content"] == "tiny result" + + +@pytest.mark.asyncio +async def test_multimodal_text_replaced_non_text_preserved( + guardrail: CompresrGuardrail, +): + image_part = {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}} + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "tool", + "tool_call_id": "c1", + "content": [{"type": "text", "text": TOOL_OUTPUT}, image_part], + }, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + content = result["structured_messages"][1]["content"] + assert isinstance(content, list) + assert content[0]["type"] == "text" + assert content[0]["text"].startswith("compressed summary") + assert content[1] == image_part + + +# ── passthrough / bypass ───────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_bypass_header_skips_compression_when_allowed(): + guardrail = _make_guardrail(allow_bypass_header=True) + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock() + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"headers": {"x-compresr-bypass": "true"}}, + } + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") + + mock_post.assert_not_called() + assert result is inputs + + +@pytest.mark.asyncio +async def test_bypass_header_ignored_by_default(guardrail: CompresrGuardrail): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"headers": {"x-compresr-bypass": "true"}}, + } + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") + + mock_post.assert_called_once() + + +@pytest.mark.asyncio +async def test_response_input_type_passthrough(guardrail: CompresrGuardrail): + inputs = _apply_inputs(AGENT_MESSAGES) + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + assert result is inputs + + +@pytest.mark.asyncio +async def test_missing_structured_messages_passthrough(guardrail: CompresrGuardrail): + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + assert result is inputs + + +# ── failure policy ──────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_transport_error_raises_when_fail_closed(guardrail: CompresrGuardrail): + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(side_effect=httpx.ConnectError("boom")), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_transport_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(side_effect=httpx.ConnectError("boom")), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_non_json_response_raises_when_fail_closed(guardrail: CompresrGuardrail): + mock = MagicMock() + mock.status_code = 200 + mock.json.side_effect = ValueError("not json") + mock.text = "gateway error" + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock)): + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_http_exception_does_not_reflect_upstream_body(guardrail: CompresrGuardrail): + mock = MagicMock() + mock.status_code = 500 + mock.json.side_effect = ValueError("not json") + mock.text = "SECRET_INSTANCE_METADATA_TOKEN=aws-imds-response" + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert "SECRET_INSTANCE_METADATA_TOKEN" not in json.dumps(exc_info.value.detail) + + +def test_init_rejects_non_http_api_base(): + with pytest.raises(ValueError, match="scheme"): + CompresrGuardrail( + api_base="file:///etc/passwd", + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +def test_init_rejects_cloud_metadata_api_base(): + with pytest.raises(ValueError, match="metadata"): + CompresrGuardrail( + api_base="http://169.254.169.254", + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +@pytest.mark.parametrize( + "api_base", + [ + "http://2852039166", # decimal encoding of 169.254.169.254 + "http://0xa9fea9fe", # hex encoding + "http://[::ffff:169.254.169.254]", # IPv4-mapped IPv6 + ], +) +def test_init_rejects_encoded_cloud_metadata_api_base(api_base): + with pytest.raises(ValueError, match="metadata"): + CompresrGuardrail( + api_base=api_base, + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_ignores_user_supplied_call_id(guardrail: CompresrGuardrail): + mock_post = AsyncMock(return_value=_make_single_compress_response()) + attacker_call_id = "victim-tenant-call-id" + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o", "litellm_call_id": attacker_call_id}, + input_type="request", + logging_obj=_logging_obj("real-framework-call-id"), + ) + + assert not any(attacker_call_id in k for k in guardrail._originals_by_call_id) + assert any(k.endswith("real-framework-call-id") for k in guardrail._originals_by_call_id) + + +@pytest.mark.asyncio +async def test_agentic_plan_ignores_user_supplied_call_id(guardrail: CompresrGuardrail): + hash_value = "d" * 24 + guardrail._store_originals("victim-tenant-call-id", {hash_value: "victim-original"}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("attacker-call-id"), + stream=False, + kwargs={"litellm_call_id": "victim-tenant-call-id"}, + ) + + # Attacker's scope resolves nothing, so the loop is vetoed and the victim + # original never surfaces. + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_recovery_store_partitioned_by_caller_identity(guardrail: CompresrGuardrail): + """Two tenants that set the SAME client-forgeable x-litellm-call-id must not + read each other's stored originals, and each still reads its own.""" + shared_call_id = "shared-call-id" + expected_hash = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + + async def _plan_for(user_api_key: str, tool_id: str): + return await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(expected_hash, tool_id)]}, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": expected_hash}, tool_id=tool_id + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj_with_key(shared_call_id, user_api_key), + stream=False, + kwargs={}, + ) + + # Tenant A compresses and stores its original. + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=_make_single_compress_response())): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj_with_key(shared_call_id, "hash-tenant-A"), + ) + + # Tenant B, same call id, different virtual-key hash → different bucket, so + # nothing resolves and the loop is vetoed (Tenant A's original never leaks). + plan_b = await _plan_for("hash-tenant-B", "call_b") + assert plan_b.run_agentic_loop is False + assert plan_b.request_patch is None + + # Tenant A retrieves its own content successfully. + plan_a = await _plan_for("hash-tenant-A", "call_a") + assert TOOL_OUTPUT in plan_a.request_patch.messages[-1]["content"] + + +@pytest.mark.asyncio +async def test_caller_scope_read_from_litellm_metadata(guardrail: CompresrGuardrail): + """/v1/messages and /v1/responses carry the auth object under + litellm_metadata rather than metadata; the store key must be scoped by it + there too, without relying on upstream's metadata backfill.""" + logging_obj = _logging_obj_with_key("call-lm", "hash-tenant-lm", meta_key="litellm_metadata") + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + + assert "hash-tenant-lm\x00call-lm" in guardrail._originals_by_call_id + assert "call-lm" not in guardrail._originals_by_call_id + + +@pytest.mark.asyncio +async def test_caller_scope_rejects_forged_user_api_key_string(guardrail: CompresrGuardrail): + """A client-supplied metadata.user_api_key STRING (no server-set + UserAPIKeyAuth object) must not be trusted as a tenant scope — otherwise a + caller could forge another tenant's recovery bucket on /v1/messages.""" + logging_obj = SimpleNamespace( + litellm_call_id="call-forge", + model_call_details={"litellm_params": {"metadata": {"user_api_key": "victim-tenant-hash"}}}, + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + + # Forged string is ignored: scope resolves to empty, so recovery is + # disabled entirely (no bucket keyed on victim-tenant-hash, no unscoped + # bucket that another caller could reuse). + assert not guardrail._originals_by_call_id + + +@pytest.mark.asyncio +async def test_compress_post_called_with_real_handler_signature(): + """AsyncHTTPHandler.post has a fixed signature; an autospec mock enforces it + (unlike AsyncMock(spec=...), which silently accepts any kwarg) so a kwarg the + real handler rejects — which would raise TypeError past the fail policy — + fails the test instead.""" + guardrail = _make_guardrail() + autospec_post = create_autospec(guardrail.async_handler.post, return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", autospec_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + assert result["structured_messages"][3]["content"].startswith("compressed summary") + + +@pytest.mark.asyncio +async def test_batch_result_count_mismatch_raises_when_fail_closed(): + guardrail = _make_guardrail(compress_system=True) + messages = [ + {"role": "system", "content": "Rules. " * 200}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_batch_compress_response(["only one"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +# ── recovery (compresr_retrieve) ────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_recovery_marker_tool_injection_and_original_stored( + guardrail: CompresrGuardrail, +): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + compressed_content = result["structured_messages"][3]["content"] + expected_hash = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + assert f"compresr hash={expected_hash}" in compressed_content + + tools = result.get("tools") + assert tools is not None and has_compresr_retrieve_tool(tools) + + scoped_key = next(k for k in guardrail._originals_by_call_id if k.endswith("call-id-1")) + originals, _expiry = guardrail._originals_by_call_id[scoped_key] + assert originals[expected_hash] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_enable_retrieval_false_no_marker_no_tool(): + guardrail = _make_guardrail(enable_retrieval=False) + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result["structured_messages"][3]["content"] == "compressed summary" + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +@pytest.mark.asyncio +async def test_existing_tools_preserved_when_injecting(guardrail: CompresrGuardrail): + existing_tool = {"type": "function", "function": {"name": "my_tool", "parameters": {}}} + inputs = GenericGuardrailAPIInputs( + structured_messages=[dict(m) for m in AGENT_MESSAGES], + tools=[existing_tool], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + tools = result["tools"] + assert existing_tool in tools + assert has_compresr_retrieve_tool(tools) + assert len(tools) == 2 + + +# ── agentic loop ────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_true_for_retrieve_call( + guardrail: CompresrGuardrail, +): + response = _make_openai_response_with_tool_call(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": "a" * 24}) + tools = [dict(t) for t in [_retrieve_tool_stub()]] + + should_run, gate_tools = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=tools, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + assert should_run is True + assert gate_tools["tool_calls"][0]["name"] == COMPRESR_RETRIEVE_TOOL_NAME + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_false_without_retrieve_tool( + guardrail: CompresrGuardrail, +): + response = _make_openai_response_with_tool_call("other_tool", {"x": 1}) + should_run, _ = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + assert should_run is False + + +@pytest.mark.asyncio +async def test_agentic_plan_returns_stored_original(guardrail: CompresrGuardrail): + hash_value = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assert plan.run_agentic_loop is True + follow_up = plan.request_patch.messages + tool_result = follow_up[-1] + assert tool_result["role"] == "tool" + assert tool_result["tool_call_id"] == "call_abc" + assert tool_result["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_preserves_list_shaped_assistant_text(guardrail: CompresrGuardrail): + """Some providers return chat assistant content as list-of-parts; the + retrieval follow-up must keep that text, not drop it to None.""" + hash_value = "a" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: "original"}) + response = _make_openai_response_with_tool_calls( + [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, "call_r")], + content=[{"type": "text", "text": "Let me fetch the original."}], + ) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_r")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assistant_message = plan.request_patch.messages[-2] + assert assistant_message["role"] == "assistant" + assert assistant_message["content"] == "Let me fetch the original." + + +@pytest.mark.asyncio +async def test_agentic_plan_rejects_hash_from_other_request( + guardrail: CompresrGuardrail, +): + hash_value = "b" * 24 + guardrail._store_originals("someone-elses-call", {hash_value: "secret"}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("my-call"), + stream=False, + kwargs={}, + ) + + # Hash belongs to another caller's scope; the loop is vetoed and the secret + # never surfaces. + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_agentic_loop_vetoed_when_no_recovery_state(guardrail: CompresrGuardrail): + # A caller-defined compresr_retrieve tool with no stored original must not + # trigger an extra provider round-trip. + hash_value = "f" * 24 + response = _make_openai_response_with_tool_call(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_x") + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_x")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_agentic_loop_dedupes_repeated_retrievals(guardrail: CompresrGuardrail): + # Retrieving the same marker many times expands the original once; repeats + # get a short marker (no follow-up amplification). + hash_value = "a" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + calls = [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, f"call_{i}") for i in range(5)] + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, f"call_{i}") for i in range(5)]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_calls(calls), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + tool_results = [m for m in plan.request_patch.messages if m.get("role") == "tool"] + assert len(tool_results) == 5 + assert sum(1 for m in tool_results if m["content"] == TOOL_OUTPUT) == 1 + assert all("already retrieved" in m["content"] for m in tool_results if m["content"] != TOOL_OUTPUT) + + +@pytest.mark.asyncio +async def test_agentic_loop_caps_retrieval_count(guardrail: CompresrGuardrail): + # Beyond _MAX_RETRIEVALS_PER_LOOP retrievals, extra calls get a bounded marker. + n = 10 + hashes = [f"{i:024x}" for i in range(n)] + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {h: f"original-{h}" for h in hashes}) + calls = [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": h}, f"call_{i}") for i, h in enumerate(hashes)] + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(h, f"call_{i}") for i, h in enumerate(hashes)]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_calls(calls), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + tool_results = [m for m in plan.request_patch.messages if m.get("role") == "tool"] + assert len(tool_results) == n + over_limit = [m for m in tool_results if "retrieval limit reached" in m["content"]] + assert len(over_limit) == n - 8 # only the first 8 expand + + +def test_display_hash_strips_control_characters(): + """The compresr_retrieve `hash` argument is model/tool-output-influenced, so + control characters (newlines, ANSI escapes) must be stripped — not just + length-capped — before it is echoed into logs or the fallback message.""" + from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import _display_hash + + assert _display_hash("a" * 24) == "a" * 24 # a real marker hash passes through + assert "\n" not in _display_hash("abc\ndef\rFORGED LOG LINE") + assert "\x1b" not in _display_hash("hash\x1b[31mred") + capped = _display_hash("z" * 100) + assert capped.endswith("…") and len(capped) <= 33 + + +@pytest.mark.asyncio +async def test_agentic_plan_builds_anthropic_followup_shape( + guardrail: CompresrGuardrail, +): + hash_value = "c" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = None + response.content = [{"type": "tool_use", "id": "toolu_1"}] + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + follow_up = plan.request_patch.messages + assistant_msg, user_msg = follow_up[-2], follow_up[-1] + assert assistant_msg["role"] == "assistant" + assert assistant_msg["content"][0]["type"] == "tool_use" + assert user_msg["content"][0]["type"] == "tool_result" + assert user_msg["content"][0]["tool_use_id"] == "toolu_1" + assert user_msg["content"][0]["content"] == TOOL_OUTPUT + assert plan.request_patch.max_tokens == 1024 + + +@pytest.mark.asyncio +async def test_agentic_plan_builds_responses_followup_shape( + guardrail: CompresrGuardrail, +): + """The /v1/responses path echoes the function_call and pairs it with a + function_call_output keyed by the same call_id.""" + hash_value = "e" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = [{"type": "function_call", "call_id": "fc_1"}] # responses-API shape + response.content = None + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + call_item, output_item = plan.request_patch.messages[-2], plan.request_patch.messages[-1] + assert call_item["type"] == "function_call" + assert call_item["call_id"] == "fc_1" + assert output_item["type"] == "function_call_output" + assert output_item["call_id"] == "fc_1" + assert output_item["output"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_chat_parallel_tool_calls_echoes_only_retrieve( + guardrail: CompresrGuardrail, +): + """When the model calls a real tool alongside compresr_retrieve in one turn, + only the retrieve call may be echoed in the reconstructed assistant message: + every echoed tool_call must have a matching tool result or the provider 400s. + The real call is re-planned by the follow-up; the assistant text is kept.""" + hash_value = "f" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = _make_openai_response_with_tool_calls( + [ + ("get_weather", {"city": "Paris"}, "call_weather"), + (COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, "call_retrieve"), + ], + content="Let me expand that note and check the weather.", + ) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_retrieve")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + follow_up = plan.request_patch.messages + assistant_msg = follow_up[-2] + echoed_ids = {tc["id"] for tc in assistant_msg["tool_calls"]} + result_ids = {m["tool_call_id"] for m in follow_up if m.get("role") == "tool"} + # get_weather is not echoed; every echoed tool_call is answered. + assert echoed_ids == {"call_retrieve"} + assert echoed_ids == result_ids + assert assistant_msg["content"] == "Let me expand that note and check the weather." + assert follow_up[-1]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_anthropic_parallel_preserves_text_and_balances( + guardrail: CompresrGuardrail, +): + """Anthropic parallel-tool-call turn: the assistant text is preserved, the + real tool_use is dropped (re-planned), and the reconstructed turn stays + balanced — one tool_result per echoed tool_use.""" + hash_value = "a" * 23 + "9" + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = None + response.content = [ + {"type": "text", "text": "Checking the weather and expanding the note."}, + {"type": "tool_use", "id": "toolu_weather", "name": "get_weather", "input": {"city": "Paris"}}, + { + "type": "tool_use", + "id": "toolu_retrieve", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "input": {"hash": hash_value}, + }, + ] + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "toolu_retrieve")]}, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assistant_msg, user_msg = plan.request_patch.messages[-2], plan.request_patch.messages[-1] + assert assistant_msg["content"][0] == { + "type": "text", + "text": "Checking the weather and expanding the note.", + } + echoed_ids = [b["id"] for b in assistant_msg["content"] if b["type"] == "tool_use"] + answered_ids = [b["tool_use_id"] for b in user_msg["content"]] + # get_weather dropped; balanced tool_use/tool_result pairing. + assert echoed_ids == ["toolu_retrieve"] + assert answered_ids == echoed_ids + + +# ── store hygiene ───────────────────────────────────────────────────── + + +def test_originals_store_prunes_expired(guardrail: CompresrGuardrail): + guardrail._originals_by_call_id["old"] = ({"a" * 24: "x"}, 0.0) # already expired + guardrail._store_originals("new", {"b" * 24: "y"}) + assert "old" not in guardrail._originals_by_call_id + assert "new" in guardrail._originals_by_call_id + + +def test_originals_store_caps_tracked_calls(guardrail: CompresrGuardrail): + for i in range(300): + guardrail._store_originals(f"call-{i}", {("%024x" % i): "x"}) + assert len(guardrail._originals_by_call_id) <= 256 + # Most recent entries survive. + assert "call-299" in guardrail._originals_by_call_id + + +def test_originals_store_caps_bytes_per_call(): + guardrail = _make_guardrail(max_bytes_per_call=1000) + hashes = tuple(f"{i:024x}" for i in range(5)) + values = tuple("x" * 400 for _ in range(5)) + guardrail._store_originals("c", dict(zip(hashes, values))) + + stored, _expiry = guardrail._originals_by_call_id["c"] + assert sum(len(v.encode("utf-8")) for v in stored.values()) <= 1000 + # Oldest entries are evicted first; newest survives. + assert hashes[-1] in stored + assert hashes[0] not in stored + + +def test_originals_store_byte_cap_survives_lone_surrogates(): + # Regression: eviction path must use surrogatepass to match the hash + # function; a bare encode("utf-8") crashed on lone surrogates. + guardrail = _make_guardrail(max_bytes_per_call=500) + surrogate_value = "\ud800" * 60 + hashes = tuple(f"{i:024x}" for i in range(3)) + guardrail._store_originals("c", dict(zip(hashes, (surrogate_value, surrogate_value, surrogate_value)))) + + stored, _expiry = guardrail._originals_by_call_id["c"] + assert hashes[-1] in stored + assert hashes[0] not in stored + + +def test_originals_store_caps_total_bytes_across_calls(monkeypatch: pytest.MonkeyPatch): + # Global byte budget: many distinct call ids must not retain unbounded memory. + monkeypatch.setattr( + "litellm.proxy.guardrails.guardrail_hooks.compresr.compresr._MAX_TOTAL_STORE_BYTES", + 10_000, + ) + guardrail = _make_guardrail(max_bytes_per_call=4_000) + for i in range(20): + guardrail._store_originals(f"call-{i}", {f"{i:024x}": "x" * 3_000}) + + total = sum( + len(v.encode("utf-8")) + for originals, _expiry in guardrail._originals_by_call_id.values() + for v in originals.values() + ) + assert total <= 10_000 + assert guardrail._store_total_bytes == total # running counter stays exact + # Oldest calls evicted; the most-recent call's originals survive. + assert "call-0" not in guardrail._originals_by_call_id + assert "call-19" in guardrail._originals_by_call_id + + +def test_originals_store_global_cap_keeps_current_when_single_call_is_large( + monkeypatch: pytest.MonkeyPatch, +): + # One call over the global cap is still kept (only max_bytes_per_call trims it); + # global eviction never empties the store. + monkeypatch.setattr( + "litellm.proxy.guardrails.guardrail_hooks.compresr.compresr._MAX_TOTAL_STORE_BYTES", + 1_000, + ) + guardrail = _make_guardrail(max_bytes_per_call=5_000) + guardrail._store_originals("solo", {f"{0:024x}": "x" * 4_000}) + assert "solo" in guardrail._originals_by_call_id + + +def test_recovery_markers_respect_per_call_byte_cap(): + # Regression: markers were built from every original before _store_originals + # applied the byte cap, so an evicted original left a dangling marker the + # model could never retrieve. Recovery must be attached only for originals + # that fit the cap, so every shipped marker stays retrievable. + guardrail = _make_guardrail(max_bytes_per_call=1000) + contexts = ["a" * 400, "b" * 400, "c" * 400] + messages = [{"role": "tool", "content": text} for text in contexts] + results = [{"compressed_context": f"small-{i}"} for i in range(3)] + + applied = guardrail._apply_compression_results( + messages, [0, 1, 2], contexts, results, recovery_enabled=True + ) + + # 400 + 400 fit under 1000; the third (which would reach 1200) is skipped. + assert applied.messages_compressed == 3 + assert len(applied.originals) == 2 + third_hash = _content_hash("c" * 400) + assert third_hash not in applied.originals + assert f"compresr hash={third_hash}" not in applied.compressed_messages[2]["content"] + + # Every marker still shipped must resolve to a stored original. + guardrail._store_originals("c", applied.originals) + for hash_value in applied.originals: + assert guardrail._retrieve_original("c", hash_value) is not None + assert f"compresr hash={hash_value}" in "".join( + str(m["content"]) for m in applied.compressed_messages + ) + + +def test_recovery_markers_respect_byte_cap_across_reused_store_key(): + # Regression: the per-call budget must also count bytes already stored under + # the same store key (a later turn reusing the call id). Otherwise merging + # this call's originals with the existing entry overflows the cap and + # _store_originals evicts an original this call just shipped a marker for. + guardrail = _make_guardrail(max_bytes_per_call=100) + old_hash = _content_hash("A" * 50) + guardrail._store_originals("k", {old_hash: "A" * 50}) + guardrail._store_originals("k", {_content_hash("C" * 40): "C" * 40}) + existing = guardrail._originals_by_call_id["k"][0] + + # This turn recompresses the same "A" (already stored) plus a new "D". + contexts = ["D" * 40, "A" * 50] + messages = [{"role": "tool", "content": text} for text in contexts] + results = [{"compressed_context": "dd"}, {"compressed_context": "aa"}] + applied = guardrail._apply_compression_results( + messages, [0, 1], contexts, results, recovery_enabled=True, existing_originals=existing + ) + + guardrail._store_originals("k", applied.originals) + # No marker shipped this turn may dangle after the store enforces the cap. + for hash_value in applied.originals: + assert guardrail._retrieve_original("k", hash_value) is not None + assert f"compresr hash={hash_value}" in "".join( + str(m["content"]) for m in applied.compressed_messages + ) + # The zero-cost repeat of an already-stored original stays retrievable. + assert old_hash in applied.originals + assert guardrail._retrieve_original("k", old_hash) is not None + + +# ── dynamic (adaptive) compression — latte_v2 Kneedle ───────────────── + + +@pytest.mark.asyncio +async def test_dynamic_flag_in_payload(): + """dynamic=True must appear in the compress payload; unset bounds omitted.""" + guardrail = _make_guardrail(dynamic=True) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["dynamic"] is True + assert "dynamic_min_ratio" not in payload + assert "dynamic_max_ratio" not in payload + + +@pytest.mark.asyncio +async def test_dynamic_bounds_in_payload_when_set(): + guardrail = _make_guardrail(dynamic=True, dynamic_min_ratio=2.0, dynamic_max_ratio=8.0) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["dynamic"] is True + assert payload["dynamic_min_ratio"] == 2.0 + assert payload["dynamic_max_ratio"] == 8.0 + + +@pytest.mark.asyncio +async def test_dynamic_on_by_default(): + guardrail = _make_guardrail() # dynamic defaults on (latte_v2) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert mock_post.call_args.kwargs["json"]["dynamic"] is True + + +# ── generic passthrough compression params ──────────────────────────── + + +@pytest.mark.asyncio +async def test_compression_params_passthrough_in_payload(): + """Extra params in compression_params are forwarded verbatim; named fields + still win on collision.""" + guardrail = _make_guardrail(compression_params={"heuristic_chunking": True, "coarse": False}) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["heuristic_chunking"] is True + # named `coarse` (default True) wins over the passthrough's coarse=False + assert payload["coarse"] is True + + +@pytest.mark.asyncio +async def test_compression_params_cannot_override_request_content_fields(): + """context/query/inputs carry the actual content being compressed; a + passthrough collision on them must be dropped, not silently win.""" + guardrail = _make_guardrail( + compression_params={ + "context": "injected", + "query": "injected", + "inputs": [], + "heuristic_chunking": True, + } + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["context"] == TOOL_OUTPUT + assert payload["query"] == 'web_search: {"query": "2026 EV range"}' + assert "inputs" not in payload + assert payload["heuristic_chunking"] is True + + +# ── compress_last_user ──────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_compress_last_user_compresses_with_verbatim_query(): + """compress_last_user=True compresses the last user message, but the query + sent to Compresr is still the original verbatim user text.""" + guardrail = _make_guardrail(compress_last_user=True) + long_question = "Which 2026 EV has the longest range? " * 20 # > 500 chars + messages = [{"role": "user", "content": long_question}] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["context"] == long_question + assert payload["query"] == long_question # verbatim, not the compressed text + assert result["structured_messages"][0]["content"] == "compressed summary" + + +# ── malformed-but-200 token stats (must not defeat fail policy) ─────── + + +@pytest.mark.asyncio +async def test_non_numeric_token_stats_do_not_raise(guardrail: CompresrGuardrail): + """A 200 response with non-numeric token counts must not raise: _call_compress + already succeeded, so a bare int() here would 500 even under fail policy.""" + resp = _make_single_compress_response() + resp.json.return_value["data"]["original_tokens"] = "not-a-number" + resp.json.return_value["data"]["compressed_tokens"] = None + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result["structured_messages"][3]["content"].startswith("compressed summary") + + +# ── HTTP status errors (non-2xx from the shared handler) ────────────── + + +def _http_status_error(status: int = 500, text: str = "upstream error body") -> httpx.HTTPStatusError: + # The shared AsyncHTTPHandler.post() raises HTTPStatusError on any non-2xx, + # carrying the upstream body and request headers; this simulates that. + request = httpx.Request("POST", f"{FAKE_API_BASE}/api/compress/question-specific/") + response = httpx.Response(status, text=text, request=request) + return httpx.HTTPStatusError(str(status), request=request, response=response) + + +@pytest.mark.asyncio +async def test_http_status_error_raises_when_fail_closed(guardrail: CompresrGuardrail): + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(500))): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_http_status_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(429))): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_http_status_error_does_not_leak_upstream_body(guardrail: CompresrGuardrail): + secret = "SECRET_INSTANCE_METADATA_TOKEN=aws-imds-response" + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(500, text=secret))): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert secret not in json.dumps(exc_info.value.detail) + + +# ── non-transport httpx errors must still honor the fail policy ──────── +# TooManyRedirects and DecodingError are httpx.RequestError but NOT +# httpx.TransportError, so a narrow except would let them escape as a 500 +# even under fail_open. These lock in that they are routed through the policy. + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + httpx.TooManyRedirects("redirect loop"), + httpx.DecodingError("bad content-encoding"), + ], +) +async def test_request_errors_raise_when_fail_closed(guardrail: CompresrGuardrail, error): + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=error)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + httpx.TooManyRedirects("redirect loop"), + httpx.DecodingError("bad content-encoding"), + ], +) +async def test_request_errors_fail_open_forwards_uncompressed(error): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=error)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_undecodable_body_on_200_forwards_uncompressed_when_fail_open(): + """A 200 whose body raises DecodingError on .json()/.text must not 500.""" + guardrail = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock() + resp.status_code = 200 + resp.json.side_effect = httpx.DecodingError("bad content-encoding") + type(resp).text = property(lambda self: (_ for _ in ()).throw(httpx.DecodingError("bad content-encoding"))) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_recursion_error_on_json_forwards_when_fail_open(): + """A deeply nested JSON body can raise RecursionError while parsing; it must + route through the fail policy, not escape as a 500.""" + guardrail = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock() + resp.status_code = 200 + resp.json.side_effect = RecursionError("maximum recursion depth exceeded") + resp.text = "" + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_lone_surrogate_in_content_does_not_crash(guardrail: CompresrGuardrail): + """A lone Unicode surrogate (reachable via a JSON \\uXXXX escape) in content + must not crash hashing/byte-accounting after the fail-policy decision.""" + surrogate_output = ("x" * 600) + "\ud800" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "web_search", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": surrogate_output}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + assert result["structured_messages"][2]["content"].startswith("compressed summary") + # The original (surrogate included) is recoverable by its hash. + stored = next(iter(guardrail._originals_by_call_id.values()))[0] + assert surrogate_output in stored.values() + + +@pytest.mark.asyncio +async def test_identical_compressed_text_treated_as_noop(guardrail: CompresrGuardrail): + """If the service returns text byte-identical to the original, nothing + changed: the exact inputs object is returned so no write-back is forced.""" + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response(compressed_context=TOOL_OUTPUT)) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + assert result is inputs + + +# ── recovery requires a framework-issued call id ────────────────────── + + +@pytest.mark.asyncio +async def test_recovery_disabled_without_call_id(): + # enable_retrieval defaults True, but with no framework litellm_call_id we + # cannot scope stored originals to the request, so compression proceeds + # without markers, the retrieve tool, or any stored originals. + guardrail = _make_guardrail() + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result["structured_messages"][3]["content"] == "compressed summary" + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +def test_config_model_exposes_unreachable_fallback(): + from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, + ) + + field = CompresrGuardrailConfigModel.model_fields.get("unreachable_fallback") + assert field is not None + assert field.default == "fail_closed" + + +# ── audit fixes ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_cancelled_error_propagates_not_swallowed(): + # Regression: CancelledError is a BaseException, not caught by + # (RequestError, Timeout). It must re-raise so cooperative cancellation + # (asyncio.wait_for, client disconnect) still fires. + import asyncio as _asyncio + + guardrail = _make_guardrail(unreachable_fallback="fail_open") + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_asyncio.CancelledError())): + with pytest.raises(_asyncio.CancelledError): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +def test_max_bytes_per_call_negative_rejected(): + # Regression: a negative value silently disabled the byte cap (< 0 behaves + # like 0 in _bound_call_bytes). Validate at construction so the footgun + # surfaces as a ValueError at startup, not silent unbounded storage. + with pytest.raises(ValueError, match="max_bytes_per_call"): + _make_guardrail(max_bytes_per_call=-1) + + +@pytest.mark.asyncio +async def test_max_tokens_zero_from_optional_params_wins_over_kwargs(): + # Regression: `or` short-circuits on falsy values, so an explicit + # max_tokens=0 from optional_params fell through to kwargs["max_tokens"]. + # Must use `is not None`. + guardrail = _make_guardrail() + hash_value = "deadbeef" + guardrail._store_originals(_scoped_store_key(_logging_obj("call-1")), {hash_value: TOOL_OUTPUT}) + response = MagicMock() + response.content = [{"type": "tool_use", "id": "toolu_1"}] + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 0}, + logging_obj=_logging_obj("call-1"), + stream=False, + kwargs={"max_tokens": 999}, + ) + assert plan.request_patch.max_tokens == 0 + + +@pytest.mark.asyncio +async def test_recovery_disabled_when_no_caller_scope(): + # Regression: on a no-auth deployment (no UserAPIKeyAuth in metadata) the + # store key would fall back to the client-settable call id alone, letting + # any caller retrieve any other caller's originals. Recovery must be off. + guardrail = _make_guardrail() + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-abc" + logging_obj.model_call_details = {"litellm_params": {"metadata": {}}} + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +@pytest.mark.asyncio +async def test_warns_once_when_recovery_skipped_without_scope(): + # enable_retrieval is on but the request has no per-key auth scope: recovery + # is silently skipped, so a call-time warning must surface it (once). + guardrail = _make_guardrail() + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-abc" + logging_obj.model_call_details = {"litellm_params": {"metadata": {}}} + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + with patch("litellm.proxy.guardrails.guardrail_hooks.compresr.compresr.verbose_proxy_logger") as mock_log: + for _ in range(3): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + + no_scope_warnings = [c for c in mock_log.warning.call_args_list if "no per-key auth scope" in str(c)] + assert len(no_scope_warnings) == 1