diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index b97259c3f6b..13c998fc163 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -72,6 +72,16 @@ DEFAULT_MIN_CHARS_TO_SELECT: Final = 500 # to hold an inbound LLM request behind an optional optimisation. A stall past # this bound is a decline, and the original message is forwarded. _SELECT_TIMEOUT_SECONDS: Final = 30.0 +# A single proxy request can carry arbitrarily many eligible messages, and each +# one becomes an outbound selection call. These two bounds keep a pathological +# request (say, thousands of minimally qualifying tool outputs) from +# monopolising the proxy's shared HTTP pool or burning selection quota: at most +# _MAX_TARGETS_PER_REQUEST messages are selected per request -- the largest +# first, where selection pays off most -- and at most +# _MAX_CONCURRENT_SELECTIONS calls are in flight at once. Messages past the +# cap are forwarded untouched, the same fail-open outcome as any decline. +_MAX_TARGETS_PER_REQUEST: Final = 16 +_MAX_CONCURRENT_SELECTIONS: Final = 4 # The service reports a deliberate no-op through the gate. Any reason under this # prefix means "the engine chose not to select"; the original content is what # the caller should send. @@ -516,31 +526,56 @@ class NeedlepathGuardrail(CustomGuardrail): # One (idx, query) pair per eligible message whose query is non-blank. # Selection is conditioned on a query; without one there is nothing to # select against, so that message is left exactly as it arrived. - selected_pairs: Final = tuple( + eligible_pairs: Final = tuple( (idx, query) for idx in self._select_targets(messages, query_idx) if (query := _query_for_target(messages, idx, fallback_query)).strip() ) - if not selected_pairs: + if not eligible_pairs: verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") return inputs - start_time: Final = time.monotonic() - # One request per message: each carries its own query, and the service - # renders one block per call. Running them concurrently keeps the added - # latency at roughly one round trip rather than one per message. - selections: Final = await asyncio.gather( - *( - self._selected_text( + # Cap the fan-out: the largest messages are kept because that is where + # selection saves the most, and the choice must not depend on message + # order in the request. Everything past the cap is forwarded untouched. + selected_pairs: Final = ( + eligible_pairs + if len(eligible_pairs) <= _MAX_TARGETS_PER_REQUEST + else tuple( + sorted( + eligible_pairs, + key=lambda pair: len(content_to_text(messages[pair[0]].get("content"))), + reverse=True, + )[:_MAX_TARGETS_PER_REQUEST] + ) + ) + if len(selected_pairs) < len(eligible_pairs): + verbose_proxy_logger.debug( + "Needlepath: %d eligible messages, selecting only the %d largest", + len(eligible_pairs), + _MAX_TARGETS_PER_REQUEST, + ) + + semaphore: Final = asyncio.Semaphore(_MAX_CONCURRENT_SELECTIONS) + + async def _bounded_selected_text(idx: int, query: str) -> str | None: + # The semaphore bounds how many selection calls this one proxy + # request holds open at a time; see _MAX_CONCURRENT_SELECTIONS. + async with semaphore: + return await self._selected_text( text=content_to_text(messages[idx].get("content")), query=query, title=_title_for(messages, idx), source=messages[idx].get("tool_call_id"), kind=_record_kind(messages[idx].get("role")), ) - for idx, query in selected_pairs - ) - ) + + start_time: Final = time.monotonic() + # One request per message: each carries its own query, and the service + # renders one block per call. Running them concurrently (up to the + # semaphore's bound) keeps the added latency near one round trip + # rather than one per message. + selections: Final = await asyncio.gather(*(_bounded_selected_text(idx, query) for idx, query in selected_pairs)) end_time: Final = time.monotonic() # Needs in-place index assignment below to build the edited copy without diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py index b8621d2c43b..1a019595042 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -13,6 +13,7 @@ Tests cover: the messages come back byte-identical. """ +import asyncio import copy from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +21,8 @@ import httpx import pytest from litellm.proxy.guardrails.guardrail_hooks.needlepath.needlepath import ( + _MAX_CONCURRENT_SELECTIONS, + _MAX_TARGETS_PER_REQUEST, DEFAULT_MAX_CONTEXT_TOKENS, DEFAULT_MIN_CHARS_TO_SELECT, DEFAULT_OPERATING_POINT, @@ -395,3 +398,68 @@ async def test_one_message_declining_does_not_block_another(): assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK # The system message stood down, so it is byte-identical. assert out[0] == AGENT_MESSAGES[0] + + +# ── fan-out bounds ──────────────────────────────────────────────────── + + +def _many_tool_messages(count: int) -> list: + """A user question plus `count` eligible tool outputs of strictly increasing size.""" + messages = [{"role": "user", "content": USER_QUESTION}] + for i in range(count): + messages.append( + { + "role": "tool", + "tool_call_id": f"call_{i}", + # All above the min-chars threshold; index i is the (i+1)-th smallest. + "content": "a" * (DEFAULT_MIN_CHARS_TO_SELECT + 20 + i), + } + ) + return messages + + +@pytest.mark.asyncio +async def test_target_cap_selects_only_the_largest_messages(guardrail: NeedlepathGuardrail): + """Past the per-request cap, only the largest messages are selected. + + The smallest overflow messages come back byte-identical and no service + call is made for them. + """ + overflow = 4 + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST + overflow) + + post = AsyncMock(return_value=_select_response()) + result = await _run(guardrail, post, messages=messages) + + assert post.await_count == _MAX_TARGETS_PER_REQUEST + out = result["structured_messages"] + # Tool messages start at index 1 and grow with index: the first `overflow` + # are the smallest, so they are the ones left untouched. + for idx in range(1, 1 + overflow): + assert out[idx] == messages[idx] + for idx in range(1 + overflow, len(messages)): + assert out[idx]["content"] == SELECTED_BLOCK + + +@pytest.mark.asyncio +async def test_concurrent_selections_are_bounded(guardrail: NeedlepathGuardrail): + """No more than _MAX_CONCURRENT_SELECTIONS service calls are in flight at once.""" + in_flight = 0 + high_water = 0 + + async def _tracking_post(*args, **kwargs): + nonlocal in_flight, high_water + in_flight += 1 + high_water = max(high_water, in_flight) + # Yield twice so every scheduled call gets a chance to start before + # this one finishes; without the semaphore all of them would overlap. + await asyncio.sleep(0) + await asyncio.sleep(0) + in_flight -= 1 + return _select_response() + + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST) + result = await _run(guardrail, _tracking_post, messages=messages) + + assert high_water <= _MAX_CONCURRENT_SELECTIONS + assert all(m["content"] == SELECTED_BLOCK for m in result["structured_messages"][1:])