scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

feat(guardrails): add needlepath context selection guardrail

Browse source 
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.
This commit is contained in:
kirankashalkar 2026-08-10 22:59:20 -04:00
parent 79d412efc2
commit 845ee4b2d1
5 changed files with 1105 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
litellm/proxy/guardrails/guardrail_hooks/needlepath/__init__.py Normal file
View file

@ -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,
}

578
litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py Normal file
View file

@ -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 "<undecodable response body>"
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

5
litellm/types/guardrails.py
View file

@ -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,

54
litellm/types/proxy/guardrails/guardrail_hooks/needlepath.py Normal file
View file

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

397
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py Normal file
View file

@ -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 = "<html>gateway error</html>"
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]