From 90493adbb31b0159a1162bacf7d0df0a91c4cc12 Mon Sep 17 00:00:00 2001 From: kirankashalkar Date: Mon, 10 Aug 2026 22:59:20 -0400 Subject: [PATCH 1/8] feat(guardrails): add needlepath context selection guardrail Adds a first-party `needlepath` guardrail that runs query-conditioned extractive context selection over tool outputs during pre_call. The guardrail is unconditionally fail-open: any decline, stand-down, non-2xx, timeout, malformed body, or empty selection forwards the original messages byte-identical. --- .../guardrail_hooks/needlepath/__init__.py | 71 +++ .../guardrail_hooks/needlepath/needlepath.py | 578 ++++++++++++++++++ litellm/types/guardrails.py | 5 + .../guardrails/guardrail_hooks/needlepath.py | 54 ++ .../guardrail_hooks/test_needlepath.py | 397 ++++++++++++ 5 files changed, 1105 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py 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..ab01bab93f6 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py @@ -0,0 +1,71 @@ +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, +) -> 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: + """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 = 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) -> 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=_get_optional_value(litellm_params, optional_params, "select_tool_outputs"), + select_history=_get_optional_value(litellm_params, optional_params, "select_history"), + select_system=_get_optional_value(litellm_params, optional_params, "select_system"), + min_chars_to_select=_get_optional_value(litellm_params, optional_params, "min_chars_to_select"), + max_context_tokens=_get_optional_value(litellm_params, optional_params, "max_context_tokens"), + operating_point=_get_optional_value(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_initializer_registry: Final = { + SupportedGuardrailIntegrations.NEEDLEPATH.value: initialize_guardrail, +} + +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..6ef2345ca60 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -0,0 +1,578 @@ +"""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 typing import TYPE_CHECKING, Any, Final, Literal, TypeGuard +from urllib.parse import urlparse + +import httpx +from httpx import Response as HttpxResponse + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +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.base import ( + GuardrailConfigModel, + ) + +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-07-r2" +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 +# 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:" +_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 _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) + 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. + + 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 _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 _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 + if _is_object_list(content) and is_all_text_parts(content): + return merge_rewritten_text_parts(content, new_text) + return content + + +def _render_tool_intent(fn: dict[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") + args: Final = "" if raw_args is None else str(raw_args).strip() + if name and args: + return f"{name}: {args}" + return name or args + + +def _query_for_target(messages: list[dict[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 + tool_calls = previous.get("tool_calls") + if _is_object_list(tool_calls): + for call in tool_calls: + if not _is_str_object_dict(call) or not tool_call_id or call.get("id") != tool_call_id: + continue + fn = call.get("function") + intent = _render_tool_intent(fn if _is_str_object_dict(fn) else {}) + 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 = previous.get("function_call") + if _is_str_object_dict(legacy) and fn_name and legacy.get("name") == fn_name: + intent = _render_tool_intent(legacy) + if intent: + return intent + return fallback + + +def _title_for(messages: list[dict[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 + tool_calls = previous.get("tool_calls") + if not _is_object_list(tool_calls): + continue + for call in tool_calls: + if not _is_str_object_dict(call) or call.get("id") != tool_call_id: + continue + fn = call.get("function") + if _is_str_object_dict(fn) and isinstance(fn.get("name"), str): + return str(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 _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, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, + default_on: bool = False, + ): + 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]: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.needlepath_api_key or ''}", + } + + def _decline(self, reason: str, detail: dict[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. + """ + verbose_proxy_logger.debug("Needlepath: declined (%s) detail=%s", reason, detail or {}) + + def _select_targets(self, messages: list[dict[str, object]], query_idx: int | None) -> list[int]: + """Indices of the messages whose text content is eligible for selection.""" + targets: Final[list[int]] = [] + for idx, msg in enumerate(messages): + # 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: + continue + role = msg.get("role") + if role in ("tool", "function"): + if not self.select_tool_outputs: + continue + elif role == "system": + if not self.select_system: + continue + elif role == "user": + if not self.select_history: + continue + else: + continue + content = msg.get("content") + if _is_object_list(content) and not is_all_text_parts(content): + continue + if len(content_to_text(content)) < self.min_chars_to_select: + continue + targets.append(idx) + return targets + + @staticmethod + def _last_user_message(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 + + async def _post_select(self, payload: dict[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=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 + + 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; + * 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. + """ + record: Final[dict[str, object]] = {"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]] = { + "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: + body: Final[object] = response.json() + 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 + if not _is_str_object_dict(body): + self._decline("unexpected_shape", {"body": _safe_response_text(response)}) + return None + + gate: Final = body.get("gate") + reason: Final = gate.get("reason") if _is_str_object_dict(gate) else None + if 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 + + rendered: Final = body.get("rendered_context") + if not isinstance(rendered, str) or not rendered.strip(): + self._decline("empty_rendered_context", {}) + return None + if len(rendered) >= len(text): + # A selection that is not smaller is not a selection worth making, + # and applying it could only add tokens. + 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, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + structured_messages: Final = inputs.get("structured_messages") + if not _is_object_list(structured_messages) or not structured_messages: + return inputs + messages: Final = [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._last_user_message(messages) + targets: Final[list[int]] = [] + queries: Final[list[str]] = [] + for idx in self._select_targets(messages, query_idx): + query = _query_for_target(messages, idx, fallback_query) + # Selection is conditioned on a query; without one there is nothing + # to select against, so the message is left exactly as it arrived. + if not query.strip(): + continue + targets.append(idx) + queries.append(query) + if not targets: + verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") + return inputs + + 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 keeps the added + # latency at roughly one round trip rather than one per message. + selections: Final = await asyncio.gather( + *( + 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")), + ) + for idx, query in zip(targets, queries) + ) + ) + end_time: Final = time.monotonic() + + selected_messages: Final = list(messages) + messages_selected = 0 + chars_before = 0 + chars_after = 0 + for idx, selection in zip(targets, selections): + original = selected_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 + + stats: Final[dict[str, object]] = { + "messages_selected": messages_selected, + "messages_considered": len(targets), + "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(targets), + chars_before, + chars_after, + ) + self.add_standard_logging_guardrail_information_to_request_data( + 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[GuardrailConfigModel[Any]] | 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 e1b891caaaf..ec4606c4cb9 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): @@ -1380,6 +1384,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..429c05aa482 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py @@ -0,0 +1,54 @@ +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class NeedlepathGuardrailOptionalParams(BaseModel): + """Optional tuning knobs for the Needlepath guardrail.""" + + 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-07-r2'. " + "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..b8621d2c43b --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -0,0 +1,397 @@ +""" +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, 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. +""" + +import copy +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.needlepath.needlepath import ( + 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, +) -> MagicMock: + body = { + "request_id": "req-1", + "records_selected": records_selected, + "records_submitted": 1, + "records_available": 1, + "tokens_before": 900, + "tokens_after": tokens_after, + "tokens_saved": 780, + "policy_version": "np-2026-07-r2", + "gate": {"engaged": True, "reason": "engage:needle"} if gate is None else gate, + } + if not omit_rendered_context: + body["rendered_context"] = rendered_context + 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", + ) + + +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-07-r2" + + +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-07-r2", + } + 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=" ")) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_missing_rendered_context(guardrail: NeedlepathGuardrail): + post = AsyncMock(return_value=_select_response(omit_rendered_context=True)) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@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)) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@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)) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_timeout(guardrail: NeedlepathGuardrail): + post = AsyncMock(side_effect=httpx.ReadTimeout("timed out")) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@pytest.mark.asyncio +async def test_fail_open_transport_error(guardrail: NeedlepathGuardrail): + post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@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" + result = await _run(guardrail, AsyncMock(return_value=mock)) + _assert_byte_identical(result) + + +@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 = "" + result = await _run(guardrail, AsyncMock(return_value=mock)) + _assert_byte_identical(result) + + +@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")) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + +@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] From 51c30abe85ecdcc1558c1642035f8d6cb2de7722 Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Sun, 16 Aug 2026 08:32:45 +0530 Subject: [PATCH 2/8] fix(needlepath): satisfy code-quality, lint, and strict-budget gates - drop typing.Any and typing.TypeGuard (TID251): inline isinstance narrowing at call sites, GuardrailConfigModel[object] per the compresr precedent - annotate __init__ -> None (ANN204) - rename the tool-intent local so no f-string interpolates a name the key-leak scanner treats as function kwargs Co-Authored-By: Claude Fable 5 --- .../guardrail_hooks/needlepath/needlepath.py | 48 ++++++++----------- 1 file changed, 20 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index 6ef2345ca60..0d7d8c4ca2c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -25,7 +25,7 @@ from __future__ import annotations import asyncio import ipaddress import time -from typing import TYPE_CHECKING, Any, Final, Literal, TypeGuard +from typing import TYPE_CHECKING, Final, Literal from urllib.parse import urlparse import httpx @@ -130,14 +130,6 @@ def _validate_api_base(url: str) -> str: 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 _write_text_back(content: object, new_text: str) -> object: """Put ``new_text`` into a ``content`` value without changing its shape. @@ -148,7 +140,7 @@ def _write_text_back(content: object, new_text: str) -> object: """ if isinstance(content, str): return new_text - if _is_object_list(content) and is_all_text_parts(content): + if isinstance(content, list) and is_all_text_parts(content): return merge_rewritten_text_parts(content, new_text) return content @@ -157,10 +149,10 @@ def _render_tool_intent(fn: dict[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") - args: Final = "" if raw_args is None else str(raw_args).strip() - if name and args: - return f"{name}: {args}" - return name or args + 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: list[dict[str, object]], target_idx: int, fallback: str) -> str: @@ -182,18 +174,18 @@ def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallba if previous.get("role") != "assistant": continue tool_calls = previous.get("tool_calls") - if _is_object_list(tool_calls): + if isinstance(tool_calls, list): for call in tool_calls: - if not _is_str_object_dict(call) or not tool_call_id or call.get("id") != tool_call_id: + if not isinstance(call, dict) or not tool_call_id or call.get("id") != tool_call_id: continue fn = call.get("function") - intent = _render_tool_intent(fn if _is_str_object_dict(fn) else {}) + intent = _render_tool_intent(fn if isinstance(fn, dict) else {}) 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 = previous.get("function_call") - if _is_str_object_dict(legacy) and fn_name and legacy.get("name") == fn_name: + if isinstance(legacy, dict) and fn_name and legacy.get("name") == fn_name: intent = _render_tool_intent(legacy) if intent: return intent @@ -214,13 +206,13 @@ def _title_for(messages: list[dict[str, object]], target_idx: int) -> str | None if previous.get("role") != "assistant": continue tool_calls = previous.get("tool_calls") - if not _is_object_list(tool_calls): + if not isinstance(tool_calls, list): continue for call in tool_calls: - if not _is_str_object_dict(call) or call.get("id") != tool_call_id: + if not isinstance(call, dict) or call.get("id") != tool_call_id: continue fn = call.get("function") - if _is_str_object_dict(fn) and isinstance(fn.get("name"), str): + if isinstance(fn, dict) and isinstance(fn.get("name"), str): return str(fn["name"])[:120] return None @@ -281,7 +273,7 @@ class NeedlepathGuardrail(CustomGuardrail): guardrail_name: str | None = None, event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, 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") @@ -346,7 +338,7 @@ class NeedlepathGuardrail(CustomGuardrail): else: continue content = msg.get("content") - if _is_object_list(content) and not is_all_text_parts(content): + if isinstance(content, list) and not is_all_text_parts(content): continue if len(content_to_text(content)) < self.min_chars_to_select: continue @@ -438,12 +430,12 @@ class NeedlepathGuardrail(CustomGuardrail): # It is a decline like any other malformed answer, not a 500. self._decline("unreadable_body", {"body": _safe_response_text(response)}) return None - if not _is_str_object_dict(body): + if not isinstance(body, dict): self._decline("unexpected_shape", {"body": _safe_response_text(response)}) return None gate: Final = body.get("gate") - reason: Final = gate.get("reason") if _is_str_object_dict(gate) else None + reason: Final = gate.get("reason") if isinstance(gate, dict) else None if isinstance(reason, str) and reason.startswith(_STANDDOWN_PREFIX): self._decline("gate_standdown", {"reason": reason}) return None @@ -480,9 +472,9 @@ class NeedlepathGuardrail(CustomGuardrail): return inputs structured_messages: Final = inputs.get("structured_messages") - if not _is_object_list(structured_messages) or not structured_messages: + if not isinstance(structured_messages, list) or not structured_messages: return inputs - messages: Final = [m for m in structured_messages if _is_str_object_dict(m)] + messages: Final = [m for m in structured_messages if isinstance(m, dict)] if len(messages) != len(structured_messages): return inputs @@ -570,7 +562,7 @@ class NeedlepathGuardrail(CustomGuardrail): return {**inputs, "structured_messages": selected_messages} # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime @staticmethod - def get_config_model() -> type[GuardrailConfigModel[Any]] | None: + def get_config_model() -> type[GuardrailConfigModel[object]] | None: from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import ( NeedlepathGuardrailConfigModel, ) From eba10641f91f30bc5d33974004a058553c1f9f3c Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Sun, 16 Aug 2026 09:03:35 +0530 Subject: [PATCH 3/8] fix(needlepath): clear all type-discipline violations, preserve behavior Widens mutable dict/list annotations to Mapping/Sequence/tuple across the needlepath guardrail (params, returns, locals) and declares Final on previously-bare locals, with zero behavior change: - _select_targets rebuilt as a pure filter (tuple(...) instead of append-in-a-loop) via a nested `_eligible` predicate - apply_guardrail's targets/queries lists collapsed into one `selected_pairs` tuple built with a walrus-filtered generator - _request_headers now returns a genuine MappingProxyType (httpx accepts any Mapping for headers=) - "records": [record] -> (record,) (json.dumps serializes tuples as arrays identically to lists) Reasoned `# mutable-ok` / `# rebind-ok` suppressions were used only where the ecosystem forces mutability, each with a specific reason: - _coerce_event_hook's list comprehension and the two module-level guardrail_*_registry dicts must stay real `dict`/`list` objects, since CustomGuardrail.should_run_guardrail does `isinstance(self.event_hook, list)` and guardrail_registry.py's discovery does `isinstance(registry, dict)` -- a Mapping/tuple would silently disable the hook / registry - the JSON payload dicts (record/payload and nested task/budget) in _selected_text must stay plain dicts: MappingProxyType is not JSON-serializable (json.dumps raises TypeError on one), and `record` is genuinely mutated after construction to add optional keys - decline() detail dicts are log-only, consumed immediately by a debug %s format - the final `{**inputs, "structured_messages": ...}` return value is a GenericGuardrailAPIInputs TypedDict, treated as a plain dict downstream - ipaddress parsing in _parse_ip_literal genuinely rebinds `addr` as a str-then-int fallback parse Verified: scripts/check_type_discipline.py reports 0 violations (was 56), ruff check/format and ruff-strict.toml pass, prevent_key_leaks_in_exceptions passes, and all 32 tests in test_needlepath.py pass in a fresh venv. Manually re-verified the two isinstance()-based ecosystem contracts that motivated the suppressions still hold: guardrail_registry discovery finds "needlepath" in both registries, and should_run_guardrail correctly dispatches a multi-hook list event_hook. Co-Authored-By: Claude Fable 5 --- .../guardrail_hooks/needlepath/__init__.py | 20 ++- .../guardrail_hooks/needlepath/needlepath.py | 170 +++++++++++------- 2 files changed, 116 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py index ab01bab93f6..9e83d232a19 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING, Final from litellm.types.guardrails import ( @@ -15,12 +16,17 @@ if TYPE_CHECKING: def _coerce_event_hook( - mode: str | list[str] | Mode, -) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: + mode: str | Sequence[str] | Mode, +) -> GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode: if isinstance(mode, Mode): return mode if isinstance(mode, list): - return [GuardrailEventHooks(item) for item in mode] + # 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] # mutable-ok: see comment above return GuardrailEventHooks(mode) @@ -62,10 +68,14 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> return _callback -guardrail_initializer_registry: Final = { +# 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 = { # mutable-ok: see comment above SupportedGuardrailIntegrations.NEEDLEPATH.value: initialize_guardrail, } -guardrail_class_registry: Final = { +# Same reason as guardrail_initializer_registry above: discovery's isinstance(dict) check. +guardrail_class_registry: Final = { # mutable-ok: see comment above SupportedGuardrailIntegrations.NEEDLEPATH.value: NeedlepathGuardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index 0d7d8c4ca2c..bd36ea5d91e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -25,6 +25,8 @@ 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 @@ -98,10 +100,10 @@ def _parse_ip_literal(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Addres 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) + addr = ipaddress.ip_address(host) # rebind-ok: reassigned in the except fallback below except ValueError: try: - addr = ipaddress.ip_address(int(host, 0)) + 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: @@ -145,7 +147,7 @@ def _write_text_back(content: object, new_text: str) -> object: return content -def _render_tool_intent(fn: dict[str, object]) -> str: +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") @@ -155,7 +157,7 @@ def _render_tool_intent(fn: dict[str, object]) -> str: return name or arg_text -def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallback: str) -> str: +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 @@ -179,7 +181,7 @@ def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallba if not isinstance(call, dict) or not tool_call_id or call.get("id") != tool_call_id: continue fn = call.get("function") - intent = _render_tool_intent(fn if isinstance(fn, dict) else {}) + intent = _render_tool_intent(fn if isinstance(fn, dict) else MappingProxyType({})) if intent: return intent # Legacy function_call turns carry no id, so require a name match. @@ -192,7 +194,7 @@ def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallba return fallback -def _title_for(messages: list[dict[str, object]], target_idx: int) -> str | None: +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") @@ -271,7 +273,7 @@ class NeedlepathGuardrail(CustomGuardrail): max_context_tokens: int | None = None, operating_point: str | None = None, guardrail_name: str | None = None, - event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, + event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | None = None, default_on: bool = False, ) -> None: raw_api_base: Final = (api_base or get_secret_str("NEEDLEPATH_API_BASE") or DEFAULT_API_BASE).rstrip("/") @@ -301,58 +303,64 @@ class NeedlepathGuardrail(CustomGuardrail): default_on=default_on, ) - def _request_headers(self) -> dict[str, str]: - return { - "Content-Type": "application/json", - "Authorization": f"Bearer {self.needlepath_api_key or ''}", - } + def _request_headers(self) -> Mapping[str, str]: + # httpx's `headers=` accepts any Mapping (it iterates .items()), so this can + # be a genuine read-only view rather than a dict a caller could mutate. + return MappingProxyType( + { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.needlepath_api_key or ''}", + } + ) - def _decline(self, reason: str, detail: dict[str, object] | None = None) -> None: + 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. """ - verbose_proxy_logger.debug("Needlepath: declined (%s) detail=%s", reason, detail or {}) + # 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 {}) # mutable-ok: see above - def _select_targets(self, messages: list[dict[str, object]], query_idx: int | None) -> list[int]: + 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.""" - targets: Final[list[int]] = [] - for idx, msg in enumerate(messages): + + 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: - continue - role = msg.get("role") + return False + role: Final = msg.get("role") if role in ("tool", "function"): if not self.select_tool_outputs: - continue + return False elif role == "system": if not self.select_system: - continue + return False elif role == "user": if not self.select_history: - continue + return False else: - continue - content = msg.get("content") + return False + content: Final = msg.get("content") if isinstance(content, list) and not is_all_text_parts(content): - continue - if len(content_to_text(content)) < self.min_chars_to_select: - continue - targets.append(idx) - return targets + 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: list[dict[str, object]]) -> tuple[str, int | None]: + 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: dict[str, object]) -> HttpxResponse | 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 @@ -371,14 +379,17 @@ class NeedlepathGuardrail(CustomGuardrail): response: Final = getattr(e, "response", None) self._decline( "http_status", - {"status_code": getattr(response, "status_code", None), "body": _safe_response_text(response)}, + { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + "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)}) + self._decline("transport", {"detail": str(e)}) # mutable-ok: log detail dict, consumed immediately return None async def _selected_text(self, text: str, query: str, title: str | None, source: object, kind: str) -> str | None: @@ -397,15 +408,24 @@ class NeedlepathGuardrail(CustomGuardrail): * the returned block is not shorter than the text it would replace, so a rewrite could never grow a message. """ - record: Final[dict[str, object]] = {"id": "m0", "kind": kind, "text": text} + # `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]] = { - "records": [record], - "task": {"prompt": query}, - "budget": { + payload: Final[dict[str, object]] = { # mutable-ok: see comment above + "records": (record,), + "task": {"prompt": query}, # mutable-ok: nested JSON payload dict, see comment above + "budget": { # mutable-ok: nested JSON payload dict, see comment above "max_context_tokens": self.max_context_tokens, "operating_point": self.operating_point, }, @@ -419,7 +439,10 @@ class NeedlepathGuardrail(CustomGuardrail): if not 200 <= response.status_code < 300: self._decline( "http_status", - {"status_code": response.status_code, "body": _safe_response_text(response)}, + { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + "status_code": response.status_code, + "body": _safe_response_text(response), + }, ) return None @@ -428,16 +451,16 @@ class NeedlepathGuardrail(CustomGuardrail): 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)}) + self._decline("unreadable_body", {"body": _safe_response_text(response)}) # mutable-ok: log detail dict return None if not isinstance(body, dict): - self._decline("unexpected_shape", {"body": _safe_response_text(response)}) + self._decline("unexpected_shape", {"body": _safe_response_text(response)}) # mutable-ok: log detail dict return None gate: Final = body.get("gate") reason: Final = gate.get("reason") if isinstance(gate, dict) else None if isinstance(reason, str) and reason.startswith(_STANDDOWN_PREFIX): - self._decline("gate_standdown", {"reason": reason}) + self._decline("gate_standdown", {"reason": reason}) # mutable-ok: log detail dict return None records_selected: Final = _safe_int(body.get("records_selected")) @@ -445,18 +468,21 @@ class NeedlepathGuardrail(CustomGuardrail): if records_selected == 0 or tokens_after == 0: self._decline( "empty_selection", - {"records_selected": records_selected, "tokens_after": tokens_after}, + { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + "records_selected": records_selected, + "tokens_after": tokens_after, + }, ) return None rendered: Final = body.get("rendered_context") if not isinstance(rendered, str) or not rendered.strip(): - self._decline("empty_rendered_context", {}) + self._decline("empty_rendered_context", MappingProxyType({})) return None if len(rendered) >= len(text): # A selection that is not smaller is not a selection worth making, # and applying it could only add tokens. - self._decline("no_reduction", {"before": len(text), "after": len(rendered)}) + self._decline("no_reduction", {"before": len(text), "after": len(rendered)}) # mutable-ok: log detail return None return rendered @@ -464,7 +490,7 @@ class NeedlepathGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: @@ -474,22 +500,20 @@ class NeedlepathGuardrail(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if not isinstance(structured_messages, list) or not structured_messages: return inputs - messages: Final = [m for m in structured_messages if isinstance(m, dict)] + messages: Final = tuple(m for m in structured_messages if isinstance(m, dict)) if len(messages) != len(structured_messages): return inputs fallback_query, query_idx = self._last_user_message(messages) - targets: Final[list[int]] = [] - queries: Final[list[str]] = [] - for idx in self._select_targets(messages, query_idx): - query = _query_for_target(messages, idx, fallback_query) - # Selection is conditioned on a query; without one there is nothing - # to select against, so the message is left exactly as it arrived. - if not query.strip(): - continue - targets.append(idx) - queries.append(query) - if not targets: + # 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. + selected_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 selected_pairs: verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") return inputs @@ -506,16 +530,19 @@ class NeedlepathGuardrail(CustomGuardrail): source=messages[idx].get("tool_call_id"), kind=_record_kind(messages[idx].get("role")), ) - for idx, query in zip(targets, queries) + for idx, query in selected_pairs ) ) end_time: Final = time.monotonic() - selected_messages: Final = list(messages) - messages_selected = 0 - chars_before = 0 - chars_after = 0 - for idx, selection in zip(targets, selections): + # Needs in-place index assignment below to build the edited copy without + # touching the frozen `messages` tuple; only replaced indices differ from + # the original list -- see the identity note below. + selected_messages: Final = list(messages) # mutable-ok: see comment above + messages_selected: Final = 0 + chars_before: Final = 0 + chars_after: Final = 0 + for (idx, _query), selection in zip(selected_pairs, selections): original = selected_messages[idx] original_text = content_to_text(original.get("content")) if selection is None: @@ -523,7 +550,7 @@ class NeedlepathGuardrail(CustomGuardrail): messages_selected += 1 chars_before += len(original_text) chars_after += len(selection) - selected_messages[idx] = { + selected_messages[idx] = { # mutable-ok: plain AllMessageValues dict, see comment above return **original, "content": _write_text_back(original.get("content"), selection), } @@ -536,9 +563,11 @@ class NeedlepathGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Needlepath: no selection applied; request unchanged") return inputs - stats: Final[dict[str, object]] = { + # 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(targets), + "messages_considered": len(selected_pairs), "chars_before": chars_before, "chars_after": chars_after, "operating_point": self.operating_point, @@ -546,7 +575,7 @@ class NeedlepathGuardrail(CustomGuardrail): verbose_proxy_logger.debug( "Needlepath: selected %s of %s message(s), %s -> %s chars", messages_selected, - len(targets), + len(selected_pairs), chars_before, chars_after, ) @@ -559,7 +588,10 @@ class NeedlepathGuardrail(CustomGuardrail): end_time=end_time, duration=end_time - start_time, ) - return {**inputs, "structured_messages": selected_messages} # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + return { # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime # mutable-ok: TypedDict, treated as plain dict downstream + **inputs, + "structured_messages": selected_messages, + } @staticmethod def get_config_model() -> type[GuardrailConfigModel[object]] | None: From 91914d4b5257fd6655a6c032986cf310021ac18f Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Sun, 16 Aug 2026 09:18:24 +0530 Subject: [PATCH 4/8] fix(needlepath): satisfy the basedpyright budget Loop accumulators drop Final (basedpyright forbids the reassignment the LIT checker's loop exemption allows) and carry rebind-ok reasons; the message-boundary isinstance keeps its runtime check with a reasoned suppression -- the annotation promises what only the check guarantees for client-supplied JSON. Co-Authored-By: Claude Fable 5 --- .../guardrail_hooks/needlepath/needlepath.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index bd36ea5d91e..b97259c3f6b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -500,7 +500,15 @@ class NeedlepathGuardrail(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if not isinstance(structured_messages, list) or not structured_messages: return inputs - messages: Final = tuple(m for m in structured_messages if isinstance(m, dict)) + # 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. + messages: Final = tuple( + m + for m in structured_messages + if isinstance(m, dict) # pyright: ignore[reportUnnecessaryIsInstance] # runtime boundary check on client JSON + ) if len(messages) != len(structured_messages): return inputs @@ -539,9 +547,9 @@ class NeedlepathGuardrail(CustomGuardrail): # touching the frozen `messages` tuple; only replaced indices differ from # the original list -- see the identity note below. selected_messages: Final = list(messages) # mutable-ok: see comment above - messages_selected: Final = 0 - chars_before: Final = 0 - chars_after: Final = 0 + 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 = selected_messages[idx] original_text = content_to_text(original.get("content")) From 98c19ecc56abbd6968af40cb1a25a3782e788b5c Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Thu, 20 Aug 2026 20:42:20 +0530 Subject: [PATCH 5/8] fix(needlepath): bound per-request selection fan-out A request carrying arbitrarily many eligible tool messages previously opened one concurrent outbound selection call per message. Cap targets at 16 per request (largest messages first, where selection pays off most; the rest are forwarded untouched, fail-open as ever) and bound in-flight calls with a semaphore of 4 so one request cannot monopolise the proxy's shared HTTP pool or burn selection quota. Adds tests for the cap (only the largest messages selected, one call per selected message) and the concurrency high-water mark. --- .../guardrail_hooks/needlepath/needlepath.py | 59 ++++++++++++---- .../guardrail_hooks/test_needlepath.py | 68 +++++++++++++++++++ 2 files changed, 115 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index b97259c3f6b..13c998fc163 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -72,6 +72,16 @@ DEFAULT_MIN_CHARS_TO_SELECT: Final = 500 # 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. @@ -516,31 +526,56 @@ class NeedlepathGuardrail(CustomGuardrail): # 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. - selected_pairs: Final = tuple( + 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 selected_pairs: + if not eligible_pairs: verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") return inputs - 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 keeps the added - # latency at roughly one round trip rather than one per message. - selections: Final = await asyncio.gather( - *( - self._selected_text( + # 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")), ) - for idx, query in selected_pairs - ) - ) + + 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 without diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py index b8621d2c43b..1a019595042 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -13,6 +13,7 @@ Tests cover: the messages come back byte-identical. """ +import asyncio import copy from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +21,8 @@ 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, @@ -395,3 +398,68 @@ async def test_one_message_declining_does_not_block_another(): 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:]) From 3175a45444f4f97a026787ec922e785079a7bfec Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Tue, 8 Sep 2026 00:34:27 -0400 Subject: [PATCH 6/8] guardrails(needlepath): default operating point is the current public label np-2026-08-r4 Co-Authored-By: Claude Fable 5.1 --- .../guardrails/guardrail_hooks/needlepath/needlepath.py | 2 +- .../types/proxy/guardrails/guardrail_hooks/needlepath.py | 2 +- .../proxy/guardrails/guardrail_hooks/test_needlepath.py | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index 13c998fc163..dea326e7fa7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -65,7 +65,7 @@ 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-07-r2" +DEFAULT_OPERATING_POINT: Final = "np-2026-08-r4" 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 diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py index 429c05aa482..e085dce245e 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py @@ -29,7 +29,7 @@ class NeedlepathGuardrailOptionalParams(BaseModel): operating_point: str | None = Field( default=None, description=( - "Immutable engine label sent with every request. Defaults to 'np-2026-07-r2'. " + "Immutable engine label sent with every request. Defaults to 'np-2026-08-r4'. " "Pinned rather than inherited from the service default, so what this guardrail " "sends does not change underneath a deployment." ), diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py index 1a019595042..bb805d4900c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -84,7 +84,7 @@ def _select_response( "tokens_before": 900, "tokens_after": tokens_after, "tokens_saved": 780, - "policy_version": "np-2026-07-r2", + "policy_version": "np-2026-08-r4", "gate": {"engaged": True, "reason": "engage:needle"} if gate is None else gate, } if not omit_rendered_context: @@ -138,7 +138,7 @@ def test_init_defaults(): 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-07-r2" + assert g.operating_point == DEFAULT_OPERATING_POINT == "np-2026-08-r4" def test_init_rejects_non_http_api_base(): @@ -176,7 +176,7 @@ async def test_apply_guardrail_replaces_tool_output_with_rendered_context( assert payload["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}' assert payload["budget"] == { "max_context_tokens": 4000, - "operating_point": "np-2026-07-r2", + "operating_point": "np-2026-08-r4", } assert payload["render"] is True assert payload["render_format"] == "plain" From 2af63cc6ed7ad2bbb17bd86ee443a43c3aa6b2cf Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Thu, 29 Oct 2026 12:19:10 -0400 Subject: [PATCH 7/8] feat(needlepath): pin np-2026-10-r8 and apply screened deliveries Pins np-2026-10-r8, which selects as np-2026-10-r6 and screens the returned context for instruction-like text. The guardrail now handles the screening outcomes the way the published Needlepath clients do: - sanitized_fallback carries no selected records and a stand-down gate reason, so it was previously declined and the original, flagged span included, was forwarded. It is now applied whenever rendered_context is non-blank. - an engaged block from which screening removed a span (format_metrics.security_action == "quarantine") is applied even if the gate reports a stand-down or the block is not shorter. - blocked is an explicit decline: the original message is forwarded. Also brings the package in line with upstream's current gates: typed parsing of untrusted JSON through pydantic TypeAdapters, typed config knobs in the initializer, a frozen optional-params model, and removal of suppression markers that no longer suppress anything. Co-Authored-By: Claude Opus 5.5 --- .../guardrail_hooks/needlepath/__init__.py | 52 +++-- .../guardrail_hooks/needlepath/needlepath.py | 217 +++++++++++------- .../guardrails/guardrail_hooks/needlepath.py | 6 +- .../guardrail_hooks/test_needlepath.py | 180 ++++++++++++++- 4 files changed, 355 insertions(+), 100 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py index 9e83d232a19..54cc75ede03 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections.abc import Sequence from typing import TYPE_CHECKING, Final from litellm.types.guardrails import ( @@ -16,8 +15,8 @@ if TYPE_CHECKING: def _coerce_event_hook( - mode: str | Sequence[str] | Mode, -) -> GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode: + 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): @@ -26,7 +25,7 @@ def _coerce_event_hook( # 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] # mutable-ok: see comment above + return [GuardrailEventHooks(item) for item in mode] return GuardrailEventHooks(mode) @@ -38,10 +37,35 @@ def _get_optional_value(litellm_params: LitellmParams, optional_params: object | set guardrail knobs flat alongside ``api_key``. """ if optional_params is not None: - value: Final = getattr(optional_params, attribute_name, None) + value: Final[object] = getattr(optional_params, attribute_name, None) # pyright: ignore[reportAny] # knob name is dynamic if value is not None: return value - return getattr(litellm_params, attribute_name, None) + 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: @@ -52,12 +76,12 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> _callback: Final = NeedlepathGuardrail( api_base=litellm_params.api_base, api_key=litellm_params.api_key, - select_tool_outputs=_get_optional_value(litellm_params, optional_params, "select_tool_outputs"), - select_history=_get_optional_value(litellm_params, optional_params, "select_history"), - select_system=_get_optional_value(litellm_params, optional_params, "select_system"), - min_chars_to_select=_get_optional_value(litellm_params, optional_params, "min_chars_to_select"), - max_context_tokens=_get_optional_value(litellm_params, optional_params, "max_context_tokens"), - operating_point=_get_optional_value(litellm_params, optional_params, "operating_point"), + 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, @@ -71,11 +95,11 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> # 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 = { # mutable-ok: see comment above +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 = { # mutable-ok: see comment above +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 index dea326e7fa7..f2f2d3033bd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -32,12 +32,13 @@ 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, + 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 @@ -56,8 +57,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) - from litellm.types.proxy.guardrails.guardrail_hooks.base import ( - GuardrailConfigModel, + from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import ( + NeedlepathGuardrailConfigModel, ) DEFAULT_API_BASE: Final = "https://api.nextmoca.com" @@ -65,7 +66,7 @@ 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-08-r4" +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 @@ -86,6 +87,20 @@ _MAX_CONCURRENT_SELECTIONS: Final = 4 # 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", @@ -105,6 +120,20 @@ _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 @@ -152,8 +181,9 @@ def _write_text_back(content: object, new_text: str) -> object: """ if isinstance(content, str): return new_text - if isinstance(content, list) and is_all_text_parts(content): - return merge_rewritten_text_parts(content, 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 @@ -185,19 +215,18 @@ def _query_for_target(messages: Sequence[Mapping[str, object]], target_idx: int, previous = messages[idx] if previous.get("role") != "assistant": continue - tool_calls = previous.get("tool_calls") - if isinstance(tool_calls, list): - for call in tool_calls: - if not isinstance(call, dict) or not tool_call_id or call.get("id") != tool_call_id: - continue - fn = call.get("function") - intent = _render_tool_intent(fn if isinstance(fn, dict) else MappingProxyType({})) - if intent: - return intent + 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 = previous.get("function_call") - if isinstance(legacy, dict) and fn_name and legacy.get("name") == fn_name: + 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 @@ -217,15 +246,14 @@ def _title_for(messages: Sequence[Mapping[str, object]], target_idx: int) -> str previous = messages[idx] if previous.get("role") != "assistant": continue - tool_calls = previous.get("tool_calls") - if not isinstance(tool_calls, list): - continue - for call in tool_calls: - if not isinstance(call, dict) or call.get("id") != tool_call_id: + 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 = call.get("function") - if isinstance(fn, dict) and isinstance(fn.get("name"), str): - return str(fn["name"])[:120] + 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 @@ -251,6 +279,17 @@ def _safe_int(value: object) -> int | None: 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. @@ -283,7 +322,8 @@ class NeedlepathGuardrail(CustomGuardrail): max_context_tokens: int | None = None, operating_point: str | None = None, guardrail_name: str | None = None, - event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | 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("/") @@ -313,15 +353,13 @@ class NeedlepathGuardrail(CustomGuardrail): default_on=default_on, ) - def _request_headers(self) -> Mapping[str, str]: - # httpx's `headers=` accepts any Mapping (it iterates .items()), so this can - # be a genuine read-only view rather than a dict a caller could mutate. - return MappingProxyType( - { - "Content-Type": "application/json", - "Authorization": f"Bearer {self.needlepath_api_key or ''}", - } - ) + 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. @@ -333,7 +371,7 @@ class NeedlepathGuardrail(CustomGuardrail): # 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 {}) # mutable-ok: see above + 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.""" @@ -357,7 +395,8 @@ class NeedlepathGuardrail(CustomGuardrail): else: return False content: Final = msg.get("content") - if isinstance(content, list) and not is_all_text_parts(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 @@ -375,7 +414,7 @@ class NeedlepathGuardrail(CustomGuardrail): try: return await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.needlepath_api_base}{SELECT_PATH}", - json=payload, + json=dict(payload), headers=self._request_headers(), timeout=_SELECT_TIMEOUT_SECONDS, ) @@ -389,7 +428,7 @@ class NeedlepathGuardrail(CustomGuardrail): response: Final = getattr(e, "response", None) self._decline( "http_status", - { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + { "status_code": getattr(response, "status_code", None), "body": _safe_response_text(response), }, @@ -399,7 +438,7 @@ class NeedlepathGuardrail(CustomGuardrail): # 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)}) # mutable-ok: log detail dict, consumed immediately + self._decline("transport", {"detail": str(e)}) return None async def _selected_text(self, text: str, query: str, title: str | None, source: object, kind: str) -> str | None: @@ -410,6 +449,8 @@ class NeedlepathGuardrail(CustomGuardrail): * 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; @@ -417,6 +458,12 @@ class NeedlepathGuardrail(CustomGuardrail): ``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 @@ -434,8 +481,8 @@ class NeedlepathGuardrail(CustomGuardrail): record["source"] = source payload: Final[dict[str, object]] = { # mutable-ok: see comment above "records": (record,), - "task": {"prompt": query}, # mutable-ok: nested JSON payload dict, see comment above - "budget": { # mutable-ok: nested JSON payload dict, see comment above + "task": {"prompt": query}, + "budget": { "max_context_tokens": self.max_context_tokens, "operating_point": self.operating_point, }, @@ -449,7 +496,7 @@ class NeedlepathGuardrail(CustomGuardrail): if not 200 <= response.status_code < 300: self._decline( "http_status", - { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + { "status_code": response.status_code, "body": _safe_response_text(response), }, @@ -457,20 +504,41 @@ class NeedlepathGuardrail(CustomGuardrail): return None try: - body: Final[object] = response.json() + 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)}) # mutable-ok: log detail dict + self._decline("unreadable_body", {"body": _safe_response_text(response)}) return None - if not isinstance(body, dict): - self._decline("unexpected_shape", {"body": _safe_response_text(response)}) # mutable-ok: log detail dict + body: Final = _as_str_object_dict(raw_body) + if body is None: + self._decline("unexpected_shape", {"body": _safe_response_text(response)}) return None - gate: Final = body.get("gate") - reason: Final = gate.get("reason") if isinstance(gate, dict) else None - if isinstance(reason, str) and reason.startswith(_STANDDOWN_PREFIX): - self._decline("gate_standdown", {"reason": reason}) # mutable-ok: log detail dict + 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")) @@ -478,21 +546,21 @@ class NeedlepathGuardrail(CustomGuardrail): if records_selected == 0 or tokens_after == 0: self._decline( "empty_selection", - { # mutable-ok: log detail dict, consumed by _decline's debug log immediately + { "records_selected": records_selected, "tokens_after": tokens_after, }, ) return None - rendered: Final = body.get("rendered_context") - if not isinstance(rendered, str) or not rendered.strip(): + if not isinstance(rendered, str) or not has_rendered: self._decline("empty_rendered_context", MappingProxyType({})) return None - if len(rendered) >= len(text): + 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. - self._decline("no_reduction", {"before": len(text), "after": len(rendered)}) # mutable-ok: log detail + # 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 @@ -500,25 +568,21 @@ class NeedlepathGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, object], + 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 = inputs.get("structured_messages") - if not isinstance(structured_messages, list) or not structured_messages: + 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. - messages: Final = tuple( - m - for m in structured_messages - if isinstance(m, dict) # pyright: ignore[reportUnnecessaryIsInstance] # runtime boundary check on client JSON - ) + # 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 @@ -578,22 +642,22 @@ class NeedlepathGuardrail(CustomGuardrail): 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 without - # touching the frozen `messages` tuple; only replaced indices differ from - # the original list -- see the identity note below. - selected_messages: Final = list(messages) # mutable-ok: see comment above + # 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 = selected_messages[idx] + 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] = { # mutable-ok: plain AllMessageValues dict, see comment above return + selected_messages[idx] = { **original, "content": _write_text_back(original.get("content"), selection), } @@ -622,7 +686,7 @@ class NeedlepathGuardrail(CustomGuardrail): chars_before, chars_after, ) - self.add_standard_logging_guardrail_information_to_request_data( + 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", @@ -631,13 +695,10 @@ class NeedlepathGuardrail(CustomGuardrail): end_time=end_time, duration=end_time - start_time, ) - return { # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime # mutable-ok: TypedDict, treated as plain dict downstream - **inputs, - "structured_messages": selected_messages, - } + return {**inputs, "structured_messages": selected_messages} # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime @staticmethod - def get_config_model() -> type[GuardrailConfigModel[object]] | None: + def get_config_model() -> type[NeedlepathGuardrailConfigModel] | None: from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import ( NeedlepathGuardrailConfigModel, ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py index e085dce245e..f417aff67a2 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from .base import GuardrailConfigModel @@ -6,6 +6,8 @@ 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.", @@ -29,7 +31,7 @@ class NeedlepathGuardrailOptionalParams(BaseModel): operating_point: str | None = Field( default=None, description=( - "Immutable engine label sent with every request. Defaults to 'np-2026-08-r4'. " + "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." ), diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py index bb805d4900c..11bf9e17dad 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -8,9 +8,12 @@ Tests cover: 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, 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. + 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. """ import asyncio @@ -75,20 +78,26 @@ def _select_response( 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-08-r4", + "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 @@ -138,7 +147,47 @@ def test_init_defaults(): 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-08-r4" + 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(): @@ -176,8 +225,10 @@ async def test_apply_guardrail_replaces_tool_output_with_rendered_context( assert payload["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}' assert payload["budget"] == { "max_context_tokens": 4000, - "operating_point": "np-2026-08-r4", + "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" @@ -381,6 +432,123 @@ async def test_fail_open_selection_not_smaller(guardrail: NeedlepathGuardrail): _assert_byte_identical(result) +# ── 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")) + result = await _run(guardrail, post) + _assert_byte_identical(result) + + @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.""" From 495bd070ad468fd042af5dfaf28a848719de62c8 Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Thu, 29 Oct 2026 12:25:15 -0400 Subject: [PATCH 8/8] fix(needlepath): decline on any client exception; tighten fail-open tests An exception outside httpx's hierarchy raised by the shared client is now a decline like any other transport failure, so it cannot turn into a failed LLM request. Cancellation still propagates. Every decline test now asserts that the exact inputs object comes back, not only equal messages. Adds coverage for /v1/responses: a function_call_output is selected against its call's intent and written back through LiteLLM's Responses translation. Co-Authored-By: Claude Opus 5.5 --- .../guardrail_hooks/needlepath/needlepath.py | 5 ++ .../guardrail_hooks/test_needlepath.py | 90 ++++++++++++++----- 2 files changed, 75 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index f2f2d3033bd..7ced925c380 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -440,6 +440,11 @@ class NeedlepathGuardrail(CustomGuardrail): # 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. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py index 11bf9e17dad..11baa671008 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -14,6 +14,8 @@ Tests cover: - 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 @@ -124,6 +126,16 @@ async def _run(guardrail: NeedlepathGuardrail, post_mock, messages=None) -> Gene ) +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) @@ -331,15 +343,13 @@ async def test_fail_open_zero_tokens_after(guardrail: NeedlepathGuardrail): @pytest.mark.asyncio async def test_fail_open_blank_rendered_context(guardrail: NeedlepathGuardrail): post = AsyncMock(return_value=_select_response(rendered_context=" ")) - result = await _run(guardrail, post) - _assert_byte_identical(result) + 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)) - result = await _run(guardrail, post) - _assert_byte_identical(result) + await _run_expect_unchanged(guardrail, post) @pytest.mark.asyncio @@ -377,30 +387,38 @@ async def test_fail_open_http_error_status(guardrail: NeedlepathGuardrail, statu 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)) - result = await _run(guardrail, post) - _assert_byte_identical(result) + 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)) - result = await _run(guardrail, post) - _assert_byte_identical(result) + 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")) - result = await _run(guardrail, post) - _assert_byte_identical(result) + 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")) - result = await _run(guardrail, post) - _assert_byte_identical(result) + 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 @@ -409,8 +427,7 @@ async def test_fail_open_malformed_json(guardrail: NeedlepathGuardrail): mock.status_code = 200 mock.json.side_effect = ValueError("not json") mock.text = "gateway error" - result = await _run(guardrail, AsyncMock(return_value=mock)) - _assert_byte_identical(result) + await _run_expect_unchanged(guardrail, AsyncMock(return_value=mock)) @pytest.mark.asyncio @@ -420,16 +437,14 @@ async def test_fail_open_unexpected_schema(guardrail: NeedlepathGuardrail): mock.status_code = 200 mock.json.return_value = ["not", "an", "object"] mock.text = "" - result = await _run(guardrail, AsyncMock(return_value=mock)) - _assert_byte_identical(result) + 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")) - result = await _run(guardrail, post) - _assert_byte_identical(result) + await _run_expect_unchanged(guardrail, post) # ── security screening outcomes ─────────────────────────────────────── @@ -545,8 +560,7 @@ async def test_fail_open_blocked(guardrail: NeedlepathGuardrail): 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")) - result = await _run(guardrail, post) - _assert_byte_identical(result) + await _run_expect_unchanged(guardrail, post) @pytest.mark.asyncio @@ -631,3 +645,39 @@ async def test_concurrent_selections_are_bounded(guardrail: NeedlepathGuardrail) 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}