From 845ee4b2d1b4e8c86a281c182af915dabb4fee71 Mon Sep 17 00:00:00 2001 From: kirankashalkar Date: Mon, 10 Aug 2026 22:59:20 -0400 Subject: [PATCH] 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/test_litellm/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 bbb6d758814..54848d37a01 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -36,6 +36,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, ) @@ -134,6 +137,7 @@ class SupportedGuardrailIntegrations(Enum): HEADROOM = "headroom" COMPRESR = "compresr" STRAIKER = "straiker" + NEEDLEPATH = "needlepath" class Role(Enum): @@ -974,6 +978,7 @@ class LitellmParams( LakeraV2GuardrailConfigModel, HeadroomGuardrailConfigModel, CompresrGuardrailConfigModel, + 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/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py new file mode 100644 index 00000000000..b8621d2c43b --- /dev/null +++ b/tests/test_litellm/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]