From 1f52df40d61a22925ce2b1db7e38b51fd4d922a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:44:07 -0700 Subject: [PATCH] feat(complexity_router): classify agent loop turns on their tool calls and results Add classifier_context_include_tool_turns. When on, the tool calls and tool results that followed the newest human ask go to the LLM or Jev classifier under the ask being classified, so an every_request loop whose tool output reveals a complex task routes the next turn on that task --- .../prompt_templates/factory.py | 105 +++++++++++++++++- .../complexity_router/complexity_router.py | 57 +++++++++- .../complexity_router/config.py | 18 +++ 3 files changed, 178 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3d865bd1819..113c2beadd4 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -9,7 +9,7 @@ import xml.etree.ElementTree as ET from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from types import MappingProxyType -from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload +from typing import Annotated, Any, Final, Literal, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator @@ -6289,3 +6289,106 @@ def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape ) return tuple(anthropic_tools) + + +ToolExchangeKind: TypeAlias = Literal["tool call", "tool result"] + +_ExchangeValue: Final = TypeVar("_ExchangeValue") + + +class _ExchangeTextBlock(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["text"] + text: str + + +class _ExchangeFunction(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + arguments: str = "" + + +class _ExchangeChatToolCall(BaseModel): + model_config = ConfigDict(frozen=True) + + function: _ExchangeFunction + + +class _ExchangeToolUse(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["tool_use"] + name: str + input: JsonValue = None + + +class _ExchangeToolResult(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["tool_result"] + content: str | tuple[object, ...] = "" + + +_EXCHANGE_PARTS: Final = TypeAdapter(tuple[object, ...]) +_EXCHANGE_TEXT_BLOCK: Final = TypeAdapter(_ExchangeTextBlock) +_EXCHANGE_CHAT_TOOL_CALL: Final = TypeAdapter(_ExchangeChatToolCall) +_EXCHANGE_BLOCK: Final[TypeAdapter[_ExchangeToolUse | _ExchangeToolResult]] = TypeAdapter( + _ExchangeToolUse | _ExchangeToolResult +) + + +def _validated_or_none(adapter: TypeAdapter[_ExchangeValue], value: object) -> _ExchangeValue | None: + try: + return adapter.validate_python(value) + except ValidationError: + return None + + +def _exchange_parts(value: object) -> tuple[object, ...]: + return () if isinstance(value, str) else _validated_or_none(_EXCHANGE_PARTS, value) or () + + +def _exchange_content_text(content: object) -> str: + if isinstance(content, str): + return content.strip() + texts: Final = tuple( + block.text + for part in _exchange_parts(content) + if (block := _validated_or_none(_EXCHANGE_TEXT_BLOCK, part)) is not None + ) + return " ".join(texts).strip() + + +def _block_exchange(part: object) -> tuple[ToolExchangeKind, str] | None: + match _validated_or_none(_EXCHANGE_BLOCK, part): + case _ExchangeToolUse(name=name, input=arguments): + return ("tool call", f"{name}({json.dumps(arguments)})") + case _ExchangeToolResult(content=content): + result: Final = _exchange_content_text(content) + return ("tool result", result) if result else None + case None: + return None + + +def tool_exchange_texts(message: Mapping[str, object]) -> tuple[tuple[ToolExchangeKind, str], ...]: + """The tool calls and tool results one chat-completions or Anthropic Messages turn carries, as text, in order. + + Chat completions puts calls on an assistant turn's `tool_calls` and each result on its own `tool` role, + while Anthropic Messages puts both inside content blocks, `tool_use` on the assistant turn and + `tool_result` on the next user turn, so this reads both places. Responses API input reaches callers + already resolved to the chat-completions shape by `resolve_structured_messages`. + """ + if message.get("role") == "tool": + result: Final = _exchange_content_text(message.get("content")) + return (("tool result", result),) if result else () + block_exchanges: Final = tuple( + exchange for part in _exchange_parts(message.get("content")) if (exchange := _block_exchange(part)) is not None + ) + chat_calls: Final[tuple[tuple[ToolExchangeKind, str], ...]] = tuple( + ("tool call", f"{call.function.name}({call.function.arguments})") + for part in _exchange_parts(message.get("tool_calls")) + if (call := _validated_or_none(_EXCHANGE_CHAT_TOOL_CALL, part)) is not None + ) + return (*block_exchanges, *chat_calls) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a475358e1f0..819e9a2f960 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -983,6 +983,35 @@ def _extract_prior_turns( return tuple(reversed(_turns_within_budget(clamped, budget_chars))) +def _extract_agent_progress( + messages: Sequence[Mapping[str, object]], + budget_chars: int, + per_turn_chars: int | None, + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, +) -> tuple[tuple[str, str], ...]: + """The tool calls and tool results after the current ask, oldest first, as (kind, text). + + An agent loop learns what the ask really involves from its tools, so these are what the classifier + needs on a continuation turn. Only what follows the newest human ask counts, since tool traffic + before it served an earlier ask. Newest first is what the budget keeps, so the latest result always + reaches the classifier. + """ + from litellm.litellm_core_utils.prompt_templates.factory import tool_exchange_texts + + ask_index: Final = _last_human_ask_index(messages, marker_pairs) + if ask_index is None: + return () + newest_first: Final = tuple( + reversed(tuple(chain.from_iterable(tool_exchange_texts(msg) for msg in islice(messages, ask_index + 1, None)))) + ) + clamped: Final = ( + newest_first + if per_turn_chars is None + else tuple((kind, _truncate(text, per_turn_chars)) for kind, text in newest_first) + ) + return tuple(reversed(_turns_within_budget(clamped, budget_chars))) + + def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bool: """Whether a first-turn decision is worth pinning for the rest of the session. @@ -2595,6 +2624,16 @@ class ComplexityRouter(CustomLogger): and len(tuple(islice(_iter_context_turns_newest_first(messages or (), include_assistant, marker_pairs), 2))) > 1 ) + agent_progress: Final = ( + _extract_agent_progress( + messages, + budget_chars=self.config.classifier_context_budget_chars, + per_turn_chars=self.config.classifier_context_per_turn_chars, + marker_pairs=marker_pairs, + ) + if messages and context_enabled and self.config.classifier_context_include_tool_turns and not encrypted_task + else () + ) return self._build_classifier_user_payload( prompt="The delegated task in the following agent_message." if encrypted_task else prompt, system_prompt=self._classifier_caller_constraints(system_prompt, request_kwargs), @@ -2602,6 +2641,7 @@ class ComplexityRouter(CustomLogger): messages=messages, has_prior_conversation=has_prior_conversation, label_roles=include_assistant, + agent_progress=agent_progress, ) async def _classify_with_llm( @@ -2901,8 +2941,9 @@ class ComplexityRouter(CustomLogger): messages: Sequence[Mapping[str, object]] | None = None, has_prior_conversation: bool = False, label_roles: bool = False, + agent_progress: Sequence[tuple[str, str]] = (), ) -> str: - """Build the classifier's user message: caller constraints, prior turns, depth, current ask. + """Build the classifier's user message: caller constraints, prior turns, depth, current ask, agent progress. Everything here is caller-controlled, which is why none of it is interpolated into the system role: that role carries only the operator's rubric, matching how the LLM-as-a-judge guardrail @@ -2922,6 +2963,10 @@ class ComplexityRouter(CustomLogger): Turns are labelled by role only when assistant turns can appear, since otherwise the section header already says whose turns these are and labelling them would reword the prompt of every deployment that never asked for assistant context. + + The agent's progress sits under the ask as part of what is classified, never among the + context-only turns: quoted there, a tool result describing a complex task left the JEV + classifier grading the loop's opening "What's in TASK.md?" as SIMPLE. """ caller_prompt_block: Final = ( ("\nCaller system prompt, quoted as task context:", system_prompt) if system_prompt else () @@ -2946,11 +2991,21 @@ class ComplexityRouter(CustomLogger): else () ) + agent_progress_block: Final = ( + ( + "\nWork the agent has done on this message so far, as tool calls and their results:", + *(f"[{i}] {kind}: {text}" for i, (kind, text) in enumerate(agent_progress, start=1)), + ) + if agent_progress + else () + ) + parts: Final = ( caller_prompt_block, prior_turns_block, trajectory_block, (f"\nClassify this message:\n{prompt}",), + agent_progress_block, ) return "\n".join(part for group in parts for part in group) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 4a2c56daf00..5e8ee3ec3cc 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1358,6 +1358,24 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "spend, for an already-deployed router. Applies to LLM and JEV classification." ), ) + classifier_context_include_tool_turns: bool = Field( + default=False, + description=( + "Send the tool calls and tool results that followed the current ask to the classifier as " + "the work done on it so far, so an agent loop's continuation turn is classified on what " + "its tools turned up rather than on the ask alone: 'What is in TASK.md?' stays SIMPLE " + "until the file it reads describes a complex task, and the turns after that classify " + "as that task. Takes effect on every request that is classified, which with " + "classification_mode='every_request' is every turn of the loop. The block is bounded " + "on its own by classifier_context_budget_chars and classifier_context_per_turn_chars, " + "newest first, but not by classifier_context_window_size, since one ask can run dozens " + "of tool calls; setting the window size to 0 still turns all context off. Tool output " + "such as file contents is sent to the classifier model, which may be a different " + "deployment or provider than the routed completion model. Off by default because " + "enabling it shifts tier decisions, and therefore spend, for an already-deployed " + "router. Applies to LLM and JEV classification." + ), + ) adaptive: bool = Field( default=False,