diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py new file mode 100644 index 00000000000..54cc75ede03 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import ( + GuardrailEventHooks, + Mode, + SupportedGuardrailIntegrations, +) + +from .needlepath import NeedlepathGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def _coerce_event_hook( + mode: str | list[str] | Mode, # mutable-ok: the type of LitellmParams.mode +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: # mutable-ok: dispatch needs a list + if isinstance(mode, Mode): + return mode + if isinstance(mode, list): + # A real `list` is required here, not just any Sequence: CustomGuardrail's + # dispatch (should_run_guardrail) does `isinstance(self.event_hook, list)` + # to decide whether this hook fires for a given event. A tuple would fail + # that check silently and the guardrail would never run for a multi-hook + # config. + 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: + """Read a knob from ``optional_params`` first, then from the top level. + + Both spellings are accepted because both appear in the wild: the nested + ``optional_params`` block is the documented form, and several deployments + set guardrail knobs flat alongside ``api_key``. + """ + if optional_params is not None: + value: Final[object] = getattr(optional_params, attribute_name, None) # pyright: ignore[reportAny] # knob name is dynamic + if value is not None: + return value + flat: Final[object] = getattr(litellm_params, attribute_name, None) # pyright: ignore[reportAny] # flat knobs are pydantic extras + return flat + + +def _bool_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> bool | None: + value: Final = _get_optional_value(litellm_params, optional_params, attribute_name) + return value if isinstance(value, bool) else None + + +def _int_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> int | None: + value: Final = _get_optional_value(litellm_params, optional_params, attribute_name) + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, str): + try: + return int(value) + except ValueError: + return None + return None + + +def _str_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> str | None: + value: Final = _get_optional_value(litellm_params, optional_params, attribute_name) + return value if isinstance(value, str) else None + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> NeedlepathGuardrail: + import litellm + + optional_params: Final = getattr(litellm_params, "optional_params", None) + + _callback: Final = NeedlepathGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + select_tool_outputs=_bool_knob(litellm_params, optional_params, "select_tool_outputs"), + select_history=_bool_knob(litellm_params, optional_params, "select_history"), + select_system=_bool_knob(litellm_params, optional_params, "select_system"), + min_chars_to_select=_int_knob(litellm_params, optional_params, "min_chars_to_select"), + max_context_tokens=_int_knob(litellm_params, optional_params, "max_context_tokens"), + operating_point=_str_knob(litellm_params, optional_params, "operating_point"), + guardrail_name=guardrail["guardrail_name"], + event_hook=_coerce_event_hook(litellm_params.mode), + default_on=litellm_params.default_on or False, + ) + litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped + _callback + ) + return _callback + + +# guardrail_registry.py's module discovery does `isinstance(registry, dict)` before +# merging this into the central registry; a Mapping/MappingProxyType would fail that +# check silently and this guardrail would never be discovered at proxy startup. +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.NEEDLEPATH.value: initialize_guardrail, +} + +# Same reason as guardrail_initializer_registry above: discovery's isinstance(dict) check. +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.NEEDLEPATH.value: NeedlepathGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py new file mode 100644 index 00000000000..7ced925c380 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -0,0 +1,711 @@ +"""Needlepath guardrail: query-conditioned extractive context selection. + +Bulky message content (tool outputs by default) is sent to the Needlepath +selection service before the request reaches the LLM. The service returns the +spans of that content which carry the answer to a query, and the guardrail +writes those spans back over the message they came from. Nothing is +paraphrased or rewritten: the returned block is made of extracts of the text +that was submitted. + +Selection is per message and query-conditioned. The query for a tool output is +the intent of the tool call that produced it (``name`` plus ``arguments``, +found through ``tool_call_id``); anything else uses the last user message. Each +message is selected independently, so one message's outcome never changes +another's. + +**This guardrail is unconditionally fail-open.** Every path that does not +produce a usable selection returns the caller's messages untouched. A proxy +that silently blanks a tool output is far worse than a proxy that does nothing, +so there is no configuration in which a selection failure becomes a request +failure. See ``_selected_text`` for the full list of declines. +""" + +from __future__ import annotations + +import asyncio +import ipaddress +import time +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal +from urllib.parse import urlparse + +import httpx +from httpx import Response as HttpxResponse +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped in custom_guardrail +) +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.guardrails.guardrail_hooks.content_text import ( + content_to_text, + is_all_text_parts, + merge_rewritten_text_parts, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import GuardrailEventHooks, Mode +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.needlepath import ( + NeedlepathGuardrailConfigModel, + ) + +DEFAULT_API_BASE: Final = "https://api.nextmoca.com" +SELECT_PATH: Final = "/v1/context/select" +# Immutable engine label. Pinned rather than inherited from the service default +# so an upgrade on the service side cannot change what this proxy sends without +# an operator changing this config. +DEFAULT_OPERATING_POINT: Final = "np-2026-10-r8" +DEFAULT_MAX_CONTEXT_TOKENS: Final = 4000 +DEFAULT_MIN_CHARS_TO_SELECT: Final = 500 +# The shared client's read timeout is measured in minutes, which is far too long +# to hold an inbound LLM request behind an optional optimisation. A stall past +# this bound is a decline, and the original message is forwarded. +_SELECT_TIMEOUT_SECONDS: Final = 30.0 +# A single proxy request can carry arbitrarily many eligible messages, and each +# one becomes an outbound selection call. These two bounds keep a pathological +# request (say, thousands of minimally qualifying tool outputs) from +# monopolising the proxy's shared HTTP pool or burning selection quota: at most +# _MAX_TARGETS_PER_REQUEST messages are selected per request -- the largest +# first, where selection pays off most -- and at most +# _MAX_CONCURRENT_SELECTIONS calls are in flight at once. Messages past the +# cap are forwarded untouched, the same fail-open outcome as any decline. +_MAX_TARGETS_PER_REQUEST: Final = 16 +_MAX_CONCURRENT_SELECTIONS: Final = 4 +# The service reports a deliberate no-op through the gate. Any reason under this +# prefix means "the engine chose not to select"; the original content is what +# the caller should send. +_STANDDOWN_PREFIX: Final = "standdown:" +# Security screening outcomes (top-level ``outcome``). A ``sanitized_fallback`` +# is the caller's full context with every flagged span replaced by a marker; it +# is applied, because forwarding the original would put the flagged text back. +# ``blocked`` delivers none of the caller's content and is a decline, the same +# handling the published Needlepath clients give it. +_OUTCOME_SANITIZED_FALLBACK: Final = "sanitized_fallback" +_OUTCOME_BLOCKED: Final = "blocked" +# ``format_metrics.security_action`` when screening removed flagged spans from +# what it delivered. +_SECURITY_ACTION_QUARANTINE: Final = "quarantine" +# Request and response JSON is untrusted. These adapters turn a value of +# unknown shape into a typed dict or list, or into None when it is not one. +_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_BLOCKED_METADATA_HOSTS: Final = frozenset( + { + "metadata.google.internal", + "metadata.goog", + "metadata.azure.com", + "metadata.azure.internal", + } +) +_BLOCKED_METADATA_IPS: Final = frozenset( + ipaddress.ip_address(ip) for ip in ("169.254.169.254", "fd00:ec2::254", "100.100.100.200", "168.63.129.16") +) +# Record kinds the service publishes. A system prompt has no dedicated kind, so +# it is submitted as external_data rather than invented as a new one: an +# unrecognised kind is a 400 for the whole call. +_KIND_TOOL_RESULT: Final = "tool_result" +_KIND_USER_INPUT: Final = "user_input" +_KIND_EXTERNAL_DATA: Final = "external_data" + + +def _as_str_object_dict(value: object) -> Mapping[str, object] | None: + try: + return _STR_OBJECT_DICT_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _as_object_list(value: object) -> Sequence[object] | None: + try: + return _OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _parse_ip_literal(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None: + """Parse ``host`` as an IP literal, including the alternate spellings the + socket layer accepts (single-integer IPv4, IPv4-mapped IPv6), so a blocked + address cannot be smuggled past a plain string comparison.""" + try: + addr = ipaddress.ip_address(host) # rebind-ok: reassigned in the except fallback below + except ValueError: + try: + addr = ipaddress.ip_address(int(host, 0)) # rebind-ok: fallback re-parse as an int literal + 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. + + Defense in depth against a mistyped or hostile ``api_base``: non-http(s) + schemes and cloud-metadata hosts/IPs are refused. Private ranges stay + allowed so on-prem deployments work. This is not a complete SSRF control: + there is no DNS resolution here and the shared client follows redirects. + ``api_base`` is operator config, so that is an accepted limit. + """ + parsed: Final = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Needlepath guardrail api_base must be http or https, got scheme={parsed.scheme!r}") + host: Final = (parsed.hostname or "").lower() + if not host: + raise ValueError("Needlepath guardrail api_base has no host") + ip_literal: Final = _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"Needlepath guardrail api_base {host!r} is a blocked cloud-metadata host") + return url + + +def _write_text_back(content: object, new_text: str) -> object: + """Put ``new_text`` into a ``content`` value without changing its shape. + + A string is replaced directly. An all-text part list collapses to a single + part carrying the last declared cache_control breakpoint. Anything else is + returned untouched: breakpoints are positional, so one selected block cannot + be written across a non-text part without moving text past it. + """ + if isinstance(content, str): + return new_text + parts: Final = _as_object_list(content) + if parts is not None and is_all_text_parts(parts): + return merge_rewritten_text_parts(parts, new_text) + return content + + +def _render_tool_intent(fn: Mapping[str, object]) -> str: + """A tool call rendered as the query its output should be selected against.""" + name: Final = str(fn.get("name") or "").strip() + raw_args: Final = fn.get("arguments") + arg_text: Final = "" if raw_args is None else str(raw_args).strip() + if name and arg_text: + return f"{name}: {arg_text}" + return name or arg_text + + +def _query_for_target(messages: Sequence[Mapping[str, object]], target_idx: int, fallback: str) -> str: + """The query ``messages[target_idx]`` should be selected against. + + A tool or function result is selected against the intent of the call that + produced it, located by ``tool_call_id`` on an earlier assistant message. + Everything else, and any tool result whose call cannot be found, uses the + last user message. + """ + msg: Final = messages[target_idx] + if msg.get("role") not in ("tool", "function"): + return fallback + + tool_call_id: Final = msg.get("tool_call_id") + fn_name: Final = msg.get("name") + for idx in range(target_idx - 1, -1, -1): + previous = messages[idx] + if previous.get("role") != "assistant": + continue + for raw_call in _as_object_list(previous.get("tool_calls")) or (): + call = _as_str_object_dict(raw_call) + if call is None or not tool_call_id or call.get("id") != tool_call_id: + continue + fn = _as_str_object_dict(call.get("function")) + intent = _render_tool_intent(fn if fn is not None else MappingProxyType({})) + if intent: + return intent + # Legacy function_call turns carry no id, so require a name match. + # Without it an older, unrelated call would supply the wrong intent. + legacy = _as_str_object_dict(previous.get("function_call")) + if legacy is not None and fn_name and legacy.get("name") == fn_name: + intent = _render_tool_intent(legacy) + if intent: + return intent + return fallback + + +def _title_for(messages: Sequence[Mapping[str, object]], target_idx: int) -> str | None: + """The tool name behind a message, used as the record title.""" + msg: Final = messages[target_idx] + name: Final = msg.get("name") + if isinstance(name, str) and name.strip(): + return name.strip()[:120] + tool_call_id: Final = msg.get("tool_call_id") + if not tool_call_id: + return None + for idx in range(target_idx - 1, -1, -1): + previous = messages[idx] + if previous.get("role") != "assistant": + continue + for raw_call in _as_object_list(previous.get("tool_calls")) or (): + call = _as_str_object_dict(raw_call) + if call is None or call.get("id") != tool_call_id: + continue + fn = _as_str_object_dict(call.get("function")) + fn_name = fn.get("name") if fn is not None else None + if isinstance(fn_name, str): + return fn_name[:120] + return None + + +def _record_kind(role: object) -> str: + if role in ("tool", "function"): + return _KIND_TOOL_RESULT + if role == "user": + return _KIND_USER_INPUT + return _KIND_EXTERNAL_DATA + + +def _safe_int(value: object) -> int | None: + """Read an integer counter from an untrusted body without raising. + + A field that is missing or not a number is reported as ``None`` so the + caller can treat it as "unknown" rather than as zero, which is a decline. + """ + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _screening_removed_content(body: Mapping[str, object]) -> bool: + """True when security screening removed flagged spans from the delivery. + + Such a block must be applied even where an ordinary selection would be + declined: the declined path forwards the original content, flagged spans + included. + """ + format_metrics: Final = _as_str_object_dict(body.get("format_metrics")) + return format_metrics is not None and format_metrics.get("security_action") == _SECURITY_ACTION_QUARANTINE + + +def _safe_response_text(response: object, limit: int = 500) -> str: + """Read a response body for a log line without letting the read itself raise. + + A corrupt ``Content-Encoding`` makes httpx's ``.text`` raise, which would + turn an already-handled decline into an unhandled 500. + """ + try: + text: Final = getattr(response, "text", "") + except httpx.DecodingError: + return "" + return (text or "")[:limit] + + +class NeedlepathGuardrail(CustomGuardrail): + """Select the spans of a message that answer the current query. + + Every knob is optional and the defaults are deliberately narrow: tool + outputs only, nothing under ``min_chars_to_select`` characters, and a pinned + operating point. + """ + + def __init__( + self, + api_base: str | None = None, + api_key: str | None = None, + select_tool_outputs: bool | None = None, + select_history: bool | None = None, + select_system: bool | None = None, + min_chars_to_select: int | None = None, + max_context_tokens: int | None = None, + operating_point: str | None = None, + guardrail_name: str | None = None, + # A list, not a Sequence: CustomGuardrail's dispatch checks isinstance(list). + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, # mutable-ok: see above + default_on: bool = False, + ) -> None: + raw_api_base: Final = (api_base or get_secret_str("NEEDLEPATH_API_BASE") or DEFAULT_API_BASE).rstrip("/") + self.needlepath_api_base = _validate_api_base(raw_api_base) + self.needlepath_api_key = api_key or get_secret_str("NEEDLEPATH_API_KEY") + if not self.needlepath_api_key: + raise ValueError( + "Needlepath guardrail requires an API key. Set `api_key` in the " + "guardrail config or the NEEDLEPATH_API_KEY env var." + ) + self.select_tool_outputs = True if select_tool_outputs is None else select_tool_outputs + self.select_history = False if select_history is None else select_history + self.select_system = False if select_system is None else select_system + self.min_chars_to_select = ( + DEFAULT_MIN_CHARS_TO_SELECT if min_chars_to_select is None else int(min_chars_to_select) + ) + self.max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS if max_context_tokens is None else int(max_context_tokens) + # Pinned, not inherited. The labels are immutable, so a pinned one makes + # what this guardrail sends reproducible across service releases. + self.operating_point = operating_point or DEFAULT_OPERATING_POINT + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + def _request_headers( + self, + ) -> dict[str, str]: # mutable-ok: the shared handler's `headers=` is typed as a dict; a fresh one per call + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.needlepath_api_key or ''}", + } + + def _decline(self, reason: str, detail: Mapping[str, object] | None = None) -> None: + """Record a decline. The caller then forwards the original content. + + Declining is the normal outcome for a stand-down or a service problem, + so it is logged at debug and never raised. Details can include upstream + response bytes and stay in the proxy's own logs. + """ + # The dict fallback below is never a value a caller could mutate: `or {}` only + # feeds the debug-log %s formatting, and swapping it for a frozen mapping would + # change the logged repr from `{}` to `mappingproxy({})`. + verbose_proxy_logger.debug("Needlepath: declined (%s) detail=%s", reason, detail or {}) + + def _select_targets(self, messages: Sequence[Mapping[str, object]], query_idx: int | None) -> tuple[int, ...]: + """Indices of the messages whose text content is eligible for selection.""" + + def _eligible(idx: int, msg: Mapping[str, object]) -> bool: + # The message carrying the query is never rewritten: selecting the + # question against itself is meaningless and it is what the rest of + # the selection is conditioned on. + if idx == query_idx: + return False + role: Final = msg.get("role") + if role in ("tool", "function"): + if not self.select_tool_outputs: + return False + elif role == "system": + if not self.select_system: + return False + elif role == "user": + if not self.select_history: + return False + else: + return False + content: Final = msg.get("content") + parts: Final = _as_object_list(content) + if parts is not None and not is_all_text_parts(parts): + return False + return len(content_to_text(content)) >= self.min_chars_to_select + + return tuple(idx for idx, msg in enumerate(messages) if _eligible(idx, msg)) + + @staticmethod + def _last_user_message(messages: Sequence[Mapping[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 + + async def _post_select(self, payload: Mapping[str, object]) -> HttpxResponse | None: + """POST one selection request. Returns None on any transport failure.""" + try: + return await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped + url=f"{self.needlepath_api_base}{SELECT_PATH}", + json=dict(payload), + headers=self._request_headers(), + timeout=_SELECT_TIMEOUT_SECONDS, + ) + except asyncio.CancelledError: + raise + except httpx.HTTPStatusError as e: + # The shared handler calls raise_for_status(), so 402, 403, 429 and + # every other non-2xx arrives here carrying the upstream body and + # our Authorization header. It is a decline and nothing about it + # reaches the client. + response: Final = getattr(e, "response", None) + self._decline( + "http_status", + { + "status_code": getattr(response, "status_code", None), + "body": _safe_response_text(response), + }, + ) + return None + except (httpx.RequestError, litellm.Timeout) as e: + # Every request-side httpx failure, timeouts included, is a + # RequestError. Catching the whole class is what keeps a network + # problem from escaping as a 500. + self._decline("transport", {"detail": str(e)}) + return None + except Exception as e: # noqa: BLE001 # fail-open: any client error is a decline + # Anything else the shared client raises is still a decline. A + # selection problem must never become a failed LLM request. + self._decline("unexpected_error", {"error_type": type(e).__name__, "detail": str(e)}) + return None + + async def _selected_text(self, text: str, query: str, title: str | None, source: object, kind: str) -> str | None: + """The selected block for one message, or None to keep the original. + + None is returned, and the original content is forwarded, when: + + * the service could not be reached, timed out, or answered non-2xx + (402, 403, 429 and everything else); + * the body is not JSON, or is JSON of an unexpected shape; + * ``outcome`` is ``blocked``: security screening left nothing + deliverable; + * the gate reports a stand-down (any ``gate.reason`` under + ``standdown:``), which is the service saying the full content is what + the caller should send; + * ``records_selected`` is 0, ``tokens_after`` is 0, or + ``rendered_context`` is missing or blank; + * the returned block is not shorter than the text it would replace, so a + rewrite could never grow a message. + + A block from which security screening removed flagged spans is applied + whenever it is non-blank: ``outcome: sanitized_fallback``, or + ``format_metrics.security_action: quarantine``. The stand-down and + size checks above do not apply to it, because declining it would + forward the original content with the flagged spans still in it. + """ + # `record` and `payload` are both real, mutable dicts by necessity: `record` + # gains optional keys after construction (below), and both are handed straight + # to httpx's `json=`, which needs a genuine dict/list -- a MappingProxyType + # is not JSON-serializable (json.dumps raises TypeError on one), so the + # freezing wrappers used elsewhere in this file cannot be used here. + record: Final[dict[str, object]] = { # mutable-ok: see comment above + "id": "m0", + "kind": kind, + "text": text, + } + if title: + record["title"] = title + if isinstance(source, str) and source: + record["source"] = source + payload: Final[dict[str, object]] = { # mutable-ok: see comment above + "records": (record,), + "task": {"prompt": query}, + "budget": { + "max_context_tokens": self.max_context_tokens, + "operating_point": self.operating_point, + }, + "render": True, + "render_format": "plain", + } + + response: Final = await self._post_select(payload) + if response is None: + return None + if not 200 <= response.status_code < 300: + self._decline( + "http_status", + { + "status_code": response.status_code, + "body": _safe_response_text(response), + }, + ) + return None + + try: + raw_body: Final[object] = response.json() # pyright: ignore[reportAny] # httpx Response.json() is untyped + except (ValueError, httpx.DecodingError, RecursionError): + # RecursionError: a deeply nested body overflows the JSON parser. + # It is a decline like any other malformed answer, not a 500. + self._decline("unreadable_body", {"body": _safe_response_text(response)}) + return None + body: Final = _as_str_object_dict(raw_body) + if body is None: + self._decline("unexpected_shape", {"body": _safe_response_text(response)}) + return None + + outcome: Final = body.get("outcome") + rendered: Final = body.get("rendered_context") + has_rendered: Final = isinstance(rendered, str) and bool(rendered.strip()) + if outcome == _OUTCOME_BLOCKED: + # Screening found nothing deliverable. The original content is + # forwarded, as the published clients do; the outcome is in the + # service's receipt for operators who want to alarm on it. + self._decline("blocked", MappingProxyType({})) + return None + if outcome == _OUTCOME_SANITIZED_FALLBACK: + # The engine stood down, so this carries no selected records and + # its gate reports a stand-down, and it can be about as long as the + # original. It is still the block to send: the original is the same + # content with the flagged spans still in it. + if not isinstance(rendered, str) or not has_rendered: + self._decline("empty_rendered_context", MappingProxyType({})) + return None + return rendered + screened: Final = _screening_removed_content(body) + + gate: Final = _as_str_object_dict(body.get("gate")) + reason: Final = gate.get("reason") if gate is not None else None + if not screened and isinstance(reason, str) and reason.startswith(_STANDDOWN_PREFIX): + self._decline("gate_standdown", {"reason": reason}) + return None + + records_selected: Final = _safe_int(body.get("records_selected")) + tokens_after: Final = _safe_int(body.get("tokens_after")) + if records_selected == 0 or tokens_after == 0: + self._decline( + "empty_selection", + { + "records_selected": records_selected, + "tokens_after": tokens_after, + }, + ) + return None + + if not isinstance(rendered, str) or not has_rendered: + self._decline("empty_rendered_context", MappingProxyType({})) + return None + if not screened and len(rendered) >= len(text): + # A selection that is not smaller is not a selection worth making, + # and applying it could only add tokens. A screened block is exempt: + # declining it would forward the flagged spans it removed. + self._decline("no_reduction", {"before": len(text), "after": len(rendered)}) + return None + return rendered + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: base-class signature; logging info is written into it + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + structured_messages: Final = _as_object_list(inputs.get("structured_messages")) + if not structured_messages: + return inputs + # The static type says every entry is an AllMessageValues TypedDict, but this + # is the proxy's client-facing boundary: the list arrives from arbitrary + # request JSON, so the runtime check is the guarantee the annotation only + # promises. Any entry that is not a dict leaves the request untouched. + messages: Final = tuple(m for raw in structured_messages if (m := _as_str_object_dict(raw)) is not None) + if len(messages) != len(structured_messages): + return inputs + + fallback_query, query_idx = self._last_user_message(messages) + # One (idx, query) pair per eligible message whose query is non-blank. + # Selection is conditioned on a query; without one there is nothing to + # select against, so that message is left exactly as it arrived. + eligible_pairs: Final = tuple( + (idx, query) + for idx in self._select_targets(messages, query_idx) + if (query := _query_for_target(messages, idx, fallback_query)).strip() + ) + if not eligible_pairs: + verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") + return inputs + + # Cap the fan-out: the largest messages are kept because that is where + # selection saves the most, and the choice must not depend on message + # order in the request. Everything past the cap is forwarded untouched. + selected_pairs: Final = ( + eligible_pairs + if len(eligible_pairs) <= _MAX_TARGETS_PER_REQUEST + else tuple( + sorted( + eligible_pairs, + key=lambda pair: len(content_to_text(messages[pair[0]].get("content"))), + reverse=True, + )[:_MAX_TARGETS_PER_REQUEST] + ) + ) + if len(selected_pairs) < len(eligible_pairs): + verbose_proxy_logger.debug( + "Needlepath: %d eligible messages, selecting only the %d largest", + len(eligible_pairs), + _MAX_TARGETS_PER_REQUEST, + ) + + semaphore: Final = asyncio.Semaphore(_MAX_CONCURRENT_SELECTIONS) + + async def _bounded_selected_text(idx: int, query: str) -> str | None: + # The semaphore bounds how many selection calls this one proxy + # request holds open at a time; see _MAX_CONCURRENT_SELECTIONS. + async with semaphore: + return await self._selected_text( + text=content_to_text(messages[idx].get("content")), + query=query, + title=_title_for(messages, idx), + source=messages[idx].get("tool_call_id"), + kind=_record_kind(messages[idx].get("role")), + ) + + start_time: Final = time.monotonic() + # One request per message: each carries its own query, and the service + # renders one block per call. Running them concurrently (up to the + # semaphore's bound) keeps the added latency near one round trip + # rather than one per message. + selections: Final = await asyncio.gather(*(_bounded_selected_text(idx, query) for idx, query in selected_pairs)) + end_time: Final = time.monotonic() + + # Needs in-place index assignment below to build the edited copy. It starts + # from the caller's own message objects, so only replaced indices differ + # from the original list -- see the identity note below. + selected_messages: Final[list[object]] = list(structured_messages) # mutable-ok: see comment above + messages_selected = 0 # rebind-ok: loop accumulator, incremented per selection below + chars_before = 0 # rebind-ok: loop accumulator, incremented per selection below + chars_after = 0 # rebind-ok: loop accumulator, incremented per selection below + for (idx, _query), selection in zip(selected_pairs, selections): + original = messages[idx] + original_text = content_to_text(original.get("content")) + if selection is None: + continue + messages_selected += 1 + chars_before += len(original_text) + chars_after += len(selection) + selected_messages[idx] = { + **original, + "content": _write_text_back(original.get("content"), selection), + } + + if messages_selected == 0: + # Nothing was replaced, so hand back the exact inputs object. The + # handlers detect a guardrail edit by identity, and a fresh list + # would force a write-back of an untouched request (on Anthropic + # that reconversion strips cache_control from thinking blocks). + verbose_proxy_logger.debug("Needlepath: no selection applied; request unchanged") + return inputs + + # Serialized as guardrail_json_response for standard logging (DataDog, + # Langfuse, ...); must stay a plain dict for json.dumps compatibility. + stats: Final[dict[str, object]] = { # mutable-ok: see comment above + "messages_selected": messages_selected, + "messages_considered": len(selected_pairs), + "chars_before": chars_before, + "chars_after": chars_after, + "operating_point": self.operating_point, + } + verbose_proxy_logger.debug( + "Needlepath: selected %s of %s message(s), %s -> %s chars", + messages_selected, + len(selected_pairs), + chars_before, + chars_after, + ) + self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper + guardrail_json_response=stats, + request_data=request_data, + guardrail_status="success", + guardrail_provider="needlepath", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + return {**inputs, "structured_messages": selected_messages} # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + + @staticmethod + def get_config_model() -> type[NeedlepathGuardrailConfigModel] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import ( + NeedlepathGuardrailConfigModel, + ) + + return NeedlepathGuardrailConfigModel diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 718a651bdee..9e1b9cebc12 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -47,6 +47,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import ( + NeedlepathGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import ( OvalixGuardrailConfigModel, ) @@ -163,6 +166,7 @@ class SupportedGuardrailIntegrations(Enum): AIRIA = "airia" WINGBACK = "wingback" DECISION_MODEL = "decision_model" + NEEDLEPATH = "needlepath" class Role(Enum): @@ -1397,6 +1401,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o HeadroomGuardrailConfigModel, CompresrGuardrailConfigModel, TypeSafeGuardrailConfigModel, + NeedlepathGuardrailConfigModel, RepelloAIGuardrailConfigModel, LassoGuardrailConfigModel, DeepKeepGuardrailConfigModel, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py new file mode 100644 index 00000000000..f417aff67a2 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py @@ -0,0 +1,56 @@ +from pydantic import BaseModel, ConfigDict, Field + +from .base import GuardrailConfigModel + + +class NeedlepathGuardrailOptionalParams(BaseModel): + """Optional tuning knobs for the Needlepath guardrail.""" + + model_config = ConfigDict(frozen=True) + + select_tool_outputs: bool | None = Field( + default=None, + description="Select over tool/function result messages (search hits, RAG chunks, API dumps). Defaults to True.", + ) + select_history: bool | None = Field( + default=None, + description="Also select over prior (non-last) user messages. Defaults to False.", + ) + select_system: bool | None = Field( + default=None, + description="Also select over system messages. Defaults to False.", + ) + min_chars_to_select: int | None = Field( + default=None, + description="Skip messages whose text is shorter than this many characters. Defaults to 500.", + ) + max_context_tokens: int | None = Field( + default=None, + description="Token budget requested for the selected block of a single message. Defaults to 4000.", + ) + operating_point: str | None = Field( + default=None, + description=( + "Immutable engine label sent with every request. Defaults to 'np-2026-10-r8'. " + "Pinned rather than inherited from the service default, so what this guardrail " + "sends does not change underneath a deployment." + ), + ) + + +class NeedlepathGuardrailConfigModel(GuardrailConfigModel[NeedlepathGuardrailOptionalParams]): + api_key: str | None = Field( + default=None, + description="Needlepath API key. Falls back to the NEEDLEPATH_API_KEY env var.", + ) + api_base: str | None = Field( + default=None, + description=( + "Base URL of the Needlepath API. Falls back to the NEEDLEPATH_API_BASE env var, " + "then https://api.nextmoca.com." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Needlepath (context selection)" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py new file mode 100644 index 00000000000..11baa671008 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -0,0 +1,683 @@ +""" +Unit tests for the Needlepath guardrail. + +Tests cover: +- apply_guardrail replaces an eligible message with the selected block, query- + conditioned on the intent of the tool call that produced it +- target selection: tool outputs by default, system/history opt-in, min-chars + threshold, the query message itself is never rewritten +- the fail-open contract, one test per decline: empty selection + (records_selected 0, tokens_after 0, blank or missing rendered_context), a + gate stand-down, a ``blocked`` security outcome, HTTP 402/403/429 and other + non-2xx, timeouts and transport errors, malformed JSON, and an unexpected + response schema. Every one asserts the messages come back byte-identical. +- security screening: a ``sanitized_fallback``, or an engaged block that + screening quarantined, is applied even where an ordinary selection would be + declined, because declining would forward the flagged spans. +- /v1/responses: a function_call_output is selected and written back through + LiteLLM's Responses translation. +""" + +import asyncio +import copy +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.needlepath.needlepath import ( + _MAX_CONCURRENT_SELECTIONS, + _MAX_TARGETS_PER_REQUEST, + DEFAULT_MAX_CONTEXT_TOKENS, + DEFAULT_MIN_CHARS_TO_SELECT, + DEFAULT_OPERATING_POINT, + NeedlepathGuardrail, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +FAKE_API_BASE = "https://needlepath.example.com" +FAKE_API_KEY = "np_live_0123456789abcdefghijklmnopqr" + +TOOL_OUTPUT = "Result row: quarterly filing extract. " * 40 # > 500 chars +SELECTED_BLOCK = "Q3 revenue was 4.2 billion." +USER_QUESTION = "What was Q3 revenue?" + +AGENT_MESSAGES = [ + {"role": "system", "content": "You are a filings analyst. " * 40}, + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "edgar_search", "arguments": '{"cik": "0000320193"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT}, +] + +TOOL_INDEX = 3 + + +def _make_guardrail(**kwargs) -> NeedlepathGuardrail: + defaults = dict( + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + guardrail_name="needlepath-selection", + default_on=True, + ) + defaults.update(kwargs) + return NeedlepathGuardrail(**defaults) + + +def _select_response( + rendered_context: object = SELECTED_BLOCK, + records_selected: object = 1, + tokens_after: object = 120, + gate: object = None, + status: int = 200, + omit_rendered_context: bool = False, + outcome: object = "engaged", + format_metrics: object = None, +) -> MagicMock: + body = { + "request_id": "req-1", + "outcome": outcome, + "fallback_used": False, + "records_selected": records_selected, + "records_submitted": 1, + "records_available": 1, + "tokens_before": 900, + "tokens_after": tokens_after, + "tokens_saved": 780, + "policy_version": "np-2026-10-r8", + "gate": {"engaged": True, "reason": "engage:needle"} if gate is None else gate, + } + if not omit_rendered_context: + body["rendered_context"] = rendered_context + if format_metrics is not None: + body["format_metrics"] = format_metrics + mock = MagicMock() + mock.status_code = status + mock.json.return_value = body + mock.text = "" + return mock + + +def _apply_inputs(messages: list) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs(structured_messages=copy.deepcopy(messages)) + + +@pytest.fixture +def guardrail() -> NeedlepathGuardrail: + return _make_guardrail() + + +async def _run(guardrail: NeedlepathGuardrail, post_mock, messages=None) -> GenericGuardrailAPIInputs: + inputs = _apply_inputs(AGENT_MESSAGES if messages is None else messages) + with patch.object(guardrail.async_handler, "post", post_mock): + return await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +async def _run_expect_unchanged(guardrail: NeedlepathGuardrail, post_mock, messages=None) -> None: + """A decline hands back the exact inputs object, with every message byte-identical.""" + msgs = AGENT_MESSAGES if messages is None else messages + inputs = _apply_inputs(msgs) + with patch.object(guardrail.async_handler, "post", post_mock): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + assert result["structured_messages"] == msgs + + +def _assert_byte_identical(result: GenericGuardrailAPIInputs, messages=None) -> None: + """Every message came back exactly as it went in.""" + assert result["structured_messages"] == (AGENT_MESSAGES if messages is None else messages) + + +# ── init ────────────────────────────────────────────────────────────── + + +def test_init_raises_without_api_key(monkeypatch): + monkeypatch.delenv("NEEDLEPATH_API_KEY", raising=False) + with pytest.raises(ValueError, match="API key"): + NeedlepathGuardrail(guardrail_name="needlepath-selection") + + +def test_init_defaults(): + g = _make_guardrail() + assert g.needlepath_api_base == FAKE_API_BASE + assert g.select_tool_outputs is True + assert g.select_history is False + assert g.select_system is False + assert g.min_chars_to_select == DEFAULT_MIN_CHARS_TO_SELECT == 500 + assert g.max_context_tokens == DEFAULT_MAX_CONTEXT_TOKENS == 4000 + # Pinned, not inherited from the service default. + assert g.operating_point == DEFAULT_OPERATING_POINT == "np-2026-10-r8" + + +def test_initialize_guardrail_reads_nested_then_flat_knobs(): + from litellm.proxy.guardrails.guardrail_hooks.needlepath import ( + guardrail_class_registry, + guardrail_initializer_registry, + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams, SupportedGuardrailIntegrations + + litellm_params = LitellmParams( + guardrail="needlepath", + mode="pre_call", + api_key=FAKE_API_KEY, + api_base=FAKE_API_BASE, + optional_params={ + "select_history": True, + "min_chars_to_select": 1200, + "operating_point": "np-2026-10-r6", + }, + max_context_tokens=2000, + ) + callback = initialize_guardrail(litellm_params, {"guardrail_name": "needlepath-selection"}) + assert isinstance(callback, NeedlepathGuardrail) + assert callback.select_history is True + assert callback.select_tool_outputs is True + assert callback.min_chars_to_select == 1200 + assert callback.max_context_tokens == 2000 + assert callback.operating_point == "np-2026-10-r6" + assert guardrail_initializer_registry[SupportedGuardrailIntegrations.NEEDLEPATH.value] is initialize_guardrail + assert guardrail_class_registry[SupportedGuardrailIntegrations.NEEDLEPATH.value] is NeedlepathGuardrail + + +def test_initialize_guardrail_defaults_to_r8(): + from litellm.proxy.guardrails.guardrail_hooks.needlepath import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams(guardrail="needlepath", mode="pre_call", api_key=FAKE_API_KEY) + callback = initialize_guardrail(litellm_params, {"guardrail_name": "needlepath-selection"}) + assert callback.operating_point == "np-2026-10-r8" + + +def test_init_rejects_non_http_api_base(): + with pytest.raises(ValueError, match="http or https"): + _make_guardrail(api_base="file:///etc/passwd") + + +def test_init_rejects_cloud_metadata_api_base(): + with pytest.raises(ValueError, match="cloud-metadata"): + _make_guardrail(api_base="http://169.254.169.254") + + +# ── selection core ──────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_apply_guardrail_replaces_tool_output_with_rendered_context( + guardrail: NeedlepathGuardrail, +): + post = AsyncMock(return_value=_select_response()) + result = await _run(guardrail, post) + + _, call_kwargs = post.call_args + assert call_kwargs["url"] == f"{FAKE_API_BASE}/v1/context/select" + assert call_kwargs["headers"]["Authorization"] == f"Bearer {FAKE_API_KEY}" + payload = call_kwargs["json"] + assert len(payload["records"]) == 1 + record = payload["records"][0] + assert record["text"] == TOOL_OUTPUT + assert record["kind"] == "tool_result" + # Title from the tool name, source from the tool_call_id. + assert record["title"] == "edgar_search" + assert record["source"] == "call_1" + # The query is the tool call's intent, not the user question. + assert payload["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}' + assert payload["budget"] == { + "max_context_tokens": 4000, + "operating_point": "np-2026-10-r8", + } + # The label travels in budget.operating_point; a top-level one is ignored. + assert "operating_point" not in payload + assert payload["render"] is True + assert payload["render_format"] == "plain" + + out = result["structured_messages"] + assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK + # Only that message changed; every other one is byte-identical. + assert out[0] == AGENT_MESSAGES[0] + assert out[1] == AGENT_MESSAGES[1] + assert out[2] == AGENT_MESSAGES[2] + # The rewritten message keeps its envelope. + assert out[TOOL_INDEX]["role"] == "tool" + assert out[TOOL_INDEX]["tool_call_id"] == "call_1" + + +@pytest.mark.asyncio +async def test_tool_output_without_matching_call_uses_user_question( + guardrail: NeedlepathGuardrail, +): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_missing", "content": TOOL_OUTPUT}, + ] + post = AsyncMock(return_value=_select_response()) + await _run(guardrail, post, messages) + + _, call_kwargs = post.call_args + assert call_kwargs["json"]["task"]["prompt"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_system_and_history_not_selected_by_default( + guardrail: NeedlepathGuardrail, +): + post = AsyncMock(return_value=_select_response()) + await _run(guardrail, post) + # Only the tool output was submitted, even though the system message is + # comfortably over the character threshold. + assert post.call_count == 1 + assert post.call_args[1]["json"]["records"][0]["text"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_opt_in_system_selects_system_message(): + guardrail = _make_guardrail(select_system=True) + post = AsyncMock(return_value=_select_response()) + await _run(guardrail, post) + + assert post.call_count == 2 + submitted = {call[1]["json"]["records"][0]["kind"] for call in post.call_args_list} + assert submitted == {"tool_result", "external_data"} + + +@pytest.mark.asyncio +async def test_short_messages_skipped(guardrail: NeedlepathGuardrail): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_1", "content": "short"}, + ] + post = AsyncMock(return_value=_select_response()) + result = await _run(guardrail, post, messages) + + post.assert_not_called() + assert result["structured_messages"] == messages + + +@pytest.mark.asyncio +async def test_response_input_type_passthrough(guardrail: NeedlepathGuardrail): + inputs = _apply_inputs(AGENT_MESSAGES) + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + assert result is inputs + + +# ── fail-open contract ──────────────────────────────────────────────── +# +# One test per decline. Each asserts the exact inputs object comes back +# (handlers detect a guardrail edit by identity) and that the messages are +# byte-identical to what arrived. + + +@pytest.mark.asyncio +async def test_fail_open_zero_records_selected(guardrail: NeedlepathGuardrail): + post = AsyncMock(return_value=_select_response(records_selected=0)) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_zero_tokens_after(guardrail: NeedlepathGuardrail): + post = AsyncMock(return_value=_select_response(tokens_after=0)) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_blank_rendered_context(guardrail: NeedlepathGuardrail): + post = AsyncMock(return_value=_select_response(rendered_context=" ")) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_fail_open_missing_rendered_context(guardrail: NeedlepathGuardrail): + post = AsyncMock(return_value=_select_response(omit_rendered_context=True)) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "reason", + [ + "standdown:flat_gap", + "standdown:high_drift", + "standdown:insufficient_candidates", + "standdown:answerability_low", + "standdown:gate_error:ValueError", + ], +) +async def test_fail_open_gate_standdown(guardrail: NeedlepathGuardrail, reason: str): + """A stand-down is the service saying the full content is what to send. + + The reason string is an open enum, so this matches on the ``standdown:`` + prefix rather than on the individual values. + """ + post = AsyncMock(return_value=_select_response(gate={"engaged": False, "reason": reason})) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + # The service was called and answered 200; the decline is ours. + assert post.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 402, 403, 429, 500, 503]) +async def test_fail_open_http_error_status(guardrail: NeedlepathGuardrail, status_code: int): + """The shared handler raises for status, so every non-2xx arrives as an + HTTPStatusError carrying the upstream body. None of it reaches the client.""" + request = httpx.Request("POST", f"{FAKE_API_BASE}/v1/context/select") + response = httpx.Response(status_code, request=request, text="upstream detail") + post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=response)) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_fail_open_non_2xx_without_raise_for_status(guardrail: NeedlepathGuardrail): + """A handler configured not to raise still yields a non-2xx response object.""" + post = AsyncMock(return_value=_select_response(status=429)) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_fail_open_timeout(guardrail: NeedlepathGuardrail): + post = AsyncMock(side_effect=httpx.ReadTimeout("timed out")) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_fail_open_transport_error(guardrail: NeedlepathGuardrail): + post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_fail_open_unexpected_client_exception(guardrail: NeedlepathGuardrail): + """An exception outside httpx's hierarchy is still a decline, not a 500.""" + await _run_expect_unchanged(guardrail, AsyncMock(side_effect=RuntimeError("client closed"))) + + +@pytest.mark.asyncio +async def test_cancellation_is_not_swallowed(guardrail: NeedlepathGuardrail): + with pytest.raises(asyncio.CancelledError): + await _run(guardrail, AsyncMock(side_effect=asyncio.CancelledError())) + + +@pytest.mark.asyncio +async def test_fail_open_malformed_json(guardrail: NeedlepathGuardrail): + mock = MagicMock() + mock.status_code = 200 + mock.json.side_effect = ValueError("not json") + mock.text = "gateway error" + await _run_expect_unchanged(guardrail, AsyncMock(return_value=mock)) + + +@pytest.mark.asyncio +async def test_fail_open_unexpected_schema(guardrail: NeedlepathGuardrail): + """A 200 whose body is well-formed JSON but not the documented object.""" + mock = MagicMock() + mock.status_code = 200 + mock.json.return_value = ["not", "an", "object"] + mock.text = "" + await _run_expect_unchanged(guardrail, AsyncMock(return_value=mock)) + + +@pytest.mark.asyncio +async def test_fail_open_selection_not_smaller(guardrail: NeedlepathGuardrail): + """A block that is not shorter than the original could only add tokens.""" + post = AsyncMock(return_value=_select_response(rendered_context=TOOL_OUTPUT + " and more")) + await _run_expect_unchanged(guardrail, post) + + +# ── security screening outcomes ─────────────────────────────────────── +# Shapes follow the documented r8 responses: a sanitized_fallback carries no +# selected records, a stand-down gate reason, and the full context with each +# flagged span replaced by a marker. + +QUARANTINE_MARKER = "[QUARANTINED: instruction-like candidate text removed]" +SANITIZED_PACKET = TOOL_OUTPUT[:-60] + QUARANTINE_MARKER +QUARANTINE_METRICS = { + "security_mode": "enforce", + "security_action": "quarantine", + "security_findings_count": 1, +} + + +def _sanitized_fallback_response(rendered_context: object = SANITIZED_PACKET) -> MagicMock: + return _select_response( + rendered_context=rendered_context, + records_selected=0, + tokens_after=210, + gate={"engaged": False, "reason": "standdown:source_support_missing"}, + outcome="sanitized_fallback", + format_metrics={**QUARANTINE_METRICS, "security_delivery_outcome": "sanitized_fallback"}, + ) + + +@pytest.mark.asyncio +async def test_sanitized_fallback_is_applied(guardrail: NeedlepathGuardrail): + """Declining would forward the original, flagged span included.""" + post = AsyncMock(return_value=_sanitized_fallback_response()) + result = await _run(guardrail, post) + out = result["structured_messages"] + assert out[TOOL_INDEX]["content"] == SANITIZED_PACKET + assert out[TOOL_INDEX]["tool_call_id"] == "call_1" + assert out[:TOOL_INDEX] == AGENT_MESSAGES[:TOOL_INDEX] + + +@pytest.mark.asyncio +async def test_sanitized_fallback_is_applied_when_not_shorter(guardrail: NeedlepathGuardrail): + """The marker can make the packet longer than the span it replaced.""" + longer = TOOL_OUTPUT + QUARANTINE_MARKER + post = AsyncMock(return_value=_sanitized_fallback_response(rendered_context=longer)) + result = await _run(guardrail, post) + assert result["structured_messages"][TOOL_INDEX]["content"] == longer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rendered", ["", " ", None]) +async def test_sanitized_fallback_without_a_packet_fails_open(guardrail: NeedlepathGuardrail, rendered): + post = AsyncMock(return_value=_sanitized_fallback_response(rendered_context=rendered)) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_engaged_quarantined_block_is_applied_despite_standdown_and_size( + guardrail: NeedlepathGuardrail, +): + """An engaged selection with a span quarantined is the safe delivery.""" + longer = TOOL_OUTPUT + QUARANTINE_MARKER + post = AsyncMock( + return_value=_select_response( + rendered_context=longer, + gate={"engaged": False, "reason": "standdown:source_support_missing"}, + format_metrics=QUARANTINE_METRICS, + ) + ) + result = await _run(guardrail, post) + assert result["structured_messages"][TOOL_INDEX]["content"] == longer + + +@pytest.mark.asyncio +async def test_screening_keep_does_not_bypass_declines(guardrail: NeedlepathGuardrail): + """Screening that flagged nothing leaves the ordinary declines in force.""" + post = AsyncMock( + return_value=_select_response( + gate={"engaged": False, "reason": "standdown:source_support_missing"}, + format_metrics={"security_mode": "enforce", "security_action": "keep", "security_findings_count": 0}, + ) + ) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_blocked(guardrail: NeedlepathGuardrail): + """blocked delivers none of the caller's content: a decline, as in the + published clients. The placeholder line is never written into a message.""" + post = AsyncMock( + return_value=_select_response( + rendered_context="Relevant State: none selected.", + records_selected=0, + tokens_after=6, + outcome="blocked", + format_metrics={"security_mode": "enforce", "security_action": "block", "security_findings_count": 3}, + ) + ) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request") + assert result is inputs + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_blocked_even_with_selected_records(guardrail: NeedlepathGuardrail): + """The outcome decides, not the counters next to it.""" + post = AsyncMock(return_value=_select_response(outcome="blocked")) + await _run_expect_unchanged(guardrail, post) + + +@pytest.mark.asyncio +async def test_one_message_declining_does_not_block_another(): + """Selection is per message: a decline on one leaves the other applied.""" + guardrail = _make_guardrail(select_system=True) + + def _by_kind(**kwargs): + if kwargs["json"]["records"][0]["kind"] == "tool_result": + return _select_response() + return _select_response(gate={"engaged": False, "reason": "standdown:flat_gap"}) + + post = AsyncMock(side_effect=_by_kind) + result = await _run(guardrail, post) + + out = result["structured_messages"] + assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK + # The system message stood down, so it is byte-identical. + assert out[0] == AGENT_MESSAGES[0] + + +# ── fan-out bounds ──────────────────────────────────────────────────── + + +def _many_tool_messages(count: int) -> list: + """A user question plus `count` eligible tool outputs of strictly increasing size.""" + messages = [{"role": "user", "content": USER_QUESTION}] + for i in range(count): + messages.append( + { + "role": "tool", + "tool_call_id": f"call_{i}", + # All above the min-chars threshold; index i is the (i+1)-th smallest. + "content": "a" * (DEFAULT_MIN_CHARS_TO_SELECT + 20 + i), + } + ) + return messages + + +@pytest.mark.asyncio +async def test_target_cap_selects_only_the_largest_messages(guardrail: NeedlepathGuardrail): + """Past the per-request cap, only the largest messages are selected. + + The smallest overflow messages come back byte-identical and no service + call is made for them. + """ + overflow = 4 + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST + overflow) + + post = AsyncMock(return_value=_select_response()) + result = await _run(guardrail, post, messages=messages) + + assert post.await_count == _MAX_TARGETS_PER_REQUEST + out = result["structured_messages"] + # Tool messages start at index 1 and grow with index: the first `overflow` + # are the smallest, so they are the ones left untouched. + for idx in range(1, 1 + overflow): + assert out[idx] == messages[idx] + for idx in range(1 + overflow, len(messages)): + assert out[idx]["content"] == SELECTED_BLOCK + + +@pytest.mark.asyncio +async def test_concurrent_selections_are_bounded(guardrail: NeedlepathGuardrail): + """No more than _MAX_CONCURRENT_SELECTIONS service calls are in flight at once.""" + in_flight = 0 + high_water = 0 + + async def _tracking_post(*args, **kwargs): + nonlocal in_flight, high_water + in_flight += 1 + high_water = max(high_water, in_flight) + # Yield twice so every scheduled call gets a chance to start before + # this one finishes; without the semaphore all of them would overlap. + await asyncio.sleep(0) + await asyncio.sleep(0) + in_flight -= 1 + return _select_response() + + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST) + result = await _run(guardrail, _tracking_post, messages=messages) + + assert high_water <= _MAX_CONCURRENT_SELECTIONS + assert all(m["content"] == SELECTED_BLOCK for m in result["structured_messages"][1:]) + + +# ── /v1/responses ───────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_responses_api_function_call_output_is_selected_and_written_back( + guardrail: NeedlepathGuardrail, +): + """Through LiteLLM's Responses translation, a function_call_output is + selected against its call's intent and the block is written back.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + + data = { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": USER_QUESTION}, + { + "type": "function_call", + "call_id": "call_1", + "name": "edgar_search", + "arguments": '{"cik": "0000320193"}', + }, + {"type": "function_call_output", "call_id": "call_1", "output": TOOL_OUTPUT}, + ], + } + post = AsyncMock(return_value=_select_response()) + with patch.object(guardrail.async_handler, "post", post): + out = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert post.call_count == 1 + assert post.call_args.kwargs["json"]["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}' + assert out["input"][2]["output"] == SELECTED_BLOCK + assert out["input"][0] == {"role": "user", "content": USER_QUESTION}