mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(router): preserve held tiers for OpenCode subagent results (#51176)
* fix(router): preserve held tiers for OpenCode task results * fix(router): handle OpenCode subagent completion envelopes --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
1d97931c84
commit
0a4ccf9554
3 changed files with 138 additions and 4 deletions
|
|
@ -690,7 +690,7 @@ generic behavior; use the real client when checking Claude Code routing.
|
|||
|
||||
Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead
|
||||
|
||||
By default the router strips complete `<system-reminder>` blocks. For requests with a Codex user agent, it also strips complete `<environment_context>`, `<recommended_plugins>`, `<user_instructions>`, and `<environments_instructions>` blocks, plus repository instructions from the fixed heading prefix `# AGENTS.md instructions for ` through `</INSTRUCTIONS>`, regardless of the repository path. Other clients keep those tags and their contents
|
||||
By default the router strips complete `<system-reminder>` blocks. For requests with a Codex user agent, it also strips complete `<environment_context>`, `<recommended_plugins>`, `<user_instructions>`, and `<environments_instructions>` blocks, plus repository instructions from the fixed heading prefix `# AGENTS.md instructions for ` through `</INSTRUCTIONS>`, regardless of the repository path. For requests with an `opencode/` user agent, it also strips complete `<task_result>` blocks and `<subagent ...>` blocks through `</subagent>`, attributes included, which OpenCode uses to return subagent output as a user turn. Other clients keep those tags and their contents
|
||||
|
||||
The proxy records the incoming user agent in request metadata. SDK callers can supply `metadata.user_agent` (or `litellm_metadata.user_agent` on Responses requests), or configure `reminder_markers` explicitly when their client identity is unavailable
|
||||
|
||||
|
|
@ -713,7 +713,7 @@ model_list:
|
|||
close: "[[SUBAGENT_CONTEXT_END]]"
|
||||
```
|
||||
|
||||
Setting `reminder_markers` replaces all built-in pairs, including the Codex heading pair, so include every default your harness still needs. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten
|
||||
Setting `reminder_markers` replaces all built-in pairs, including the Codex and OpenCode pairs, so include every default your harness still needs. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten
|
||||
|
||||
### Code Detection
|
||||
|
||||
|
|
|
|||
|
|
@ -520,6 +520,11 @@ _CODEX_REMINDER_MARKERS: Final = _DEFAULT_REMINDER_MARKERS + (
|
|||
("<environments_instructions>", "</environments_instructions>"),
|
||||
("# agents.md instructions for ", "</instructions>"),
|
||||
)
|
||||
_OPENCODE_REMINDER_MARKERS: Final = _DEFAULT_REMINDER_MARKERS + (
|
||||
("<task_result>", "</task_result>"),
|
||||
("<subagent ", "</subagent>"),
|
||||
)
|
||||
_OPENCODE_USER_AGENT_PREFIX: Final = "opencode/"
|
||||
|
||||
_TRUNCATION_MARKER: Final = "..."
|
||||
_TRUNCATION_HEAD_FRACTION: Final = 0.3
|
||||
|
|
@ -530,6 +535,10 @@ _CLASSIFIER_CIRCUIT_OPEN_SIGNAL: Final = "classifier-circuit-open"
|
|||
_CJK_CHARACTER: Final = re.compile("[-ヿㇰ-ㇿ㐀-䶿一-鿿豈-ｦ-ﾝ\U00020000-\U0003ffff]")
|
||||
|
||||
|
||||
def _is_opencode_user_agent(user_agent: str) -> bool:
|
||||
return user_agent.lower().startswith(_OPENCODE_USER_AGENT_PREFIX)
|
||||
|
||||
|
||||
def _message_text(content: object) -> str:
|
||||
"""Flatten message content to plain text, joining multi-part text blocks.
|
||||
|
||||
|
|
@ -4418,12 +4427,20 @@ class ComplexityRouter(CustomLogger):
|
|||
def _reminder_markers_for_request(self, request_kwargs: Mapping[str, object]) -> tuple[tuple[str, str], ...]:
|
||||
if self.config.reminder_markers is not None:
|
||||
return self._reminder_markers
|
||||
return _CODEX_REMINDER_MARKERS if self._is_codex_request(request_kwargs) else _DEFAULT_REMINDER_MARKERS
|
||||
if self._is_codex_request(request_kwargs):
|
||||
return _CODEX_REMINDER_MARKERS
|
||||
if self._request_user_agent_matches(request_kwargs, _is_opencode_user_agent):
|
||||
return _OPENCODE_REMINDER_MARKERS
|
||||
return _DEFAULT_REMINDER_MARKERS
|
||||
|
||||
@staticmethod
|
||||
def _is_codex_request(request_kwargs: Mapping[str, object]) -> bool:
|
||||
return ComplexityRouter._request_user_agent_matches(request_kwargs, is_codex_user_agent)
|
||||
|
||||
@staticmethod
|
||||
def _request_user_agent_matches(request_kwargs: Mapping[str, object], matches: Callable[[str], bool]) -> bool:
|
||||
return any(
|
||||
is_codex_user_agent(user_agent)
|
||||
matches(user_agent)
|
||||
for metadata_key in ("litellm_metadata", "metadata")
|
||||
if isinstance(metadata := request_kwargs.get(metadata_key), Mapping)
|
||||
if isinstance(user_agent := metadata.get("user_agent"), str)
|
||||
|
|
|
|||
|
|
@ -10250,6 +10250,12 @@ _CODEX_ENVELOPES: Final = (
|
|||
"<environments_instructions>LITELLM ESCALATE environment</environments_instructions>",
|
||||
"# AGENTS.md instructions for /repo with spaces/中文\n<INSTRUCTIONS>LITELLM ESCALATE instructions</INSTRUCTIONS>",
|
||||
)
|
||||
_OPENCODE_USER_AGENT: Final = "opencode/latest/2.0.29/cli"
|
||||
_OPENCODE_SUBAGENT_RESULTS: Final = (
|
||||
"<task_result>\nLITELLM ESCALATE subagent found the race in cache.py</task_result>",
|
||||
'<subagent sessionID="ses_7f3a9c1e" state="completed" description="Explore cache.py">\n'
|
||||
"LITELLM ESCALATE subagent found the race in cache.py\n</subagent>",
|
||||
)
|
||||
|
||||
|
||||
class TestContextAwareClassifier:
|
||||
|
|
@ -10490,6 +10496,117 @@ class TestContextAwareClassifier:
|
|||
completion.assert_awaited_once()
|
||||
assert completion.call_args.kwargs["messages"][1]["content"].strip() == f"Classify this message:\n{envelope}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("subagent_result", _OPENCODE_SUBAGENT_RESULTS)
|
||||
@pytest.mark.parametrize("user_agent", (_OPENCODE_USER_AGENT, "OpenCode/latest/2.0.29/cli"))
|
||||
async def test_opencode_subagent_result_turn_keeps_held_tier_until_a_new_ask(
|
||||
self, user_agent: str, subagent_result: str
|
||||
) -> None:
|
||||
completion: Final = AsyncMock(
|
||||
side_effect=[_llm_response('{"tier":"SIMPLE"}'), _llm_response('{"tier":"COMPLEX"}')]
|
||||
)
|
||||
router: Final = ComplexityRouter(
|
||||
model_name="router",
|
||||
ctx=RoutingContext(MagicMock(acompletion=completion, **{"health.return_value.store": DualCache()})),
|
||||
complexity_router_config={
|
||||
"tiers": {"SIMPLE": "simple-model", "COMPLEX": "task-model"},
|
||||
"classifier_type": "llm",
|
||||
"classifier_llm_config": {"model": "classifier-model"},
|
||||
"classification_mode": "user_turn",
|
||||
},
|
||||
)
|
||||
metadata: Final = {"user_agent": user_agent, "session_id": "opencode-task-session"}
|
||||
asked: Final = [{"role": "user", "content": "Find the race in cache.py"}]
|
||||
task_returned: Final = [
|
||||
*asked,
|
||||
{"role": "assistant", "content": "Delegating to a subagent"},
|
||||
{"role": "user", "content": subagent_result},
|
||||
]
|
||||
new_ask: Final = [
|
||||
*task_returned,
|
||||
{"role": "assistant", "content": "Found it"},
|
||||
{"role": "user", "content": "Now design cache invalidation"},
|
||||
{"role": "user", "content": subagent_result},
|
||||
]
|
||||
original: Final = deepcopy((asked, task_returned, new_ask))
|
||||
|
||||
results: Final = [
|
||||
await router.async_pre_routing_hook(
|
||||
model="router", request_kwargs={"metadata": metadata}, messages=messages
|
||||
)
|
||||
for messages in (asked, task_returned, new_ask)
|
||||
]
|
||||
|
||||
assert [(result.model, result.routing_decision["cause"]) for result in results] == [
|
||||
("simple-model", "llm_classifier"),
|
||||
("simple-model", "user_turn_continuation"),
|
||||
("task-model", "llm_classifier"),
|
||||
]
|
||||
classified: Final = [call.kwargs["messages"][1]["content"] for call in completion.await_args_list]
|
||||
assert [payload.rsplit("Classify this message:\n", 1)[1] for payload in classified] == [
|
||||
"Find the race in cache.py",
|
||||
"Now design cache invalidation",
|
||||
]
|
||||
assert all("subagent found the race" not in payload for payload in classified)
|
||||
assert tuple(result.messages for result in results) == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("subagent_result", _OPENCODE_SUBAGENT_RESULTS)
|
||||
@pytest.mark.parametrize(
|
||||
"request_kwargs,content,custom_markers,expected_ask",
|
||||
(
|
||||
({"metadata": {"user_agent": _OPENCODE_USER_AGENT}}, "{result}\nAdd retries", False, "Add retries"),
|
||||
({"litellm_metadata": {"user_agent": _OPENCODE_USER_AGENT}}, "Add retries {result}", False, "Add retries"),
|
||||
(
|
||||
{"metadata": {"user_agent": _OPENCODE_USER_AGENT}},
|
||||
"{unclosed}\nAdd retries",
|
||||
False,
|
||||
"{unclosed}\nAdd retries",
|
||||
),
|
||||
({"metadata": {"user_agent": _OPENCODE_USER_AGENT}}, "{result}", True, "{result}"),
|
||||
({}, "{result}", False, "{result}"),
|
||||
({"metadata": {"user_agent": "curl/8.7.1"}}, "{result}", False, "{result}"),
|
||||
({"metadata": {"user_agent": "codex-tui"}}, "{result}", False, "{result}"),
|
||||
({"metadata": {"user_agent": "opencode_go/1.0"}}, "{result}", False, "{result}"),
|
||||
({"metadata": {"user_agent": "opencodex/1.0"}}, "{result}", False, "{result}"),
|
||||
),
|
||||
)
|
||||
async def test_classifier_strips_only_complete_opencode_subagent_results(
|
||||
self,
|
||||
request_kwargs: dict[str, object],
|
||||
content: str,
|
||||
custom_markers: bool,
|
||||
expected_ask: str,
|
||||
subagent_result: str,
|
||||
) -> None:
|
||||
unclosed: Final = subagent_result[: subagent_result.rindex("</")].rstrip()
|
||||
message: Final = content.format(result=subagent_result, unclosed=unclosed)
|
||||
completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}'))
|
||||
router: Final = ComplexityRouter(
|
||||
model_name="router",
|
||||
ctx=RoutingContext(MagicMock(acompletion=completion)),
|
||||
complexity_router_config={
|
||||
"tiers": {"COMPLEX": "task-model"},
|
||||
"default_model": "fallback-model",
|
||||
"classifier_type": "llm",
|
||||
"classifier_llm_config": {"model": "classifier-model"},
|
||||
"escalation_keywords": [],
|
||||
**({"reminder_markers": [{"open": "<custom>", "close": "</custom>"}]} if custom_markers else {}),
|
||||
},
|
||||
)
|
||||
messages: Final = [{"role": "user", "content": message}]
|
||||
|
||||
response: Final = await router.async_pre_routing_hook(
|
||||
model="router", request_kwargs=request_kwargs, messages=messages
|
||||
)
|
||||
|
||||
assert response is not None and response.model == "task-model"
|
||||
assert response.messages == [{"role": "user", "content": message}]
|
||||
completion.assert_awaited_once()
|
||||
classified: Final = completion.call_args.kwargs["messages"][1]["content"]
|
||||
expected: Final = expected_ask.format(result=subagent_result, unclosed=unclosed)
|
||||
assert classified.strip() == f"Classify this message:\n{expected}"
|
||||
|
||||
@pytest.mark.parametrize("envelope", _CODEX_ENVELOPES)
|
||||
def test_codex_envelopes_preserve_delegated_task_and_prior_context(self, envelope: str) -> None:
|
||||
from litellm.router_strategy.complexity_router.complexity_router import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue