mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
460ceab075 | ||
|
|
fb6606279f | ||
|
|
efb698a4d0 | ||
|
|
cc2b3351b8 | ||
|
|
af7769507a | ||
|
|
d93a99a355 |
@@ -429,3 +429,67 @@ def is_known_openai_bare_model(model_name: str) -> bool:
|
||||
return False
|
||||
entry = litellm.model_cost.get(name)
|
||||
return bool(entry and entry.get("litellm_provider") == "openai")
|
||||
|
||||
|
||||
def is_claude_model(model_name: str) -> bool:
|
||||
return "claude" in (model_name or "").strip().lower()
|
||||
|
||||
|
||||
def is_bedrock_route(model_name: str) -> bool:
|
||||
"""Whether ``model_name`` resolves to an AWS Bedrock route.
|
||||
|
||||
Matches the ``bedrock/...`` LiteLLM route prefix and bare Bedrock model ids
|
||||
(``[region.]anthropic.claude-...``).
|
||||
"""
|
||||
name = (model_name or "").strip().lower()
|
||||
return name.startswith("bedrock/") or "anthropic." in name
|
||||
|
||||
|
||||
def _prompt_cache_name_candidates(model_name: str) -> list[str]:
|
||||
"""Candidate LiteLLM model-map keys for ``model_name``, most→least specific.
|
||||
|
||||
LiteLLM keys the same model under several names (``bedrock/global.anthropic.
|
||||
claude-opus-4-1``, ``anthropic.claude-opus-4-1``, ``claude-opus-4-1``) and not
|
||||
every provider/region-prefixed variant is present for every model. Strip the
|
||||
LiteLLM route prefix, then leading dotted segments (region, then provider) so
|
||||
a prefixed name still resolves to a bare key.
|
||||
"""
|
||||
name = (model_name or "").strip().lower()
|
||||
for prefix in ("litellm/", "bedrock/"):
|
||||
if name.startswith(prefix):
|
||||
name = name[len(prefix) :]
|
||||
break
|
||||
candidates = [name]
|
||||
rest = name
|
||||
while "." in rest:
|
||||
rest = rest.split(".", 1)[1]
|
||||
candidates.append(rest)
|
||||
return candidates
|
||||
|
||||
|
||||
def bedrock_route_supports_prompt_caching(model_name: str) -> bool:
|
||||
"""Whether LiteLLM can confirm this Bedrock model supports prompt caching.
|
||||
|
||||
Bedrock's Converse API rejects unknown request fields outright
|
||||
(``ValidationException: cache_control_injection_points: Extra inputs are
|
||||
not permitted``), and LiteLLM only consumes the cache marker for models its
|
||||
(statically bundled) model map recognises as cache-capable. For a Bedrock
|
||||
model missing from that map — a just-released model, or any model when the
|
||||
remote model-map refresh fails and a stale local copy is used — the marker
|
||||
would pass straight through and fail every call, so callers must withhold
|
||||
it unless support is confirmed here.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
checker = getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None)
|
||||
for cand in _prompt_cache_name_candidates(model_name):
|
||||
if checker is not None:
|
||||
# supports_prompt_caching raises for models missing from the map;
|
||||
# keep checking the remaining name candidates.
|
||||
with contextlib.suppress(Exception):
|
||||
if checker(cand):
|
||||
return True
|
||||
entry = litellm.model_cost.get(cand)
|
||||
if entry and entry.get("supports_prompt_caching"):
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -40,6 +40,10 @@ class LlmSettings(BaseSettings):
|
||||
default=False,
|
||||
alias="STRIX_FORCE_REQUIRED_TOOL_CHOICE",
|
||||
)
|
||||
prompt_cache: bool = Field(
|
||||
default=True,
|
||||
alias="STRIX_PROMPT_CACHE",
|
||||
)
|
||||
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
|
||||
|
||||
|
||||
|
||||
@@ -10,6 +10,9 @@ from openai.types.shared import Reasoning
|
||||
|
||||
from strix.config.models import (
|
||||
DEFAULT_MODEL_RETRY,
|
||||
bedrock_route_supports_prompt_caching,
|
||||
is_bedrock_route,
|
||||
is_claude_model,
|
||||
is_known_openai_bare_model,
|
||||
model_supports_reasoning,
|
||||
request_timeout_extra_args,
|
||||
@@ -128,6 +131,7 @@ def make_model_settings(
|
||||
model_name: str,
|
||||
force_required_tool_choice: bool = False,
|
||||
request_timeout: float | None = None,
|
||||
prompt_cache: bool = True,
|
||||
) -> ModelSettings:
|
||||
model_settings = ModelSettings(
|
||||
parallel_tool_calls=False,
|
||||
@@ -145,9 +149,95 @@ def make_model_settings(
|
||||
)
|
||||
if force_required_tool_choice and _accepts_required_tool_choice(model_name):
|
||||
model_settings = model_settings.resolve(ModelSettings(tool_choice="required"))
|
||||
|
||||
cache_extra_args = _prompt_cache_extra_args(model_name) if prompt_cache else None
|
||||
if cache_extra_args:
|
||||
# Merge into any existing extra_args (e.g. the request timeout) rather
|
||||
# than relying on resolve()'s dict-merge semantics, so it is obvious at
|
||||
# the call site that unrelated LiteLLM options are preserved.
|
||||
model_settings = model_settings.resolve(
|
||||
ModelSettings(
|
||||
extra_args={**(model_settings.extra_args or {}), **cache_extra_args},
|
||||
),
|
||||
)
|
||||
return model_settings
|
||||
|
||||
|
||||
def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
|
||||
"""LiteLLM ``extra_args`` that enable Anthropic/Bedrock prompt caching.
|
||||
|
||||
A Strix scan is a long, multi-turn agentic loop that re-sends a large,
|
||||
STABLE prefix every turn — the system prompt plus the tool schemas — AND an
|
||||
append-only conversation transcript that only grows. Without caching
|
||||
breakpoints the whole request is re-tokenised and billed at the full input
|
||||
rate on every turn; on Claude that is the single biggest lever on scan cost
|
||||
(measured: ``cache-read 0% -> ~66%`` on a real scan once these points are set).
|
||||
|
||||
This mirrors the caching policy of production agent harnesses (e.g.
|
||||
anomalyco/opencode's ``cache-policy``): cache the tool schemas, the system
|
||||
prompt, and the latest conversation message, capped at Anthropic's 4
|
||||
breakpoints. We express it through LiteLLM's ``cache_control_injection_points``
|
||||
— the ``AnthropicCacheControlHook`` fires on that kwarg and emits the
|
||||
provider-appropriate breakpoint (Anthropic ``cache_control``; Bedrock
|
||||
Converse ``cachePoint``). ``LitellmModel`` forwards ``ModelSettings.extra_args``
|
||||
straight into ``litellm.acompletion()``, so passing the points there is all
|
||||
that is required; this is the LiteLLM-config-layer approach the Agents SDK
|
||||
maintainer prescribed over a native ``ModelSettings`` caching flag
|
||||
(openai/openai-agents-python#3008 / #3009).
|
||||
|
||||
Returns ``None`` (a strict no-op — the hook never fires) for every route
|
||||
that would not benefit or could break:
|
||||
- non-Claude models, and
|
||||
- Bedrock Claude routes LiteLLM can't confirm as cache-capable. Bedrock's
|
||||
Converse API rejects unknown request fields outright
|
||||
(``ValidationException: cache_control_injection_points: Extra inputs are
|
||||
not permitted``) and LiteLLM only consumes the marker for models its
|
||||
model map recognises; an unmapped Bedrock model would pass the marker
|
||||
straight through and crash the first call. Only Bedrock hard-rejects, so
|
||||
only Bedrock is guarded — gating Anthropic-native/Vertex/OpenRouter on
|
||||
confirmed support would needlessly disable caching for capable models.
|
||||
"""
|
||||
if not is_claude_model(model_name):
|
||||
return None
|
||||
if is_bedrock_route(model_name) and not bedrock_route_supports_prompt_caching(model_name):
|
||||
return None
|
||||
|
||||
points = _prompt_cache_injection_points(model_name)
|
||||
return {"cache_control_injection_points": points}
|
||||
|
||||
|
||||
def _prompt_cache_injection_points(model_name: str) -> list[dict[str, Any]]:
|
||||
"""Cache breakpoints for a Claude route (system + tools + latest message).
|
||||
|
||||
At most 3 of Anthropic's 4 allowed breakpoints, leaving headroom:
|
||||
- system prompt (``role: system``) — the largest repeated span.
|
||||
- tool schemas (``tool_config``) — Bedrock Converse ONLY. LiteLLM's
|
||||
``tool_config`` location is implemented solely by the Bedrock Converse
|
||||
transform (which appends a ``cachePoint`` to the tool list); on any other
|
||||
route it is not consumed and would leak onto the wire as an unknown
|
||||
top-level ``cache_control_injection_points`` field. It is also redundant
|
||||
elsewhere: Anthropic orders tools BEFORE the system prompt, so the system
|
||||
breakpoint already caches the tool schemas in the shared prefix.
|
||||
- latest message (``index: -1``) — a ROLLING breakpoint on the last
|
||||
message. A scan transcript is append-only (prior turns are immutable, each
|
||||
turn just appends new assistant/tool messages), so without it the growing
|
||||
body is re-sent at full input price every turn and cache-read decays as a
|
||||
denominator effect even though the prefix keeps hitting. Re-caching the
|
||||
whole prefix-so-far each turn keeps cache-read high on long scans. (This
|
||||
is the Strix analogue of opencode's ``latest-user-message``; ``index: -1``
|
||||
tracks the true tail because Strix appends tool-role, not user-role,
|
||||
messages each turn.)
|
||||
|
||||
Unrecognised locations degrade gracefully on older LiteLLM — they are simply
|
||||
not injected (no error).
|
||||
"""
|
||||
points: list[dict[str, Any]] = [{"location": "message", "role": "system"}]
|
||||
if is_bedrock_route(model_name):
|
||||
points.append({"location": "tool_config"})
|
||||
points.append({"location": "message", "index": -1})
|
||||
return points
|
||||
|
||||
|
||||
def child_initial_input(
|
||||
*,
|
||||
name: str,
|
||||
|
||||
@@ -236,6 +236,7 @@ async def run_strix_scan(
|
||||
model_name=resolved_model,
|
||||
force_required_tool_choice=settings.llm.force_required_tool_choice,
|
||||
request_timeout=settings.llm.timeout,
|
||||
prompt_cache=settings.llm.prompt_cache,
|
||||
)
|
||||
run_config = RunConfig(
|
||||
model=resolved_model,
|
||||
|
||||
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
from itertools import pairwise
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
from strix.core.inputs import build_root_task, child_initial_input, make_model_settings
|
||||
@@ -55,6 +56,138 @@ def test_child_initial_input_no_consecutive_same_role(parent_history: list[Any])
|
||||
assert all(prev != nxt for prev, nxt in pairwise(roles))
|
||||
|
||||
|
||||
def _cache_points(model_name: str) -> Any:
|
||||
extra = make_model_settings(None, model_name=model_name).extra_args or {}
|
||||
return extra.get("cache_control_injection_points")
|
||||
|
||||
|
||||
def test_make_model_settings_enables_prompt_cache_for_bedrock_claude() -> None:
|
||||
# Bedrock Converse is the only route whose transform consumes the
|
||||
# ``tool_config`` location, so it gets all three breakpoints.
|
||||
assert _cache_points("bedrock/global.anthropic.claude-opus-4-8") == [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "tool_config"},
|
||||
{"location": "message", "index": -1},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
"openrouter/anthropic/claude-3.5-sonnet",
|
||||
"vertex_ai/claude-sonnet-4-5",
|
||||
],
|
||||
)
|
||||
def test_make_model_settings_enables_prompt_cache_for_non_bedrock_claude(model_name: str) -> None:
|
||||
# Non-Bedrock routes get system + tail only: LiteLLM implements the
|
||||
# ``tool_config`` location solely in the Bedrock Converse transform, so
|
||||
# sending it elsewhere would leak an unknown top-level field (see the
|
||||
# dedicated no-leak test). Tools are already cached by the system breakpoint
|
||||
# there — Anthropic orders tools ahead of the system prompt in the prefix.
|
||||
assert _cache_points(model_name) == [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1},
|
||||
]
|
||||
|
||||
|
||||
def test_tool_config_point_not_leaked_to_non_bedrock_claude() -> None:
|
||||
# Regression guard: the ``tool_config`` injection point must NEVER be sent on
|
||||
# a non-Bedrock route. LiteLLM leaves it on the outgoing body as a top-level
|
||||
# ``cache_control_injection_points`` field there, which native Anthropic
|
||||
# would reject as an unknown field.
|
||||
for model in ("anthropic/claude-sonnet-4-5", "openrouter/anthropic/claude-3.5-sonnet"):
|
||||
points = _cache_points(model) or []
|
||||
assert all(p.get("location") != "tool_config" for p in points)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["gpt-5", "vertex_ai/gemini-2.5-pro", "openai/o3"])
|
||||
def test_make_model_settings_no_prompt_cache_for_non_claude(model_name: str) -> None:
|
||||
# No injection points for non-Claude models: the LiteLLM cache hook never
|
||||
# fires, so this stays a strict no-op (won't emit cache_control to strict
|
||||
# OpenAI-compatible endpoints).
|
||||
assert make_model_settings(None, model_name=model_name).extra_args is None
|
||||
|
||||
|
||||
def test_prompt_cache_can_be_disabled() -> None:
|
||||
# STRIX_PROMPT_CACHE=false kill switch: no injection points even for Claude.
|
||||
assert (
|
||||
make_model_settings(
|
||||
None, model_name="anthropic/claude-sonnet-4-5", prompt_cache=False
|
||||
).extra_args
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_no_prompt_cache_for_unmapped_bedrock_claude_model(monkeypatch: Any) -> None:
|
||||
"""A BEDROCK Claude route LiteLLM has NOT mapped (a new release, or any model
|
||||
when LiteLLM can't refresh its model map and falls back to a stale local
|
||||
copy) must run UNCACHED, not crash. Bedrock's Converse API rejects the
|
||||
unknown field outright (ValidationException 'cache_control_injection_points:
|
||||
Extra inputs are not permitted'); LiteLLM only strips the marker for models
|
||||
it recognises as cache-capable, so an unmapped model would 500 the first
|
||||
call and fail the whole run."""
|
||||
unmapped = "bedrock/global.anthropic.claude-brand-new-9"
|
||||
# Simulate a model LiteLLM doesn't know: no cost-map entry, checker says no.
|
||||
monkeypatch.setattr(litellm, "model_cost", {}, raising=False)
|
||||
if getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None):
|
||||
monkeypatch.setattr(litellm.utils, "supports_prompt_caching", lambda *_a, **_k: False)
|
||||
|
||||
# Bedrock Claude by name, but unmapped → no injection points, no crash.
|
||||
assert make_model_settings(None, model_name=unmapped).extra_args is None
|
||||
|
||||
|
||||
def test_prompt_cache_kept_for_non_bedrock_claude_even_if_unmapped(monkeypatch: Any) -> None:
|
||||
"""Non-Bedrock Claude routes must KEEP caching-by-family even when LiteLLM
|
||||
can't confirm support — those providers tolerate/ignore the marker (or
|
||||
LiteLLM maps them under keys we don't resolve, e.g. OpenRouter), so gating
|
||||
them on confirmed support would DISABLE caching for capable models — a
|
||||
regression. Only Bedrock hard-rejects, so only Bedrock is guarded."""
|
||||
monkeypatch.setattr(litellm, "model_cost", {}, raising=False)
|
||||
if getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None):
|
||||
monkeypatch.setattr(litellm.utils, "supports_prompt_caching", lambda *_a, **_k: False)
|
||||
|
||||
# Even with LiteLLM knowing nothing, an Anthropic-native / OpenRouter Claude
|
||||
# still gets the injection points (system + tail, no tool_config).
|
||||
for model in ("anthropic/claude-brand-new-9", "openrouter/anthropic/claude-brand-new"):
|
||||
assert _cache_points(model) == [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1},
|
||||
]
|
||||
|
||||
|
||||
def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None:
|
||||
"""The tail breakpoint's premise, end-to-end: LiteLLM's own message-injection
|
||||
logic must place the cache_control on the LAST message for both a short and a
|
||||
long transcript — i.e. it tracks the growing (append-only) tail rather than a
|
||||
fixed position — so the immutable prefix-so-far is cached and re-read next
|
||||
turn.
|
||||
|
||||
Driven through the hook's static ``_apply_message_injections`` primitive
|
||||
(stable across LiteLLM versions) rather than the prompt-manager entrypoint
|
||||
(whose signature drifts).
|
||||
"""
|
||||
hook_mod = pytest.importorskip("litellm.integrations.anthropic_cache_control_hook")
|
||||
apply = hook_mod.AnthropicCacheControlHook._apply_message_injections
|
||||
points = _cache_points("bedrock/global.anthropic.claude-opus-4-8")
|
||||
msg_points = [p for p in points if p.get("location") == "message"]
|
||||
|
||||
def last_msg_cache_control(n_turns: int) -> Any:
|
||||
messages: list[dict[str, Any]] = [{"role": "system", "content": "stable prompt"}]
|
||||
for i in range(n_turns):
|
||||
messages.append({"role": "assistant", "content": f"turn {i} action"})
|
||||
messages.append({"role": "user", "content": f"turn {i} tool result"})
|
||||
processed = apply(msg_points, messages, 4)
|
||||
last = processed[-1]
|
||||
content = last.get("content")
|
||||
if isinstance(content, list):
|
||||
return content[-1].get("cache_control")
|
||||
return last.get("cache_control")
|
||||
|
||||
assert last_msg_cache_control(2) == {"type": "ephemeral"}
|
||||
assert last_msg_cache_control(20) == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_build_root_task_empty_config() -> None:
|
||||
assert build_root_task({}) == ""
|
||||
|
||||
|
||||
@@ -38,6 +38,7 @@ async def test_persistent_rate_limit_stops_gracefully(
|
||||
reasoning_effort="high",
|
||||
force_required_tool_choice=False,
|
||||
timeout=300,
|
||||
prompt_cache=True,
|
||||
),
|
||||
runtime=types.SimpleNamespace(max_context_images=3),
|
||||
)
|
||||
|
||||
@@ -17,6 +17,7 @@ import strix.tools.notes.tools as notes_tools
|
||||
import strix.tools.todo.tools as todo_tools
|
||||
from strix.core import runner
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.runtime import session_manager
|
||||
|
||||
|
||||
def _make_rate_limit_error() -> RateLimitError:
|
||||
@@ -46,6 +47,7 @@ def _patch_engine_scaffold(
|
||||
reasoning_effort="high",
|
||||
force_required_tool_choice=False,
|
||||
timeout=300,
|
||||
prompt_cache=True,
|
||||
),
|
||||
runtime=types.SimpleNamespace(max_context_images=3),
|
||||
)
|
||||
@@ -66,8 +68,8 @@ def _patch_engine_scaffold(
|
||||
async def _cleanup(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(runner.session_manager, "create_or_reuse", _create_or_reuse)
|
||||
monkeypatch.setattr(runner.session_manager, "cleanup", _cleanup)
|
||||
monkeypatch.setattr(session_manager, "create_or_reuse", _create_or_reuse)
|
||||
monkeypatch.setattr(session_manager, "cleanup", _cleanup)
|
||||
|
||||
monkeypatch.setattr(runner, "build_root_task", lambda _scan_config: "task")
|
||||
monkeypatch.setattr(runner, "build_scope_context", lambda _scan_config: scope_context)
|
||||
|
||||
Reference in New Issue
Block a user