needlepath-litellm

Needlepath context selection as a LiteLLM proxy CustomGuardrail. Every client behind the proxy gets context selection with no client-side change at all — and when the selection service is slow or down, every one of them still gets their completion.

pip install needlepath-litellm     # pulls litellm[proxy]

Configure

# config.yaml
guardrails:
  - guardrail_name: "needlepath"
    litellm_params:
      guardrail: needlepath_litellm.NeedlepathGuardrail
      mode: "pre_call"
      default_on: true
      operating_point: "np-2026-08-r4"     # required; immutable label
      history_max_tokens: 8000

Set NEEDLEPATH_API_KEY in the proxy's environment (or pass api_key: above, which puts a secret in your config file — prefer the environment).

Everything under litellm_params other than guardrail, mode and default_on is forwarded to the constructor, so every knob below is settable from config with no code.

With default_on: false, clients opt in per request:

{"model": "gpt-4o", "messages": [...], "guardrails": ["needlepath"]}

What it does

Before the request is routed, tool replies in the message array are selected against the current user turn. A reply the current step needs is replaced by the part of it that answers; a reply it does not need collapses to a short placeholder. The system message and assistant messages carrying tool_calls are never touched. There is no "protect the last N messages" option — see below.

A message is never removed, only rewritten. Dropping an assistant message that carries tool_calls while keeping its tool reply — or the reverse — is rejected by every provider. Rewriting content in place makes that impossible by construction, at the cost of the tokens in the message envelopes.

What the engine returns is what gets applied. A preserve_recent option existed in an earlier version and was removed (see CHANGELOG.md): a client-side "never touch the last N messages" rule is a selection decision made locally, and it made .stats misreport what was actually sent — the engine's own tokens_after describes its full selection, not the narrower one this guardrail would have installed. If you pass preserve_recent today the constructor raises, naming the migration.

Two stats objects, two truths

guardrail.stats reports what this guardrail actually did to data["messages"] — the same rewrite_tokens_before/rewrite_tokens_after numbers the per-request metadata blob already carries, aggregated. guardrail.engine_stats reports what the engine's raw response claimed, unmodified — the same tokens_before/tokens_after the metadata blob's needlepath key carries. They differ because the engine cannot see the placeholder text substituted for a declined reply. Use .stats for anything that has to be true of the request; use .engine_stats to reconcile against what the service measured or billed. In shadow mode .stats.tokens_saved is always 0 (nothing is ever applied); .engine_stats.tokens_saved carries the engine's prediction.

Hook coverage — read this before you deploy

Mutation on LiteLLM is proxy-only in practice, and exactly one hook can replace a request: async_pre_call_hook. Its returned dict becomes the request data. Everything else on the base class is observability, rejection, or post-processing.

Which routes this guardrail acts on

mode: "pre_call" fires on a long list of routes, but only one of them carries an OpenAI-format messages array we can safely rewrite. The rest are left strictly alone, with reason: "unsupported" recorded and no call made:

Skipping is not an oversight. Rewriting an Anthropic content-block body or a Responses-API input as if it were a chat message array would corrupt the request; the safe move is to do nothing and say so.

If your traffic is mostly /v1/messages or /responses, this guardrail saves you nothing today. Route-specific support is additive work that does not exist yet — which is a better thing to read here than to discover from a savings number that turns out to be zero.

Other caveats worth knowing

• apply_guardrail is never overridden here, and you must not add it. If a subclass defines it, LiteLLM routes the call through its unified_guardrail singleton and async_pre_call_hook never runs. Silent, and total.
• data is mutated in place and returned. Returning a fresh dict works on the proxy but loses non-messages keys on LiteLLM's SDK path, where only result["messages"] is copied back.
• The SDK path is not purely observability. A CustomGuardrail in litellm.callbacks will fire in an SDK process on completion/acompletion when the caller passes guardrails=[...], because CustomGuardrail bridges the SDK-side async_pre_call_deployment_hook to async_pre_call_hook. This guardrail is safe there for the same reason: in-place mutation.
• Metadata goes in metadata / litellm_metadata, never a new top-level key. Unknown top-level keys are forwarded toward providers on some routes and rejected. A debug marker must not become an outage.
• This hook never raises. An exception here is raised to the client, which would turn a selection-service hiccup into a failed LLM request for everyone behind the proxy. The whole body is wrapped; worst case, the request goes through unchanged.
• Streaming egress is untouched. This guardrail is ingress-only. If you ever add egress control, use async_post_call_streaming_iterator_hook, not async_post_call_streaming_hook — the latter computes its payload only for ModelResponse-typed chunks and is a silent no-op on Anthropic and pass-through SSE, which emit raw str/bytes.

Observability

Every request the guardrail touches carries a metadata-only blob:

{"metadata": {"needlepath": {
  "applied": true, "reason": "ok", "request_id": "np-…",
  "operating_point": "np-2026-08-r4", "tokens_saved": 6800,
  "rewrite_tokens_before": 7100, "rewrite_tokens_after": 900,
  "gate_reason": "engage:needle"}}}

rewrite_tokens_* is what this adapter measured; tokens_* is the service's own accounting. They are reported separately and never blended.

guardrail.stats.as_dict() aggregates the rewrite_tokens_* numbers per process — what was actually applied. guardrail.engine_stats.as_dict() aggregates the tokens_* numbers — the engine's raw claim. See "Two stats objects, two truths" above.

Nothing derived from message text is in the blob. One consequence is non-obvious: the service's selection_error is built as f"{type(exc).__name__}: {exc}" and can carry fragments of the request, so only its type is reported.

Lifecycle

A proxy builds its guardrails once at startup, so this rarely matters. If you construct guardrails repeatedly, await guardrail.aclose() releases the client's connection pool; a client you passed in with client= is never closed for you.

Shadow mode

shadow: true makes every call and changes nothing. It is the day-one deployment: put it in front of real traffic, read the numbers, decide afterwards. Check guardrail.engine_stats for what the engine predicts a live run would save — guardrail.stats correctly shows 0 saved, since nothing was applied.

Failure behaviour

There is no configuration in which a Needlepath failure becomes an LLM failure.

Tested against

litellm==1.94.1 (litellm[proxy]), capped at <2.0.0. CI runs against the newest 1.x minor; see .github/workflows/sdk-python.yml.