mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
Merge remote-tracking branch 'origin/main' into litellm_usage_key_free_aggregate_split
This commit is contained in:
commit
b5e2e9a392
76 changed files with 6508 additions and 1261 deletions
|
|
@ -9,7 +9,7 @@ commands:
|
|||
parameters:
|
||||
category:
|
||||
type: enum
|
||||
enum: ["backend", "client"]
|
||||
enum: ["backend", "client", "provider-harness"]
|
||||
default: "backend"
|
||||
steps:
|
||||
- run:
|
||||
|
|
@ -2918,19 +2918,30 @@ jobs:
|
|||
provider_replay_harness:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
working_directory: ~/project
|
||||
resource_class: medium
|
||||
environment:
|
||||
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
|
||||
E2E_PROVIDER_CACHE: "0"
|
||||
E2E_FIXTURE_MODE: live
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes:
|
||||
category: provider-harness
|
||||
- setup_litellm_test_deps
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:6379
|
||||
- run:
|
||||
name: Test provider replay harness
|
||||
name: Test provider capture and replay harness
|
||||
command: |
|
||||
mkdir -p test-results/provider-replay-harness
|
||||
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
|
||||
--junitxml=test-results/provider-replay-harness/junit.xml \
|
||||
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
|
||||
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py \
|
||||
tests/code_coverage_tests/test_provider_cache.py
|
||||
- store_test_results:
|
||||
path: test-results/provider-replay-harness
|
||||
|
||||
|
|
|
|||
|
|
@ -1,13 +1,19 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
has_provider_harness=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
tests/e2e/*/*.py) : ;;
|
||||
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
|
||||
has_provider_harness=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
ui/* | tests/e2e/ui/*) has_client=true ;;
|
||||
docs/* | *.md | *.mdx) : ;;
|
||||
|
|
@ -17,6 +23,9 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
done
|
||||
|
||||
case "$category" in
|
||||
provider-harness)
|
||||
[ "$has_provider_harness" = true ] && echo run || echo skip
|
||||
;;
|
||||
backend)
|
||||
[ "$has_backend" = true ] && echo run || echo skip
|
||||
;;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: path_filter.sh <backend|client>}"
|
||||
category="${1:?usage: path_filter.sh <backend|client|provider-harness>}"
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
run_full() {
|
||||
|
|
@ -36,5 +36,5 @@ if [ "$decision" = run ]; then
|
|||
run_full "$category-relevant changes detected"
|
||||
fi
|
||||
|
||||
echo "path-filter[$category]: only unrelated (docs/client) changes detected; halting job as successful"
|
||||
echo "path-filter[$category]: only unrelated changes detected; halting job as successful"
|
||||
circleci-agent step halt
|
||||
|
|
|
|||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -26,6 +26,7 @@ on:
|
|||
- ui/Dockerfile
|
||||
- ui/nginx.conf
|
||||
- .github/workflows/image-scan.yml
|
||||
- .grype.yaml
|
||||
schedule:
|
||||
- cron: "41 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
|
@ -93,6 +94,7 @@ jobs:
|
|||
GRYPE_MATCH_PYTHON_USING_CPES: "true"
|
||||
run: |
|
||||
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
|
||||
--config .grype.yaml \
|
||||
--only-fixed \
|
||||
--fail-on high \
|
||||
--output table
|
||||
|
|
|
|||
13
.grype.yaml
Normal file
13
.grype.yaml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# Wolfi's security database names zlib 1.3.3-r0 as the fix for CVE-2026-85091,
|
||||
# but the newest zlib published to the Wolfi apk repo is 1.3.2-r7, so every
|
||||
# wolfi-base digest reports it and no `apk upgrade` can clear it.
|
||||
# Drop this once Wolfi ships zlib >= 1.3.3-r0; expected by 2026-10-15.
|
||||
ignore:
|
||||
- vulnerability: CVE-2026-85091
|
||||
package:
|
||||
name: zlib
|
||||
type: apk
|
||||
- vulnerability: GHSA-g5fp-32jq-cfw2
|
||||
package:
|
||||
name: zlib
|
||||
type: apk
|
||||
|
|
@ -1167,7 +1167,9 @@ class ModelResponseIterator:
|
|||
# (matches OpenAI behavior and non-streaming Anthropic implementation)
|
||||
if self.converted_response_format_tool:
|
||||
finish_reason = "stop"
|
||||
usage: Final = self._handle_usage(anthropic_usage_chunk=message_delta["usage"])
|
||||
usage: Final = (
|
||||
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
|
||||
)
|
||||
container: Final = message_delta["delta"].get("container")
|
||||
return finish_reason, usage, container
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,9 @@ if TYPE_CHECKING:
|
|||
import litellm
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) -> tuple[float, float]:
|
||||
def cost_per_token(
|
||||
model: str, usage: "Usage", service_tier: str | None = None, model_info: "ModelInfo | None" = None
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
||||
|
|
@ -27,6 +29,7 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
|
|||
- usage: LiteLLM Usage block, containing anthropic caching information
|
||||
- service_tier: the service tier the request was served at (e.g. "priority"),
|
||||
read from the Anthropic response usage and used to select tier-specific pricing
|
||||
- model_info: effective deployment prices, when they override public rates
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
|
|
@ -36,16 +39,23 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
|
|||
usage=usage,
|
||||
custom_llm_provider="anthropic",
|
||||
service_tier=service_tier,
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
# Apply provider_specific_entry multipliers for geo/speed routing
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
|
||||
provider_specific_entry: Final[dict] = model_info.get("provider_specific_entry") or {}
|
||||
effective_info: Final = (
|
||||
model_info
|
||||
if model_info is not None
|
||||
else litellm.get_model_info(model=model, custom_llm_provider="anthropic")
|
||||
)
|
||||
provider_specific_entry: Final = effective_info.get("provider_specific_entry")
|
||||
|
||||
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
|
||||
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=effective_info, usage=usage)
|
||||
speed_multiplier: Final = (
|
||||
provider_specific_entry.get("fast", 1.0) if getattr(usage, "speed", None) == "fast" else 1.0
|
||||
provider_specific_entry.get("fast", 1.0)
|
||||
if provider_specific_entry and getattr(usage, "speed", None) == "fast"
|
||||
else 1.0
|
||||
)
|
||||
|
||||
if speed_multiplier != 1.0:
|
||||
|
|
|
|||
|
|
@ -376,10 +376,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
|
||||
chunk.usage
|
||||
)
|
||||
merged_chunk["usage"] = usage_dict
|
||||
if self.applied_edits and "context_management" not in merged_chunk:
|
||||
merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
|
||||
return self._augment_message_delta_usage(merged_chunk)
|
||||
return self._augment_message_delta_usage({**merged_chunk, "usage": usage_dict})
|
||||
|
||||
def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool:
|
||||
"""Consume an OpenAI-compatible chunk that carries no ``choices``.
|
||||
|
|
@ -448,8 +447,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
}
|
||||
iterations.append(message_iteration)
|
||||
augmented_usage["iterations"] = iterations
|
||||
augmented["usage"] = augmented_usage
|
||||
return augmented
|
||||
return {**augmented, "usage": augmented_usage}
|
||||
|
||||
def _next_compaction_event(self) -> dict[str, object] | None:
|
||||
"""Return the next compaction content-block SSE event, or ``None``.
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
For calculating cost of fireworks ai serverless inference models.
|
||||
"""
|
||||
|
||||
import math
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import (
|
||||
Final,
|
||||
cast, # noqa: TID251 # the fallback entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it
|
||||
)
|
||||
|
||||
from litellm.constants import (
|
||||
FIREWORKS_AI_4_B,
|
||||
|
|
@ -12,12 +14,10 @@ from litellm.constants import (
|
|||
FIREWORKS_AI_56_B_MOE,
|
||||
FIREWORKS_AI_176_B_MOE,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
NO_CACHE_READ_RATE: Final = float("nan")
|
||||
|
||||
|
||||
# Extract the number of billion parameters from the model name
|
||||
# only used for together_computer LLMs
|
||||
|
|
@ -67,6 +67,28 @@ def _resolve_model_info(model: str) -> ModelInfo:
|
|||
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
|
||||
|
||||
|
||||
def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo:
|
||||
"""Entries without a cache-read rate keep the previous calculator's input-rate fallback for cached
|
||||
reads (LIT-7845 tracks the documented discount); the shared map is never mutated, so a copy carries it."""
|
||||
input_rate: Final = model_info.get("input_cost_per_token")
|
||||
if model_info.get("cache_read_input_token_cost") is not None or input_rate is None:
|
||||
return model_info
|
||||
off_peak: Final = model_info.get("off_peak_pricing")
|
||||
if off_peak is None or "cache_read_input_token_cost" in off_peak:
|
||||
return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": input_rate})
|
||||
return cast(
|
||||
ModelInfo,
|
||||
{
|
||||
**model_info,
|
||||
"cache_read_input_token_cost": input_rate,
|
||||
"off_peak_pricing": {
|
||||
**off_peak,
|
||||
"cache_read_input_token_cost": off_peak.get("input_cost_per_token", input_rate),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
|
||||
|
|
@ -80,29 +102,11 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
|
|||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
"""
|
||||
model_info: Final = _resolve_model_info(model)
|
||||
standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
|
||||
rates: Final = apply_off_peak_pricing(
|
||||
model_info,
|
||||
current_time,
|
||||
TokenRates(
|
||||
input_rate=model_info["input_cost_per_token"] or 0.0,
|
||||
output_rate=model_info["output_cost_per_token"] or 0.0,
|
||||
cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
|
||||
cache_creation_rate=0.0,
|
||||
reasoning_rate=None,
|
||||
),
|
||||
model_info: Final = _with_cache_read_fallback(_resolve_model_info(model))
|
||||
return generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider="fireworks_ai",
|
||||
model_info=model_info,
|
||||
current_time=current_time,
|
||||
)
|
||||
cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
|
||||
|
||||
prompt_tokens_details: Final = usage.prompt_tokens_details
|
||||
cached_tokens: Final[int] = (
|
||||
prompt_tokens_details.cached_tokens
|
||||
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
|
||||
else 0
|
||||
)
|
||||
non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
|
||||
prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
|
||||
completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
|
|
|
|||
|
|
@ -124,6 +124,12 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
|
|||
)
|
||||
|
||||
|
||||
def _served_model_name(model_version: object) -> str | None:
|
||||
if not isinstance(model_version, str) or not model_version:
|
||||
return None
|
||||
return model_version.split("@", 1)[0]
|
||||
|
||||
|
||||
class VertexAIBaseConfig:
|
||||
def get_mapped_special_auth_params(self) -> dict:
|
||||
"""
|
||||
|
|
@ -1951,6 +1957,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
def _check_prompt_level_content_filter(
|
||||
processed_chunk: GenerateContentResponseBody,
|
||||
response_id: str | None,
|
||||
model: str | None = None,
|
||||
) -> Optional["ModelResponseStream"]:
|
||||
"""
|
||||
Check if prompt is blocked due to content filtering at the prompt level.
|
||||
|
|
@ -1990,7 +1997,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
enhancements=None,
|
||||
)
|
||||
|
||||
model_response: Final = ModelResponseStream(choices=[choice], id=response_id)
|
||||
model_response: Final = ModelResponseStream(choices=[choice], id=response_id, model=model)
|
||||
return model_response
|
||||
|
||||
return None
|
||||
|
|
@ -2434,7 +2441,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
completion_response = GenerateContentResponseBody(**completion_response)
|
||||
|
||||
## GET MODEL ##
|
||||
model_response.model = model
|
||||
served: Final = _served_model_name(completion_response.get("modelVersion"))
|
||||
model_response.model = served if served is not None else model
|
||||
|
||||
## CHECK IF RESPONSE FLAGGED
|
||||
if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]:
|
||||
|
|
@ -3264,12 +3272,18 @@ class ModelResponseIterator:
|
|||
|
||||
processed_chunk: Final = GenerateContentResponseBody(**chunk)
|
||||
response_id: Final = processed_chunk.get("responseId")
|
||||
model_response = ModelResponseStream(choices=[], id=response_id)
|
||||
served: Final = _served_model_name(processed_chunk.get("modelVersion"))
|
||||
model_response = ModelResponseStream(
|
||||
choices=[],
|
||||
id=response_id,
|
||||
model=served,
|
||||
)
|
||||
|
||||
# Check if prompt is blocked due to content filtering
|
||||
blocked_response: Final = VertexGeminiConfig._check_prompt_level_content_filter(
|
||||
processed_chunk=processed_chunk,
|
||||
response_id=response_id,
|
||||
model=served,
|
||||
)
|
||||
if blocked_response is not None:
|
||||
model_response = blocked_response
|
||||
|
|
|
|||
|
|
@ -49,7 +49,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
|
||||
Inherits from OpenAIResponsesAPIConfig since XAI's Responses API is largely
|
||||
compatible with OpenAI's, with a few differences:
|
||||
- Does not support the 'instructions' parameter
|
||||
- Requires code_interpreter tools to have 'container' field removed
|
||||
- Recommends store=false when sending images
|
||||
|
||||
|
|
@ -60,20 +59,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.XAI
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
Get supported parameters for XAI Responses API.
|
||||
|
||||
XAI supports most OpenAI Responses API params except 'instructions'.
|
||||
"""
|
||||
supported_params: Final = super().get_supported_openai_params(model)
|
||||
|
||||
# Remove 'instructions' as it's not supported by XAI
|
||||
if "instructions" in supported_params:
|
||||
supported_params.remove("instructions")
|
||||
|
||||
return supported_params
|
||||
|
||||
def _transform_web_search_tool(self, tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""
|
||||
Transform web_search tool to XAI format.
|
||||
|
|
@ -158,19 +143,13 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
Map parameters for XAI Responses API.
|
||||
|
||||
Handles XAI-specific transformations:
|
||||
1. Drops 'instructions' parameter (not supported)
|
||||
2. Transforms code_interpreter tools to remove 'container' field
|
||||
3. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
|
||||
4. Transforms x_search tools to XAI format
|
||||
5. Sets store=false when images are detected (recommended by XAI)
|
||||
1. Transforms code_interpreter tools to remove 'container' field
|
||||
2. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
|
||||
3. Transforms x_search tools to XAI format
|
||||
4. Sets store=false when images are detected (recommended by XAI)
|
||||
"""
|
||||
params: Final = dict(response_api_optional_params)
|
||||
|
||||
# Drop instructions parameter (not supported by XAI)
|
||||
if "instructions" in params:
|
||||
verbose_logger.debug("XAI Responses API does not support 'instructions' parameter. Dropping it.")
|
||||
params.pop("instructions")
|
||||
|
||||
if "metadata" in params:
|
||||
verbose_logger.debug("XAI Responses API does not support 'metadata' parameter. Dropping it.")
|
||||
params.pop("metadata")
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria."""
|
||||
"""LLM-as-a-Judge guardrail: uses an LLM to score requests or responses against weighted criteria."""
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.llm_judge import (
|
||||
default_router_provider,
|
||||
|
|
@ -16,8 +19,9 @@ from litellm.litellm_core_utils.llm_judge import (
|
|||
judge_acompletion,
|
||||
parse_json_verdict,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN, GenericGuardrailAPIInputs, GuardrailStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
|
@ -26,18 +30,65 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import StandardLoggingEvalInformation
|
||||
|
||||
JUDGE_SYSTEM_PROMPT = """You are a quality judge. Evaluate the assistant's response against the criteria provided.
|
||||
For each criterion, assign a score from 0 to 100 and provide concise reasoning.
|
||||
JudgeInputType = Literal["request", "response"]
|
||||
JudgeEventHook = GuardrailEventHooks | list[GuardrailEventHooks] | Mode
|
||||
JudgeModeParam = str | list[str] | Mode | GuardrailEventHooks | list[GuardrailEventHooks] | None
|
||||
|
||||
_JUDGE_SYSTEM_PROMPT_TEMPLATE: Final = """You are a quality judge. Evaluate the {subject} against the criteria provided.
|
||||
{focus}For each criterion, assign a score from 0 to 100 and provide concise reasoning.
|
||||
Return ONLY valid JSON in this exact format:
|
||||
{
|
||||
{{
|
||||
"verdicts": [
|
||||
{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}
|
||||
{{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}}
|
||||
],
|
||||
"overall_score": <weighted average 0-100>
|
||||
}"""
|
||||
}}"""
|
||||
|
||||
JUDGE_SYSTEM_PROMPTS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
|
||||
{
|
||||
"request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(
|
||||
subject="request",
|
||||
focus="Judge the most recent user turn; treat earlier turns in the conversation only as context.\n",
|
||||
),
|
||||
"response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response", focus=""),
|
||||
}
|
||||
)
|
||||
|
||||
_JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
|
||||
{"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"}
|
||||
)
|
||||
|
||||
_LIFECYCLE_HOOKS: Final[MappingProxyType[JudgeInputType, tuple[GuardrailEventHooks, ...]]] = MappingProxyType(
|
||||
{
|
||||
"request": (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.logging_only),
|
||||
"response": (GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only),
|
||||
}
|
||||
)
|
||||
|
||||
_VALID_ON_FAILURE: Final = frozenset({"block", "log"})
|
||||
|
||||
_JUDGE_CALL_METADATA: Final = MappingProxyType(
|
||||
{INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}
|
||||
)
|
||||
|
||||
|
||||
class _LoggedCallParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metadata: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def _is_logged_judge_call(data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
"""logging_only is the only event whose ``data`` is the SDK's model_call_details rather than the client body."""
|
||||
if event_type is not GuardrailEventHooks.logging_only:
|
||||
return False
|
||||
try:
|
||||
params: Final = _LoggedCallParams.model_validate(data.get("litellm_params") or {})
|
||||
except ValidationError:
|
||||
return False
|
||||
return (params.metadata or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
|
||||
|
||||
|
||||
_default_router_provider: Final = default_router_provider
|
||||
_parse_judge_verdict: Final = parse_json_verdict
|
||||
_extract_text_from_content: Final = extract_text_from_content
|
||||
|
|
@ -86,10 +137,29 @@ def _get_litellm_param(
|
|||
return default
|
||||
|
||||
|
||||
def _coerce_event_hook(mode: JudgeModeParam) -> JudgeEventHook:
|
||||
if mode is None:
|
||||
return GuardrailEventHooks.post_call
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [GuardrailEventHooks(hook) for hook in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _text_under_review(inputs: GenericGuardrailAPIInputs, input_type: JudgeInputType) -> str:
|
||||
all_text: Final = "\n".join(inputs.get("texts") or [])
|
||||
if input_type == "response":
|
||||
return all_text
|
||||
latest_user_turn: Final = get_last_user_message(inputs.get("structured_messages") or [])
|
||||
return latest_user_turn if latest_user_turn is not None else all_text
|
||||
|
||||
|
||||
def _build_judge_prompt(
|
||||
criteria: Sequence[JudgeCriterion],
|
||||
messages: Sequence[JudgeMessage],
|
||||
response_text: str,
|
||||
text_under_review: str,
|
||||
input_type: JudgeInputType = "response",
|
||||
) -> str:
|
||||
criteria_block: Final = "\n".join(
|
||||
f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria
|
||||
|
|
@ -99,15 +169,16 @@ def _build_judge_prompt(
|
|||
for m in messages
|
||||
if m.get("content") is not None
|
||||
)
|
||||
conversation_block: Final = f"Conversation:\n{conversation}\n\n" if conversation or input_type == "response" else ""
|
||||
return (
|
||||
f"Criteria to evaluate:\n{criteria_block}\n\n"
|
||||
f"Conversation:\n{conversation}\n\n"
|
||||
f"Assistant response to evaluate:\n{response_text}"
|
||||
f"{conversation_block}"
|
||||
f"{_JUDGE_SUBJECT_LABELS[input_type]}:\n{text_under_review}"
|
||||
)
|
||||
|
||||
|
||||
class LLMAsAJudgeGuardrail(CustomGuardrail):
|
||||
"""Post-call guardrail that judges response quality via an LLM."""
|
||||
"""Guardrail that judges request (pre_call/during_call) or response (post_call) quality via an LLM."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -116,22 +187,15 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
criteria: Sequence[JudgeCriterion],
|
||||
overall_threshold: float = 80.0,
|
||||
on_failure: Literal["block", "log"] = "block",
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None,
|
||||
event_hook: JudgeModeParam = None,
|
||||
default_on: bool = False,
|
||||
router_provider: "Callable[[], Router | None] | None" = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
_event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None
|
||||
if event_hook is not None:
|
||||
if isinstance(event_hook, list):
|
||||
_event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook]
|
||||
else:
|
||||
_event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, str) else event_hook
|
||||
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=list(self.get_supported_event_hooks()),
|
||||
event_hook=_event_hook or GuardrailEventHooks.post_call,
|
||||
event_hook=_coerce_event_hook(event_hook),
|
||||
default_on=default_on,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -143,18 +207,24 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
|
||||
return [GuardrailEventHooks.post_call]
|
||||
return [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
|
||||
|
||||
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
if _is_logged_judge_call(data, event_type):
|
||||
return False
|
||||
return super().should_run_guardrail(data, event_type)
|
||||
|
||||
async def _run_judge(
|
||||
self,
|
||||
messages: Sequence[JudgeMessage],
|
||||
response_text: str,
|
||||
text_under_review: str,
|
||||
input_type: JudgeInputType = "response",
|
||||
) -> dict[str, object]:
|
||||
judge_messages: Final[list[AllMessageValues]] = [
|
||||
{"role": "system", "content": JUDGE_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": JUDGE_SYSTEM_PROMPTS[input_type]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": _build_judge_prompt(self.criteria, messages, response_text),
|
||||
"content": _build_judge_prompt(self.criteria, messages, text_under_review, input_type),
|
||||
},
|
||||
]
|
||||
response: Final = await judge_acompletion(
|
||||
|
|
@ -163,6 +233,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_messages,
|
||||
response_format={"type": "json_object"},
|
||||
temperature=0,
|
||||
metadata=dict(_JUDGE_CALL_METADATA),
|
||||
)
|
||||
raw: Final = response.choices[0].message.content or "{}"
|
||||
return _parse_judge_verdict(raw)
|
||||
|
|
@ -174,13 +245,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Only evaluate post-call (response text). Fail open on pre-call.
|
||||
if input_type != "response":
|
||||
return inputs
|
||||
|
||||
texts: Final = inputs.get("texts") or []
|
||||
response_text: Final = " ".join(texts)
|
||||
if not response_text:
|
||||
text_under_review: Final = _text_under_review(inputs, input_type)
|
||||
if not text_under_review:
|
||||
return inputs
|
||||
|
||||
start_time: Final = datetime.now()
|
||||
|
|
@ -188,10 +254,12 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_result: dict[str, object] = {}
|
||||
|
||||
try:
|
||||
messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or []
|
||||
messages: Final[Sequence[JudgeMessage]] = (
|
||||
inputs.get("structured_messages") or request_data.get("messages") or []
|
||||
)
|
||||
|
||||
try:
|
||||
judge_result = await self._run_judge(messages, response_text)
|
||||
judge_result = await self._run_judge(messages, text_under_review, input_type)
|
||||
except Exception as judge_err:
|
||||
verbose_logger.warning(
|
||||
"llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err
|
||||
|
|
@ -230,7 +298,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail={
|
||||
"error": "LLM judge rejected response: score below threshold",
|
||||
"error": f"LLM judge rejected {input_type}: score below threshold",
|
||||
"overall_score": overall_score,
|
||||
"threshold": self.overall_threshold,
|
||||
"verdicts": judge_result.get("verdicts", []),
|
||||
|
|
@ -252,9 +320,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
guardrail_status=status,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now().timestamp(),
|
||||
event_type=GuardrailEventHooks.post_call,
|
||||
event_type=self._event_type_for(input_type),
|
||||
)
|
||||
|
||||
def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks | None:
|
||||
configured: Final = tuple(hook for hook in _LIFECYCLE_HOOKS[input_type] if self._event_hook_is_event_type(hook))
|
||||
return configured[0] if len(configured) == 1 else None
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams",
|
||||
|
|
@ -282,10 +354,7 @@ def initialize_guardrail(
|
|||
|
||||
overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0))
|
||||
|
||||
mode: Final[str | None] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
event_hook: GuardrailEventHooks | None = None
|
||||
if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}:
|
||||
event_hook = GuardrailEventHooks(mode)
|
||||
mode: Final[JudgeModeParam] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
|
||||
instance: Final = LLMAsAJudgeGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
|
|
@ -293,7 +362,7 @@ def initialize_guardrail(
|
|||
criteria=criteria,
|
||||
overall_threshold=overall_threshold,
|
||||
on_failure=on_failure,
|
||||
event_hook=event_hook,
|
||||
event_hook=mode,
|
||||
default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -19,20 +22,26 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import (
|
||||
GuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
|
||||
AssistantMessage,
|
||||
SingulrGuardrailPayload,
|
||||
SingulrGuardrailRequest,
|
||||
SingulrGuardrailResponse,
|
||||
SingulrMcpGuardrailPayload,
|
||||
ToolCall,
|
||||
ToolCallFunction,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
||||
|
||||
_DEFAULT_API_BASE: Final = "http://localhost:8003"
|
||||
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm"
|
||||
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2"
|
||||
_DEFAULT_TIMEOUT: Final = 30.0
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_MCP_MODEL_PREFIX: Final = "MCP:"
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
|
|
@ -51,8 +60,8 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
) -> None:
|
||||
self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY")
|
||||
self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip(
|
||||
"/"
|
||||
self.singulr_api_base = (
|
||||
(singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).strip().rstrip("/")
|
||||
)
|
||||
parsed: Final = urlparse(self.singulr_api_base)
|
||||
if parsed.scheme == "http" and parsed.hostname not in (
|
||||
|
|
@ -85,6 +94,9 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
kwargs["supported_event_hooks"] = [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.logging_only,
|
||||
GuardrailEventHooks.pre_mcp_call,
|
||||
GuardrailEventHooks.post_mcp_call,
|
||||
]
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -97,52 +109,70 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
|
||||
return SingulrGuardrailConfigModel
|
||||
|
||||
def _build_payload(
|
||||
self,
|
||||
request_data: dict[str, Any],
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
input_type: str,
|
||||
) -> dict[str, object]:
|
||||
if not request_data:
|
||||
texts: Final = inputs.get("texts", [])
|
||||
|
||||
payload = SingulrGuardrailPayload(
|
||||
input_type=input_type,
|
||||
is_playground_request=True,
|
||||
playground_text=texts[0] if texts else None,
|
||||
@staticmethod
|
||||
def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
|
||||
litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING
|
||||
return tuple(
|
||||
container
|
||||
for container in (
|
||||
request_data.get("litellm_metadata"),
|
||||
request_data.get("metadata"),
|
||||
litellm_params.get("litellm_metadata") if litellm_params else None,
|
||||
litellm_params.get("metadata") if litellm_params else None,
|
||||
)
|
||||
else:
|
||||
response: Final = request_data.get("response")
|
||||
singulr_req_object: Final = SingulrGuardrailRequest(
|
||||
model=request_data.get("model"),
|
||||
messages=request_data.get("messages"),
|
||||
tools=request_data.get("tools"),
|
||||
model_response=response.model_dump(mode="json") if input_type == "response" and response else None,
|
||||
litellm_metadata=request_data.get("litellm_metadata"),
|
||||
)
|
||||
payload = SingulrGuardrailPayload(
|
||||
litellm_call_id=request_data.get("litellm_call_id"),
|
||||
request_data=singulr_req_object,
|
||||
input_type=input_type,
|
||||
)
|
||||
|
||||
return payload.model_dump(mode="json")
|
||||
|
||||
def _build_headers(self) -> dict[str, str]:
|
||||
return dict(
|
||||
(header, value)
|
||||
for header, value in (
|
||||
("Content-Type", "application/json"),
|
||||
("X-Singulr-Gateway-Token", self.singulr_api_key),
|
||||
(
|
||||
"X-Singulr-Enforcement-Entity-Id",
|
||||
self.singulr_application_id or "",
|
||||
),
|
||||
("X-Singulr-Guardrail-Id", self.singulr_guardrail_id or ""),
|
||||
)
|
||||
if value
|
||||
if container
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _resolve_metadata_value(cls, request_data: Mapping[str, Any], key: str) -> str | None:
|
||||
for container in cls._metadata_containers(request_data=request_data):
|
||||
value = container.get(key)
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _resolve_user_role_from_request_data(cls, request_data: Mapping[str, Any]) -> str | None:
|
||||
for container in cls._metadata_containers(request_data=request_data):
|
||||
auth = container.get("user_api_key_auth")
|
||||
if isinstance(auth, UserAPIKeyAuth) and auth.user_role:
|
||||
return auth.user_role.value
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, str] | None:
|
||||
fields: Final = (
|
||||
"user_api_key_alias",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_user_email",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_org_alias",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_team_alias",
|
||||
)
|
||||
resolved: Final = (
|
||||
*((field, cls._resolve_metadata_value(request_data=request_data, key=field)) for field in fields),
|
||||
("user_api_key_user_role", cls._resolve_user_role_from_request_data(request_data=request_data)),
|
||||
)
|
||||
if not any(value for _, value in resolved):
|
||||
return None
|
||||
return {key: value for key, value in resolved if value} # mutable-ok: short-lived JSON payload dict
|
||||
|
||||
@staticmethod
|
||||
def _build_user_message(text: str) -> Mapping[str, Any]:
|
||||
return {"role": "user", "content": text} # mutable-ok: short-lived JSON payload dict
|
||||
|
||||
def _build_headers(self) -> Mapping[str, str]:
|
||||
all_headers: Final = MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/json",
|
||||
"X-Singulr-Gateway-Token": self.singulr_api_key,
|
||||
"X-Singulr-Enforcement-Entity-Id": self.singulr_application_id,
|
||||
"X-Singulr-Guardrail-Id": self.singulr_guardrail_id,
|
||||
}
|
||||
)
|
||||
return MappingProxyType({header: value for header, value in all_headers.items() if value})
|
||||
|
||||
async def _call_api(self, payload: dict[str, object]) -> SingulrGuardrailResponse | None:
|
||||
endpoint: Final = f"{self.singulr_api_base}{_GUARD_ENDPOINT}"
|
||||
verbose_proxy_logger.debug("Singulr: %s", endpoint)
|
||||
|
|
@ -168,7 +198,7 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
if self.block_on_error:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=(f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}"),
|
||||
message=f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}",
|
||||
) from exc
|
||||
return None
|
||||
|
||||
|
|
@ -190,33 +220,218 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
) from exc
|
||||
return None
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
async def _apply_guardrail_on_request(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: str,
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
texts: Sequence[str],
|
||||
structured_messages: Sequence[Any],
|
||||
request_data: Mapping[str, Any],
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
payload: Final = self._build_payload(request_data, inputs, input_type)
|
||||
if not payload:
|
||||
return inputs
|
||||
|
||||
result: Final = await self._call_api(payload)
|
||||
if result is None:
|
||||
return inputs
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Singulr: should_block=%s blocking_due_to=%s",
|
||||
result.should_block,
|
||||
result.blocking_due_to,
|
||||
messages: Final = (
|
||||
tuple(structured_messages)
|
||||
if structured_messages
|
||||
else tuple(self._build_user_message(text) for text in texts)
|
||||
)
|
||||
|
||||
if result.should_block:
|
||||
images: Final = inputs.get("images")
|
||||
tools: Final = inputs.get("tools")
|
||||
|
||||
if not messages and not images and not tools:
|
||||
verbose_proxy_logger.debug("Singulr: No messages, images, or tools to check after filtering")
|
||||
return inputs
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_req_obj = SingulrGuardrailPayload(
|
||||
correlation_id=request_data.get("litellm_call_id"),
|
||||
model_name=inputs.get("model"),
|
||||
guardrail_scope="request",
|
||||
messages=messages,
|
||||
images=images,
|
||||
tools=tools,
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_req_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}",
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
return inputs
|
||||
|
||||
@staticmethod
|
||||
def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None:
|
||||
return request_data.get("mcp_tool_name") or request_data.get("name")
|
||||
|
||||
@staticmethod
|
||||
def _mcp_arguments(request_data: Mapping[str, Any]) -> object:
|
||||
arguments: Final = request_data.get("mcp_arguments")
|
||||
return arguments if arguments is not None else request_data.get("arguments")
|
||||
|
||||
@staticmethod
|
||||
def _is_mcp_call(request_data: Mapping[str, Any], logging_obj: LiteLLMLoggingObj | None) -> bool:
|
||||
call_type: Final = logging_obj.call_type if logging_obj is not None else request_data.get("call_type")
|
||||
if call_type is not None:
|
||||
return call_type == CallTypes.call_mcp_tool.value
|
||||
model: Final = request_data.get("model")
|
||||
return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX))
|
||||
|
||||
async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None:
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_mcp_obj = SingulrMcpGuardrailPayload(
|
||||
guardrail_scope="mcp_request",
|
||||
tool_name=self._mcp_tool_name(request_data),
|
||||
tool_arguments=self._mcp_arguments(request_data),
|
||||
mcp_server_name=request_data.get("mcp_server_name"),
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_mcp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
async def _apply_guardrail_on_mcp_response(
|
||||
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if not texts:
|
||||
return inputs
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_mcp_obj = SingulrMcpGuardrailPayload(
|
||||
model_name=request_data.get("model"),
|
||||
guardrail_scope="mcp_response",
|
||||
tool_result=texts,
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_mcp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
@staticmethod
|
||||
def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None":
|
||||
tool_call_id: Final = tool_call.get("id")
|
||||
fun: Final = tool_call.get("function")
|
||||
if not tool_call_id or not fun:
|
||||
return None
|
||||
func_name: Final = fun.get("name")
|
||||
args: Final = fun.get("arguments")
|
||||
if not func_name or args is None:
|
||||
return None
|
||||
call_type: Final = tool_call.get("type")
|
||||
return ToolCall(
|
||||
id=tool_call_id,
|
||||
type=call_type if isinstance(call_type, str) and call_type else "function",
|
||||
function=ToolCallFunction(
|
||||
name=func_name,
|
||||
arguments=args if isinstance(args, str) else json.dumps(args, default=str),
|
||||
),
|
||||
)
|
||||
|
||||
async def _apply_guardrail_on_response(
|
||||
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
combined_texts: Final = "\n".join(texts) if texts else None
|
||||
|
||||
tool_calls: Final = inputs.get("tool_calls", ())
|
||||
tool_calls_res: Final = tuple(
|
||||
tool_call_res
|
||||
for tool_call_res in (self._build_tool_call(tool_call) for tool_call in tool_calls)
|
||||
if tool_call_res is not None
|
||||
)
|
||||
|
||||
assistant_message: Final = AssistantMessage(
|
||||
role="assistant",
|
||||
content=combined_texts,
|
||||
tool_calls=tool_calls_res,
|
||||
)
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_resp_obj = SingulrGuardrailPayload(
|
||||
correlation_id=request_data.get("litellm_call_id"),
|
||||
guardrail_scope="response",
|
||||
model_name=request_data.get("model"),
|
||||
messages=request_data.get("messages"),
|
||||
images=inputs.get("images"),
|
||||
response=assistant_message,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
payload = singulr_resp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
return inputs
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict, # mutable-ok: required by CustomGuardrail.apply_guardrail override signature
|
||||
input_type: str,
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
texts: Final = inputs.get("texts", ())
|
||||
structured_messages: Final = inputs.get("structured_messages", ())
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Singulr Guardrail: apply_guardrail called with input_type=%s, texts=%d, structured_messages=%d",
|
||||
input_type,
|
||||
len(texts),
|
||||
len(structured_messages),
|
||||
)
|
||||
|
||||
is_mcp_call: Final = self._is_mcp_call(request_data, logging_obj)
|
||||
if input_type == "request":
|
||||
if is_mcp_call:
|
||||
await self._apply_guardrail_on_mcp_request(request_data=request_data)
|
||||
return inputs
|
||||
return await self._apply_guardrail_on_request(
|
||||
inputs=inputs, texts=texts, structured_messages=structured_messages, request_data=request_data
|
||||
)
|
||||
elif input_type == "response":
|
||||
if is_mcp_call:
|
||||
return await self._apply_guardrail_on_mcp_response(
|
||||
inputs=inputs, texts=texts, request_data=request_data
|
||||
)
|
||||
return await self._apply_guardrail_on_response(inputs=inputs, texts=texts, request_data=request_data)
|
||||
return inputs
|
||||
|
|
|
|||
|
|
@ -171,15 +171,25 @@ def _cost_of_usage(
|
|||
) -> float | None:
|
||||
"""What ``usage`` costs on ``model``, or ``None`` when the model has no pricing."""
|
||||
try:
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
custom_llm_provider=model.provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
model_info=model_info,
|
||||
vertex_location=basis.vertex_location,
|
||||
)
|
||||
if model.provider == "anthropic":
|
||||
from litellm.llms.anthropic.cost_calculation import cost_per_token
|
||||
|
||||
prompt_cost, completion_cost = cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
service_tier=basis.service_tier,
|
||||
model_info=model_info,
|
||||
)
|
||||
else:
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
custom_llm_provider=model.provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
model_info=model_info,
|
||||
vertex_location=basis.vertex_location,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings
|
||||
verbose_proxy_logger.debug(
|
||||
"savings: cannot price usage for provider=%s model=%s (%s)", model.provider, model.model, e
|
||||
|
|
@ -198,11 +208,6 @@ def _cache_token_split(usage: Usage) -> tuple[int, int]:
|
|||
return int(read), int(created)
|
||||
|
||||
|
||||
_CACHE_SPLIT_FIELDS: Final = frozenset(
|
||||
("cached_tokens", "cache_creation_tokens", "cache_write_tokens", "cache_creation_token_details", "text_tokens")
|
||||
)
|
||||
|
||||
|
||||
def _baseline_cache_rate_keys(baseline_info: ModelInfo | None) -> tuple[bool, bool]:
|
||||
"""Whether the baseline model has a ``(cache read, cache write)`` rate of its own.
|
||||
|
||||
|
|
@ -274,19 +279,22 @@ def _baseline_usage(usage: Usage, conversation_continuing: bool, baseline_info:
|
|||
(getattr(details, field, 0) or 0) for field in ("audio_tokens", "image_tokens", "video_tokens")
|
||||
)
|
||||
return Usage(
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
completion_tokens_details=usage.completion_tokens_details,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
**details.model_dump(exclude=_CACHE_SPLIT_FIELDS),
|
||||
cached_tokens=reads,
|
||||
cache_creation_tokens=writes,
|
||||
cache_write_tokens=writes,
|
||||
cache_creation_token_details=details.cache_creation_token_details if writes else None,
|
||||
# Whatever no longer sits in a cache bucket is plain input on the baseline.
|
||||
text_tokens=max(usage.prompt_tokens - reads - writes - other_modalities, 0),
|
||||
),
|
||||
**{
|
||||
**usage.model_dump(),
|
||||
# Rebuild through Usage so private fallback counts agree with the public buckets.
|
||||
"cache_read_input_tokens": reads,
|
||||
"cache_creation_input_tokens": writes,
|
||||
"prompt_tokens_details": PromptTokensDetailsWrapper(
|
||||
**{
|
||||
**details.model_dump(),
|
||||
"cached_tokens": reads,
|
||||
"cache_creation_tokens": writes,
|
||||
"cache_write_tokens": writes,
|
||||
"cache_creation_token_details": details.cache_creation_token_details if writes else None,
|
||||
"text_tokens": max(usage.prompt_tokens - reads - writes - other_modalities, 0),
|
||||
}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2669,34 +2669,15 @@ class ProxyLogging:
|
|||
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict)
|
||||
else:
|
||||
user_api_key_auth_dict = user_api_key_dict
|
||||
# Add task to list for parallel execution
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not getattr(callback, "use_native_during_call_hook", False)
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
guardrail_tasks.append(
|
||||
self._run_during_call_guardrail(
|
||||
callback=callback,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_api_key_auth_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
else:
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
guardrail_tasks.append(guardrail_task)
|
||||
)
|
||||
|
||||
# Step 2: Run all guardrail tasks in parallel
|
||||
if guardrail_tasks:
|
||||
|
|
@ -2708,6 +2689,41 @@ class ProxyLogging:
|
|||
|
||||
return data
|
||||
|
||||
async def _run_during_call_guardrail(
|
||||
self,
|
||||
callback: CustomGuardrail,
|
||||
data: dict[str, object], # mutable-ok: request payload dict, guardrail_to_apply is written in place
|
||||
user_api_key_dict: UserAPIKeyAuth | None,
|
||||
user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> None:
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not callback.use_native_during_call_hook
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
return
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
|
||||
async def failed_tracking_alert(
|
||||
self,
|
||||
error_message: str,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
)
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
|
||||
from litellm.types.integrations.custom_logger import converted_stream_requested
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._failure_handled = False # Track if failure handler has been called
|
||||
self._yielded_first_chunk = False
|
||||
self._generated_content = ""
|
||||
self._generated_tool_arguments = ""
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: bool | None = None
|
||||
|
|
@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_delta, str):
|
||||
self._generated_content += _delta
|
||||
elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS:
|
||||
_args_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_args_delta, str):
|
||||
self._generated_tool_arguments += _args_delta
|
||||
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
|
||||
if _event_type in (
|
||||
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
|
|
@ -419,14 +427,41 @@ class BaseResponsesAPIStreamingIterator:
|
|||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
):
|
||||
self.completed_response = openai_responses_api_chunk
|
||||
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
|
||||
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
_estimate_wanted: Final[bool] = _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
)
|
||||
_billed_response: Final[ResponsesAPIResponse | None] = _billed_terminal_response(
|
||||
_response_obj,
|
||||
(
|
||||
lambda: (
|
||||
_estimate_usage_safely(
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
self._generated_content + self._generated_tool_arguments,
|
||||
)
|
||||
if _estimate_wanted
|
||||
else None
|
||||
)
|
||||
),
|
||||
)
|
||||
_terminal_chunk: Final = (
|
||||
openai_responses_api_chunk
|
||||
if _billed_response is None or _billed_response is _response_obj
|
||||
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
|
||||
)
|
||||
self.completed_response = _terminal_chunk
|
||||
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
self._handle_logging_failed_response()
|
||||
else:
|
||||
self._handle_logging_completed_response()
|
||||
|
||||
return _terminal_chunk
|
||||
|
||||
return openai_responses_api_chunk
|
||||
|
||||
return None
|
||||
|
|
@ -655,7 +690,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if cache is None:
|
||||
return
|
||||
|
||||
cached_response: Final = response_obj.model_dump_json()
|
||||
cached_response: Final = _dump_json_safely(response_obj)
|
||||
if cached_response is None:
|
||||
return
|
||||
if is_async:
|
||||
from litellm.caching.caching_handler import create_cache_write_task
|
||||
|
||||
|
|
@ -1301,6 +1338,31 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
def _billed_terminal_response(
|
||||
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
|
||||
) -> ResponsesAPIResponse | None:
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
return (
|
||||
response_obj
|
||||
if response_obj.usage is not None or estimate is None
|
||||
else response_obj.model_copy(update={"usage": estimate()})
|
||||
)
|
||||
if not isinstance(response_obj, dict):
|
||||
return None
|
||||
usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict
|
||||
return ResponsesAPIResponse.model_construct(
|
||||
**{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread
|
||||
)
|
||||
|
||||
|
||||
def _dump_json_safely(response: BaseModel) -> str | None:
|
||||
try:
|
||||
return response.model_dump_json()
|
||||
except Exception as exc:
|
||||
verbose_logger.debug("could not serialize completed response for cache: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _logging_copy(event: object) -> object:
|
||||
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
|
||||
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
|
||||
|
|
@ -1332,6 +1394,56 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
|
|||
return None
|
||||
|
||||
|
||||
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
|
||||
{
|
||||
ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
|
||||
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_from_text(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage:
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
|
||||
responses_api_request=dict(responses_api_request),
|
||||
)
|
||||
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, messages=messages
|
||||
)
|
||||
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, text=generated_text, count_response_tokens=True
|
||||
)
|
||||
return ResponseAPIUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_safely(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage | None:
|
||||
try:
|
||||
return _estimate_usage_from_text(
|
||||
model=model,
|
||||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
generated_text=generated_text,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _stamp_responses_usage_cost(
|
||||
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -586,7 +586,7 @@ class MessageBlockDelta(TypedDict):
|
|||
|
||||
type: Literal["message_delta"]
|
||||
delta: MessageDelta
|
||||
usage: UsageDelta
|
||||
usage: NotRequired[ReadOnly[UsageDelta]]
|
||||
context_management: NotRequired[ContextManagementResponse]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,24 +1,53 @@
|
|||
from typing import Any
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class SingulrGuardrailRequest(BaseModel):
|
||||
model: str | None = None
|
||||
messages: list[dict[str, Any]] | None = None
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
model_response: dict[str, Any] | None = None
|
||||
litellm_metadata: dict[str, Any] | None = None
|
||||
class ContentBlock(BaseModel):
|
||||
type: str | None = None
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class ToolCallFunction(BaseModel):
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
id: str
|
||||
type: str = "function"
|
||||
function: ToolCallFunction
|
||||
|
||||
|
||||
class AssistantMessage(BaseModel):
|
||||
role: Literal["assistant"] = "assistant"
|
||||
content: str | Sequence[ContentBlock] | None = None
|
||||
tool_calls: Sequence[ToolCall] | None = None
|
||||
|
||||
|
||||
class SingulrGuardrailPayload(BaseModel):
|
||||
litellm_call_id: str | None = None
|
||||
request_data: SingulrGuardrailRequest | None = None
|
||||
input_type: str
|
||||
is_playground_request: bool | None = None
|
||||
playground_text: str | None = None
|
||||
correlation_id: str | None = None
|
||||
model_name: str | None = None
|
||||
model_provider_name: str | None = None
|
||||
guardrail_scope: str | None = None
|
||||
messages: Sequence[Mapping[str, object]] | None = None
|
||||
images: Sequence[str] | None = None
|
||||
tools: Sequence[Mapping[str, object]] | None = None
|
||||
response: AssistantMessage | None = None
|
||||
metadata: Mapping[str, str] | None = None
|
||||
|
||||
|
||||
class SingulrMcpGuardrailPayload(BaseModel):
|
||||
model_name: str | None = None
|
||||
guardrail_scope: str | None = None
|
||||
tool_name: str | None = None
|
||||
tool_arguments: object = None
|
||||
mcp_server_name: str | None = None
|
||||
tool_result: Sequence[str] | None = None
|
||||
metadata: Mapping[str, str] | None = None
|
||||
|
||||
|
||||
class SingulrGuardrailResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -2957,6 +2957,7 @@ InternalCallOrigin = Literal[
|
|||
"autorouter_classifier",
|
||||
"shadow_eval_router",
|
||||
"shadow_eval_judge",
|
||||
"llm_as_a_judge_guardrail",
|
||||
"background_response_cost_poll",
|
||||
]
|
||||
"""Which internal litellm feature originated a billed sub-call, so a spend log row
|
||||
|
|
@ -2965,6 +2966,7 @@ records that it is not traffic the caller sent."""
|
|||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier"
|
||||
SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
|
||||
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
|
||||
LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail"
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"
|
||||
|
||||
|
||||
|
|
|
|||
472
tests/code_coverage_tests/test_provider_cache.py
Normal file
472
tests/code_coverage_tests/test_provider_cache.py
Normal file
|
|
@ -0,0 +1,472 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from http.client import HTTPConnection
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward
|
||||
from models import LiteLLMParamsBody
|
||||
from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response
|
||||
from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model
|
||||
from provider_edge import configured_cache_backend, start_provider_edge
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
SECRET: Final = b"synthetic-cache-hmac-key-for-tests"
|
||||
BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}'
|
||||
SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}'
|
||||
HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"}
|
||||
|
||||
|
||||
class Provider(ThreadingHTTPServer):
|
||||
hits: tuple[tuple[str, bytes], ...] = ()
|
||||
response: bytes = SUCCESS
|
||||
status: int = 200
|
||||
delay: float = 0
|
||||
stream: bool = False
|
||||
truncated: bool = False
|
||||
cookie: str = ""
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def do_POST(self) -> None:
|
||||
server: Final = self.server
|
||||
assert isinstance(server, Provider)
|
||||
body: Final = self.rfile.read(int(self.headers.get("content-length", "0")))
|
||||
server.hits += ((self.path, body),)
|
||||
time.sleep(server.delay)
|
||||
self.send_response(server.status)
|
||||
if server.stream:
|
||||
self.send_header("content-type", "text/event-stream")
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"%x\r\n%s\r\n" % (len(server.response), server.response))
|
||||
if server.truncated:
|
||||
self.close_connection = True
|
||||
return
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
return
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", str(len(server.response)))
|
||||
if server.cookie:
|
||||
self.send_header("set-cookie", server.cookie)
|
||||
self.end_headers()
|
||||
self.wfile.write(server.response)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def provider() -> Generator[Provider, None, None]:
|
||||
server: Final = Provider(("127.0.0.1", 0), Handler)
|
||||
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield server
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def redis_url(tmp_path_factory: pytest.TempPathFactory) -> Generator[str, None, None]:
|
||||
configured: Final = os.environ.get("E2E_CACHE_TEST_REDIS_URL")
|
||||
if configured:
|
||||
yield configured
|
||||
return
|
||||
binary: Final = shutil.which("redis-server")
|
||||
assert binary is not None, "Set E2E_CACHE_TEST_REDIS_URL or install Redis for cache integration checks"
|
||||
root: Final = tmp_path_factory.mktemp("provider-cache-redis")
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port: Final = probe.getsockname()[1]
|
||||
with (root / "redis.log").open("wb") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[binary, "--bind", "127.0.0.1", "--port", str(port), "--save", "", "--appendonly", "no", "--dir", str(root)],
|
||||
stdout=log, stderr=subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
deadline: Final = time.monotonic() + 5
|
||||
while True:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.1):
|
||||
break
|
||||
except OSError:
|
||||
assert process.poll() is None and time.monotonic() < deadline
|
||||
time.sleep(0.02)
|
||||
yield f"redis://127.0.0.1:{port}/0"
|
||||
finally:
|
||||
process.terminate()
|
||||
process.wait(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(redis_url: str) -> RedisResponseStore:
|
||||
return redis_store(redis_url, "test-" + uuid.uuid4().hex)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]:
|
||||
upstream: Final = f"http://127.0.0.1:{provider.server_port}"
|
||||
running: Final = start_provider_edge(cache, mounts={"openai": upstream})
|
||||
try:
|
||||
yield running.edge.api_base("openai") + "/v1/chat/completions"
|
||||
finally:
|
||||
running.shutdown()
|
||||
|
||||
|
||||
def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse:
|
||||
result: Final = forward("POST", url, headers=headers, body=body, timeout=5)
|
||||
assert isinstance(result, RawResponse), result
|
||||
return result
|
||||
|
||||
|
||||
def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
with edge(CacheEdge(store, SECRET), provider) as other:
|
||||
assert call(other).body == SUCCESS
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")])
|
||||
def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url, body)
|
||||
call(url, body)
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")])
|
||||
def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url, headers=HEADERS | {name: value})
|
||||
call(url + "?x=1")
|
||||
assert len(provider.hits) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status,response", [(429, b'{"error":"rate limited"}'), (500, b'failed'), (200, b'{"error":"bad"}'), (200, b'not json')])
|
||||
def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None:
|
||||
provider.status = status
|
||||
provider.response = response
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).status_code == status
|
||||
assert call(url).body == response
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None:
|
||||
provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure"
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
replies: Final = tuple(call(url) for _ in range(2))
|
||||
assert len(provider.hits) == 1
|
||||
assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies)
|
||||
|
||||
|
||||
def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None:
|
||||
short: Final = replace(store, lifetime_ms=250)
|
||||
with edge(CacheEdge(short, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url)
|
||||
time.sleep(0.3)
|
||||
call(url)
|
||||
call(url)
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None:
|
||||
provider.delay = 0.15
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
with ThreadPoolExecutor(max_workers=5) as executor:
|
||||
replies: Final = tuple(executor.map(lambda _: call(url).body, range(5)))
|
||||
assert replies == (SUCCESS,) * 5
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("age_past_expiry_ms", [0, 1])
|
||||
def test_expired_response_is_rejected_without_physical_eviction(
|
||||
store: RedisResponseStore, age_past_expiry_ms: int,
|
||||
) -> None:
|
||||
response_key: Final = store.keys("expired")[0]
|
||||
retained: Final = store.client.eval(
|
||||
"""
|
||||
local clock = redis.call('TIME')
|
||||
local expires = clock[1] * 1000 + math.floor(clock[2] / 1000) - tonumber(ARGV[1])
|
||||
redis.call('HSET', KEYS[1], 'captured', expires - 86400000, 'expires', expires, 'payload', 'old-response')
|
||||
return redis.call('PTTL', KEYS[1])
|
||||
""",
|
||||
1, response_key, age_past_expiry_ms,
|
||||
)
|
||||
assert retained == -1
|
||||
replacement: Final = store.lookup("expired")
|
||||
assert isinstance(replacement, CaptureLease)
|
||||
assert replacement.expires_at_ms - replacement.captured_at_ms == 86_400_000
|
||||
assert store.publish("expired", replacement, b"fresh-response")
|
||||
hit: Final = store.lookup("expired")
|
||||
assert isinstance(hit, CacheHit) and hit.payload == b"fresh-response"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("truncated", [False, True])
|
||||
def test_stream_completion_controls_publication(store: RedisResponseStore, provider: Provider, truncated: bool) -> None:
|
||||
provider.stream = True
|
||||
provider.truncated = truncated
|
||||
provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n'
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
for _ in range(2):
|
||||
result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5)
|
||||
if truncated:
|
||||
assert isinstance(result, NetworkError)
|
||||
else:
|
||||
assert isinstance(result, RawResponse) and result.body == provider.response
|
||||
assert len(provider.hits) == (2 if truncated else 1)
|
||||
|
||||
|
||||
def test_store_outage_preserves_provider_success(provider: Provider) -> None:
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port: Final = probe.getsockname()[1]
|
||||
unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable")
|
||||
with edge(CacheEdge(unavailable, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None:
|
||||
short: Final = replace(store, lease_ms=50)
|
||||
old: Final = short.lookup("key")
|
||||
assert isinstance(old, CaptureLease)
|
||||
time.sleep(0.08)
|
||||
current: Final = short.lookup("key")
|
||||
assert isinstance(current, CaptureLease)
|
||||
assert not short.publish("key", old, b"old")
|
||||
assert short.publish("key", current, b"new")
|
||||
hit: Final = short.lookup("key")
|
||||
assert isinstance(hit, CacheHit) and hit.payload == b"new"
|
||||
|
||||
|
||||
def test_identity_preserves_values_and_never_contains_credentials() -> None:
|
||||
variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'')
|
||||
keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants)
|
||||
assert len(set(keys)) == len(variants)
|
||||
assert all(len(key) == 64 and "synthetic-account" not in key for key in keys)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()])
|
||||
def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
|
||||
upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY)
|
||||
lease: Final = store.lookup(key)
|
||||
assert isinstance(lease, CaptureLease)
|
||||
assert store.publish(key, lease, payload)
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
for _ in range(2):
|
||||
head = cache.forward("POST", upstream, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
|
||||
assert len(provider.hits) == 1
|
||||
assert dict(cache.counters.counts) == {
|
||||
"corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [
|
||||
b'data: {}\n\ndata: [DONE]\n\n',
|
||||
b'data: {"choices":[{"index":0,"delta":{}}]}\n\ndata: [DONE]\n\n',
|
||||
b'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]',
|
||||
b'data: {"error":{"message":"failed"}}\n\ndata: [DONE]\n\n',
|
||||
])
|
||||
def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
|
||||
provider.stream = True
|
||||
provider.response = payload
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == payload
|
||||
assert call(url).body == payload
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_anthropic_stream_requires_start_finish_and_stop() -> None:
|
||||
start: Final = b'data: {"type":"message_start","message":{}}\n\n'
|
||||
finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n'
|
||||
stop: Final = b'data: {"type":"message_stop"}\n\n'
|
||||
url: Final = "https://example.invalid/v1/messages"
|
||||
headers: Final = {"content-type": "text/event-stream"}
|
||||
assert successful_response(url, 200, headers, start + finish + stop)
|
||||
assert not successful_response(url, 200, headers, start + stop)
|
||||
assert not successful_response(url, 200, headers, finish + stop)
|
||||
assert not successful_response(url, 200, headers, start + finish)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")])
|
||||
def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None:
|
||||
params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12)
|
||||
routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True)
|
||||
assert routed.api_base == f"http://edge.invalid/{provider}{suffix}"
|
||||
assert routed.model_dump(exclude={"api_base"}) == params.model_dump(exclude={"api_base"})
|
||||
assert params.api_base is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("params", [
|
||||
LiteLLMParamsBody(model="bedrock/test"),
|
||||
LiteLLMParamsBody(model="azure/test"),
|
||||
LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"),
|
||||
LiteLLMParamsBody(model="openai/test", api_base=""),
|
||||
LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"),
|
||||
LiteLLMParamsBody(model="openai/test", mock_response="synthetic"),
|
||||
])
|
||||
def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMParamsBody) -> None:
|
||||
def unexpected_edge(mount: str) -> str:
|
||||
pytest.fail(f"should not start edge for {mount}")
|
||||
assert route_cache_model(params, unexpected_edge, enabled=True) is params
|
||||
|
||||
|
||||
def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None:
|
||||
params: Final = LiteLLMParamsBody(model="openai/test")
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True, mode="realtime") is params
|
||||
token: Final = LIVE_PROVIDER_REQUIRED.set(True)
|
||||
try:
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True) is params
|
||||
finally:
|
||||
LIVE_PROVIDER_REQUIRED.reset(token)
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True).api_base == "http://edge.invalid/v1"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PublishOutage:
|
||||
client: RedisCommands
|
||||
|
||||
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object:
|
||||
if script == PUBLISH:
|
||||
raise RedisConnectionError("synthetic publication outage")
|
||||
return self.client.eval(script, numkeys, *args)
|
||||
|
||||
|
||||
def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None:
|
||||
unavailable: Final = replace(store, client=PublishOutage(store.client))
|
||||
cache: Final = CacheEdge(unavailable, SECRET)
|
||||
with edge(cache, provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
assert dict(cache.counters.counts)["write_failures"] == 2
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 3
|
||||
|
||||
|
||||
def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> None:
|
||||
with socket.socket() as unavailable:
|
||||
unavailable.bind(("127.0.0.1", 0))
|
||||
url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError)
|
||||
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
|
||||
slot: Final = store.lookup(key)
|
||||
assert isinstance(slot, CaptureLease)
|
||||
assert store.release(key, slot)
|
||||
assert dict(cache.counters.counts)["rejected"] == 1
|
||||
|
||||
|
||||
def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None:
|
||||
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
head: Final = cache.forward("POST", url, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
head.steps.close()
|
||||
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
|
||||
slot: Final = store.lookup(key)
|
||||
assert isinstance(slot, CaptureLease)
|
||||
assert store.release(key, slot)
|
||||
|
||||
|
||||
def test_effective_account_change_cannot_reuse_cache(
|
||||
store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path,
|
||||
) -> None:
|
||||
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
for account in ("account-a", "account-b", "account-b"):
|
||||
netrc = tmp_path / account
|
||||
netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n")
|
||||
monkeypatch.setenv("NETRC", str(netrc))
|
||||
head = cache.forward("POST", url, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
assert dict(cache.counters.counts)["hits"] == 1
|
||||
|
||||
|
||||
def test_enabled_environment_reuses_store_across_fresh_backends(
|
||||
redis_url: str, provider: Provider, monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE", "1")
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_REDIS_URL", redis_url)
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_HMAC_KEY", SECRET.decode())
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_NAMESPACE", "environment-" + uuid.uuid4().hex)
|
||||
configured_cache.cache_clear()
|
||||
try:
|
||||
for _ in range(2):
|
||||
backend = configured_cache_backend()
|
||||
assert isinstance(backend, CacheEdge)
|
||||
with edge(backend, provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
configured_cache.cache_clear()
|
||||
assert len(provider.hits) == 1
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE", "0")
|
||||
assert configured_cache_backend() is None
|
||||
finally:
|
||||
configured_cache.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("known_mount", (True, False))
|
||||
def test_duplicate_headers_bypass_cache_and_count_live_calls(
|
||||
store: RedisResponseStore, provider: Provider, known_mount: bool,
|
||||
) -> None:
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
with edge(cache, provider) as url:
|
||||
parsed: Final = urlsplit(url)
|
||||
for _ in range(2):
|
||||
connection = HTTPConnection(str(parsed.hostname), parsed.port, timeout=5)
|
||||
try:
|
||||
connection.putrequest("POST", parsed.path if known_mount else "/unknown/v1/chat/completions")
|
||||
connection.putheader("content-length", str(len(BODY)))
|
||||
connection.putheader("content-type", "application/json")
|
||||
connection.putheader("x-duplicate", "first")
|
||||
connection.putheader("x-duplicate", "second")
|
||||
connection.endheaders(BODY)
|
||||
response = connection.getresponse()
|
||||
assert response.status == (200 if known_mount else 404)
|
||||
payload = response.read()
|
||||
assert payload == SUCCESS if known_mount else b"unknown provider mount" in payload
|
||||
finally:
|
||||
connection.close()
|
||||
assert len(provider.hits) == (2 if known_mount else 0)
|
||||
assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2
|
||||
assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0)
|
||||
33
tests/e2e/PROVIDER_CACHE.md
Normal file
33
tests/e2e/PROVIDER_CACHE.md
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
# Shared provider-response cache
|
||||
|
||||
`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
|
||||
|
||||
The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
|
||||
|
||||
An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
|
||||
|
||||
Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires
|
||||
|
||||
## Configuration
|
||||
|
||||
The trusted runner receives:
|
||||
|
||||
- `E2E_PROVIDER_CACHE`: `1` to enable, `0` to use the normal live path
|
||||
- `E2E_PROVIDER_CACHE_REDIS_URL`: authenticated dedicated Redis URL
|
||||
- `E2E_PROVIDER_CACHE_HMAC_KEY`: dedicated secret containing at least 32 bytes
|
||||
- `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision
|
||||
- `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory
|
||||
|
||||
Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
|
||||
|
||||
Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
|
||||
|
||||
## Recorded response semantics
|
||||
|
||||
Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching
|
||||
|
||||
Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers
|
||||
|
||||
## Qualification
|
||||
|
||||
`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence
|
||||
|
|
@ -22,7 +22,6 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from e2e_config import (
|
||||
CONTROL_PLANE_BASE_URL,
|
||||
FIXTURE_DIR,
|
||||
|
|
@ -41,6 +40,7 @@ from idp import Identity, Keycloak, keycloak_from_env
|
|||
from junit_properties import attach_result_properties
|
||||
from lifecycle import ProxyClientProvider, ResourceManager
|
||||
from models import TeamNewBody, UserNewBody, UserNewResponse
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
|
||||
from provider_edge import replay_leftover_error
|
||||
from proxy_client import ProxyClient, build_proxy_client
|
||||
|
||||
|
|
@ -85,6 +85,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
|
|||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache")
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"e2e: live test that requires a running proxy and real provider keys",
|
||||
|
|
@ -192,11 +193,13 @@ def _proxy_fail_reason() -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_runtest_setup(item: pytest.Item) -> None:
|
||||
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
|
||||
Unmarked tests (unit coverage of the harness) don't touch the proxy, so they
|
||||
run even when none is up. Never skip for a missing proxy. Replay mode needs
|
||||
the proxy too: only provider-bound traffic replays from the bundle."""
|
||||
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
|
||||
if item.get_closest_marker("e2e") is None:
|
||||
return
|
||||
reason = _proxy_fail_reason()
|
||||
|
|
@ -235,6 +238,7 @@ def pytest_runtest_teardown(item: pytest.Item) -> Generator[None, None, None]:
|
|||
yield so fixture finalizers replay their recorded calls first. Failed tests
|
||||
are left alone - their own failure already explains any unconsumed tail."""
|
||||
result = yield
|
||||
LIVE_PROVIDER_REQUIRED.set(False)
|
||||
if not item.stash.get(_CALL_PASSED, False):
|
||||
return result
|
||||
reason = replay_leftover_error(
|
||||
|
|
|
|||
|
|
@ -853,6 +853,7 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
|
|||
the chunks already delivered are exactly what makes a mid-stream failure
|
||||
different from a request that never streamed at all."""
|
||||
try:
|
||||
yield StreamChunk(b"")
|
||||
for piece in cast("Iterator[bytes]", resp.iter_content(chunk_size=None)):
|
||||
if piece:
|
||||
yield StreamChunk(data=piece)
|
||||
|
|
@ -862,6 +863,44 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
|
|||
resp.close()
|
||||
|
||||
|
||||
def primed_steps(steps: Generator[StreamStep, None, None]) -> Generator[StreamStep, None, None]:
|
||||
first: Final = next(steps)
|
||||
assert isinstance(first, StreamChunk) and first.data == b""
|
||||
return steps
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
class PreparedForward:
|
||||
request: requests.PreparedRequest
|
||||
url: str
|
||||
headers: dict[str, str]
|
||||
|
||||
|
||||
def prepare_forward(
|
||||
method: str, url: str, headers: dict[str, str], body: bytes | None,
|
||||
) -> PreparedForward | NetworkError:
|
||||
try:
|
||||
with requests.Session() as session:
|
||||
request: Final = session.prepare_request(requests.Request(method, url, headers=headers, data=body))
|
||||
except requests.RequestException as exc:
|
||||
return NetworkError(message=str(exc))
|
||||
assert request.url is not None
|
||||
return PreparedForward(request, request.url, dict(request.headers))
|
||||
|
||||
|
||||
def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> StreamHead | NetworkError:
|
||||
try:
|
||||
with requests.Session() as session:
|
||||
settings: Final = session.merge_environment_settings(prepared.url, {}, True, None, None)
|
||||
resp: Final = session.send(prepared.request, timeout=timeout, allow_redirects=False, **settings)
|
||||
except requests.RequestException as exc:
|
||||
return NetworkError(message=str(exc))
|
||||
return StreamHead(
|
||||
resp.status_code, {name.lower(): value for name, value in resp.headers.items()},
|
||||
primed_steps(_stream_steps(resp)),
|
||||
)
|
||||
|
||||
|
||||
def open_stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamHead | NetworkError:
|
||||
"""POST a streaming request and return the moment its response head arrives,
|
||||
leaving the body unread behind ``StreamHead.steps``. For a test that must keep
|
||||
|
|
@ -907,5 +946,5 @@ def forward_stream(
|
|||
return StreamHead(
|
||||
status_code=resp.status_code,
|
||||
headers={name.lower(): value for name, value in resp.headers.items()},
|
||||
steps=_stream_steps(resp),
|
||||
steps=primed_steps(_stream_steps(resp)),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage
|
|||
from passthrough_client import PassthroughClient
|
||||
import os
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
|
||||
|
||||
BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
VERTEX_MODEL = "vertex_ai/gemini-2.5-flash"
|
||||
|
|
|
|||
|
|
@ -24,10 +24,10 @@ from models import (
|
|||
AnthropicAssistantTurn,
|
||||
AnthropicContentBlock,
|
||||
AnthropicCustomTool,
|
||||
AnthropicMessagesBody,
|
||||
AnthropicToolChoice,
|
||||
AnthropicToolResultBlock,
|
||||
AnthropicToolResultTurn,
|
||||
AnthropicMessagesBody,
|
||||
ChatMessage,
|
||||
JsonSchemaProperty,
|
||||
LiteLLMParamsBody,
|
||||
|
|
@ -165,6 +165,7 @@ class TestAnthropicMessages:
|
|||
)
|
||||
|
||||
@pytest.mark.covers("llm.messages.anthropic.basic.stream.works")
|
||||
@pytest.mark.provider_live
|
||||
def test_messages_streams_completion(
|
||||
self, endpoints_client: EndpointsClient, resources: ResourceManager
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -744,6 +744,7 @@ class TestTogetherMessages:
|
|||
assert "22" in text, f"the model never saw the tool result: {response.content}"
|
||||
|
||||
@pytest.mark.covers("llm.messages.together_ai.basic.stream.works")
|
||||
@pytest.mark.provider_live
|
||||
def test_streams_text_deltas(
|
||||
self, client: PassthroughClient, resources: ResourceManager, reasoning_tool_backend: str
|
||||
) -> None:
|
||||
|
|
|
|||
299
tests/e2e/provider_cache.py
Normal file
299
tests/e2e/provider_cache.py
Normal file
|
|
@ -0,0 +1,299 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal, Protocol
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from e2e_http import (
|
||||
NetworkError,
|
||||
StreamChunk,
|
||||
StreamHead,
|
||||
StreamStep,
|
||||
StreamTruncation,
|
||||
forward_prepared_stream,
|
||||
forward_stream,
|
||||
prepare_forward,
|
||||
primed_steps,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
LIFETIME_SECONDS: Final = 86_400
|
||||
MAX_REQUEST_BYTES: Final = 256 * 1024
|
||||
MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
|
||||
UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
|
||||
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheHit:
|
||||
payload: bytes
|
||||
valid_until: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaptureLease:
|
||||
token: str
|
||||
captured_at_ms: int
|
||||
expires_at_ms: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheBusy:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheUnavailable:
|
||||
pass
|
||||
|
||||
|
||||
type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable
|
||||
|
||||
|
||||
class ResponseStore(Protocol):
|
||||
def lookup(self, key: str) -> CacheLookup: ...
|
||||
|
||||
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool: ...
|
||||
|
||||
def release(self, key: str, lease: CaptureLease) -> bool: ...
|
||||
|
||||
def discard(self, key: str, payload: bytes) -> bool: ...
|
||||
|
||||
|
||||
class CachedResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
format_version: Literal[1] = 1
|
||||
request_key: str
|
||||
status_code: int
|
||||
headers: dict[str, str]
|
||||
chunks: tuple[str, ...]
|
||||
|
||||
|
||||
class SignedResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
response: str
|
||||
signature: str
|
||||
|
||||
|
||||
def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str:
|
||||
fields: Final = (
|
||||
b"provider-cache-exact-v1", method.encode(), url.encode(),
|
||||
*(part.encode() for pair in sorted(headers.items()) for part in pair),
|
||||
b"no-body" if body is None else b"body", b"" if body is None else body,
|
||||
)
|
||||
encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields)
|
||||
return hmac.new(secret, encoded, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
|
||||
return (
|
||||
method == "POST"
|
||||
and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"}
|
||||
and body is not None
|
||||
and len(body) <= MAX_REQUEST_BYTES
|
||||
)
|
||||
|
||||
|
||||
def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
|
||||
if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
|
||||
return False
|
||||
streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
|
||||
if streaming:
|
||||
try:
|
||||
text: Final = body.decode("utf-8").replace("\r\n", "\n")
|
||||
if not text.endswith("\n\n"):
|
||||
return False
|
||||
events: Final = tuple(
|
||||
"\n".join(line[5:].removeprefix(" ") for line in event.split("\n") if line.startswith("data:"))
|
||||
for event in text.split("\n\n") if any(line.startswith("data:") for line in event.split("\n"))
|
||||
)
|
||||
values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]")
|
||||
except (UnicodeDecodeError, ValidationError):
|
||||
return False
|
||||
if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values):
|
||||
return False
|
||||
if urlsplit(url).path == "/v1/chat/completions":
|
||||
return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values)
|
||||
return (
|
||||
"[DONE]" not in events
|
||||
and isinstance(values[0], dict) and values[0].get("type") == "message_start"
|
||||
and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop"
|
||||
and any(
|
||||
isinstance(value, dict) and value.get("type") == "message_delta"
|
||||
and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
|
||||
for value in values
|
||||
)
|
||||
)
|
||||
try:
|
||||
value: Final = JSON_VALUE.validate_json(body)
|
||||
except ValidationError:
|
||||
return False
|
||||
if not isinstance(value, dict) or "error" in value:
|
||||
return False
|
||||
if urlsplit(url).path == "/v1/messages":
|
||||
return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str)
|
||||
choices: Final = value.get("choices")
|
||||
return isinstance(choices, list) and bool(choices) and all(
|
||||
isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str)
|
||||
for choice in choices
|
||||
)
|
||||
|
||||
|
||||
def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool:
|
||||
if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values):
|
||||
return False
|
||||
choices: Final = tuple(
|
||||
choice for value in values if isinstance(value, dict)
|
||||
if isinstance(items := value.get("choices"), list) for choice in items
|
||||
)
|
||||
if not choices or any(
|
||||
not isinstance(choice, dict) or type(choice.get("index")) is not int
|
||||
or not isinstance(choice.get("delta"), dict)
|
||||
for choice in choices
|
||||
):
|
||||
return False
|
||||
indices: Final = frozenset(choice["index"] for choice in choices if isinstance(choice, dict))
|
||||
return all(
|
||||
isinstance(tuple(choice for choice in choices if isinstance(choice, dict) and choice["index"] == index)[-1].get("finish_reason"), str)
|
||||
for index in indices
|
||||
)
|
||||
|
||||
|
||||
def encode_response(secret: bytes, response: CachedResponse) -> bytes:
|
||||
raw: Final = response.model_dump_json()
|
||||
return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode()
|
||||
|
||||
|
||||
def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None:
|
||||
if len(payload) > 2 * MAX_RESPONSE_BYTES:
|
||||
return None
|
||||
try:
|
||||
signed: Final = SignedResponse.model_validate_json(payload)
|
||||
if not hmac.compare_digest(signed.signature.encode(), hmac.new(secret, signed.response.encode(), hashlib.sha256).hexdigest().encode()):
|
||||
return None
|
||||
response: Final = CachedResponse.model_validate_json(signed.response)
|
||||
chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks)
|
||||
except (ValidationError, ValueError):
|
||||
return None
|
||||
if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)):
|
||||
return None
|
||||
return response
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CacheCounters:
|
||||
counts: tuple[tuple[str, int], ...] = ()
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def increment(self, name: str) -> None:
|
||||
with self.lock:
|
||||
current: Final = dict(self.counts)
|
||||
self.counts = tuple((current | {name: current.get(name, 0) + 1}).items())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ResponseCapture:
|
||||
buffer: io.BytesIO = field(default_factory=io.BytesIO)
|
||||
size: int = 0
|
||||
eligible: bool = True
|
||||
|
||||
def observe(self, step: StreamStep) -> None:
|
||||
if not self.eligible:
|
||||
return
|
||||
if isinstance(step, StreamTruncation) or self.size + len(step.data) + 8 > MAX_RESPONSE_BYTES:
|
||||
self.eligible = False
|
||||
self.buffer.close()
|
||||
return
|
||||
self.buffer.write(len(step.data).to_bytes(8, "big"))
|
||||
self.buffer.write(step.data)
|
||||
self.size += len(step.data) + 8
|
||||
|
||||
def chunks(self) -> tuple[bytes, ...]:
|
||||
self.buffer.seek(0)
|
||||
return tuple(self.buffer.read(int.from_bytes(size, "big")) for size in iter(lambda: self.buffer.read(8), b""))
|
||||
|
||||
|
||||
def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None]:
|
||||
for chunk in response.chunks:
|
||||
yield StreamChunk(base64.b64decode(chunk, validate=True))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheEdge:
|
||||
store: ResponseStore
|
||||
secret: bytes = field(repr=False)
|
||||
counters: CacheCounters = field(default_factory=CacheCounters)
|
||||
wait_seconds: float = 2.0
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
sleep: Callable[[float], None] = time.sleep
|
||||
|
||||
def lookup(self, key: str) -> CacheLookup:
|
||||
deadline: Final = self.clock() + self.wait_seconds
|
||||
while isinstance(result := self.store.lookup(key), CacheBusy) and self.clock() < deadline:
|
||||
self.sleep(min(0.05, max(0, deadline - self.clock())))
|
||||
return result
|
||||
|
||||
def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError:
|
||||
if not cacheable_endpoint(method, url, body):
|
||||
self.counters.increment("bypass")
|
||||
self.counters.increment("upstream_attempts")
|
||||
return forward_stream(method, url, headers=headers, body=body, timeout=timeout)
|
||||
prepared: Final = prepare_forward(method, url, headers, body)
|
||||
if isinstance(prepared, NetworkError):
|
||||
self.counters.increment("rejected")
|
||||
return prepared
|
||||
key: Final = exact_key(self.secret, method, url, prepared.headers, body)
|
||||
found: Final = self.lookup(key)
|
||||
if isinstance(found, CacheHit):
|
||||
response: Final = decode_response(self.secret, key, found.payload, url)
|
||||
if response is not None and self.clock() < found.valid_until:
|
||||
self.counters.increment("hits")
|
||||
return StreamHead(response.status_code, response.headers, response_steps(response))
|
||||
self.counters.increment("corrupt" if response is None else "expired")
|
||||
self.store.discard(key, found.payload)
|
||||
capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found
|
||||
self.counters.increment("misses")
|
||||
if isinstance(capture_slot, CacheUnavailable):
|
||||
self.counters.increment("cache_errors")
|
||||
self.counters.increment("upstream_attempts")
|
||||
head: Final = forward_prepared_stream(prepared, timeout)
|
||||
if not isinstance(capture_slot, CaptureLease):
|
||||
return head
|
||||
if isinstance(head, NetworkError):
|
||||
self.store.release(key, capture_slot)
|
||||
self.counters.increment("rejected")
|
||||
return head
|
||||
return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head)))
|
||||
|
||||
def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]:
|
||||
capture: Final = ResponseCapture()
|
||||
try:
|
||||
with closing(head.steps):
|
||||
yield StreamChunk(b"")
|
||||
for step in head.steps:
|
||||
yield step
|
||||
capture.observe(step)
|
||||
chunks: Final = capture.chunks() if capture.eligible else ()
|
||||
headers: Final = {
|
||||
name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
|
||||
}
|
||||
if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
|
||||
self.counters.increment("rejected")
|
||||
return
|
||||
response: Final = CachedResponse(
|
||||
request_key=key, status_code=head.status_code, headers=headers,
|
||||
chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
|
||||
)
|
||||
published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
|
||||
self.counters.increment("writes" if published else "write_failures")
|
||||
finally:
|
||||
self.store.release(key, lease)
|
||||
capture.buffer.close()
|
||||
154
tests/e2e/provider_cache_redis.py
Normal file
154
tests/e2e/provider_cache_redis.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Protocol, cast
|
||||
|
||||
from provider_cache import LIFETIME_SECONDS, CacheBusy, CacheEdge, CacheHit, CacheLookup, CacheUnavailable, CaptureLease
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from redis import Redis
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
REDIS_ARRAY: Final[TypeAdapter[list[bytes]]] = TypeAdapter(list[bytes])
|
||||
|
||||
LOOKUP: Final = """
|
||||
local clock = redis.call('TIME')
|
||||
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
|
||||
local row = redis.call('HMGET', KEYS[1], 'captured', 'expires', 'payload')
|
||||
if row[3] then
|
||||
local captured = tonumber(row[1])
|
||||
local expires = tonumber(row[2])
|
||||
if captured and expires and captured <= now and expires > now
|
||||
and expires - captured == tonumber(ARGV[2]) then
|
||||
return {'hit', row[3], tostring(expires - now)}
|
||||
end
|
||||
redis.call('DEL', KEYS[1])
|
||||
end
|
||||
if redis.call('SET', KEYS[2], ARGV[1], 'NX', 'PX', ARGV[3]) then
|
||||
return {'lease', tostring(now), tostring(now + tonumber(ARGV[2]))}
|
||||
end
|
||||
return {'busy'}
|
||||
"""
|
||||
|
||||
PUBLISH: Final = """
|
||||
if redis.call('GET', KEYS[2]) ~= ARGV[1] then return 0 end
|
||||
local clock = redis.call('TIME')
|
||||
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
|
||||
local captured = tonumber(ARGV[2])
|
||||
local expires = tonumber(ARGV[3])
|
||||
if captured > now or expires <= now or expires - captured ~= tonumber(ARGV[5]) then return 0 end
|
||||
if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end
|
||||
redis.call('HSET', KEYS[1], 'captured', ARGV[2], 'expires', ARGV[3], 'payload', ARGV[4])
|
||||
redis.call('PEXPIREAT', KEYS[1], expires)
|
||||
redis.call('DEL', KEYS[2])
|
||||
return 1
|
||||
"""
|
||||
|
||||
RELEASE: Final = """
|
||||
if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end
|
||||
return redis.call('DEL', KEYS[1])
|
||||
"""
|
||||
|
||||
DISCARD: Final = """
|
||||
if redis.call('HGET', KEYS[1], 'payload') ~= ARGV[1] then return 0 end
|
||||
return redis.call('DEL', KEYS[1])
|
||||
"""
|
||||
|
||||
|
||||
class RedisCommands(Protocol):
|
||||
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RedisResponseStore:
|
||||
client: RedisCommands
|
||||
namespace: str
|
||||
lifetime_ms: int = LIFETIME_SECONDS * 1000
|
||||
lease_ms: int = 120_000
|
||||
|
||||
def keys(self, key: str) -> tuple[str, str]:
|
||||
prefix: Final = f"e2e-provider-cache:v1:{self.namespace}:{{{key}}}"
|
||||
return prefix + ":response", prefix + ":lease"
|
||||
|
||||
def lookup(self, key: str) -> CacheLookup:
|
||||
token: Final = uuid.uuid4().hex
|
||||
started: Final = time.monotonic()
|
||||
try:
|
||||
result: Final = self.client.eval(LOOKUP, 2, *self.keys(key), token, self.lifetime_ms, self.lease_ms)
|
||||
except (RedisError, OSError):
|
||||
return CacheUnavailable()
|
||||
try:
|
||||
parts: Final = tuple(REDIS_ARRAY.validate_python(result, strict=True))
|
||||
except ValidationError:
|
||||
return CacheUnavailable()
|
||||
if len(parts) == 3 and parts[0] == b"hit" and parts[2].isdigit():
|
||||
return CacheHit(parts[1], started + int(parts[2]) / 1000)
|
||||
if len(parts) == 3 and parts[0] == b"lease" and parts[1].isdigit() and parts[2].isdigit():
|
||||
return CaptureLease(token, int(parts[1]), int(parts[2]))
|
||||
if parts == (b"busy",):
|
||||
return CacheBusy()
|
||||
return CacheUnavailable()
|
||||
|
||||
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(
|
||||
PUBLISH, 2, *self.keys(key), lease.token, lease.captured_at_ms, lease.expires_at_ms, payload, self.lifetime_ms,
|
||||
)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
def release(self, key: str, lease: CaptureLease) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(RELEASE, 1, self.keys(key)[1], lease.token)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
def discard(self, key: str, payload: bytes) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(DISCARD, 1, self.keys(key)[0], payload)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
|
||||
def redis_store(url: str, namespace: str) -> RedisResponseStore:
|
||||
client: Final = Redis.from_url(url, socket_timeout=0.25, socket_connect_timeout=0.25, decode_responses=False)
|
||||
return RedisResponseStore(cast(RedisCommands, client), namespace)
|
||||
|
||||
|
||||
def write_metrics(cache: CacheEdge) -> None:
|
||||
report: Final = json.dumps({"provider_cache": dict(cache.counters.counts)})
|
||||
directory: Final = os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR")
|
||||
if directory:
|
||||
try:
|
||||
root: Final = Path(directory)
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / f"{os.getpid()}.json").write_text(report + "\n")
|
||||
except OSError:
|
||||
logging.getLogger(__name__).warning("provider cache metrics artifact unavailable")
|
||||
logging.getLogger(__name__).info("%s", report)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def configured_cache() -> CacheEdge | None:
|
||||
if os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
|
||||
return None
|
||||
if os.environ.get("E2E_PROVIDER_CACHE") != "1":
|
||||
raise ValueError("E2E_PROVIDER_CACHE must be 0 or 1")
|
||||
secret: Final = os.environ.get("E2E_PROVIDER_CACHE_HMAC_KEY", "").encode()
|
||||
namespace: Final = os.environ.get("E2E_PROVIDER_CACHE_NAMESPACE", "")
|
||||
if len(secret) < 32 or re.fullmatch(r"[a-zA-Z0-9_-]{1,64}", namespace) is None:
|
||||
raise ValueError("provider cache requires a dedicated key and namespace")
|
||||
cache: Final = CacheEdge(redis_store(os.environ["E2E_PROVIDER_CACHE_REDIS_URL"], namespace), secret)
|
||||
atexit.register(write_metrics, cache)
|
||||
return cache
|
||||
23
tests/e2e/provider_cache_routing.py
Normal file
23
tests/e2e/provider_cache_routing.py
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextvars import ContextVar
|
||||
from typing import Final
|
||||
|
||||
from models import LiteLLMParamsBody, ModelMode
|
||||
|
||||
LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
|
||||
|
||||
|
||||
def route_cache_model(
|
||||
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
|
||||
) -> LiteLLMParamsBody:
|
||||
if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None:
|
||||
return params
|
||||
provider: Final = params.model.partition("/")[0]
|
||||
if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None:
|
||||
return params
|
||||
base: Final = base_for(provider)
|
||||
if base is None:
|
||||
return params
|
||||
return params.model_copy(update={"api_base": f"{base}/v1" if provider == "openai" else base})
|
||||
|
|
@ -42,6 +42,7 @@ import base64
|
|||
import difflib
|
||||
import functools
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
from collections import deque
|
||||
|
|
@ -93,6 +94,8 @@ from fixture_mode import (
|
|||
parse_fixture_mode,
|
||||
)
|
||||
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
|
||||
from provider_cache import CacheEdge
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
|
|
@ -506,7 +509,7 @@ class LiveEdge:
|
|||
pass
|
||||
|
||||
|
||||
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge
|
||||
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
|
@ -750,12 +753,16 @@ def _handle_record(
|
|||
|
||||
|
||||
def _handle_live(
|
||||
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float
|
||||
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
|
||||
cache: CacheEdge | None = None,
|
||||
) -> EdgeOutcome:
|
||||
forwarded: Final = {
|
||||
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
|
||||
}
|
||||
head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
|
||||
head: Final = (
|
||||
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
|
||||
if cache is None else cache.forward(method, url, forwarded, body, timeout)
|
||||
)
|
||||
match head:
|
||||
case NetworkError(message=message):
|
||||
return _recorded_outcome(_network_error_response(message))
|
||||
|
|
@ -821,6 +828,10 @@ def handle_edge_request(
|
|||
else edge_request(method, split.path, split.query, body, _header_value(headers, "content-type"))
|
||||
)
|
||||
match backend:
|
||||
case CacheEdge():
|
||||
return _handle_live(
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend,
|
||||
)
|
||||
case LiveEdge():
|
||||
return _handle_live(
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
|
||||
|
|
@ -871,11 +882,19 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
or isinstance(edge_server.backend, ReplayEdge)
|
||||
and edge_server.backend.source.bundle.manifest.match_profile == "stateless_v1"
|
||||
)
|
||||
if strict and len({name.lower() for name in self.headers.keys()}) != len(self.headers):
|
||||
if strict and len({name.lower() for name in self.headers}) != len(self.headers):
|
||||
self._write_reply(_text_reply(REPLAY_MISS_STATUS, "stateless_v1 eligibility error: duplicate headers"))
|
||||
return
|
||||
duplicate_headers: Final = len({name.lower() for name in self.headers}) != len(self.headers)
|
||||
selected_backend: Final = (
|
||||
LiveEdge() if isinstance(edge_server.backend, CacheEdge) and duplicate_headers else edge_server.backend
|
||||
)
|
||||
if isinstance(edge_server.backend, CacheEdge) and duplicate_headers:
|
||||
edge_server.backend.counters.increment("duplicate_header_bypass")
|
||||
if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts:
|
||||
edge_server.backend.counters.increment("upstream_attempts")
|
||||
outcome: Final = handle_edge_request(
|
||||
edge_server.backend,
|
||||
selected_backend,
|
||||
edge_server.mounts,
|
||||
self.command,
|
||||
self.path,
|
||||
|
|
@ -908,12 +927,12 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
shuts down write-side first: the proxy sees a graceful close mid-message,
|
||||
which is the incomplete chunked read a provider hanging up produces, and not
|
||||
the reset that could discard the chunks already in flight."""
|
||||
self.send_response(stream.status_code)
|
||||
for name, value in stream.headers.items():
|
||||
self.send_header(name, value)
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
with closing(stream.steps) as steps:
|
||||
self.send_response(stream.status_code)
|
||||
for name, value in stream.headers.items():
|
||||
self.send_header(name, value)
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
for step in steps:
|
||||
match step:
|
||||
case StreamChunk(data=data):
|
||||
|
|
@ -923,7 +942,7 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
return
|
||||
case _:
|
||||
assert_never(step)
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
"""Silence the per-request stderr line BaseHTTPRequestHandler emits."""
|
||||
|
|
@ -1056,6 +1075,8 @@ def provider_edge_api_base(
|
|||
case InvalidFixtureMode(value=value):
|
||||
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
|
||||
case "live":
|
||||
if configured_cache_backend() is not None:
|
||||
return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount)
|
||||
return None
|
||||
case "record" | "replay":
|
||||
if mount not in EDGE_MOUNTS:
|
||||
|
|
@ -1073,7 +1094,7 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
|
|||
case InvalidFixtureMode(value=value):
|
||||
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
|
||||
case "live":
|
||||
return LiveEdge()
|
||||
return configured_cache_backend() or LiveEdge()
|
||||
case "record":
|
||||
return RecordEdge(_shared_recorder(bundle_dir, match_profile()), threading.Lock())
|
||||
case "replay":
|
||||
|
|
@ -1082,6 +1103,24 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
|
|||
assert_never(mode)
|
||||
|
||||
|
||||
def configured_cache_backend() -> CacheEdge | None:
|
||||
if LIVE_PROVIDER_REQUIRED.get() or os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
|
||||
return None
|
||||
from provider_cache_redis import configured_cache
|
||||
|
||||
return configured_cache()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _shared_cache_edge(bind_host: str, advertise_host: str, forward_timeout: float) -> ProviderEdge:
|
||||
backend: Final = configured_cache_backend()
|
||||
assert backend is not None
|
||||
return start_provider_edge(
|
||||
backend, mounts=EDGE_MOUNTS, bind_host=bind_host,
|
||||
advertise_host=advertise_host, forward_timeout=forward_timeout,
|
||||
).edge
|
||||
|
||||
|
||||
@contextmanager
|
||||
def observed_provider_edge(
|
||||
observation: ProviderRequestObservation,
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ ProxyClient's key/customer methods for cleanup. Read-backs are eventually consis
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -26,6 +27,7 @@ from e2e_config import (
|
|||
PROXY_REPLICA_URLS,
|
||||
REQUEST_TIMEOUT,
|
||||
SLOW_PROVIDER_TIMEOUT_SECONDS,
|
||||
provider_edge_base,
|
||||
settle_propagation,
|
||||
)
|
||||
from e2e_http import (
|
||||
|
|
@ -93,6 +95,7 @@ from models import (
|
|||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
)
|
||||
from provider_cache_routing import route_cache_model
|
||||
from pydantic import BaseModel
|
||||
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
|
||||
|
||||
|
|
@ -645,7 +648,10 @@ class ProxyClient:
|
|||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
json=body.model_copy(update={"litellm_params": route_cache_model(
|
||||
body.litellm_params, provider_edge_base,
|
||||
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1", mode=body.model_info.mode,
|
||||
)}),
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
).model_id
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ from models import (
|
|||
)
|
||||
from quota_client import QuotaClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
|
||||
|
||||
# Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed").
|
||||
ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001"
|
||||
|
|
|
|||
|
|
@ -1254,7 +1254,8 @@ class TestHandleEdgeRequestPure:
|
|||
|
||||
|
||||
class TestApiBaseSeam:
|
||||
def test_live_mode_returns_none(self, tmp_path: Path) -> None:
|
||||
def test_live_mode_returns_none(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("E2E_PROVIDER_CACHE", raising=False)
|
||||
for mode_raw in ("live", ""):
|
||||
assert (
|
||||
provider_edge_api_base(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
import {
|
||||
E2E_TEAM_CRUD_ALIAS,
|
||||
E2E_TEAM_ORG_ALIAS,
|
||||
E2E_TEAM_ORG_ID,
|
||||
INTERNAL_USER_STORAGE_PATH,
|
||||
} from "../../constants";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
|
|
@ -179,18 +180,28 @@ test.describe("Models and Endpoints for an internal user", () => {
|
|||
`switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`,
|
||||
).toHaveCount(1, { timeout: 15_000 });
|
||||
|
||||
await expect(page).toHaveURL((url) =>
|
||||
url.searchParams.get("filter_team") === E2E_TEAM_ORG_ID &&
|
||||
url.searchParams.get("view_mode") === "all",
|
||||
);
|
||||
await page.reload();
|
||||
await expect(
|
||||
teamSelector(page),
|
||||
"the team selection is not persisted across a reload, so the table returns to the personal view",
|
||||
).toContainText(PERSONAL_TEAM, { timeout: 15_000 });
|
||||
"the selected team is restored from the URL after a reload",
|
||||
).toContainText(E2E_TEAM_ORG_ALIAS, { timeout: 15_000 });
|
||||
await expect(
|
||||
viewSelector(page),
|
||||
"the view selection is not persisted across a reload either",
|
||||
).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 });
|
||||
"the selected view is restored from the URL after a reload",
|
||||
).toContainText(ALL_MODELS_VIEW, { timeout: 15_000 });
|
||||
await expect(modelRow(page, CHAT_MODEL_A)).toHaveCount(1, { timeout: 15_000 });
|
||||
await expect(page.getByTestId("pagination-range")).toHaveText("Showing 1-1 of 1");
|
||||
await expect(modelRow(page, CHAT_MODEL_B)).toHaveCount(0);
|
||||
await expect(modelRow(page, ungrantedModelName)).toHaveCount(0);
|
||||
|
||||
await chooseOption(page, teamSelector(page), PERSONAL_TEAM);
|
||||
await expect(
|
||||
modelRow(page, ungrantedModelName),
|
||||
"the personal view still renders models after a reload rather than coming back empty",
|
||||
"switching back to the personal team restores models outside the selected team",
|
||||
).toHaveCount(1, { timeout: 30_000 });
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -163,6 +163,9 @@ async def test_reset_budget_keys_partial_failure():
|
|||
key1, key2, key3, key4, key5, key6 = (
|
||||
_attrify(k) for k in [key1, key2, key3, key4, key5, key6]
|
||||
)
|
||||
pre_reset_spend = {
|
||||
k["token"]: k["spend"] for k in [key2, key3, key4, key5, key6]
|
||||
}
|
||||
prisma_client.get_data = AsyncMock(
|
||||
return_value=[key1, key2, key3, key4, key5, key6]
|
||||
)
|
||||
|
|
@ -201,7 +204,7 @@ async def test_reset_budget_keys_partial_failure():
|
|||
# And every write must carry only {spend, budget_reset_at} — never the full row.
|
||||
for c in key_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {"decrement": pre_reset_spend[c["where"]["token"]]}
|
||||
|
||||
# Verify that the failure logging hook was scheduled (due to the failure for key1)
|
||||
failure_hook_calls = (
|
||||
|
|
@ -252,6 +255,9 @@ async def test_reset_budget_users_partial_failure():
|
|||
user1, user2, user3, user4, user5, user6 = (
|
||||
_attrify(u) for u in [user1, user2, user3, user4, user5, user6]
|
||||
)
|
||||
pre_reset_spend = {
|
||||
u["user_id"]: u["spend"] for u in [user2, user3, user4, user5, user6]
|
||||
}
|
||||
prisma_client.get_data = AsyncMock(
|
||||
return_value=[user1, user2, user3, user4, user5, user6]
|
||||
)
|
||||
|
|
@ -280,7 +286,9 @@ async def test_reset_budget_users_partial_failure():
|
|||
assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
|
||||
for c in user_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {
|
||||
"decrement": pre_reset_spend[c["where"]["user_id"]]
|
||||
}
|
||||
|
||||
failure_hook_calls = (
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
||||
|
|
@ -441,6 +449,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
for t in [team1, team2]:
|
||||
t.setdefault("team_id", t["id"])
|
||||
team1, team2 = _attrify(team1), _attrify(team2)
|
||||
pre_reset_spend = team2["spend"]
|
||||
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
|
||||
|
||||
async def fake_reset_team(team, current_time, reset_settings=None):
|
||||
|
|
@ -465,7 +474,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
assert len(team_writes) == 1
|
||||
assert team_writes[0]["where"] == {"team_id": "team2"}
|
||||
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert team_writes[0]["data"]["spend"] == 0
|
||||
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
|
||||
|
||||
failure_hook_calls = (
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
||||
|
|
@ -542,6 +551,11 @@ async def test_reset_budget_continues_other_categories_on_failure():
|
|||
user1, user2 = _attrify(user1), _attrify(user2)
|
||||
team1, team2 = _attrify(team1), _attrify(team2)
|
||||
enduser1 = _attrify(enduser1)
|
||||
pre_reset_spend = {
|
||||
**{k["token"]: k["spend"] for k in [key1, key2]},
|
||||
**{u["user_id"]: u["spend"] for u in [user2]},
|
||||
**{t["team_id"]: t["spend"] for t in [team1, team2]},
|
||||
}
|
||||
_wire_cascade_reads_for_test(prisma_client)
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
|
|
@ -618,7 +632,9 @@ async def test_reset_budget_continues_other_categories_on_failure():
|
|||
# Every batched write must carry only the two reset fields, never the full row.
|
||||
for c in key_writes + user_writes + team_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {
|
||||
"decrement": pre_reset_spend[next(iter(c["where"].values()))]
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ def test_async_fallbacks(caplog):
|
|||
if "Task exception was never retrieved" not in log
|
||||
and "Task was destroyed but it is pending" not in log
|
||||
and "get_available_deployment" not in log
|
||||
and "Selected deployment for model" not in log
|
||||
and "in the Langfuse queue" not in log
|
||||
and "Unclosed client session" not in log
|
||||
and "Unclosed connector" not in log
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -579,6 +580,20 @@ def test_text_only_streaming_has_index_zero():
|
|||
), f"Expected index=0, got {parsed.choices[0].index}"
|
||||
|
||||
|
||||
def test_message_delta_without_usage_returns_chunk_with_no_usage():
|
||||
iterator: Final = ModelResponseIterator(None, sync_stream=True)
|
||||
|
||||
model_response: Final = iterator.chunk_parser(
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
}
|
||||
)
|
||||
|
||||
assert model_response.choices[0].finish_reason == "stop"
|
||||
assert model_response.usage is None
|
||||
|
||||
|
||||
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
|
||||
"""Anthropic streaming usage should account for emitted thinking deltas."""
|
||||
chunks = [
|
||||
|
|
|
|||
|
|
@ -1,14 +1,20 @@
|
|||
import math
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.fireworks_ai.cost_calculator import cost_per_token
|
||||
from litellm.types.utils import OffPeakPricing, PromptTokensDetailsWrapper, Usage
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
OffPeakPricing,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
)
|
||||
|
||||
MODEL = "accounts/fireworks/models/glm-5p2"
|
||||
INPUT_COST = 1.4e-06
|
||||
# Read the cached rate from the price map so this test tracks the shipped value
|
||||
# (glm-5p2 is $0.14/1M) instead of hardcoding a number that breaks when it changes.
|
||||
CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="fireworks_ai")["cache_read_input_token_cost"]
|
||||
OUTPUT_COST = 4.4e-06
|
||||
|
||||
|
|
@ -44,13 +50,16 @@ STANDARD_CACHE_READ_COST = 1.5e-08
|
|||
def _register_off_peak_model(
|
||||
off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST
|
||||
) -> None:
|
||||
litellm.model_cost[f"fireworks_ai/{OFF_PEAK_MODEL}"] = {
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": STANDARD_INPUT_COST,
|
||||
"output_cost_per_token": STANDARD_OUTPUT_COST,
|
||||
"off_peak_pricing": off_peak_pricing,
|
||||
**({} if cache_read_cost is None else {"cache_read_input_token_cost": cache_read_cost}),
|
||||
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
|
||||
**litellm.model_cost,
|
||||
f"fireworks_ai/{OFF_PEAK_MODEL}": {
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": STANDARD_INPUT_COST,
|
||||
"output_cost_per_token": STANDARD_OUTPUT_COST,
|
||||
"off_peak_pricing": off_peak_pricing,
|
||||
**({} if cache_read_cost is None else {"cache_read_input_token_cost": cache_read_cost}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -125,3 +134,75 @@ def test_off_peak_defaults_to_the_current_time():
|
|||
|
||||
assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10)
|
||||
assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)
|
||||
|
||||
|
||||
COMPONENT_MODEL = "accounts/fireworks/models/cost-components-test"
|
||||
COMPONENT_INPUT_COST = 1e-06
|
||||
COMPONENT_OUTPUT_COST = 2e-06
|
||||
COMPONENT_CACHE_READ_COST = 1e-07
|
||||
COMPONENT_CACHE_CREATION_COST = 3e-06
|
||||
COMPONENT_REASONING_COST = 4e-06
|
||||
COMPONENT_AUDIO_IN_COST = 5e-06
|
||||
COMPONENT_AUDIO_OUT_COST = 6e-06
|
||||
|
||||
|
||||
def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates():
|
||||
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
|
||||
**litellm.model_cost,
|
||||
f"fireworks_ai/{COMPONENT_MODEL}": {
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": COMPONENT_INPUT_COST,
|
||||
"output_cost_per_token": COMPONENT_OUTPUT_COST,
|
||||
"cache_read_input_token_cost": COMPONENT_CACHE_READ_COST,
|
||||
"cache_creation_input_token_cost": COMPONENT_CACHE_CREATION_COST,
|
||||
"output_cost_per_reasoning_token": COMPONENT_REASONING_COST,
|
||||
"input_cost_per_audio_token": COMPONENT_AUDIO_IN_COST,
|
||||
"output_cost_per_audio_token": COMPONENT_AUDIO_OUT_COST,
|
||||
},
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=500,
|
||||
total_tokens=1500,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=300,
|
||||
cache_creation_tokens=200,
|
||||
audio_tokens=100,
|
||||
),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=200,
|
||||
audio_tokens=50,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = cost_per_token(model=COMPONENT_MODEL, usage=usage)
|
||||
|
||||
expected_prompt_cost = (
|
||||
400 * COMPONENT_INPUT_COST
|
||||
+ 300 * COMPONENT_CACHE_READ_COST
|
||||
+ 200 * COMPONENT_CACHE_CREATION_COST
|
||||
+ 100 * COMPONENT_AUDIO_IN_COST
|
||||
)
|
||||
expected_completion_cost = (
|
||||
250 * COMPONENT_OUTPUT_COST + 200 * COMPONENT_REASONING_COST + 50 * COMPONENT_AUDIO_OUT_COST
|
||||
)
|
||||
assert prompt_cost == pytest.approx(expected_prompt_cost)
|
||||
assert completion_cost == pytest.approx(expected_completion_cost)
|
||||
|
||||
|
||||
def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback():
|
||||
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
|
||||
**litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # the SDK types model_cost as dict[Unknown, Unknown]
|
||||
"fireworks_ai/accounts/fireworks/models/no-input-rate-test": {
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2e-06,
|
||||
},
|
||||
}
|
||||
usage: Final = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200)
|
||||
|
||||
prompt_cost, completion_cost = cost_per_token(model="accounts/fireworks/models/no-input-rate-test", usage=usage)
|
||||
|
||||
assert prompt_cost == 0
|
||||
assert completion_cost == 200 * 2e-06
|
||||
|
|
|
|||
|
|
@ -5836,3 +5836,126 @@ def test_supported_reasoning_efforts_still_map(model):
|
|||
drop_params=False,
|
||||
)
|
||||
assert "thinkingConfig" in result
|
||||
|
||||
|
||||
def _generate_content_body() -> dict:
|
||||
return {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "hi"}]},
|
||||
"finishReason": "STOP",
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 5,
|
||||
"candidatesTokenCount": 7,
|
||||
"totalTokenCount": 12,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_generate_content_transform_uses_reported_model_version():
|
||||
"""The served modelVersion must win over the requested name so downstream
|
||||
pricing sees what actually ran."""
|
||||
import httpx
|
||||
|
||||
body = {**_generate_content_body(), "modelVersion": "gemini-x-served"}
|
||||
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
|
||||
completion_response=body,
|
||||
model_response=ModelResponse(),
|
||||
model="gemini-x",
|
||||
logging_obj=MagicMock(),
|
||||
raw_response=httpx.Response(200, headers={}),
|
||||
)
|
||||
|
||||
assert response.model == "gemini-x-served"
|
||||
|
||||
|
||||
def test_generate_content_transform_falls_back_to_requested_model():
|
||||
import httpx
|
||||
|
||||
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
|
||||
completion_response=_generate_content_body(),
|
||||
model_response=ModelResponse(),
|
||||
model="gemini-x",
|
||||
logging_obj=MagicMock(),
|
||||
raw_response=httpx.Response(200, headers={}),
|
||||
)
|
||||
|
||||
assert response.model == "gemini-x"
|
||||
|
||||
|
||||
def test_streaming_chunk_carries_model_version():
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
chunk = {**_generate_content_body(), "modelVersion": "gemini-x-served"}
|
||||
iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
|
||||
streaming_chunk: Final = iterator.chunk_parser(chunk)
|
||||
|
||||
assert streaming_chunk.model == "gemini-x-served"
|
||||
|
||||
|
||||
def test_served_model_version_reaches_assembled_stream_through_custom_stream_wrapper():
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
served_model: Final = "gemini-3.8-flash-001"
|
||||
iterator: Final = ModelResponseIterator(
|
||||
streaming_response=iter(
|
||||
[json.dumps({**_generate_content_body(), "modelVersion": served_model}) for _ in range(3)]
|
||||
),
|
||||
sync_stream=True,
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
wrapper: Final = CustomStreamWrapper(
|
||||
completion_stream=iter(iterator),
|
||||
model="gemini/gemini-3.8-flash",
|
||||
custom_llm_provider="gemini",
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
chunks: Final = list(wrapper)
|
||||
|
||||
assert len(chunks) >= 3
|
||||
for chunk in chunks[:-1]:
|
||||
assert chunk._hidden_params["provider_response_model"] == served_model
|
||||
assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}])
|
||||
assert assembled._hidden_params["provider_response_model"] == served_model
|
||||
|
||||
|
||||
def test_generate_content_transform_strips_version_suffix_from_model_version():
|
||||
import httpx
|
||||
|
||||
body: Final = {**_generate_content_body(), "modelVersion": "gemini-3.8-flash-001@default"}
|
||||
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
|
||||
completion_response=body,
|
||||
model_response=ModelResponse(),
|
||||
model="gemini-3.8-flash",
|
||||
logging_obj=MagicMock(),
|
||||
raw_response=httpx.Response(200, headers={}),
|
||||
)
|
||||
|
||||
assert response.model == "gemini-3.8-flash-001"
|
||||
|
||||
|
||||
def test_prompt_blocked_chunk_keeps_served_model_version():
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
chunk: Final = {
|
||||
"promptFeedback": {"blockReason": "SAFETY", "blockReasonMessage": "prompt was blocked"},
|
||||
"modelVersion": "gemini-3.8-flash-001",
|
||||
"responseId": "resp-1",
|
||||
}
|
||||
iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
|
||||
|
||||
streaming_chunk: Final = iterator.chunk_parser(chunk)
|
||||
|
||||
assert streaming_chunk.model == "gemini-3.8-flash-001"
|
||||
assert streaming_chunk.choices[0].finish_reason == "content_filter"
|
||||
|
|
|
|||
|
|
@ -51,23 +51,23 @@ class TestXAIResponsesAPITransformation:
|
|||
assert result["tools"][0]["type"] == "code_interpreter"
|
||||
assert "container" not in result["tools"][0], "Container field should be removed"
|
||||
|
||||
def test_instructions_parameter_dropped(self):
|
||||
"""Test that instructions parameter is dropped for XAI"""
|
||||
def test_instructions_parameter_forwarded(self):
|
||||
"""xAI supports 'instructions' on /v1/responses, so it must survive param mapping"""
|
||||
config = XAIResponsesAPIConfig()
|
||||
|
||||
params = ResponsesAPIOptionalRequestParams(instructions="You are a helpful assistant.", temperature=0.7)
|
||||
|
||||
result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False)
|
||||
|
||||
assert "instructions" not in result, "Instructions should be dropped"
|
||||
assert result.get("instructions") == "You are a helpful assistant."
|
||||
assert result.get("temperature") == 0.7, "Other params should be preserved"
|
||||
|
||||
def test_supported_params_excludes_instructions(self):
|
||||
"""Test that get_supported_openai_params excludes instructions"""
|
||||
def test_supported_params_includes_instructions(self):
|
||||
"""A system message bridged to 'instructions' must not be rejected for xAI"""
|
||||
config = XAIResponsesAPIConfig()
|
||||
supported = config.get_supported_openai_params("grok-4-fast")
|
||||
|
||||
assert "instructions" not in supported, "instructions should not be supported"
|
||||
assert "instructions" in supported, "instructions should be supported"
|
||||
assert "tools" in supported, "tools should be supported"
|
||||
assert "temperature" in supported, "temperature should be supported"
|
||||
assert "model" in supported, "model should be supported"
|
||||
|
|
|
|||
|
|
@ -53,8 +53,8 @@ class TestXAIResponsesAPITransformation:
|
|||
"container" not in result["tools"][0]
|
||||
), "Container field should be removed"
|
||||
|
||||
def test_instructions_parameter_dropped(self):
|
||||
"""Test that instructions parameter is dropped for XAI"""
|
||||
def test_instructions_parameter_forwarded(self):
|
||||
"""xAI supports 'instructions' on /v1/responses, so it must survive param mapping"""
|
||||
config = XAIResponsesAPIConfig()
|
||||
|
||||
params = ResponsesAPIOptionalRequestParams(
|
||||
|
|
@ -65,15 +65,15 @@ class TestXAIResponsesAPITransformation:
|
|||
response_api_optional_params=params, model="grok-4-fast", drop_params=False
|
||||
)
|
||||
|
||||
assert "instructions" not in result, "Instructions should be dropped"
|
||||
assert result.get("instructions") == "You are a helpful assistant."
|
||||
assert result.get("temperature") == 0.7, "Other params should be preserved"
|
||||
|
||||
def test_supported_params_excludes_instructions(self):
|
||||
"""Test that get_supported_openai_params excludes instructions"""
|
||||
def test_supported_params_includes_instructions(self):
|
||||
"""A system message bridged to 'instructions' must not be rejected for xAI"""
|
||||
config = XAIResponsesAPIConfig()
|
||||
supported = config.get_supported_openai_params("grok-4-fast")
|
||||
|
||||
assert "instructions" not in supported, "instructions should not be supported"
|
||||
assert "instructions" in supported, "instructions should be supported"
|
||||
assert "tools" in supported, "tools should be supported"
|
||||
assert "temperature" in supported, "temperature should be supported"
|
||||
assert "model" in supported, "model should be supported"
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,11 +1,14 @@
|
|||
"""Unit tests for the LLM-as-a-Judge guardrail hook."""
|
||||
|
||||
import json
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
||||
LLMAsAJudgeGuardrail,
|
||||
_build_judge_prompt,
|
||||
|
|
@ -13,7 +16,8 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
|||
_parse_judge_verdict,
|
||||
initialize_guardrail,
|
||||
)
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
|
|
@ -136,17 +140,314 @@ def test_initialize_guardrail_invalid_on_failure():
|
|||
initialize_guardrail(lp, g)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "runs_pre_call", "runs_post_call"),
|
||||
[
|
||||
("pre_call", True, False),
|
||||
(["pre_call", "post_call"], True, True),
|
||||
(Mode(tags={"judge": ["pre_call"]}, default="post_call"), True, False),
|
||||
(None, False, True),
|
||||
],
|
||||
ids=["scalar", "list", "tagged", "missing"],
|
||||
)
|
||||
def test_initialize_guardrail_preserves_every_mode_shape(
|
||||
mode: str | list[str] | Mode | None,
|
||||
runs_pre_call: bool,
|
||||
runs_post_call: bool,
|
||||
):
|
||||
lp: Final = _make_litellm_params(mode=mode)
|
||||
instance: Final = initialize_guardrail(lp, _make_guardrail_dict())
|
||||
request_data: Final[dict[str, object]] = {"metadata": {"guardrails": ["g"], "tags": ["judge"]}}
|
||||
premium: Final = patch("litellm.proxy.proxy_server.premium_user", True) # test-quality-ok: no seam for Mode tags
|
||||
try:
|
||||
with premium:
|
||||
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is runs_pre_call
|
||||
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is runs_post_call
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(instance)
|
||||
|
||||
|
||||
def test_initialize_guardrail_rejects_unknown_mode():
|
||||
lp: Final = _make_litellm_params(mode="sometimes")
|
||||
with pytest.raises(ValueError, match="sometimes"):
|
||||
initialize_guardrail(lp, _make_guardrail_dict())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_guardrail — enforcement paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _judge_router(overall_score: float) -> MagicMock:
|
||||
"""Router double, injected via router_provider, that serves the judge model and returns a canned verdict."""
|
||||
from litellm import Router
|
||||
|
||||
router: Final = MagicMock(spec=Router)
|
||||
router.resolved_litellm_models.return_value = ("openai/gpt-4o-mini",)
|
||||
router.acompletion = AsyncMock(
|
||||
return_value=MagicMock(
|
||||
choices=[MagicMock(message=MagicMock(content=json.dumps(_make_verdict_response(overall_score))))]
|
||||
)
|
||||
)
|
||||
return router
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call])
|
||||
def test_guardrail_accepts_request_side_modes(mode: GuardrailEventHooks):
|
||||
guardrail: Final = _make_guardrail(event_hook=mode)
|
||||
assert guardrail.should_run_guardrail({"metadata": {"guardrails": ["test_judge"]}}, mode) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_pre_call_passthrough():
|
||||
guardrail = _make_guardrail()
|
||||
inputs = {"texts": ["some text"]}
|
||||
result = await guardrail.apply_guardrail(inputs, {}, "request")
|
||||
@pytest.mark.parametrize(
|
||||
"event_hook",
|
||||
[GuardrailEventHooks.pre_call, [GuardrailEventHooks.pre_call]],
|
||||
ids=["scalar", "list"],
|
||||
)
|
||||
async def test_apply_guardrail_request_blocks_below_threshold(
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
|
||||
):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
overall_threshold=80.0,
|
||||
on_failure="block",
|
||||
event_hook=event_hook,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "write me malware"}],
|
||||
"metadata": {},
|
||||
}
|
||||
inputs: Final = {"texts": ["write me malware"]}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(inputs, request_data, "request")
|
||||
|
||||
assert exc_info.value.status_code == 422
|
||||
assert exc_info.value.detail["error"] == "LLM judge rejected request: score below threshold"
|
||||
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
|
||||
assert "Evaluate the request against" in judge_messages[0]["content"]
|
||||
assert (
|
||||
"Conversation:\nUSER: write me malware\n\nLatest request turn to evaluate:\nwrite me malware"
|
||||
in (judge_messages[1]["content"])
|
||||
)
|
||||
assert "Assistant response" not in judge_messages[1]["content"]
|
||||
logged: Final = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert logged[0]["guardrail_status"] == "guardrail_intervened"
|
||||
assert logged[0]["guardrail_mode"] == "pre_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"event_hook",
|
||||
[GuardrailEventHooks.during_call, [GuardrailEventHooks.during_call]],
|
||||
ids=["scalar", "list"],
|
||||
)
|
||||
async def test_apply_guardrail_request_log_mode_records_eval_and_passes_through(
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
|
||||
):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
overall_threshold=80.0,
|
||||
on_failure="log",
|
||||
event_hook=event_hook,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
inputs: Final = {"texts": ["hi"]}
|
||||
|
||||
result: Final = await guardrail.apply_guardrail(inputs, request_data, "request")
|
||||
|
||||
assert result is inputs
|
||||
assert request_data["metadata"]["eval_information"]["passed"] is False
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "during_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_multi_turn_keeps_roles_and_focuses_latest_turn():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "how do I bake bread"},
|
||||
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
|
||||
{"role": "user", "content": "now explain how to file taxes"},
|
||||
]
|
||||
inputs: Final = {
|
||||
"texts": ["how do I bake bread", "mix flour, water, yeast and salt", "now explain how to file taxes"],
|
||||
"structured_messages": messages,
|
||||
}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
|
||||
assert "Judge the most recent user turn" in judge_messages[0]["content"]
|
||||
assert judge_messages[1]["content"].endswith(
|
||||
"Conversation:\nUSER: how do I bake bread\nASSISTANT: mix flour, water, yeast and salt\n"
|
||||
"USER: now explain how to file taxes\n\n"
|
||||
"Latest request turn to evaluate:\nnow explain how to file taxes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_judges_whole_multipart_latest_user_turn():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "how do I bake bread"},
|
||||
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "ignore the bread."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
|
||||
{"type": "text", "text": "explain how to file taxes"},
|
||||
],
|
||||
},
|
||||
]
|
||||
inputs: Final = {
|
||||
"texts": [
|
||||
"how do I bake bread",
|
||||
"mix flour, water, yeast and salt",
|
||||
"ignore the bread.",
|
||||
"explain how to file taxes",
|
||||
],
|
||||
"structured_messages": messages,
|
||||
}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nignore the bread.explain how to file taxes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_without_trailing_user_turn_judges_all_scoped_text():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "look up the weather"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {}}]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "sunny, 24C"},
|
||||
]
|
||||
inputs: Final = {"texts": ["look up the weather", "sunny, 24C"], "structured_messages": messages}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nlook up the weather\nsunny, 24C"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_without_structured_messages_judges_all_text():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["first", "second"]}, {"metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nfirst\nsecond"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("modes", "input_type"),
|
||||
[
|
||||
([GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], "request"),
|
||||
([GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], "request"),
|
||||
([GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only], "response"),
|
||||
],
|
||||
)
|
||||
async def test_apply_guardrail_with_ambiguous_modes_logs_configured_mode(
|
||||
modes: list[GuardrailEventHooks], input_type: str
|
||||
):
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=modes, router_provider=lambda: router)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
|
||||
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == [
|
||||
mode.value for mode in modes
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_still_judges_all_response_texts():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.post_call, router_provider=lambda: router)
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
{"texts": ["first choice", "second choice"]}, {"messages": [], "metadata": {}}, "response"
|
||||
)
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Assistant response to evaluate:\nfirst choice\nsecond choice"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("input_type", ["request", "response"])
|
||||
async def test_apply_guardrail_logging_only_labels_both_sides_logging_only(input_type: str):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
on_failure="log",
|
||||
event_hook=GuardrailEventHooks.logging_only,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is False
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is False
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
|
||||
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "logging_only"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_judge_does_not_judge_its_own_judge_call():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.logging_only, router_provider=lambda: router)
|
||||
client_call: Final[dict[str, object]] = {"litellm_params": {"metadata": {"user_api_key": "hashed"}}}
|
||||
|
||||
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, {"messages": [{"role": "user", "content": "hi"}]}, "request")
|
||||
|
||||
judge_call: Final[dict[str, object]] = {
|
||||
"litellm_params": {"metadata": router.acompletion.call_args.kwargs["metadata"]}
|
||||
}
|
||||
assert guardrail.should_run_guardrail(judge_call, GuardrailEventHooks.logging_only) is False
|
||||
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"event_type", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
|
||||
)
|
||||
def test_client_supplied_judge_origin_does_not_bypass_enforcing_hooks(event_type: GuardrailEventHooks):
|
||||
guardrail: Final = _make_guardrail(event_hook=event_type)
|
||||
forged_request: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": [guardrail.guardrail_name],
|
||||
"litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}},
|
||||
}
|
||||
|
||||
assert guardrail.should_run_guardrail(forged_request, event_type) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_prompt_unchanged():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(router_provider=lambda: router)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["hello there"]}, request_data, "response")
|
||||
|
||||
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
|
||||
assert "assistant's response" in judge_messages[0]["content"]
|
||||
assert "Conversation:\nUSER: hi\n\nAssistant response to evaluate:\nhello there" in judge_messages[1]["content"]
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "post_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -230,7 +531,7 @@ def test_parse_judge_verdict_reraises_when_no_json():
|
|||
|
||||
def test_parse_judge_verdict_rejects_json_non_object():
|
||||
"""Valid JSON that is not an object (e.g. a bare list) raises ValueError."""
|
||||
with pytest.raises(ValueError, match='judge response is not a JSON object'):
|
||||
with pytest.raises(ValueError, match="judge response is not a JSON object"):
|
||||
_parse_judge_verdict("[1, 2, 3]")
|
||||
|
||||
|
||||
|
|
@ -252,9 +553,7 @@ async def test_apply_guardrail_enforces_fenced_verdict(mock_completion):
|
|||
@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion")
|
||||
async def test_apply_guardrail_non_object_verdict_fails_open_with_status(mock_completion):
|
||||
"""A non-object verdict fails open and logs guardrail_failed_to_respond."""
|
||||
mock_completion.return_value = MagicMock(
|
||||
choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))]
|
||||
)
|
||||
mock_completion.return_value = MagicMock(choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))])
|
||||
guardrail = _make_guardrail(overall_threshold=80.0, on_failure="block", router_provider=lambda: None)
|
||||
inputs = {"texts": ["response"]}
|
||||
request_data: dict = {"messages": [], "metadata": {}}
|
||||
|
|
@ -314,7 +613,12 @@ def _real_router(model_list, **router_kwargs):
|
|||
"model_list, router_kwargs, judge_model",
|
||||
[
|
||||
(
|
||||
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "my-judge-alias",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{},
|
||||
"my-judge-alias",
|
||||
),
|
||||
|
|
@ -324,12 +628,22 @@ def _real_router(model_list, **router_kwargs):
|
|||
"anthropic/claude-sonnet-4-6",
|
||||
),
|
||||
(
|
||||
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "backing-group",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{"model_group_alias": {"my-judge-alias": "backing-group"}},
|
||||
"my-judge-alias",
|
||||
),
|
||||
(
|
||||
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "backing-group",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{"model_group_alias": {"my-judge-alias": {"model": "backing-group", "hidden": True}}},
|
||||
"my-judge-alias",
|
||||
),
|
||||
|
|
@ -412,7 +726,12 @@ async def test_judge_resolves_router_lazily_per_call(mock_sdk_completion):
|
|||
mock_sdk_completion.assert_awaited_once()
|
||||
|
||||
holder["router"] = _real_router(
|
||||
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}]
|
||||
[
|
||||
{
|
||||
"model_name": "my-judge-alias",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
]
|
||||
)
|
||||
await guardrail.apply_guardrail({"texts": ["r"]}, {"messages": [], "metadata": {}}, "response")
|
||||
holder["router"].acompletion.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.llms.anthropic.cost_calculation import cost_per_token as anthropic_cost_per_token
|
||||
from litellm.proxy.spend_tracking.savings import (
|
||||
_baseline_usage,
|
||||
_resolve_model,
|
||||
|
|
@ -17,6 +18,34 @@ from litellm.types.utils import Usage
|
|||
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("modifier", [{"speed": "fast"}, {"inference_geo": "us"}])
|
||||
@pytest.mark.parametrize("continuing", [False, True])
|
||||
def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None:
|
||||
usage: Final = _usage(1000, 0, 1000, 100).model_copy(update=modifier)
|
||||
expected: Final = (_usage(1000, 1000, 0, 100) if continuing else usage).model_copy(update=modifier)
|
||||
normalized: Final = _baseline_usage(usage, continuing)
|
||||
cache_fields: Final = {"prompt_tokens_details", "cache_read_input_tokens", "cache_creation_input_tokens"}
|
||||
assert normalized.model_dump(exclude=cache_fields) == usage.model_dump(exclude=cache_fields)
|
||||
assert usage.prompt_tokens_details.cached_tokens == 0
|
||||
selected_cost: Final = 0.013
|
||||
assert compute_autorouter_savings(
|
||||
"claude-opus-5", "claude-sonnet-5", "anthropic", usage, conversation_continuing=continuing,
|
||||
cost_breakdown={"input_cost": 0.01, "output_cost": 0.003},
|
||||
) == pytest.approx(sum(anthropic_cost_per_token("claude-opus-5", expected)) - selected_cost)
|
||||
|
||||
|
||||
def test_anthropic_baseline_keeps_negotiated_prices_with_provider_multiplier() -> None:
|
||||
info: Final = {
|
||||
**litellm.get_model_info("claude-opus-5", "anthropic"),
|
||||
"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, "cache_read_input_token_cost": 3e-7,
|
||||
}
|
||||
usage: Final = _usage(1000, 1000, 0, 100).model_copy(update={"speed": "fast"})
|
||||
assert compute_autorouter_savings(
|
||||
"claude-opus-5", "claude-sonnet-5", "anthropic", usage, baseline_info=info,
|
||||
cost_breakdown={"input_cost": 0.01, "output_cost": 0.003},
|
||||
) == pytest.approx(0.0015 * 2 - 0.013)
|
||||
|
||||
|
||||
def _anthropic_costs(model: str) -> tuple[float, float]:
|
||||
info = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
|
||||
input_cost = info["input_cost_per_token"] or 0.0
|
||||
|
|
|
|||
|
|
@ -689,6 +689,36 @@ async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_
|
|||
assert recorded["status"] == "success"
|
||||
|
||||
|
||||
class _RecordingApplyGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str, applied: list[str]) -> None:
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
default_on=True,
|
||||
)
|
||||
self._applied = applied
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
await asyncio.sleep(0)
|
||||
self._applied.append(self.guardrail_name or "")
|
||||
return inputs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_hook_runs_every_unified_guardrail(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
applied: list[str] = []
|
||||
guardrails = [_RecordingApplyGuardrail(f"judge-{i}", applied) for i in range(3)]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
|
||||
await proxy_logging.during_call_hook(
|
||||
data={"model": "m", "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert sorted(applied) == ["judge-0", "judge-1", "judge-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
cb = _moderation_guardrail()
|
||||
|
|
|
|||
|
|
@ -5,17 +5,20 @@ completion_start_time = end_time."""
|
|||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic_core import PydanticSerializationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import (
|
||||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
_estimate_usage_from_text,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseAPIUsage,
|
||||
|
|
@ -31,16 +34,23 @@ def _sse_event(payload: dict) -> bytes:
|
|||
|
||||
def _mock_config() -> Mock:
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
|
||||
mock_responses_api_response.id = "resp_ttft"
|
||||
mock_responses_api_response = ResponsesAPIResponse(
|
||||
id="resp_ttft",
|
||||
created_at=0,
|
||||
status="completed",
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2),
|
||||
)
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
evt_type = parsed_chunk.get("type")
|
||||
if evt_type == "response.completed":
|
||||
completed = Mock(spec=ResponseCompletedEvent)
|
||||
completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
completed.response = mock_responses_api_response
|
||||
return completed
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=mock_responses_api_response,
|
||||
)
|
||||
stub = Mock()
|
||||
stub.type = evt_type
|
||||
return stub
|
||||
|
|
@ -54,6 +64,8 @@ def _make_iterator(
|
|||
sse_events: list[bytes],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
trailing_error: Optional[Exception] = None,
|
||||
config: Mock | None = None,
|
||||
request_data: dict | None = None,
|
||||
) -> ResponsesAPIStreamingIterator:
|
||||
async def aiter_bytes():
|
||||
for evt in sse_events:
|
||||
|
|
@ -68,10 +80,11 @@ def _make_iterator(
|
|||
return ResponsesAPIStreamingIterator(
|
||||
response=mock_response,
|
||||
model="gpt-4o-mini",
|
||||
responses_api_provider_config=_mock_config(),
|
||||
responses_api_provider_config=config or _mock_config(),
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata={},
|
||||
custom_llm_provider="openai",
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -329,6 +342,88 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
|
|||
assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params
|
||||
|
||||
|
||||
def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock:
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
evt_type = parsed_chunk.get("type")
|
||||
if evt_type == "response.completed":
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=response,
|
||||
)
|
||||
stub = Mock()
|
||||
stub.type = evt_type
|
||||
if "delta" in parsed_chunk:
|
||||
stub.delta = parsed_chunk.get("delta")
|
||||
if "item" in parsed_chunk:
|
||||
stub.item = parsed_chunk.get("item")
|
||||
return stub
|
||||
|
||||
mock_config.transform_streaming_response.side_effect = _transform
|
||||
return mock_config
|
||||
|
||||
|
||||
def _responses_api_response_without_usage() -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(
|
||||
id="resp_no_usage",
|
||||
created_at=int(datetime(2025, 1, 1).timestamp()),
|
||||
status="completed",
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_gets_text_estimate():
|
||||
"""A response.completed event carrying usage: null still bills: the
|
||||
iterator estimates usage from the request input and generated text."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens > 0
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_with_usage_is_left_untouched():
|
||||
"""Provider-reported usage on response.completed wins over the estimate."""
|
||||
response = _responses_api_response_with_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage.input_tokens == 20
|
||||
assert usage.output_tokens == 60
|
||||
assert usage.total_tokens == 80
|
||||
|
||||
|
||||
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(
|
||||
id="resp_lit6427",
|
||||
|
|
@ -628,3 +723,222 @@ async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_val
|
|||
assert isinstance(client_usage, ResponseAPIUsage)
|
||||
assert client_usage.input_tokens == 29
|
||||
assert client_usage.cost == pytest.approx(0.0001)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_counts_tool_call_arguments():
|
||||
"""A function-call-only stream still bills output tokens: streamed
|
||||
function_call_arguments deltas feed the text estimate."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event(
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"item": {"type": "function_call", "name": "get_weather", "call_id": "call_1"},
|
||||
}
|
||||
),
|
||||
_sse_event(
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"delta": '{"location": "San Francisco", "unit": "celsius"}',
|
||||
}
|
||||
),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "what is the weather in san francisco"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_without_usage_counts_multimodal_input_as_messages():
|
||||
"""Multimodal request input is counted as chat messages, not as a JSON blob:
|
||||
a huge base64 image must not inflate the estimated input tokens."""
|
||||
image_input: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "what is in this image"},
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": "data:image/png;base64," + "A" * 4000,
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
json_count: Final = litellm.token_counter(model="gpt-4o-mini", text=json.dumps(image_input))
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "it is a cat"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": image_input},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens < json_count / 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_survives_a_failing_usage_estimate():
|
||||
"""A malformed request input that makes the message transformer raise must not
|
||||
break a stream that previously completed: the estimate is best-effort and
|
||||
falls back to usage None."""
|
||||
malformed_input: Final = [{"type": "message", "role": "user", "content": 42}]
|
||||
with pytest.raises(ValueError, match="Invalid content type"):
|
||||
_estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world")
|
||||
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": malformed_input},
|
||||
)
|
||||
|
||||
yielded: list = []
|
||||
async for chunk in iterator:
|
||||
yielded.append(chunk)
|
||||
|
||||
assert yielded
|
||||
assert iterator.completed_response.response.usage is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"tool_delta_event_type",
|
||||
["response.custom_tool_call_input.delta", "response.mcp_call_arguments.delta"],
|
||||
)
|
||||
async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta_event_type):
|
||||
"""Custom-tool and MCP argument deltas feed the streamed usage fallback the
|
||||
same way function_call_arguments deltas do."""
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": tool_delta_event_type, "delta": '{"query": "weather in sf"}'}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=_logging_obj_stub(),
|
||||
config=_mock_config_with_completed_response(response),
|
||||
request_data={"input": "what is the weather in san francisco"},
|
||||
)
|
||||
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
usage = iterator.completed_response.response.usage
|
||||
assert usage is not None
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_event_with_a_dict_response_is_typed_and_billed():
|
||||
"""transform_streaming_response can model_construct a terminal event whose
|
||||
response stays a plain dict; the iterator must type it so the estimated
|
||||
usage reaches the cost stamping path."""
|
||||
dict_response: Final = {
|
||||
"id": "resp_dict",
|
||||
"model": "gpt-4o-mini",
|
||||
"object": "response",
|
||||
"output": [],
|
||||
"usage": None,
|
||||
}
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
if parsed_chunk.get("type") == "response.completed":
|
||||
return ResponseCompletedEvent.model_construct(type="response.completed", response=dict_response)
|
||||
stub: Final = Mock()
|
||||
stub.type = parsed_chunk.get("type")
|
||||
if "delta" in parsed_chunk:
|
||||
stub.delta = parsed_chunk.get("delta")
|
||||
return stub
|
||||
|
||||
config: Final = Mock(spec=BaseResponsesAPIConfig)
|
||||
config.transform_streaming_response.side_effect = _transform
|
||||
logging_obj: Final = _logging_obj_stub()
|
||||
logging_obj._response_cost_calculator.return_value = 0.000704
|
||||
iterator: Final = _make_iterator(
|
||||
sse_events=[
|
||||
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
|
||||
_sse_event({"type": "response.completed", "response": {}}),
|
||||
],
|
||||
logging_obj=logging_obj,
|
||||
config=config,
|
||||
request_data={"input": "count these input tokens please"},
|
||||
)
|
||||
|
||||
yielded: Final = [chunk async for chunk in iterator]
|
||||
|
||||
terminal_event: Final = iterator.completed_response
|
||||
assert yielded[-1] is terminal_event
|
||||
completed_response: Final = terminal_event.response
|
||||
assert isinstance(completed_response, ResponsesAPIResponse)
|
||||
usage: Final = completed_response.usage
|
||||
assert usage is not None
|
||||
assert usage.input_tokens > 0
|
||||
assert usage.output_tokens > 0
|
||||
assert usage.cost == pytest.approx(0.000704)
|
||||
logging_obj._response_cost_calculator.assert_any_call(result=completed_response)
|
||||
|
||||
|
||||
def test_billed_terminal_response_keeps_a_response_that_already_has_usage():
|
||||
from litellm.responses.streaming_iterator import _billed_terminal_response
|
||||
|
||||
response: Final = _responses_api_response_with_usage()
|
||||
|
||||
assert _billed_terminal_response(response, None) is response
|
||||
|
||||
|
||||
def test_billed_terminal_response_copies_when_estimating_and_leaves_the_original_untouched():
|
||||
from litellm.responses.streaming_iterator import _billed_terminal_response
|
||||
|
||||
response: Final = _responses_api_response_without_usage()
|
||||
estimated: Final = ResponseAPIUsage(input_tokens=3, output_tokens=4, total_tokens=7)
|
||||
|
||||
billed: Final = _billed_terminal_response(response, lambda: estimated)
|
||||
|
||||
assert billed is not response
|
||||
assert billed.usage is estimated
|
||||
assert response.usage is None
|
||||
|
||||
|
||||
def test_persist_completed_response_to_cache_survives_an_unserializable_response(monkeypatch):
|
||||
bad_response: Final = ResponsesAPIResponse.model_construct(id="r", output=[object()], usage=None)
|
||||
with pytest.raises(PydanticSerializationError):
|
||||
bad_response.model_dump_json()
|
||||
|
||||
logging_obj: Final = _logging_obj_stub()
|
||||
caching_handler: Final = Mock()
|
||||
caching_handler.request_kwargs = {"stream": True}
|
||||
logging_obj._llm_caching_handler = caching_handler
|
||||
iterator: Final = _make_iterator(sse_events=[], logging_obj=logging_obj)
|
||||
iterator.completed_response = ResponseCompletedEvent.model_construct(
|
||||
type="response.completed", response=bad_response
|
||||
)
|
||||
cache: Final = Mock()
|
||||
monkeypatch.setattr(litellm, "cache", cache)
|
||||
|
||||
iterator._persist_completed_response_to_cache(is_async=False)
|
||||
|
||||
cache.add_cache.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -49,6 +49,20 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"]
|
|||
@pytest.mark.parametrize(
|
||||
"category,changed,expected",
|
||||
[
|
||||
("provider-harness", ["tests/e2e/provider_cache.py"], "run"),
|
||||
("provider-harness", ["tests/e2e/conftest.py"], "run"),
|
||||
("provider-harness", ["tests/e2e/e2e_http.py"], "run"),
|
||||
("provider-harness", ["tests/code_coverage_tests/test_provider_cache.py"], "run"),
|
||||
("provider-harness", ["tests/code_coverage_tests/test_provider_replay_harness.py"], "run"),
|
||||
("provider-harness", [".circleci/config.yml"], "run"),
|
||||
("provider-harness", [".circleci/scripts/classify_changes.sh"], "run"),
|
||||
("provider-harness", ["pyproject.toml"], "run"),
|
||||
("provider-harness", ["uv.lock"], "run"),
|
||||
("provider-harness", ["tests/e2e/PROVIDER_CACHE.md"], "skip"),
|
||||
("provider-harness", ["tests/e2e/ui/test_example.py"], "skip"),
|
||||
("provider-harness", ["tests/e2e/quota_management/test_quota.py"], "skip"),
|
||||
("provider-harness", ["litellm/main.py"], "skip"),
|
||||
("provider-harness", ["ui/litellm-dashboard/src/App.tsx"], "skip"),
|
||||
# docs-only: skip everything
|
||||
("backend", DOCS, "skip"),
|
||||
("client", DOCS, "skip"),
|
||||
|
|
|
|||
|
|
@ -2,14 +2,30 @@
|
|||
Test automatic routing to xAI Responses API when tools are present
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
||||
|
||||
class _RecordingResponsesHandler:
|
||||
"""MockTransport handler that serves a canned /responses reply and keeps the body xAI would have received"""
|
||||
|
||||
def __init__(self, reply: Mapping[str, object]) -> None:
|
||||
self.reply: Final = reply
|
||||
self.request_body: Mapping[str, object] | None = None
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.request_body = json.loads(request.content)
|
||||
return httpx.Response(200, json=dict(self.reply), request=request)
|
||||
|
||||
|
||||
class TestXAIResponsesAutoRouting:
|
||||
"""Test that xAI requests with tools automatically route to Responses API"""
|
||||
|
||||
|
|
@ -254,6 +270,44 @@ class TestXAIResponsesAutoRouting:
|
|||
# Note: This test may need adjustment based on actual mock_response behavior
|
||||
# The key is that the responses_api_bridge_check logic routes correctly
|
||||
|
||||
def test_system_message_survives_web_search_bridge(self):
|
||||
"""A system message becomes 'instructions' on the bridged /responses call, and xAI accepts it"""
|
||||
handler: Final = _RecordingResponsesHandler(
|
||||
reply={
|
||||
"id": "resp_test",
|
||||
"object": "response",
|
||||
"created_at": 0,
|
||||
"status": "completed",
|
||||
"model": "grok-4.6",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_test",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "1.0.0", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
)
|
||||
|
||||
response: Final = litellm.completion(
|
||||
model="xai/grok-4.6",
|
||||
messages=[
|
||||
{"role": "system", "content": "Answer briefly."},
|
||||
{"role": "user", "content": "newest litellm version?"},
|
||||
],
|
||||
web_search_options={"search_context_size": "medium"},
|
||||
api_key="fake-key",
|
||||
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))),
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "1.0.0"
|
||||
assert handler.request_body is not None
|
||||
assert handler.request_body["instructions"] == "Answer briefly."
|
||||
assert handler.request_body["tools"] == [{"type": "web_search"}]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
|
|
|||
|
|
@ -87,8 +87,8 @@ const LLMJudgeFields: React.FC<LLMJudgeFieldsProps> = ({ availableModels, contro
|
|||
return (
|
||||
<FieldGroup>
|
||||
<div className="rounded-md border border-success/20 bg-success/10 px-3.5 py-2.5 text-[13px] text-success">
|
||||
After each LLM response, the <strong>Judge Model</strong> scores it 0–100 against your criteria. If the weighted
|
||||
average falls below the threshold, the response is blocked (or logged).
|
||||
The <strong>Judge Model</strong> scores the user request (pre_call, during_call) or the LLM response (post_call)
|
||||
0–100 against your criteria. If the weighted average falls below the threshold, it is blocked (or logged).
|
||||
</div>
|
||||
|
||||
<GuardrailField
|
||||
|
|
|
|||
|
|
@ -104,8 +104,8 @@ export function AutoRoutersPanel({
|
|||
<DialogHeader>
|
||||
<DialogTitle>Add Auto Router</DialogTitle>
|
||||
<DialogDescription>
|
||||
Routes each request to a model by classifying its complexity. Called like any other model, so clients keep
|
||||
using a single model name.
|
||||
Choose a classifier to route each request to a model. Called like any other model, so clients keep using a
|
||||
single model name.
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
<AddAutoRouterTab
|
||||
|
|
|
|||
|
|
@ -57,6 +57,8 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models));
|
|||
|
||||
const COMPLEXITY_TYPE_LABELS: Record<string, string> = {
|
||||
llm: "LLM Classifier",
|
||||
capability: "Capability",
|
||||
llm_v2: "Fuse v2",
|
||||
heuristic_first: "Heuristic first",
|
||||
hybrid: "Hybrid",
|
||||
custom: "Custom classifier",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,75 @@
|
|||
import React, { useState } from "react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
classifier_type: "llm",
|
||||
tiers: { SIMPLE: ["efficient"], MEDIUM: [], COMPLEX: [], REASONING: ["capable"] },
|
||||
};
|
||||
|
||||
function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) {
|
||||
const [value, setValue] = useState(initialValue);
|
||||
return (
|
||||
<AutoRouterClassifierTabs value={value} onChange={setValue}>
|
||||
<output aria-label="Classifier type">{value.classifier_type}</output>
|
||||
</AutoRouterClassifierTabs>
|
||||
);
|
||||
}
|
||||
|
||||
describe("AutoRouterClassifierTabs", () => {
|
||||
it.each(["heuristic", "heuristic_v2", "llm", "heuristic_first", "hybrid"] as const)(
|
||||
"groups %s under Complexity without resetting its configuration",
|
||||
(classifier_type) => {
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(
|
||||
<AutoRouterClassifierTabs value={{ ...initial, classifier_type }} onChange={onChange}>
|
||||
Existing classifier settings
|
||||
</AutoRouterClassifierTabs>,
|
||||
);
|
||||
expect(screen.getByRole("tab", { name: "Complexity" })).toHaveAttribute("aria-selected", "true");
|
||||
expect(screen.getByRole("tabpanel", { name: "Complexity" })).toHaveTextContent("Existing classifier settings");
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
["capability", "Capability"],
|
||||
["llm_v2", "Fuse v2"],
|
||||
] as const)("opens saved %s settings and switches back to local Complexity", (classifier_type, label) => {
|
||||
renderWithProviders(<Form initialValue={{ ...initial, classifier_type }} />);
|
||||
expect(screen.getByRole("tab", { name: label })).toHaveAttribute("aria-selected", "true");
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
|
||||
expect(screen.getByRole("tab", { name: "Complexity" })).toHaveAttribute("aria-selected", "true");
|
||||
expect(screen.getByRole("status", { name: "Classifier type" })).toHaveTextContent("heuristic");
|
||||
});
|
||||
|
||||
it("keeps custom tiers editable under Complexity and explains why forecast tabs are disabled", () => {
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(
|
||||
<AutoRouterClassifierTabs
|
||||
value={{
|
||||
...initial,
|
||||
custom_tier_set: {
|
||||
tiers: [{ id: "review", name: "REVIEW", definition: "Code reviews", models: ["capable"] }],
|
||||
fallback_tier_id: "review",
|
||||
},
|
||||
}}
|
||||
onChange={onChange}
|
||||
>
|
||||
Custom tiers
|
||||
</AutoRouterClassifierTabs>,
|
||||
);
|
||||
expect(screen.getByRole("tabpanel", { name: "Complexity" })).toHaveTextContent("Custom tiers");
|
||||
for (const name of ["Capability", "Fuse v2"]) {
|
||||
const tab = screen.getByRole("tab", { name });
|
||||
expect(tab).toHaveAttribute("aria-disabled", "true");
|
||||
expect(tab).toHaveAccessibleDescription("Restore standard tiers to use Capability or Fuse v2.");
|
||||
fireEvent.click(tab);
|
||||
}
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(screen.getByText("Restore standard tiers to use Capability or Fuse v2.")).toBeVisible();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
import React, { useId } from "react";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { effectiveClassifierType, type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { transitionClassifierType } from "./classifier_type_transition";
|
||||
import { isForecastClassifier } from "./forecast_classifier_config";
|
||||
|
||||
interface AutoRouterClassifierTabsProps {
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
children: React.ReactNode;
|
||||
}
|
||||
|
||||
const AutoRouterClassifierTabs: React.FC<AutoRouterClassifierTabsProps> = ({ value, onChange, children }) => {
|
||||
const restrictionId = useId();
|
||||
const classifierType = effectiveClassifierType(value);
|
||||
const selected = isForecastClassifier(classifierType) ? classifierType : "complexity";
|
||||
const hasCustomTiers = Boolean(value.custom_tier_set);
|
||||
|
||||
const handleChange = (tab: unknown) => {
|
||||
if (tab === selected) return;
|
||||
if (tab === "complexity") {
|
||||
onChange(transitionClassifierType(value, isForecastClassifier(classifierType) ? "heuristic" : classifierType));
|
||||
} else if (!hasCustomTiers && (tab === "capability" || tab === "llm_v2")) {
|
||||
onChange(transitionClassifierType(value, tab));
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Tabs value={selected} onValueChange={handleChange}>
|
||||
<p className="text-sm font-medium">Classifier type</p>
|
||||
<TabsList aria-label="Classifier type" className="w-full">
|
||||
<TabsTrigger value="complexity">Complexity</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="capability"
|
||||
disabled={hasCustomTiers}
|
||||
aria-describedby={hasCustomTiers ? restrictionId : undefined}
|
||||
>
|
||||
Capability
|
||||
</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="llm_v2"
|
||||
disabled={hasCustomTiers}
|
||||
aria-describedby={hasCustomTiers ? restrictionId : undefined}
|
||||
>
|
||||
Fuse v2
|
||||
</TabsTrigger>
|
||||
</TabsList>
|
||||
{hasCustomTiers && (
|
||||
<p id={restrictionId} className="text-sm text-muted-foreground">
|
||||
Restore standard tiers to use Capability or Fuse v2.
|
||||
</p>
|
||||
)}
|
||||
<TabsContent value={selected}>{children}</TabsContent>
|
||||
</Tabs>
|
||||
);
|
||||
};
|
||||
|
||||
export default AutoRouterClassifierTabs;
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import { transitionClassifierType } from "./classifier_type_transition";
|
||||
import { Info } from "lucide-react";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { MultiSelect } from "@/components/shared/MultiSelect";
|
||||
|
|
@ -17,7 +18,6 @@ import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect";
|
|||
import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig";
|
||||
import ClassifierVisionConfig from "./ClassifierVisionConfig";
|
||||
import type { ReasoningEffort } from "./complexity_router_tiers";
|
||||
import { nonReasoningTierFields } from "./nonReasoningTierFields";
|
||||
import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults";
|
||||
import {
|
||||
ClassificationFrequency,
|
||||
|
|
@ -33,12 +33,10 @@ import {
|
|||
DEFAULT_CLASSIFIER_FALLBACK,
|
||||
DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
DEFAULT_CLASSIFICATION_RUBRIC,
|
||||
NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
|
||||
ClassificationRubric,
|
||||
effectiveTierLabel,
|
||||
heuristicScoringRole,
|
||||
usesLlmClassifier,
|
||||
DEFAULT_HEURISTIC_FIRST_MAX_TIER,
|
||||
DEFAULT_HYBRID_BOUNDARY_MARGIN,
|
||||
HEURISTIC_FIRST_MAX_TIER_KEYS,
|
||||
effectiveClassifierType,
|
||||
|
|
@ -263,35 +261,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
const explicitlySupportedClassifierEfforts = effortOptionsByModel[classifierModel];
|
||||
|
||||
const handleClassifierTypeChange = (classifierType: ClassifierType) => {
|
||||
const nextValue: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
classifier_type: classifierType,
|
||||
classifier_llm_config: usesLlmClassifier(classifierType)
|
||||
? value.classifier_llm_config ?? {
|
||||
model: "",
|
||||
timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
classification_rubric: NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
|
||||
}
|
||||
: undefined,
|
||||
classifier_context_window_size: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
|
||||
: undefined,
|
||||
classifier_context_budget_chars: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_include_assistant_turns
|
||||
: undefined,
|
||||
classifier_fallback: usesLlmClassifier(classifierType) ? value.classifier_fallback : undefined,
|
||||
heuristic_first_max_tier:
|
||||
classifierType === "heuristic_first"
|
||||
? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER
|
||||
: undefined,
|
||||
hybrid_boundary_margin:
|
||||
classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined,
|
||||
...nonReasoningTierFields(classifierType, value),
|
||||
};
|
||||
onChange(nextValue);
|
||||
onChange(transitionClassifierType(value, classifierType));
|
||||
};
|
||||
|
||||
const handleHeuristicFirstMaxTierChange = (tier: string) => {
|
||||
|
|
@ -433,27 +403,6 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
});
|
||||
};
|
||||
|
||||
if (classifierType === "capability") {
|
||||
return (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
This router uses capability forecasting. Configure its classifier, threshold, and calibration through YAML or
|
||||
the API. Saving preserves those settings
|
||||
</p>
|
||||
);
|
||||
}
|
||||
|
||||
if (classifierType === "llm_v2") {
|
||||
return (
|
||||
<div className="rounded-md border p-4 text-sm">
|
||||
<strong>LLM V2 classifier (experimental)</strong>
|
||||
<p className="mt-2 text-muted-foreground">
|
||||
Combines task demands and model capability in one forecast. Its solver profiles and quality allowance are
|
||||
configured through the API. Saving this router preserves those settings
|
||||
</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<ClassifierTypeRadios value={value} classifierType={classifierType} onTypeChange={handleClassifierTypeChange} />
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
import RoutingOptions from "./RoutingOptions";
|
||||
import PlanModeOverrideControls from "./PlanModeOverrideControls";
|
||||
import ForecastClassifierConfig, { ForecastSolverModels } from "./ForecastClassifierConfig";
|
||||
import { isForecastClassifier, type CapabilitySettings, type FuseSettings } from "./forecast_classifier_config";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { MultiSelect } from "@/components/shared/MultiSelect";
|
||||
import { SearchSelect } from "@/components/shared/SearchSelect";
|
||||
import DefaultModelField from "./DefaultModelField";
|
||||
import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
|
||||
import { AffinityControls } from "./AffinityControls";
|
||||
import NonReasoningTierToggle from "./NonReasoningTierToggle";
|
||||
|
|
@ -204,11 +207,6 @@ const rowOrigin = (row: TierRow, editing: boolean): string => {
|
|||
return isBuiltInTierName(row.name) ? "built-in" : "custom";
|
||||
};
|
||||
|
||||
const defaultModelPlaceholderFor = (derivedDefaultModel: string | undefined, isCustomSet: boolean): string => {
|
||||
if (derivedDefaultModel) return `Derived from tiers: ${derivedDefaultModel}`;
|
||||
return isCustomSet ? "Add a model to your fallback tier" : "Add a model to the Simple or Medium tier";
|
||||
};
|
||||
|
||||
const builtInTierInfo = (rowId: string): { label: string; description: string; examples: string } | undefined => {
|
||||
const builtIn = ALL_BUILT_IN_TIERS.find((tier) => tier === rowId);
|
||||
return builtIn ? TIER_DESCRIPTIONS[builtIn] : undefined;
|
||||
|
|
@ -376,6 +374,8 @@ export interface ComplexityRouterConfigValue {
|
|||
/** An explicit pin. Unset means the default tracks the tiers - see resolveComplexityDefaultModel. */
|
||||
default_model?: string;
|
||||
classifier_type: ClassifierType;
|
||||
capability_classifier_config?: CapabilitySettings;
|
||||
llm_v2_config?: FuseSettings;
|
||||
classifier_llm_config?: ClassifierLLMConfig;
|
||||
classifier_context_window_size?: number;
|
||||
classifier_context_budget_chars?: number;
|
||||
|
|
@ -535,44 +535,6 @@ export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03;
|
|||
*/
|
||||
export const HEURISTIC_FIRST_MAX_TIER_KEYS = TIER_ORDER.slice(0, -1);
|
||||
|
||||
const PlanModeOverrideControls: React.FC<{
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
planModeTierOptions: { value: string; label: string }[];
|
||||
}> = ({ value, onChange, planModeTierOptions }) => (
|
||||
<>
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<Switch
|
||||
checked={value.plan_mode_min_tier !== undefined}
|
||||
disabled={planModeTierOptions.length === 0}
|
||||
onCheckedChange={(enabled) =>
|
||||
onChange({
|
||||
...value,
|
||||
plan_mode_min_tier: enabled ? planModeTierOptions.at(-1)?.value : undefined,
|
||||
})
|
||||
}
|
||||
aria-label="Route plan-mode requests to a minimum tier"
|
||||
/>
|
||||
<strong className="font-semibold">Route plan-mode requests to a minimum tier</strong>
|
||||
</div>
|
||||
<span className="block text-xs mb-3 text-muted-foreground">
|
||||
Requests from coding agents in plan mode (Claude Code, GitHub Copilot) route to at least this tier. The classifier
|
||||
still wins when it picks higher, and the override only lasts while plan mode is active.
|
||||
{planModeTierOptions.length === 0 && " Add models to a tier to enable this."}
|
||||
</span>
|
||||
{value.plan_mode_min_tier !== undefined && (
|
||||
<div style={{ maxWidth: 320 }}>
|
||||
<TierRowSelect
|
||||
label="Plan-mode minimum tier"
|
||||
options={planModeTierOptions}
|
||||
value={value.plan_mode_min_tier ?? null}
|
||||
onValueChange={(tier) => onChange({ ...value, plan_mode_min_tier: tier })}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
);
|
||||
|
||||
const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
||||
modelInfo,
|
||||
value,
|
||||
|
|
@ -596,6 +558,7 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
onAutoRouterCompressionChange,
|
||||
showValidationErrors = false,
|
||||
}) => {
|
||||
const forecast = isForecastClassifier(value.classifier_type);
|
||||
const customTierSet = value.custom_tier_set;
|
||||
const tierRows = activeTierRows(value);
|
||||
const tierRowsError = customTierSet ? getCustomTierRowsError(customTierSet) : null;
|
||||
|
|
@ -605,8 +568,6 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
value: row.id,
|
||||
label: tierRowLabel(row, value.tier_labels),
|
||||
}));
|
||||
const derivedDefaultModel = resolveComplexityDefaultModel(value);
|
||||
const defaultModelPlaceholder = defaultModelPlaceholderFor(derivedDefaultModel, Boolean(customTierSet));
|
||||
const defaultModel = resolveComplexityDefaultModel(value, value.default_model);
|
||||
|
||||
const dispatch = (action: TierSetAction) => {
|
||||
|
|
@ -641,298 +602,319 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, change),
|
||||
});
|
||||
|
||||
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
|
||||
// "track the tiers" everywhere downstream instead of as a blank model name.
|
||||
const handleDefaultModelChange = (model: string | null | undefined) => {
|
||||
onChange({ ...value, default_model: model || undefined });
|
||||
};
|
||||
|
||||
const handleTierLabelChange = (tier: keyof ComplexityTiers, label: string) => {
|
||||
onChange({
|
||||
...value,
|
||||
tier_labels: { ...value.tier_labels, [tier]: label },
|
||||
});
|
||||
};
|
||||
const handleTierLabelChange = (tier: keyof ComplexityTiers, label: string) =>
|
||||
onChange({ ...value, tier_labels: { ...value.tier_labels, [tier]: label } });
|
||||
|
||||
return (
|
||||
<div className="w-full max-w-none">
|
||||
<div className="inline-flex items-center gap-2 mb-4">
|
||||
<h4 className="m-0 text-xl font-semibold text-foreground">Complexity Tier Configuration</h4>
|
||||
<SimpleTooltip content="Map each complexity tier to one or more models. Simple queries use cheaper/faster models, complex queries use more capable models.">
|
||||
<Info className="size-4 text-muted-foreground" />
|
||||
</SimpleTooltip>
|
||||
<h4 className="m-0 text-xl font-semibold text-foreground">
|
||||
{forecast ? "Solver models" : "Complexity Tier Configuration"}
|
||||
</h4>
|
||||
{!forecast && (
|
||||
<SimpleTooltip content="Map each complexity tier to one or more models. Simple queries use cheaper/faster models, complex queries use more capable models.">
|
||||
<Info className="size-4 text-muted-foreground" />
|
||||
</SimpleTooltip>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<TierConfigIntro value={value} />
|
||||
|
||||
<Card>
|
||||
<CardContent>
|
||||
{!customTierSet && (
|
||||
<NonReasoningTierToggle value={value} onChange={onChange} available={value.classifier_type === "llm"} />
|
||||
)}
|
||||
|
||||
{tierRows.map((row, index) => {
|
||||
const tierInfo = builtInTierInfo(row.id);
|
||||
const label = tierRowLabel(row, value.tier_labels);
|
||||
const tierMissing = showValidationErrors && row.models.length === 0;
|
||||
const needsDefinition = Boolean(customTierSet) && !row.definition.trim() && !isBuiltInTierName(row.name);
|
||||
const definitionMissing = showValidationErrors && needsDefinition;
|
||||
const showsDisplayName = !customTierSet && !editingTiers;
|
||||
return (
|
||||
<div key={row.id}>
|
||||
{index > 0 && <Separator className="my-4" />}
|
||||
<div className="mb-4">
|
||||
<TierRowHeader
|
||||
row={row}
|
||||
index={index}
|
||||
rowCount={tierRows.length}
|
||||
label={label}
|
||||
description={tierInfo?.description}
|
||||
editing={editingTiers}
|
||||
isCustomSet={Boolean(customTierSet)}
|
||||
onRemove={() => removeTierRow(row.id)}
|
||||
/>
|
||||
{tierInfo && !customTierSet && (
|
||||
<span className="block mb-2 text-xs text-muted-foreground">Examples: {tierInfo.examples}</span>
|
||||
)}
|
||||
{editingTiers && (
|
||||
<TierRowEditFields
|
||||
row={row}
|
||||
index={index}
|
||||
definitionMissing={definitionMissing}
|
||||
onPatch={(patch) => updateTierRow(row.id, patch)}
|
||||
/>
|
||||
)}
|
||||
{showsDisplayName && tierInfo && (
|
||||
<InputGroup className="mb-2">
|
||||
<InputGroupInput
|
||||
value={value.tier_labels?.[row.id as keyof ComplexityTiers] ?? ""}
|
||||
onChange={(event) => handleTierLabelChange(row.id as keyof ComplexityTiers, event.target.value)}
|
||||
placeholder={`Display name (default: ${tierInfo.label})`}
|
||||
aria-label={`Display name for the ${tierInfo.label} tier`}
|
||||
/>
|
||||
{value.tier_labels?.[row.id as keyof ComplexityTiers] && (
|
||||
<InputGroupAddon align="inline-end">
|
||||
<InputGroupButton
|
||||
size="icon-xs"
|
||||
aria-label={`Clear display name for the ${tierInfo.label} tier`}
|
||||
onClick={() => handleTierLabelChange(row.id as keyof ComplexityTiers, "")}
|
||||
>
|
||||
<X />
|
||||
</InputGroupButton>
|
||||
</InputGroupAddon>
|
||||
)}
|
||||
</InputGroup>
|
||||
)}
|
||||
<MultiSelect
|
||||
options={modelOptions}
|
||||
value={row.models}
|
||||
onValueChange={(models: string[]) => setRowModels(row, models)}
|
||||
placeholder={`Select model(s) for ${label.toLowerCase()} queries`}
|
||||
emptyText="No models found"
|
||||
className={tierMissing ? "w-full border-destructive" : "w-full"}
|
||||
/>
|
||||
<TierModelEffortRows
|
||||
tierLabel={label}
|
||||
models={row.models}
|
||||
effortOptionsByModel={tierEffortOptionsByModel}
|
||||
paramsByModel={row.params}
|
||||
fastModeByModel={fastModeByModel}
|
||||
onEffortChange={(model, effort) =>
|
||||
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
|
||||
}
|
||||
onFastModeChange={(model, enabled) =>
|
||||
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
|
||||
}
|
||||
/>
|
||||
{row.models.length > 1 && (
|
||||
<span className="text-xs text-muted-foreground">
|
||||
Multiple models selected: the router randomly picks among them per request (or Thompson-samples
|
||||
within the pool when adaptive routing is on).
|
||||
</span>
|
||||
)}
|
||||
{tierMissing && <span className="text-xs text-destructive">The {label} tier is required</span>}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
|
||||
<TierSetToolbar
|
||||
editing={editingTiers}
|
||||
isCustomSet={Boolean(customTierSet)}
|
||||
rowCount={tierRows.length}
|
||||
rowsError={tierRowsError}
|
||||
keywordRulesError={keywordRulesError}
|
||||
onEditingChange={onEditingTiersChange}
|
||||
onAdd={addCustomTier}
|
||||
onRestore={exitToBuiltInTiers}
|
||||
{forecast ? (
|
||||
<>
|
||||
<ForecastSolverModels
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
modelOptions={modelOptions}
|
||||
effortOptionsByModel={tierEffortOptionsByModel}
|
||||
fastModeByModel={fastModeByModel}
|
||||
/>
|
||||
<ForecastClassifierConfig
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
modelOptions={modelOptions}
|
||||
effortOptionsByModel={classifierEffortOptionsByModel}
|
||||
/>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<TierConfigIntro value={value} />
|
||||
|
||||
{customTierSet && (
|
||||
<FallbackTierField
|
||||
rows={tierRows}
|
||||
fallbackTierId={customTierSet.fallback_tier_id}
|
||||
onValueChange={(fallbackTierId) => onChange(setFallbackTier(value, fallbackTierId))}
|
||||
/>
|
||||
)}
|
||||
<Card>
|
||||
<CardContent>
|
||||
{!customTierSet && (
|
||||
<NonReasoningTierToggle value={value} onChange={onChange} available={value.classifier_type === "llm"} />
|
||||
)}
|
||||
|
||||
<Separator className="my-4" />
|
||||
{tierRows.map((row, index) => {
|
||||
const tierInfo = builtInTierInfo(row.id);
|
||||
const label = tierRowLabel(row, value.tier_labels);
|
||||
const tierMissing = showValidationErrors && row.models.length === 0;
|
||||
const needsDefinition =
|
||||
Boolean(customTierSet) && !row.definition.trim() && !isBuiltInTierName(row.name);
|
||||
const definitionMissing = showValidationErrors && needsDefinition;
|
||||
const showsDisplayName = !customTierSet && !editingTiers;
|
||||
return (
|
||||
<div key={row.id}>
|
||||
{index > 0 && <Separator className="my-4" />}
|
||||
<div className="mb-4">
|
||||
<TierRowHeader
|
||||
row={row}
|
||||
index={index}
|
||||
rowCount={tierRows.length}
|
||||
label={label}
|
||||
description={tierInfo?.description}
|
||||
editing={editingTiers}
|
||||
isCustomSet={Boolean(customTierSet)}
|
||||
onRemove={() => removeTierRow(row.id)}
|
||||
/>
|
||||
{tierInfo && !customTierSet && (
|
||||
<span className="block mb-2 text-xs text-muted-foreground">Examples: {tierInfo.examples}</span>
|
||||
)}
|
||||
{editingTiers && (
|
||||
<TierRowEditFields
|
||||
row={row}
|
||||
index={index}
|
||||
definitionMissing={definitionMissing}
|
||||
onPatch={(patch) => updateTierRow(row.id, patch)}
|
||||
/>
|
||||
)}
|
||||
{showsDisplayName && tierInfo && (
|
||||
<InputGroup className="mb-2">
|
||||
<InputGroupInput
|
||||
value={value.tier_labels?.[row.id as keyof ComplexityTiers] ?? ""}
|
||||
onChange={(event) =>
|
||||
handleTierLabelChange(row.id as keyof ComplexityTiers, event.target.value)
|
||||
}
|
||||
placeholder={`Display name (default: ${tierInfo.label})`}
|
||||
aria-label={`Display name for the ${tierInfo.label} tier`}
|
||||
/>
|
||||
{value.tier_labels?.[row.id as keyof ComplexityTiers] && (
|
||||
<InputGroupAddon align="inline-end">
|
||||
<InputGroupButton
|
||||
size="icon-xs"
|
||||
aria-label={`Clear display name for the ${tierInfo.label} tier`}
|
||||
onClick={() => handleTierLabelChange(row.id as keyof ComplexityTiers, "")}
|
||||
>
|
||||
<X />
|
||||
</InputGroupButton>
|
||||
</InputGroupAddon>
|
||||
)}
|
||||
</InputGroup>
|
||||
)}
|
||||
<MultiSelect
|
||||
options={modelOptions}
|
||||
value={row.models}
|
||||
onValueChange={(models: string[]) => setRowModels(row, models)}
|
||||
placeholder={`Select model(s) for ${label.toLowerCase()} queries`}
|
||||
emptyText="No models found"
|
||||
className={tierMissing ? "w-full border-destructive" : "w-full"}
|
||||
/>
|
||||
<TierModelEffortRows
|
||||
tierLabel={label}
|
||||
models={row.models}
|
||||
effortOptionsByModel={tierEffortOptionsByModel}
|
||||
paramsByModel={row.params}
|
||||
fastModeByModel={fastModeByModel}
|
||||
onEffortChange={(model, effort) =>
|
||||
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
|
||||
}
|
||||
onFastModeChange={(model, enabled) =>
|
||||
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
|
||||
}
|
||||
/>
|
||||
{row.models.length > 1 && (
|
||||
<span className="text-xs text-muted-foreground">
|
||||
Multiple models selected: the router randomly picks among them per request (or
|
||||
Thompson-samples within the pool when adaptive routing is on).
|
||||
</span>
|
||||
)}
|
||||
{tierMissing && <span className="text-xs text-destructive">The {label} tier is required</span>}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
|
||||
<div className="mb-2">
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<strong className="text-base font-semibold">Default Model</strong>
|
||||
<SimpleTooltip content="Leave empty to follow the tiers. A model chosen here is pinned: it stays the default however the tiers change.">
|
||||
<Info className="size-4 text-muted-foreground" />
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
<SearchSelect
|
||||
options={modelOptions}
|
||||
value={value.default_model ?? ""}
|
||||
onValueChange={handleDefaultModelChange}
|
||||
placeholder={defaultModelPlaceholder}
|
||||
emptyText="No models found"
|
||||
aria-label="Default model"
|
||||
/>
|
||||
<span className="block mt-1 text-xs text-muted-foreground">
|
||||
Used when the tier the request lands in has no model, and when the classifier fails with "Route to
|
||||
the default model" selected.
|
||||
</span>
|
||||
</div>
|
||||
</CardContent>
|
||||
</Card>
|
||||
<TierSetToolbar
|
||||
editing={editingTiers}
|
||||
isCustomSet={Boolean(customTierSet)}
|
||||
rowCount={tierRows.length}
|
||||
rowsError={tierRowsError}
|
||||
keywordRulesError={keywordRulesError}
|
||||
onEditingChange={onEditingTiersChange}
|
||||
onAdd={addCustomTier}
|
||||
onRestore={exitToBuiltInTiers}
|
||||
/>
|
||||
|
||||
{customTierSet && (
|
||||
<FallbackTierField
|
||||
rows={tierRows}
|
||||
fallbackTierId={customTierSet.fallback_tier_id}
|
||||
onValueChange={(fallbackTierId) => onChange(setFallbackTier(value, fallbackTierId))}
|
||||
/>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
</>
|
||||
)}
|
||||
{!forecast && <DefaultModelField value={value} onChange={onChange} modelOptions={modelOptions} />}
|
||||
<Separator className="my-6" />
|
||||
|
||||
<div className="rounded-lg border border-border bg-muted">
|
||||
{[
|
||||
{
|
||||
key: "classifier",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Classification Method</strong>,
|
||||
children: (
|
||||
<ClassificationMethodConfig
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
modelOptions={modelOptions}
|
||||
effortOptionsByModel={classifierEffortOptionsByModel}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange}
|
||||
showValidationErrors={showValidationErrors}
|
||||
defaultModel={defaultModel}
|
||||
/>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "adaptive",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Adaptive Routing</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "adaptive")}>
|
||||
<AdaptiveRoutingConfig value={value} onChange={onChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "affinity",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Affinity</strong>,
|
||||
children: <AffinityControls value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "modality",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Modality Routing</strong>,
|
||||
children: <ModalityRoutingControls value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "plan-mode",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Plan-Mode Override</strong>,
|
||||
children: (
|
||||
<PlanModeOverrideControls value={value} onChange={onChange} planModeTierOptions={planModeTierOptions} />
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "context-window",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Context Window Escalation</strong>,
|
||||
children: <ContextWindowEscalationConfig value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "stall-escalation",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Stalled Task Escalation</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "stallEscalation")}>
|
||||
<StallEscalationConfig value={value} onChange={onChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "response",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Response Format</strong>,
|
||||
children: <ResponseFormatControls value={value} onChange={onChange} />,
|
||||
},
|
||||
...(onEscalationKeywordsChange
|
||||
? [
|
||||
{
|
||||
key: "escalation",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Escalation Keywords</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "escalation")}>
|
||||
<EscalationKeywords keywords={escalationKeywords} onChange={onEscalationKeywordsChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(onAutoRouterCompressionChange
|
||||
? [
|
||||
{
|
||||
key: "compression",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
|
||||
children: (
|
||||
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
|
||||
? [
|
||||
{
|
||||
key: "keyword-semantic",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Keyword/Semantic Matching</strong>,
|
||||
children: (
|
||||
<>
|
||||
{onKeywordTierRulesChange && (
|
||||
<KeywordTierRules
|
||||
rules={keywordTierRules}
|
||||
onChange={onKeywordTierRulesChange}
|
||||
tierLabels={value.tier_labels}
|
||||
tierNames={customTierSet && tierRows.map(activeTierName).filter(Boolean)}
|
||||
/>
|
||||
)}
|
||||
{onKeywordTierRulesChange && onSemanticMatchingEnabledChange && <Separator className="my-4" />}
|
||||
{onSemanticMatchingEnabledChange && (
|
||||
<SemanticKeywordMatching
|
||||
enabled={semanticMatchingEnabled}
|
||||
onEnabledChange={onSemanticMatchingEnabledChange}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={onEmbeddingModelChange}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={onMatchThresholdChange}
|
||||
modelInfo={modelInfo}
|
||||
showValidationErrors={showValidationErrors}
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
].map(({ key, label, children }) => (
|
||||
<Collapsible key={key} className="border-b border-border last:border-b-0">
|
||||
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left">
|
||||
<ChevronRight className="size-4 shrink-0 text-muted-foreground transition-transform group-data-panel-open:rotate-90" />
|
||||
{label}
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="px-4 pb-4">{children}</CollapsibleContent>
|
||||
</Collapsible>
|
||||
))}
|
||||
</div>
|
||||
<RoutingOptions forecast={forecast}>
|
||||
{forecast && (
|
||||
<>
|
||||
<DefaultModelField value={value} onChange={onChange} modelOptions={modelOptions} />
|
||||
<ForecastSolverModels
|
||||
additionalPoolsOnly
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
modelOptions={modelOptions}
|
||||
effortOptionsByModel={tierEffortOptionsByModel}
|
||||
fastModeByModel={fastModeByModel}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
<div className="rounded-lg border border-border bg-muted">
|
||||
{[
|
||||
...(!forecast
|
||||
? [
|
||||
{
|
||||
key: "classifier",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Classification Method</strong>,
|
||||
children: (
|
||||
<ClassificationMethodConfig
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
modelOptions={modelOptions}
|
||||
effortOptionsByModel={classifierEffortOptionsByModel}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange}
|
||||
showValidationErrors={showValidationErrors}
|
||||
defaultModel={defaultModel}
|
||||
/>
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
{
|
||||
key: "adaptive",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Adaptive Routing</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "adaptive")}>
|
||||
<AdaptiveRoutingConfig value={value} onChange={onChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "affinity",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Affinity</strong>,
|
||||
children: <AffinityControls value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "modality",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Modality Routing</strong>,
|
||||
children: <ModalityRoutingControls value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "plan-mode",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Plan-Mode Override</strong>,
|
||||
children: (
|
||||
<PlanModeOverrideControls value={value} onChange={onChange} planModeTierOptions={planModeTierOptions} />
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "context-window",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Context Window Escalation</strong>,
|
||||
children: <ContextWindowEscalationConfig value={value} onChange={onChange} />,
|
||||
},
|
||||
{
|
||||
key: "stall-escalation",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Stalled Task Escalation</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "stallEscalation")}>
|
||||
<StallEscalationConfig value={value} onChange={onChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "response",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Response Format</strong>,
|
||||
children: <ResponseFormatControls value={value} onChange={onChange} />,
|
||||
},
|
||||
...(onEscalationKeywordsChange
|
||||
? [
|
||||
{
|
||||
key: "escalation",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Escalation Keywords</strong>,
|
||||
children: (
|
||||
<Restricted by={restrictedBy(value, "escalation")}>
|
||||
<EscalationKeywords keywords={escalationKeywords} onChange={onEscalationKeywordsChange} />
|
||||
</Restricted>
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(onAutoRouterCompressionChange
|
||||
? [
|
||||
{
|
||||
key: "compression",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
|
||||
children: (
|
||||
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
|
||||
? [
|
||||
{
|
||||
key: "keyword-semantic",
|
||||
label: (
|
||||
<strong className="text-foreground font-semibold">Advanced: Keyword/Semantic Matching</strong>
|
||||
),
|
||||
children: (
|
||||
<>
|
||||
{onKeywordTierRulesChange && (
|
||||
<KeywordTierRules
|
||||
rules={keywordTierRules}
|
||||
onChange={onKeywordTierRulesChange}
|
||||
tierLabels={value.tier_labels}
|
||||
tierNames={
|
||||
customTierSet || isForecastClassifier(value.classifier_type)
|
||||
? tierRows.map(activeTierName).filter(Boolean)
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
)}
|
||||
{onKeywordTierRulesChange && onSemanticMatchingEnabledChange && <Separator className="my-4" />}
|
||||
{onSemanticMatchingEnabledChange && (
|
||||
<SemanticKeywordMatching
|
||||
enabled={semanticMatchingEnabled}
|
||||
onEnabledChange={onSemanticMatchingEnabledChange}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={onEmbeddingModelChange}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={onMatchThresholdChange}
|
||||
modelInfo={modelInfo}
|
||||
showValidationErrors={showValidationErrors}
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
]
|
||||
.filter(({ key }) => !forecast || !["adaptive", "context-window", "escalation"].includes(key))
|
||||
.map(({ key, label, children }) => (
|
||||
<Collapsible key={key} className="border-b border-border last:border-b-0">
|
||||
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left">
|
||||
<ChevronRight className="size-4 shrink-0 text-muted-foreground transition-transform group-data-panel-open:rotate-90" />
|
||||
{label}
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="px-4 pb-4">{children}</CollapsibleContent>
|
||||
</Collapsible>
|
||||
))}
|
||||
</div>
|
||||
</RoutingOptions>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import { renderWithProviders, screen, within } from "../../../tests/test-utils";
|
||||
import {
|
||||
buildUpdatedComplexityRouterConfig,
|
||||
hydrateComplexityRouterConfig,
|
||||
} from "../edit_auto_router/edit_auto_router_modal";
|
||||
import type { KeywordTierRule } from "./KeywordTierRules";
|
||||
import type { ModelGroup } from "../llm_calls/fetch_models";
|
||||
import ComplexityRouterConfig, { type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
|
||||
|
|
@ -50,8 +51,8 @@ it.each([false, true])("edits and round-trips independent model settings with cu
|
|||
const view = renderWithProviders(editor(initial));
|
||||
const fast = () => screen.getByRole("switch", { name: `Fast mode for primary in the ${label} tier` });
|
||||
|
||||
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(3);
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for (blocked|missing)/ })).not.toBeInTheDocument();
|
||||
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(4);
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for missing/ })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("combobox", { name: /^Reasoning effort for secondary/ })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("switch", { name: `Fast mode for secondary in the ${label} tier` })).toBeChecked();
|
||||
expect(fast()).not.toBeChecked();
|
||||
|
|
@ -125,19 +126,170 @@ it.each([false, true])("edits and round-trips independent model settings with cu
|
|||
);
|
||||
});
|
||||
|
||||
describe("Fast mode metadata", () => {
|
||||
it("offers nothing before model capabilities load and leaves stored speed untouched", () => {
|
||||
const value: ComplexityRouterConfigValue = {
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: [] },
|
||||
classifier_type: "heuristic",
|
||||
tier_model_params: { SIMPLE: { primary: { speed: "fast" } } },
|
||||
};
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={[]} value={value} onChange={onChange} />);
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for/ })).not.toBeInTheDocument();
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(buildUpdatedComplexityRouterConfig({}, value).tier_model_configs).toEqual({
|
||||
SIMPLE: [{ model_name: "primary", litellm_params: { speed: "fast" } }],
|
||||
});
|
||||
it.each(["capability", "llm_v2"] as const)("preserves Fast mode controls for %s solvers", async (classifierType) => {
|
||||
const user = userEvent.setup();
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
classifier_type: classifierType,
|
||||
classifier_llm_config: { model: "primary", timeout_ms: 3000 },
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["blocked"] },
|
||||
capability_classifier_config: { efficient_tier: "SIMPLE", capable_tier: "REASONING", base_threshold: 0.7 },
|
||||
llm_v2_config: {
|
||||
efficient_profile: "Small solver",
|
||||
capable_profile: "Large solver",
|
||||
harness: "One attempt",
|
||||
max_quality_gap: 0.05,
|
||||
},
|
||||
tier_model_params: { SIMPLE: { primary: { reasoning_effort: "high", max_tokens: 1024 } } },
|
||||
};
|
||||
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
|
||||
const editor = (value: ComplexityRouterConfigValue) => (
|
||||
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
|
||||
);
|
||||
const view = renderWithProviders(editor(initial));
|
||||
const fast = () => screen.getByRole("switch", { name: "Fast mode for primary in the Efficient solver tier" });
|
||||
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(1);
|
||||
expect(fast()).not.toBeChecked();
|
||||
await user.click(fast());
|
||||
const enabled = onChange.mock.lastCall![0];
|
||||
expect(enabled.tier_model_params?.SIMPLE.primary).toEqual({
|
||||
reasoning_effort: "high",
|
||||
max_tokens: 1024,
|
||||
speed: "fast",
|
||||
});
|
||||
const saved = buildUpdatedComplexityRouterConfig({}, enabled);
|
||||
expect(saved.tier_model_configs).toEqual({
|
||||
SIMPLE: [{ model_name: "primary", litellm_params: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" } }],
|
||||
});
|
||||
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined)));
|
||||
expect(fast()).toBeChecked();
|
||||
await user.click(fast());
|
||||
expect(onChange.mock.lastCall![0].tier_model_params?.SIMPLE.primary).toEqual({
|
||||
reasoning_effort: "high",
|
||||
max_tokens: 1024,
|
||||
});
|
||||
});
|
||||
|
||||
describe("Fast mode metadata", () => {
|
||||
it.each(["heuristic", "capability", "llm_v2"] as const)(
|
||||
"can clear stored Fast mode without current capability metadata for %s",
|
||||
async (classifier_type) => {
|
||||
const user = userEvent.setup();
|
||||
const value: ComplexityRouterConfigValue = {
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["secondary"] },
|
||||
classifier_type,
|
||||
tier_model_params: { SIMPLE: { primary: { speed: "fast", max_tokens: 512 } } },
|
||||
};
|
||||
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
|
||||
const editor = (current: ComplexityRouterConfigValue, info: ModelGroup[]) => (
|
||||
<ComplexityRouterConfig modelInfo={info} value={current} onChange={onChange} />
|
||||
);
|
||||
const view = renderWithProviders(editor(value, []));
|
||||
const fast = () => screen.getByRole("switch", { name: /^Fast mode for primary/ });
|
||||
expect(fast()).toBeChecked();
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
view.rerender(editor(value, [{ model_group: "primary", supports_fast_mode: false }]));
|
||||
expect(fast()).toBeChecked();
|
||||
await user.click(fast());
|
||||
const cleared = onChange.mock.lastCall![0];
|
||||
expect(cleared.tier_model_params?.SIMPLE.primary).toEqual({ max_tokens: 512 });
|
||||
const saved = buildUpdatedComplexityRouterConfig({}, cleared);
|
||||
expect(saved.tier_model_configs).toEqual({
|
||||
SIMPLE: [{ model_name: "primary", litellm_params: { max_tokens: 512 } }],
|
||||
});
|
||||
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined), []));
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for primary/ })).not.toBeInTheDocument();
|
||||
view.rerender(editor(cleared, modelInfo));
|
||||
expect(fast()).not.toBeChecked();
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it.each(["MEDIUM", "REASONING"])("clears a legacy Capability pool while reconciling plan floor %s", async (floor) => {
|
||||
const user = userEvent.setup();
|
||||
const stored = {
|
||||
classifier_type: "capability" as const,
|
||||
plan_mode_min_tier: floor,
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: ["secondary"], COMPLEX: [], REASONING: ["blocked"] },
|
||||
tier_model_configs: { MEDIUM: [{ model_name: "secondary", litellm_params: { speed: "fast" } }] },
|
||||
};
|
||||
const value = hydrateComplexityRouterConfig(stored, undefined);
|
||||
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />);
|
||||
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
|
||||
expect(screen.getByRole("switch", { name: "Fast mode for secondary in the Medium routing pool tier" })).toBeChecked();
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
await user.click(screen.getByRole("combobox", { name: "Select medium routing pool models" }));
|
||||
await user.click(await screen.findByRole("option", { name: "secondary" }));
|
||||
await user.keyboard("{Escape}");
|
||||
const cleared = onChange.mock.lastCall![0];
|
||||
expect(cleared.tiers.MEDIUM).toEqual([]);
|
||||
expect(cleared.plan_mode_min_tier).toBe(floor === "MEDIUM" ? undefined : floor);
|
||||
expect(cleared.tier_model_params).toBeUndefined();
|
||||
expect(buildUpdatedComplexityRouterConfig(stored, cleared).tiers).toEqual({
|
||||
SIMPLE: ["primary"],
|
||||
REASONING: ["blocked"],
|
||||
});
|
||||
});
|
||||
|
||||
it.each(["capability", "llm_v2"] as const)(
|
||||
"shows and clears a persisted default model in %s",
|
||||
async (classifier_type) => {
|
||||
const user = userEvent.setup();
|
||||
const stored = {
|
||||
classifier_type,
|
||||
default_model: "legacy-default",
|
||||
tiers: { SIMPLE: ["primary"], REASONING: ["secondary"] },
|
||||
};
|
||||
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
|
||||
const editor = (value: ComplexityRouterConfigValue) => (
|
||||
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
|
||||
);
|
||||
const view = renderWithProviders(editor(hydrateComplexityRouterConfig(stored, undefined)));
|
||||
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
|
||||
const select = () => screen.getByRole("combobox", { name: "Default model" });
|
||||
expect(select()).toHaveValue("legacy-default");
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
await user.click(select());
|
||||
await user.click(await screen.findByRole("option", { name: "blocked" }));
|
||||
const changed = onChange.mock.lastCall![0];
|
||||
expect(buildUpdatedComplexityRouterConfig(stored, changed).default_model).toBe("blocked");
|
||||
view.rerender(editor(changed));
|
||||
await user.click(
|
||||
within(screen.getByRole("group", { name: "Default model configuration" })).getByRole("button", { name: "Clear" }),
|
||||
);
|
||||
const cleared = onChange.mock.lastCall![0];
|
||||
expect(cleared.default_model).toBeUndefined();
|
||||
const saved = buildUpdatedComplexityRouterConfig(stored, cleared);
|
||||
expect(saved).not.toHaveProperty("default_model");
|
||||
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined)));
|
||||
expect(select()).toHaveValue("");
|
||||
expect(select()).toHaveAttribute("placeholder", expect.stringContaining("primary"));
|
||||
},
|
||||
);
|
||||
|
||||
it.each(["capability", "llm_v2"] as const)("offers only populated keyword targets for %s", async (classifier_type) => {
|
||||
const user = userEvent.setup();
|
||||
const value: ComplexityRouterConfigValue = {
|
||||
classifier_type,
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["secondary"] },
|
||||
};
|
||||
const onRulesChange = vi.fn<(rules: KeywordTierRule[]) => void>();
|
||||
const editor = (rules: KeywordTierRule[]) => (
|
||||
<ComplexityRouterConfig
|
||||
value={value}
|
||||
onChange={vi.fn()}
|
||||
modelInfo={modelInfo}
|
||||
keywordTierRules={rules}
|
||||
onKeywordTierRulesChange={onRulesChange}
|
||||
/>
|
||||
);
|
||||
const view = renderWithProviders(editor([]));
|
||||
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
|
||||
await user.click(screen.getByText("Advanced: Keyword/Semantic Matching"));
|
||||
await user.click(screen.getByRole("button", { name: "Add keyword rule" }));
|
||||
const rules = onRulesChange.mock.lastCall![0];
|
||||
expect(rules[0].tier).toBe("SIMPLE");
|
||||
view.rerender(editor(rules));
|
||||
await user.click(screen.getByRole("combobox", { name: "Route keyword rule 1 to tier" }));
|
||||
expect((await screen.findAllByRole("option")).map((option) => option.textContent)).toEqual(["Simple", "Reasoning"]);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -0,0 +1,56 @@
|
|||
import React from "react";
|
||||
import { Info } from "lucide-react";
|
||||
import { SearchSelect } from "@/components/shared/SearchSelect";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { isForecastClassifier } from "./forecast_classifier_config";
|
||||
import { resolveComplexityDefaultModel } from "./tier_rows";
|
||||
|
||||
interface DefaultModelFieldProps {
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
modelOptions: { value: string; label: string }[];
|
||||
}
|
||||
|
||||
const defaultModelPlaceholderFor = (derivedDefaultModel: string | undefined, isCustomSet: boolean): string => {
|
||||
if (derivedDefaultModel) return `Derived from tiers: ${derivedDefaultModel}`;
|
||||
return isCustomSet ? "Add a model to your fallback tier" : "Add a model to the Simple or Medium tier";
|
||||
};
|
||||
|
||||
const DefaultModelField = ({ value, onChange, modelOptions }: DefaultModelFieldProps) => {
|
||||
const defaultModelPlaceholder = defaultModelPlaceholderFor(
|
||||
resolveComplexityDefaultModel(value),
|
||||
Boolean(value.custom_tier_set),
|
||||
);
|
||||
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
|
||||
// "track the tiers" everywhere downstream instead of as a blank model name.
|
||||
const handleDefaultModelChange = (model: string | null | undefined) => {
|
||||
onChange({ ...value, default_model: model || undefined });
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="mt-4 mb-2" role="group" aria-label="Default model configuration">
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<strong className="text-base font-semibold">Default Model</strong>
|
||||
<SimpleTooltip content="Leave empty to follow the tiers. A model chosen here is pinned: it stays the default however the tiers change.">
|
||||
<Info className="size-4 text-muted-foreground" />
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
<SearchSelect
|
||||
options={modelOptions}
|
||||
value={value.default_model ?? ""}
|
||||
onValueChange={handleDefaultModelChange}
|
||||
placeholder={defaultModelPlaceholder}
|
||||
emptyText="No models found"
|
||||
aria-label="Default model"
|
||||
/>
|
||||
<span className="block mt-1 text-xs text-muted-foreground">
|
||||
{isForecastClassifier(value.classifier_type)
|
||||
? "Used when routing cannot find a suitable model. Classifier failures route to the capable solver."
|
||||
: 'Used when the tier the request lands in has no model, and when the classifier fails with "Route to the default model" selected.'}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default DefaultModelField;
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
import React, { useState } from "react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import ClassificationMethodConfig from "./ClassificationMethodConfig";
|
||||
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
|
||||
import ForecastClassifierConfig from "./ForecastClassifierConfig";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { getForecastConfigError, isForecastClassifier } from "./forecast_classifier_config";
|
||||
import { buildUpdatedComplexityRouterConfig } from "../edit_auto_router/edit_auto_router_modal";
|
||||
|
||||
vi.mock("@/components/networking", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/components/networking")>()),
|
||||
getComplexityScorerDefaults: vi.fn(async () => ({
|
||||
tier_boundaries: {},
|
||||
token_thresholds: {},
|
||||
dimension_weights: {},
|
||||
})),
|
||||
}));
|
||||
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
classifier_type: "capability",
|
||||
classifier_llm_config: { model: "judge", timeout_ms: 20000 },
|
||||
tiers: { SIMPLE: ["efficient"], MEDIUM: [], COMPLEX: [], REASONING: ["capable"] },
|
||||
capability_classifier_config: { efficient_tier: "SIMPLE", capable_tier: "REASONING", base_threshold: 0.7 },
|
||||
};
|
||||
const fuseInitial: ComplexityRouterConfigValue = {
|
||||
...initial,
|
||||
classifier_type: "llm_v2",
|
||||
capability_classifier_config: undefined,
|
||||
adaptive: false,
|
||||
llm_v2_config: {
|
||||
efficient_profile: "Small solver",
|
||||
capable_profile: "Larger solver",
|
||||
harness: "One attempt",
|
||||
max_quality_gap: 0.05,
|
||||
},
|
||||
};
|
||||
const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model }));
|
||||
|
||||
function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) {
|
||||
const [value, setValue] = useState(initialValue);
|
||||
const [saved, setSaved] = useState("");
|
||||
return (
|
||||
<>
|
||||
<AutoRouterClassifierTabs value={value} onChange={setValue}>
|
||||
{isForecastClassifier(value.classifier_type) ? (
|
||||
<ForecastClassifierConfig
|
||||
value={value}
|
||||
onChange={setValue}
|
||||
modelOptions={options}
|
||||
effortOptionsByModel={{}}
|
||||
/>
|
||||
) : (
|
||||
<ClassificationMethodConfig
|
||||
value={value}
|
||||
onChange={setValue}
|
||||
modelOptions={options}
|
||||
effortOptionsByModel={{}}
|
||||
/>
|
||||
)}
|
||||
</AutoRouterClassifierTabs>
|
||||
<button
|
||||
disabled={Boolean(getForecastConfigError(value))}
|
||||
onClick={() => setSaved(JSON.stringify(buildUpdatedComplexityRouterConfig({}, value)))}
|
||||
>
|
||||
Save configuration
|
||||
</button>
|
||||
<output aria-label="Saved configuration">{saved}</output>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
describe("forecast classifier form", () => {
|
||||
it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => {
|
||||
renderWithProviders(
|
||||
<Form
|
||||
initialValue={{
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "judge", timeout_ms: 20000, classification_rubric: "agentic" },
|
||||
adaptive: true,
|
||||
plan_mode_min_tier: "MEDIUM",
|
||||
tiers: {
|
||||
SIMPLE: ["efficient", "second-efficient"],
|
||||
MEDIUM: ["leftover-medium"],
|
||||
COMPLEX: ["leftover-complex"],
|
||||
REASONING: ["capable"],
|
||||
},
|
||||
tier_model_params: {
|
||||
SIMPLE: { efficient: { reasoning_effort: "low", speed: "fast", max_tokens: 1024 } },
|
||||
MEDIUM: { "leftover-medium": { speed: "fast" } },
|
||||
COMPLEX: { "leftover-complex": { max_tokens: 4096 } },
|
||||
REASONING: { capable: { reasoning_effort: "high" } },
|
||||
},
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Capability" }));
|
||||
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
|
||||
expect(screen.getByRole("button", { name: "Save configuration" })).toBeEnabled();
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
|
||||
const output = screen.getByRole("status", { name: "Saved configuration" });
|
||||
expect(output).toHaveTextContent('"classifier_type":"capability"');
|
||||
expect(output).toHaveTextContent('"SIMPLE":["efficient","second-efficient"]');
|
||||
expect(output).toHaveTextContent('"REASONING":["capable"]');
|
||||
expect(output).toHaveTextContent('"reasoning_effort":"low","speed":"fast","max_tokens":1024');
|
||||
expect(output).toHaveTextContent('"reasoning_effort":"high"');
|
||||
expect(output).toHaveTextContent('"adaptive":false');
|
||||
expect(output).not.toHaveTextContent("leftover-medium");
|
||||
expect(output).not.toHaveTextContent("leftover-complex");
|
||||
expect(output).not.toHaveTextContent('"plan_mode_min_tier"');
|
||||
});
|
||||
|
||||
it.each(["capability", "llm_v2"] as const)(
|
||||
"carries non-default solver assignments when switching away from %s",
|
||||
(source) => {
|
||||
const pair = { efficient_tier: "MEDIUM", capable_tier: "COMPLEX" };
|
||||
const previous: ComplexityRouterConfigValue = {
|
||||
...(source === "capability" ? initial : fuseInitial),
|
||||
tiers: { SIMPLE: [], MEDIUM: ["efficient"], COMPLEX: ["capable"], REASONING: [] },
|
||||
capability_classifier_config:
|
||||
source === "capability" ? { ...initial.capability_classifier_config!, ...pair } : undefined,
|
||||
llm_v2_config: source === "llm_v2" ? { ...fuseInitial.llm_v2_config!, ...pair } : undefined,
|
||||
plan_mode_min_tier: "COMPLEX",
|
||||
tier_model_params: { MEDIUM: { efficient: { max_tokens: 128 } }, COMPLEX: { capable: { speed: "fast" } } },
|
||||
};
|
||||
renderWithProviders(<Form initialValue={previous} />);
|
||||
fireEvent.click(screen.getByRole("tab", { name: source === "capability" ? "Fuse v2" : "Capability" }));
|
||||
if (source === "capability") {
|
||||
fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Small solver" } });
|
||||
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Large solver" } });
|
||||
fireEvent.change(screen.getByLabelText("Harness and budget"), { target: { value: "One attempt" } });
|
||||
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
|
||||
} else {
|
||||
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
|
||||
}
|
||||
expect(screen.getByRole("button", { name: "Save configuration" })).toBeEnabled();
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
|
||||
const output = screen.getByRole("status", { name: "Saved configuration" });
|
||||
expect(output).toHaveTextContent('"efficient_tier":"MEDIUM","capable_tier":"COMPLEX"');
|
||||
expect(output).toHaveTextContent('"tiers":{"MEDIUM":["efficient"],"COMPLEX":["capable"]}');
|
||||
expect(output).toHaveTextContent('"plan_mode_min_tier":"COMPLEX"');
|
||||
expect(output).toHaveTextContent('"max_tokens":128');
|
||||
expect(output).toHaveTextContent('"speed":"fast"');
|
||||
},
|
||||
);
|
||||
|
||||
it("keeps decimal and negative numbers when entered one character at a time", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<Form />);
|
||||
const threshold = screen.getByLabelText("Solve probability threshold");
|
||||
await user.clear(threshold);
|
||||
await user.type(threshold, "0.65");
|
||||
expect(threshold).toHaveValue(0.65);
|
||||
await user.click(screen.getByRole("button", { name: "Classifier options" }));
|
||||
await user.click(screen.getByRole("switch", { name: "Use fitted calibration" }));
|
||||
await user.type(screen.getByLabelText("Efficient intercept"), "-0.3");
|
||||
expect(screen.getByLabelText("Efficient intercept")).toHaveValue(-0.3);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["capability", "LLM Classifier"],
|
||||
["capability", "Heuristic first"],
|
||||
["capability", "Hybrid"],
|
||||
["llm_v2", "LLM Classifier"],
|
||||
["llm_v2", "Heuristic first"],
|
||||
["llm_v2", "Hybrid"],
|
||||
] as const)("restores the current rubric when switching %s through Complexity to %s", async (source, target) => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<Form initialValue={source === "capability" ? initial : fuseInitial} />);
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
|
||||
fireEvent.click(screen.getByRole("radio", { name: new RegExp(`^${target}`) }));
|
||||
await user.click(screen.getByRole("combobox", { name: "Classifier Model" }));
|
||||
await user.click(screen.getByRole("option", { name: "judge", exact: true }));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
|
||||
const output = screen.getByRole("status", { name: "Saved configuration" });
|
||||
expect(output).toHaveTextContent('"classification_rubric":"agentic"');
|
||||
expect(output).toHaveTextContent('"model":"judge"');
|
||||
expect(output).toHaveTextContent('"timeout_ms":3000');
|
||||
expect(output).not.toHaveTextContent('"capability_classifier_config"');
|
||||
expect(output).not.toHaveTextContent('"llm_v2_config"');
|
||||
});
|
||||
|
||||
it("saves capability threshold edits together with fitted calibration", () => {
|
||||
renderWithProviders(<Form />);
|
||||
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.6" } });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Classifier options" }));
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Use fitted calibration" }));
|
||||
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
|
||||
fireEvent.change(screen.getByLabelText("Calibration version"), { target: { value: "eval-a" } });
|
||||
fireEvent.change(screen.getByLabelText("Efficient slope"), { target: { value: "1.2" } });
|
||||
fireEvent.change(screen.getByLabelText("Efficient intercept"), { target: { value: "-0.3" } });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
|
||||
const output = screen.getByRole("status", { name: "Saved configuration" });
|
||||
expect(output).toHaveTextContent('"base_threshold":0.6');
|
||||
expect(output).toHaveTextContent('"calibration":{"version":"eval-a","slope":1.2,"intercept":-0.3}');
|
||||
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "" } });
|
||||
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
|
||||
});
|
||||
|
||||
it("switches to Fuse, requires solver context, and saves the filled fields", () => {
|
||||
renderWithProviders(<Form />);
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Fuse v2" }));
|
||||
expect(screen.queryByLabelText("Solve probability threshold")).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
|
||||
fireEvent.change(screen.getByLabelText("Efficient solver profile"), {
|
||||
target: { value: "Short reasoning budget" },
|
||||
});
|
||||
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Larger reasoning budget" } });
|
||||
fireEvent.change(screen.getByLabelText("Harness and budget"), {
|
||||
target: { value: "Shell and test runner, one attempt" },
|
||||
});
|
||||
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
|
||||
const output = screen.getByRole("status", { name: "Saved configuration" });
|
||||
expect(output).toHaveTextContent('"classifier_type":"llm_v2"');
|
||||
expect(output).toHaveTextContent('"efficient_profile":"Short reasoning budget"');
|
||||
expect(output).toHaveTextContent('"capable_profile":"Larger reasoning budget"');
|
||||
expect(output).toHaveTextContent('"harness":"Shell and test runner, one attempt"');
|
||||
expect(output).toHaveTextContent('"max_quality_gap":0.05');
|
||||
expect(output).toHaveTextContent('"adaptive":false');
|
||||
expect(output).not.toHaveTextContent('"capability_classifier_config"');
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,433 @@
|
|||
import React from "react";
|
||||
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
|
||||
import { ChevronRight } from "lucide-react";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Textarea } from "@/components/ui/textarea";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { SearchSelect } from "@/components/shared/SearchSelect";
|
||||
import { MultiSelect } from "@/components/shared/MultiSelect";
|
||||
import {
|
||||
type ComplexityRouterConfigValue,
|
||||
type ClassificationFrequency,
|
||||
classificationFrequency,
|
||||
withClassificationFrequency,
|
||||
DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
} from "./ComplexityRouterConfig";
|
||||
import {
|
||||
forecastTierNames,
|
||||
forecastModels,
|
||||
getForecastConfigError,
|
||||
newCapabilitySettings,
|
||||
newFuseSettings,
|
||||
type CapabilitySettings,
|
||||
type FuseSettings,
|
||||
} from "./forecast_classifier_config";
|
||||
import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect";
|
||||
import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig";
|
||||
import ClassifierVisionConfig from "./ClassifierVisionConfig";
|
||||
import TierModelEffortRows from "./TierModelEffortRows";
|
||||
import { activeTierRows } from "./tier_rows";
|
||||
import { setTierModels } from "./tier_set_actions";
|
||||
import { tierRowLabel, setTierModelParam, setTierModelReasoningEffort } from "./complexity_router_tiers";
|
||||
|
||||
interface Props {
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
modelOptions: { value: string; label: string }[];
|
||||
effortOptionsByModel: Record<string, string[] | null | undefined>;
|
||||
}
|
||||
|
||||
const NumberField = ({
|
||||
label,
|
||||
value,
|
||||
onChange,
|
||||
min,
|
||||
max,
|
||||
step = "any",
|
||||
help,
|
||||
}: {
|
||||
label: string;
|
||||
value: number;
|
||||
onChange: (value: number) => void;
|
||||
min?: number;
|
||||
max?: number;
|
||||
step?: number | "any";
|
||||
help?: string;
|
||||
}) => {
|
||||
const id = React.useId();
|
||||
return (
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor={id}>{label}</Label>
|
||||
<Input
|
||||
id={id}
|
||||
type="number"
|
||||
min={min}
|
||||
max={max}
|
||||
step={step}
|
||||
value={Number.isFinite(value) ? value : ""}
|
||||
onChange={(event) => onChange(event.target.value === "" ? Number.NaN : Number(event.target.value))}
|
||||
/>
|
||||
{help && <p className="text-xs text-muted-foreground">{help}</p>}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export const ForecastSolverModels = ({
|
||||
value,
|
||||
onChange,
|
||||
modelOptions,
|
||||
effortOptionsByModel,
|
||||
fastModeByModel,
|
||||
additionalPoolsOnly = false,
|
||||
}: Props & { fastModeByModel: Record<string, boolean>; additionalPoolsOnly?: boolean }) => {
|
||||
const id = React.useId();
|
||||
const names = forecastTierNames(value);
|
||||
const additionalRows =
|
||||
value.classifier_type === "capability"
|
||||
? activeTierRows(value)
|
||||
.filter((row) => !names.includes(row.id) && row.models.length > 0)
|
||||
.map((row) => ({ tier: row.id, label: `${tierRowLabel(row, value.tier_labels)} routing pool` }))
|
||||
: [];
|
||||
const rows = additionalPoolsOnly
|
||||
? additionalRows
|
||||
: names.map((tier, index) => ({ tier, label: index === 0 ? "Efficient solver" : "Capable solver" }));
|
||||
if (rows.length === 0) return null;
|
||||
return (
|
||||
<div className="rounded-lg border p-4 space-y-4">
|
||||
{rows.map(({ tier, label }) => {
|
||||
const models = forecastModels(value.tiers, tier);
|
||||
const setModels = (next: string[]) => onChange(setTierModels(value, tier, next));
|
||||
return (
|
||||
<div key={tier} className="space-y-2">
|
||||
<Label htmlFor={`${id}-${tier}`} className="block text-sm font-semibold">
|
||||
{label}
|
||||
</Label>
|
||||
{value.classifier_type === "llm_v2" ? (
|
||||
<SearchSelect
|
||||
options={modelOptions}
|
||||
inputId={`${id}-${tier}`}
|
||||
value={models[0] ?? ""}
|
||||
aria-label={label}
|
||||
placeholder={`Select ${label.toLowerCase()}`}
|
||||
onValueChange={(model) => setModels(model ? [model] : [])}
|
||||
/>
|
||||
) : (
|
||||
<MultiSelect
|
||||
options={modelOptions}
|
||||
id={`${id}-${tier}`}
|
||||
value={models}
|
||||
onValueChange={setModels}
|
||||
placeholder={`Select ${label.toLowerCase()} models`}
|
||||
/>
|
||||
)}
|
||||
<TierModelEffortRows
|
||||
tierLabel={label}
|
||||
models={models}
|
||||
effortOptionsByModel={Object.fromEntries(
|
||||
Object.entries(effortOptionsByModel).map(([model, efforts]) => [model, efforts ?? []]),
|
||||
)}
|
||||
paramsByModel={value.tier_model_params?.[tier] ?? {}}
|
||||
fastModeByModel={fastModeByModel}
|
||||
onFastModeChange={(model, enabled) =>
|
||||
onChange({
|
||||
...value,
|
||||
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, [
|
||||
"speed",
|
||||
enabled ? "fast" : undefined,
|
||||
]),
|
||||
})
|
||||
}
|
||||
onEffortChange={(model, effort) =>
|
||||
onChange({
|
||||
...value,
|
||||
tier_model_params: setTierModelReasoningEffort(value.tier_model_params, tier, model, effort),
|
||||
})
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
{!additionalPoolsOnly && (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Invalid forecasts and classifier failures route to the capable solver
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const CalibrationFields = ({
|
||||
label,
|
||||
value,
|
||||
onChange,
|
||||
bounded = false,
|
||||
}: {
|
||||
label: string;
|
||||
bounded?: boolean;
|
||||
value: { slope: number; intercept: number };
|
||||
onChange: (value: { slope: number; intercept: number }) => void;
|
||||
}) => (
|
||||
<div className="grid gap-3 sm:grid-cols-2">
|
||||
<NumberField
|
||||
label={`${label} slope`}
|
||||
value={value.slope}
|
||||
min={0}
|
||||
max={bounded ? 20 : undefined}
|
||||
onChange={(slope) => onChange({ ...value, slope })}
|
||||
/>
|
||||
<NumberField
|
||||
label={`${label} intercept`}
|
||||
value={value.intercept}
|
||||
min={bounded ? -20 : undefined}
|
||||
max={bounded ? 20 : undefined}
|
||||
onChange={(intercept) => onChange({ ...value, intercept })}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
|
||||
const emptyCoefficients = () => ({ slope: Number.NaN, intercept: Number.NaN });
|
||||
|
||||
const ForecastClassifierConfig = ({ value, onChange, modelOptions, effortOptionsByModel }: Props) => {
|
||||
const id = React.useId();
|
||||
const isCapability = value.classifier_type === "capability";
|
||||
const capability = value.capability_classifier_config ?? newCapabilitySettings();
|
||||
const fuse = value.llm_v2_config ?? newFuseSettings();
|
||||
const config = isCapability ? capability : fuse;
|
||||
const llm = value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS };
|
||||
const updateCapability = (next: CapabilitySettings) => onChange({ ...value, capability_classifier_config: next });
|
||||
const updateFuse = (next: FuseSettings) => onChange({ ...value, llm_v2_config: next });
|
||||
const updateTransport = (patch: { max_output_tokens?: number; response_format?: "json_schema" | "json_object" }) =>
|
||||
isCapability ? updateCapability({ ...capability, ...patch }) : updateFuse({ ...fuse, ...patch });
|
||||
const setCalibrationVersion = (version: string) => {
|
||||
if (isCapability && capability.calibration)
|
||||
updateCapability({ ...capability, calibration: { ...capability.calibration, version } });
|
||||
if (!isCapability && fuse.calibration) updateFuse({ ...fuse, calibration: { ...fuse.calibration, version } });
|
||||
};
|
||||
const error = getForecastConfigError(value);
|
||||
return (
|
||||
<div className="mt-4 space-y-4">
|
||||
<p className="text-sm text-muted-foreground">
|
||||
{isCapability
|
||||
? "Forecasts whether the efficient solver can complete the task using the bundled capability card"
|
||||
: "Forecasts success for both solvers and selects efficient when the estimated quality gap is within your allowance"}
|
||||
</p>
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor={`${id}-judge`}>Judge model</Label>
|
||||
<SearchSelect
|
||||
inputId={`${id}-judge`}
|
||||
aria-label="Judge model"
|
||||
options={modelOptions}
|
||||
value={llm.model}
|
||||
placeholder="Select the judge model"
|
||||
onValueChange={(model) => {
|
||||
if (model === llm.model) return;
|
||||
onChange({ ...value, classifier_llm_config: { ...llm, model: model ?? "", reasoning_effort: undefined } });
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
{isCapability ? (
|
||||
<>
|
||||
<NumberField
|
||||
label="Solve probability threshold"
|
||||
value={capability.base_threshold}
|
||||
min={0}
|
||||
max={1}
|
||||
help="Minimum estimated chance of whole-task success required to use the efficient solver"
|
||||
onChange={(base_threshold) => updateCapability({ ...capability, base_threshold })}
|
||||
/>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
{(["efficient_profile", "capable_profile", "harness"] as const).map((field) => {
|
||||
const label = {
|
||||
efficient_profile: "Efficient solver profile",
|
||||
capable_profile: "Capable solver profile",
|
||||
harness: "Harness and budget",
|
||||
}[field];
|
||||
return (
|
||||
<div key={field} className="space-y-1">
|
||||
<Label htmlFor={`${id}-${field}`}>{label}</Label>
|
||||
<Textarea
|
||||
id={`${id}-${field}`}
|
||||
value={fuse[field]}
|
||||
maxLength={4000}
|
||||
placeholder={
|
||||
field === "harness"
|
||||
? "Tools, execution environment, verification, and budget available to each solver"
|
||||
: "Describe this solver's strengths, limitations, and settings"
|
||||
}
|
||||
onChange={(event) => updateFuse({ ...fuse, [field]: event.target.value })}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
<NumberField
|
||||
label="Maximum quality gap"
|
||||
value={fuse.max_quality_gap}
|
||||
min={0}
|
||||
max={1}
|
||||
help="Allowed difference between capable and efficient success probabilities, from 0 to 1. This is an estimate, not a measured quality guarantee"
|
||||
onChange={(max_quality_gap) => updateFuse({ ...fuse, max_quality_gap })}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
<Collapsible className="rounded-lg border">
|
||||
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left font-medium">
|
||||
<ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
|
||||
Classifier options
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="space-y-4 px-4 pb-4">
|
||||
<ClassifierReasoningEffortSelect
|
||||
model={llm.model}
|
||||
value={llm.reasoning_effort}
|
||||
explicitlySupported={effortOptionsByModel[llm.model]}
|
||||
onChange={(reasoning_effort) => onChange({ ...value, classifier_llm_config: { ...llm, reasoning_effort } })}
|
||||
/>
|
||||
<NumberField
|
||||
label="Timeout (ms)"
|
||||
min={1}
|
||||
step={1}
|
||||
value={llm.timeout_ms}
|
||||
help="Allow enough time for the judge to produce its forecast"
|
||||
onChange={(timeout_ms) => onChange({ ...value, classifier_llm_config: { ...llm, timeout_ms } })}
|
||||
/>
|
||||
<ClassifierCircuitBreakerConfig
|
||||
value={llm}
|
||||
onChange={(classifier_llm_config) => onChange({ ...value, classifier_llm_config })}
|
||||
/>
|
||||
<ClassifierVisionConfig
|
||||
value={llm}
|
||||
onChange={(classifier_llm_config) => onChange({ ...value, classifier_llm_config })}
|
||||
/>
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor={`${id}-frequency`}>How often to classify</Label>
|
||||
<SearchSelect
|
||||
inputId={`${id}-frequency`}
|
||||
aria-label="How often to classify"
|
||||
value={classificationFrequency(value)}
|
||||
allowClear={false}
|
||||
options={[
|
||||
{ value: "every_request", label: "Every request" },
|
||||
{ value: "user_turn", label: "Every new user message" },
|
||||
{ value: "session", label: "Once per session" },
|
||||
]}
|
||||
onValueChange={(frequency) => {
|
||||
if (frequency) onChange(withClassificationFrequency(value, frequency as ClassificationFrequency));
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
{isCapability && (
|
||||
<NumberField
|
||||
label="Capability boundary step"
|
||||
value={capability.threshold_step ?? 0}
|
||||
min={0}
|
||||
max={0.5}
|
||||
help="Added once for uncertain or unmatched tasks and twice for unsupported tasks; the final threshold cannot exceed 1"
|
||||
onChange={(threshold_step) => updateCapability({ ...capability, threshold_step })}
|
||||
/>
|
||||
)}
|
||||
<NumberField
|
||||
label="Classifier output token limit"
|
||||
min={1}
|
||||
step={1}
|
||||
value={config.max_output_tokens ?? (isCapability ? 4096 : 1024)}
|
||||
onChange={(max_output_tokens) => updateTransport({ max_output_tokens })}
|
||||
/>
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor={`${id}-format`}>Forecast response format</Label>
|
||||
<SearchSelect
|
||||
inputId={`${id}-format`}
|
||||
aria-label="Forecast response format"
|
||||
value={config.response_format ?? "json_schema"}
|
||||
allowClear={false}
|
||||
options={[
|
||||
{ value: "json_schema", label: "Strict JSON schema" },
|
||||
{ value: "json_object", label: "JSON object (for judges without strict schema support)" },
|
||||
]}
|
||||
onValueChange={(response_format) => {
|
||||
if (response_format === "json_schema" || response_format === "json_object")
|
||||
updateTransport({ response_format });
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<div className="space-y-3 rounded-md border p-3">
|
||||
<Label>
|
||||
<Switch
|
||||
checked={Boolean(config.calibration)}
|
||||
onCheckedChange={(enabled) =>
|
||||
isCapability
|
||||
? updateCapability({
|
||||
...capability,
|
||||
calibration: enabled ? { version: "", ...emptyCoefficients() } : undefined,
|
||||
})
|
||||
: updateFuse({
|
||||
...fuse,
|
||||
calibration: enabled
|
||||
? {
|
||||
version: "",
|
||||
prompt_version: "llm-v2-1",
|
||||
efficient: emptyCoefficients(),
|
||||
capable: emptyCoefficients(),
|
||||
}
|
||||
: undefined,
|
||||
})
|
||||
}
|
||||
/>
|
||||
Use fitted calibration
|
||||
</Label>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Optional coefficients fitted for your judge, solvers, and harness. Leave off to use raw forecasts
|
||||
</p>
|
||||
{config.calibration && (
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor={`${id}-version`}>Calibration version</Label>
|
||||
<Input
|
||||
id={`${id}-version`}
|
||||
value={config.calibration.version}
|
||||
maxLength={isCapability ? 128 : 512}
|
||||
onChange={(event) => setCalibrationVersion(event.target.value)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
{isCapability && capability.calibration && (
|
||||
<CalibrationFields
|
||||
label="Efficient"
|
||||
bounded
|
||||
value={capability.calibration}
|
||||
onChange={(next) =>
|
||||
updateCapability({
|
||||
...capability,
|
||||
calibration: { version: capability.calibration?.version ?? "", ...next },
|
||||
})
|
||||
}
|
||||
/>
|
||||
)}
|
||||
{!isCapability &&
|
||||
fuse.calibration &&
|
||||
(["efficient", "capable"] as const).map((role) => (
|
||||
<CalibrationFields
|
||||
key={role}
|
||||
label={role === "efficient" ? "Efficient" : "Capable"}
|
||||
value={fuse.calibration![role]}
|
||||
onChange={(next) => {
|
||||
if (fuse.calibration) updateFuse({ ...fuse, calibration: { ...fuse.calibration, [role]: next } });
|
||||
}}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
</CollapsibleContent>
|
||||
</Collapsible>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
The classifier uses its bundled prompt and always falls back to the capable solver
|
||||
</p>
|
||||
{error && (
|
||||
<p role="alert" className="text-sm text-destructive">
|
||||
{error}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default ForecastClassifierConfig;
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
import React from "react";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import TierRowSelect from "./TierRowSelect";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
|
||||
const PlanModeOverrideControls: React.FC<{
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
planModeTierOptions: { value: string; label: string }[];
|
||||
}> = ({ value, onChange, planModeTierOptions }) => (
|
||||
<>
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<Switch
|
||||
checked={value.plan_mode_min_tier !== undefined}
|
||||
disabled={planModeTierOptions.length === 0}
|
||||
onCheckedChange={(enabled) =>
|
||||
onChange({
|
||||
...value,
|
||||
plan_mode_min_tier: enabled ? planModeTierOptions.at(-1)?.value : undefined,
|
||||
})
|
||||
}
|
||||
aria-label="Route plan-mode requests to a minimum tier"
|
||||
/>
|
||||
<strong className="font-semibold">Route plan-mode requests to a minimum tier</strong>
|
||||
</div>
|
||||
<span className="block text-xs mb-3 text-muted-foreground">
|
||||
Requests from coding agents in plan mode (Claude Code, GitHub Copilot) route to at least this tier. The classifier
|
||||
still wins when it picks higher, and the override only lasts while plan mode is active.
|
||||
{planModeTierOptions.length === 0 && " Add models to a tier to enable this."}
|
||||
</span>
|
||||
{value.plan_mode_min_tier !== undefined && (
|
||||
<div style={{ maxWidth: 320 }}>
|
||||
<TierRowSelect
|
||||
label="Plan-mode minimum tier"
|
||||
options={planModeTierOptions}
|
||||
value={value.plan_mode_min_tier ?? null}
|
||||
onValueChange={(tier) => onChange({ ...value, plan_mode_min_tier: tier })}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
);
|
||||
|
||||
export default PlanModeOverrideControls;
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
import React from "react";
|
||||
import { ChevronRight } from "lucide-react";
|
||||
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
|
||||
|
||||
interface RoutingOptionsProps {
|
||||
forecast: boolean;
|
||||
children: React.ReactNode;
|
||||
}
|
||||
|
||||
const RoutingOptions = ({ forecast, children }: RoutingOptionsProps) =>
|
||||
forecast ? (
|
||||
<Collapsible className="rounded-lg border">
|
||||
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left font-medium">
|
||||
<ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
|
||||
Advanced routing options
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="space-y-4 px-4 pb-4">{children}</CollapsibleContent>
|
||||
</Collapsible>
|
||||
) : (
|
||||
<>{children}</>
|
||||
);
|
||||
|
||||
export default RoutingOptions;
|
||||
|
|
@ -23,6 +23,12 @@ interface TierModelEffortRowsProps {
|
|||
onFastModeChange: (model: string, enabled: boolean) => void;
|
||||
}
|
||||
|
||||
const canEditFastMode = (
|
||||
model: string,
|
||||
fastModeByModel: TierModelEffortRowsProps["fastModeByModel"],
|
||||
paramsByModel: TierModelEffortRowsProps["paramsByModel"],
|
||||
): boolean => fastModeByModel?.[model] === true || paramsByModel?.[model]?.speed === "fast";
|
||||
|
||||
export interface TierEffortRow {
|
||||
model: string;
|
||||
effort: ReasoningEffort | undefined;
|
||||
|
|
@ -49,7 +55,7 @@ export const tierEffortRows = ({
|
|||
const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported;
|
||||
return { model, effort, options: Array.from(new Set(listed)) };
|
||||
})
|
||||
.filter(({ model, options }) => options.length > 0 || fastModeByModel?.[model] === true);
|
||||
.filter(({ model, options }) => options.length > 0 || canEditFastMode(model, fastModeByModel, paramsByModel));
|
||||
|
||||
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
|
||||
const { tierLabel, paramsByModel, onEffortChange, fastModeByModel, onFastModeChange } = props;
|
||||
|
|
@ -101,7 +107,7 @@ const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
|
|||
</SelectContent>
|
||||
</Select>
|
||||
)}
|
||||
{fastModeByModel?.[model] === true && (
|
||||
{canEditFastMode(model, fastModeByModel, paramsByModel) && (
|
||||
<SimpleTooltip content="Fast mode has higher pricing and requires an eligible provider account. Off removes this tier's speed override and inherits the request or provider default">
|
||||
<label
|
||||
className="flex items-center gap-2 text-xs"
|
||||
|
|
|
|||
|
|
@ -2,9 +2,17 @@ import React from "react";
|
|||
import { CUSTOM_TIER_RESTRICTIONS, CustomTierSet, TierRestriction } from "./tier_rows";
|
||||
|
||||
export const restrictedBy = (
|
||||
value: { custom_tier_set?: CustomTierSet },
|
||||
value: { custom_tier_set?: CustomTierSet; classifier_type?: string },
|
||||
key: keyof typeof CUSTOM_TIER_RESTRICTIONS,
|
||||
): TierRestriction | undefined => (value.custom_tier_set ? CUSTOM_TIER_RESTRICTIONS[key] : undefined);
|
||||
): TierRestriction | undefined => {
|
||||
if (value.custom_tier_set) return CUSTOM_TIER_RESTRICTIONS[key];
|
||||
if (value.classifier_type === "llm_v2" && key === "adaptive")
|
||||
return {
|
||||
omit: ["adaptive", "adaptive_weights", "adaptive_eligible", "tier_distance_penalty"],
|
||||
reason: "Fuse v2 uses its quality-gap decision directly; adaptive routing is unavailable",
|
||||
};
|
||||
return undefined;
|
||||
};
|
||||
|
||||
export const Restricted: React.FC<{ by: TierRestriction | undefined; children: React.ReactNode }> = ({
|
||||
by,
|
||||
|
|
|
|||
|
|
@ -167,6 +167,124 @@ describe("AddAutoRouterTab", () => {
|
|||
mockFetchAllModelDeployments.mockResolvedValue([]);
|
||||
});
|
||||
|
||||
it.each(["Capability", "Fuse v2"])(
|
||||
"creates %s from its dedicated tab without complexity templates",
|
||||
async (label) => {
|
||||
const user = userEvent.setup();
|
||||
mockFetchAvailableModels.mockResolvedValue([
|
||||
{ model_group: "efficient", mode: "chat" },
|
||||
{ model_group: "capable", mode: "chat" },
|
||||
{ model_group: "judge", mode: "chat" },
|
||||
]);
|
||||
renderWithProviders(<Harness />);
|
||||
await user.type(screen.getByLabelText("Auto Router Name"), "forecast-router");
|
||||
await user.click(screen.getByRole("tab", { name: label, exact: true }));
|
||||
expect(screen.getByLabelText("Auto Router Name")).toHaveValue("forecast-router");
|
||||
expect(screen.queryByTestId("template-selector")).not.toBeInTheDocument();
|
||||
expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument();
|
||||
expect(screen.queryByTestId("detailed-configuration-toggle")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Complexity Tier Configuration")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Advanced: Classification Method")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Advanced: Adaptive Routing")).not.toBeInTheDocument();
|
||||
const capability = label === "Capability";
|
||||
for (const [role, model] of [
|
||||
["Efficient", "efficient"],
|
||||
["Capable", "capable"],
|
||||
]) {
|
||||
await user.click(
|
||||
screen.getByRole("combobox", {
|
||||
name: capability ? `Select ${role.toLowerCase()} solver models` : `${role} solver`,
|
||||
}),
|
||||
);
|
||||
await user.click(await screen.findByRole("option", { name: model, exact: true }));
|
||||
if (capability) await user.keyboard("{Escape}");
|
||||
}
|
||||
await user.click(screen.getByRole("combobox", { name: "Judge model" }));
|
||||
await user.click(await screen.findByRole("option", { name: "judge", exact: true }));
|
||||
expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled();
|
||||
if (capability) {
|
||||
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
|
||||
} else {
|
||||
fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Small solver" } });
|
||||
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Large solver" } });
|
||||
fireEvent.change(screen.getByLabelText("Harness and budget"), { target: { value: "One attempt" } });
|
||||
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
|
||||
}
|
||||
expect(screen.queryByRole("alert")).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled();
|
||||
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
|
||||
for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) {
|
||||
expect(screen.queryByText(`Advanced: ${label}`)).not.toBeInTheDocument();
|
||||
}
|
||||
expect(screen.getByText("Advanced: Stalled Task Escalation")).toBeInTheDocument();
|
||||
expect(screen.getByText("Advanced: Response Format")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Advanced: Classification Method")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Advanced: Affinity")).toBeInTheDocument();
|
||||
await user.click(screen.getByRole("button", { name: "Add Auto Router" }));
|
||||
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledTimes(1));
|
||||
const expected = {
|
||||
classifier_type: capability ? "capability" : "llm_v2",
|
||||
adaptive: false,
|
||||
enable_context_window_escalation: false,
|
||||
escalation_keywords: [],
|
||||
tiers: { SIMPLE: ["efficient"], REASONING: ["capable"] },
|
||||
classifier_llm_config: { model: "judge" },
|
||||
};
|
||||
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0].complexity_router_config).toMatchObject(expected);
|
||||
},
|
||||
);
|
||||
|
||||
it("restores the automatic/template/detail flow on the Complexity tab", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS);
|
||||
renderWithProviders(<Harness />);
|
||||
await screen.findByTestId("configure-automatically-button");
|
||||
await user.click(screen.getByRole("tab", { name: "Capability", exact: true }));
|
||||
expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument();
|
||||
await user.click(screen.getByRole("tab", { name: "Complexity", exact: true }));
|
||||
expect(screen.getByTestId("configure-automatically-button")).toBeInTheDocument();
|
||||
expect(screen.getByTestId("template-selector")).toBeInTheDocument();
|
||||
expandDetailedConfiguration();
|
||||
expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument();
|
||||
for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) {
|
||||
expect(screen.getByText(`Advanced: ${label}`)).toBeInTheDocument();
|
||||
}
|
||||
await user.click(screen.getByText("Advanced: Classification Method"));
|
||||
expect(screen.queryByRole("radio", { name: /^Capability/ })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("radio", { name: /^Fuse v2/ })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("radio", { name: /^Heuristic \(default/ })).toBeChecked();
|
||||
});
|
||||
|
||||
it.each(["Capability", "Fuse v2"])(
|
||||
"retries failed model loading on %s without losing entered settings",
|
||||
async (label) => {
|
||||
const user = userEvent.setup();
|
||||
mockFetchAvailableModels.mockRejectedValueOnce(new Error("Model list unavailable")).mockResolvedValue([
|
||||
{ model_group: "efficient", mode: "chat" },
|
||||
{ model_group: "capable", mode: "chat" },
|
||||
{ model_group: "judge", mode: "chat" },
|
||||
]);
|
||||
renderWithProviders(<Harness />);
|
||||
await user.click(screen.getByRole("tab", { name: label, exact: true }));
|
||||
await user.type(screen.getByLabelText("Auto Router Name"), "forecast-retry");
|
||||
const capability = label === "Capability";
|
||||
const policyField = capability ? "Solve probability threshold" : "Efficient solver profile";
|
||||
fireEvent.change(screen.getByLabelText(policyField), {
|
||||
target: { value: capability ? "0.7" : "Small solver" },
|
||||
});
|
||||
expect(await screen.findByText("Could not load available models.")).toBeVisible();
|
||||
expect(screen.queryByTestId("template-selector")).not.toBeInTheDocument();
|
||||
await user.click(screen.getByRole("button", { name: "Retry", exact: true }));
|
||||
await waitFor(() => expect(screen.queryByText("Could not load available models.")).not.toBeInTheDocument());
|
||||
expect(screen.getByRole("tab", { name: label, exact: true })).toHaveAttribute("aria-selected", "true");
|
||||
expect(screen.getByLabelText("Auto Router Name")).toHaveValue("forecast-retry");
|
||||
expect(screen.getByLabelText(policyField)).toHaveValue(capability ? 0.7 : "Small solver");
|
||||
await user.click(screen.getByRole("combobox", { name: "Judge model" }));
|
||||
expect(await screen.findByRole("option", { name: "judge", exact: true })).toBeVisible();
|
||||
expect(mockFetchAvailableModels).toHaveBeenCalledTimes(2);
|
||||
},
|
||||
);
|
||||
|
||||
// Detailed Configuration starts collapsed so the modal opens onto just Name + Template; a caller
|
||||
// opts into the full tier/classifier form rather than always seeing it up front.
|
||||
it("keeps Detailed Configuration collapsed until a caller opens it", () => {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
|
||||
import { getForecastConfigError, isForecastClassifier } from "../add_model/forecast_classifier_config";
|
||||
import React, { useEffect, useState } from "react";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { useWatch } from "react-hook-form";
|
||||
|
|
@ -138,7 +140,9 @@ export const getSubmitBlockedReason = (
|
|||
(config.custom_tier_set
|
||||
? getCustomTierRowsError(config.custom_tier_set)
|
||||
: getTierLabelsError(config.tier_labels)) ??
|
||||
getMissingTiersError(activeTierRows(config)) ??
|
||||
(isForecastClassifier(config.classifier_type)
|
||||
? getForecastConfigError(config)
|
||||
: getMissingTiersError(activeTierRows(config))) ??
|
||||
getPlanModeTierError(config.plan_mode_min_tier, activeTierRows(config)) ??
|
||||
getKeywordTierRulesError(keywordTierRules, activeTierRows(config)) ??
|
||||
getClassifierModelError(config) ??
|
||||
|
|
@ -401,6 +405,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
classificationMode: complexityRouterConfig.classification_mode,
|
||||
tierLabels: complexityRouterConfig.tier_labels,
|
||||
classifierType: complexityRouterConfig.classifier_type,
|
||||
capabilityClassifierConfig: complexityRouterConfig.capability_classifier_config,
|
||||
llmV2Config: complexityRouterConfig.llm_v2_config,
|
||||
classifierLlmConfig: complexityRouterConfig.classifier_llm_config,
|
||||
classifierContextWindowSize: complexityRouterConfig.classifier_context_window_size,
|
||||
classifierContextBudgetChars: complexityRouterConfig.classifier_context_budget_chars,
|
||||
|
|
@ -543,81 +549,138 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
setIsTestModalVisible(true);
|
||||
};
|
||||
|
||||
const configurationForm = (
|
||||
<ComplexityRouterConfig
|
||||
editingTiers={editingTiers}
|
||||
onEditingTiersChange={setEditingTiers}
|
||||
modelInfo={modelInfo}
|
||||
value={complexityRouterConfig}
|
||||
onChange={setComplexityRouterConfig}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
|
||||
keywordTierRules={keywordTierRules}
|
||||
onKeywordTierRulesChange={setKeywordTierRules}
|
||||
keywordRulesError={getKeywordTierRulesError(keywordTierRules, activeTierRows(complexityRouterConfig))}
|
||||
semanticMatchingEnabled={semanticMatchingEnabled}
|
||||
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={setEmbeddingModel}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
|
||||
showValidationErrors={showValidationErrors}
|
||||
/>
|
||||
);
|
||||
const forecast = isForecastClassifier(complexityRouterConfig.classifier_type);
|
||||
|
||||
return (
|
||||
<TooltipProvider>
|
||||
<Card>
|
||||
<CardContent>
|
||||
<form onSubmit={form.handleSubmit(() => handleAutoRouterSubmit())} noValidate>
|
||||
<FieldGroup>
|
||||
<div>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="auto_router_name"
|
||||
label={labelWithHint("Auto Router Name", "Unique name for this auto router configuration")}
|
||||
>
|
||||
{({ ref, ...field }) => (
|
||||
<Input {...field} ref={ref} placeholder="e.g., smart_router, auto_router_1" />
|
||||
)}
|
||||
</FormField>
|
||||
|
||||
{!automaticSetupLoading && automaticRouterConfig && (
|
||||
<div className="mt-5 flex flex-wrap items-center justify-between gap-3 rounded-lg border border-border bg-muted px-4 py-3">
|
||||
<div>
|
||||
<p className="text-sm font-medium text-foreground">Not sure where to start?</p>
|
||||
<p className="text-sm text-muted-foreground">Let us pick models for each complexity tier.</p>
|
||||
</div>
|
||||
<Button type="button" data-testid="configure-automatically-button" onClick={handleAutomaticSetup}>
|
||||
Configure automatically
|
||||
</Button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-5">
|
||||
<label className="block text-sm font-medium text-foreground mb-2">Template</label>
|
||||
<Select
|
||||
items={templateItems}
|
||||
value={selectedPreset ?? null}
|
||||
onValueChange={(presetKey: string | null) => handlePresetChange(presetKey ?? undefined)}
|
||||
>
|
||||
<SelectTrigger data-testid="template-selector" className="w-full">
|
||||
<SelectValue placeholder="Choose a template or select Custom to define your own" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
|
||||
const disabledHint = presetDisabledHint(presetState);
|
||||
const hintClass = isPresetHintAlarming(presetState)
|
||||
? "text-destructive"
|
||||
: "text-muted-foreground";
|
||||
const matchedHint =
|
||||
presetState.kind === "available" && presetState.viaDeployments
|
||||
? "Matches your deployments"
|
||||
: null;
|
||||
|
||||
return (
|
||||
<SelectItem
|
||||
key={preset.key}
|
||||
value={preset.key}
|
||||
label={preset.label}
|
||||
disabled={disabledHint !== null}
|
||||
title={disabledHint ?? preset.description}
|
||||
<div className="mb-6">
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="auto_router_name"
|
||||
label={labelWithHint("Auto Router Name", "Unique name for this auto router configuration")}
|
||||
>
|
||||
{({ ref, ...field }) => <Input {...field} ref={ref} placeholder="e.g., smart_router, auto_router_1" />}
|
||||
</FormField>
|
||||
</div>
|
||||
<AutoRouterClassifierTabs
|
||||
value={complexityRouterConfig}
|
||||
onChange={(config) => {
|
||||
setSelectedPreset(undefined);
|
||||
setComplexityRouterConfig(config);
|
||||
}}
|
||||
>
|
||||
<FieldGroup>
|
||||
<div>
|
||||
{!forecast && (
|
||||
<>
|
||||
{!automaticSetupLoading && automaticRouterConfig && (
|
||||
<div className="mt-5 flex flex-wrap items-center justify-between gap-3 rounded-lg border border-border bg-muted px-4 py-3">
|
||||
<div>
|
||||
<p className="text-sm font-medium text-foreground">Not sure where to start?</p>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Let us pick models for each complexity tier.
|
||||
</p>
|
||||
</div>
|
||||
<Button
|
||||
type="button"
|
||||
data-testid="configure-automatically-button"
|
||||
onClick={handleAutomaticSetup}
|
||||
>
|
||||
<div>
|
||||
<div className="font-medium">{preset.label}</div>
|
||||
<div className="text-xs text-muted-foreground">{preset.description}</div>
|
||||
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
|
||||
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
|
||||
</div>
|
||||
</SelectItem>
|
||||
);
|
||||
})}
|
||||
<SelectItem value="custom" label="Custom Configuration">
|
||||
<div>
|
||||
<div className="font-medium">Custom Configuration</div>
|
||||
<div className="text-xs text-muted-foreground">Define your auto router from scratch</div>
|
||||
Configure automatically
|
||||
</Button>
|
||||
</div>
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
)}
|
||||
|
||||
<div className="mt-5">
|
||||
<label className="block text-sm font-medium text-foreground mb-2">Template</label>
|
||||
<Select
|
||||
items={templateItems}
|
||||
value={selectedPreset ?? null}
|
||||
onValueChange={(presetKey: string | null) => handlePresetChange(presetKey ?? undefined)}
|
||||
>
|
||||
<SelectTrigger data-testid="template-selector" className="w-full">
|
||||
<SelectValue placeholder="Choose a template or select Custom to define your own" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
|
||||
const disabledHint = presetDisabledHint(presetState);
|
||||
const hintClass = isPresetHintAlarming(presetState)
|
||||
? "text-destructive"
|
||||
: "text-muted-foreground";
|
||||
const matchedHint =
|
||||
presetState.kind === "available" && presetState.viaDeployments
|
||||
? "Matches your deployments"
|
||||
: null;
|
||||
|
||||
return (
|
||||
<SelectItem
|
||||
key={preset.key}
|
||||
value={preset.key}
|
||||
label={preset.label}
|
||||
disabled={disabledHint !== null}
|
||||
title={disabledHint ?? preset.description}
|
||||
>
|
||||
<div>
|
||||
<div className="font-medium">{preset.label}</div>
|
||||
<div className="text-xs text-muted-foreground">{preset.description}</div>
|
||||
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
|
||||
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
|
||||
</div>
|
||||
</SelectItem>
|
||||
);
|
||||
})}
|
||||
<SelectItem value="custom" label="Custom Configuration">
|
||||
<div>
|
||||
<div className="font-medium">Custom Configuration</div>
|
||||
<div className="text-xs text-muted-foreground">
|
||||
Define your auto router from scratch
|
||||
</div>
|
||||
</div>
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
{presetsPending && (
|
||||
<div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>
|
||||
)}
|
||||
{presetsUnavailable && (
|
||||
<div className="text-xs mt-1 text-destructive">
|
||||
Could not load templates, so only Custom Configuration is shown.{" "}
|
||||
<button type="button" className="underline" onClick={() => void refetchPresets()}>
|
||||
Retry
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
{modelsUnverifiable && (
|
||||
<div className="text-xs mt-1 text-destructive">
|
||||
Could not load available models.{" "}
|
||||
|
|
@ -626,163 +689,129 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
</button>
|
||||
</div>
|
||||
)}
|
||||
{presetsPending && <div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>}
|
||||
{presetsUnavailable && (
|
||||
<div className="text-xs mt-1 text-destructive">
|
||||
Could not load templates, so only Custom Configuration is shown.{" "}
|
||||
<button type="button" className="underline" onClick={() => void refetchPresets()}>
|
||||
Retry
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{requiresTeamScope && (
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="team_id"
|
||||
label={labelWithHint(
|
||||
"Select Team",
|
||||
"Select the team this auto router belongs to. Only keys for this team will be able to call it.",
|
||||
)}
|
||||
>
|
||||
{({ id, value, onChange }) => (
|
||||
<TeamDropdown
|
||||
id={id}
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
filterTeam={(team) => canCreateAutoRouterForTeam(actor, team)}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
)}
|
||||
|
||||
<div className="border border-border rounded-lg">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setDetailsExpanded((expanded) => !expanded)}
|
||||
className="w-full flex flex-col gap-1 px-4 py-3 text-left hover:bg-muted"
|
||||
data-testid="detailed-configuration-toggle"
|
||||
>
|
||||
<span className="flex items-center gap-2 font-medium text-foreground">
|
||||
{detailsExpanded ? (
|
||||
<ChevronDown className="size-3 text-muted-foreground" />
|
||||
) : (
|
||||
<ChevronRight className="size-3 text-muted-foreground" />
|
||||
{requiresTeamScope && (
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="team_id"
|
||||
label={labelWithHint(
|
||||
"Select Team",
|
||||
"Select the team this auto router belongs to. Only keys for this team will be able to call it.",
|
||||
)}
|
||||
Detailed Configuration
|
||||
</span>
|
||||
{!detailsExpanded && (
|
||||
<span className="text-xs text-muted-foreground line-clamp-2">
|
||||
{tierConfigSummary(complexityRouterConfig)}
|
||||
</span>
|
||||
)}
|
||||
</button>
|
||||
{detailsExpanded && (
|
||||
<div className="px-4 pb-4">
|
||||
<ComplexityRouterConfig
|
||||
editingTiers={editingTiers}
|
||||
onEditingTiersChange={setEditingTiers}
|
||||
modelInfo={modelInfo}
|
||||
value={complexityRouterConfig}
|
||||
onChange={setComplexityRouterConfig}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
|
||||
keywordTierRules={keywordTierRules}
|
||||
onKeywordTierRulesChange={setKeywordTierRules}
|
||||
keywordRulesError={getKeywordTierRulesError(
|
||||
keywordTierRules,
|
||||
activeTierRows(complexityRouterConfig),
|
||||
>
|
||||
{({ id, value, onChange }) => (
|
||||
<TeamDropdown
|
||||
id={id}
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
filterTeam={(team) => canCreateAutoRouterForTeam(actor, team)}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
)}
|
||||
|
||||
{forecast ? (
|
||||
configurationForm
|
||||
) : (
|
||||
<div className="border border-border rounded-lg">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setDetailsExpanded((expanded) => !expanded)}
|
||||
className="w-full flex flex-col gap-1 px-4 py-3 text-left hover:bg-muted"
|
||||
data-testid="detailed-configuration-toggle"
|
||||
>
|
||||
<span className="flex items-center gap-2 font-medium text-foreground">
|
||||
{detailsExpanded ? (
|
||||
<ChevronDown className="size-3 text-muted-foreground" />
|
||||
) : (
|
||||
<ChevronRight className="size-3 text-muted-foreground" />
|
||||
)}
|
||||
Detailed Configuration
|
||||
</span>
|
||||
{!detailsExpanded && (
|
||||
<span className="text-xs text-muted-foreground line-clamp-2">
|
||||
{tierConfigSummary(complexityRouterConfig)}
|
||||
</span>
|
||||
)}
|
||||
semanticMatchingEnabled={semanticMatchingEnabled}
|
||||
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={setEmbeddingModel}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
|
||||
showValidationErrors={showValidationErrors}
|
||||
/>
|
||||
</button>
|
||||
{detailsExpanded && <div className="px-4 pb-4">{configurationForm}</div>}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{isAdmin && (
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="model_access_group"
|
||||
label={labelWithHint(
|
||||
"Model Access Group",
|
||||
"Use model access groups to control who can access this auto router",
|
||||
)}
|
||||
>
|
||||
{({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => (
|
||||
<AccessGroupTagsCombobox
|
||||
id={id}
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
options={modelAccessGroups}
|
||||
ariaInvalid={ariaInvalid}
|
||||
ariaDescribedBy={ariaDescribedBy}
|
||||
{isAdmin && (
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="model_access_group"
|
||||
label={labelWithHint(
|
||||
"Model Access Group",
|
||||
"Use model access groups to control who can access this auto router",
|
||||
)}
|
||||
>
|
||||
{({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => (
|
||||
<AccessGroupTagsCombobox
|
||||
id={id}
|
||||
value={value}
|
||||
onChange={onChange}
|
||||
options={modelAccessGroups}
|
||||
ariaInvalid={ariaInvalid}
|
||||
ariaDescribedBy={ariaDescribedBy}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
)}
|
||||
|
||||
<div className="flex justify-between items-center">
|
||||
<Tooltip>
|
||||
<TooltipTrigger
|
||||
render={
|
||||
<a
|
||||
href="https://github.com/BerriAI/litellm/issues"
|
||||
className="text-sm text-primary underline-offset-4 hover:underline"
|
||||
>
|
||||
Need Help?
|
||||
</a>
|
||||
}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
)}
|
||||
|
||||
<div className="flex justify-between items-center">
|
||||
<Tooltip>
|
||||
<TooltipTrigger
|
||||
render={
|
||||
<a
|
||||
href="https://github.com/BerriAI/litellm/issues"
|
||||
className="text-sm text-primary underline-offset-4 hover:underline"
|
||||
<TooltipContent>Get help on our github</TooltipContent>
|
||||
</Tooltip>
|
||||
<div className="flex gap-2">
|
||||
<BlockedReasonTooltip reason={submitBlockedReason}>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
data-testid="auto-router-test-routing-btn"
|
||||
disabled={submitBlockedReason !== null || isSubmitting}
|
||||
onClick={() => setIsRoutingTestVisible(true)}
|
||||
>
|
||||
Need Help?
|
||||
</a>
|
||||
}
|
||||
/>
|
||||
<TooltipContent>Get help on our github</TooltipContent>
|
||||
</Tooltip>
|
||||
<div className="flex gap-2">
|
||||
<BlockedReasonTooltip reason={submitBlockedReason}>
|
||||
Test Routing
|
||||
</Button>
|
||||
</BlockedReasonTooltip>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
data-testid="auto-router-test-routing-btn"
|
||||
disabled={submitBlockedReason !== null || isSubmitting}
|
||||
onClick={() => setIsRoutingTestVisible(true)}
|
||||
data-testid="auto-router-test-connect-btn"
|
||||
onClick={handleTestConnection}
|
||||
disabled={isTestingConnection}
|
||||
>
|
||||
Test Routing
|
||||
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
|
||||
Test Connection
|
||||
</Button>
|
||||
</BlockedReasonTooltip>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
data-testid="auto-router-test-connect-btn"
|
||||
onClick={handleTestConnection}
|
||||
disabled={isTestingConnection}
|
||||
>
|
||||
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
|
||||
Test Connection
|
||||
</Button>
|
||||
<BlockedReasonTooltip reason={submitBlockedReason}>
|
||||
<Button
|
||||
type="button"
|
||||
disabled={submitBlockedReason !== null || isSubmitting}
|
||||
onClick={() => {
|
||||
void handleAutoRouterSubmit();
|
||||
}}
|
||||
>
|
||||
Add Auto Router
|
||||
</Button>
|
||||
</BlockedReasonTooltip>
|
||||
<BlockedReasonTooltip reason={submitBlockedReason}>
|
||||
<Button
|
||||
type="button"
|
||||
disabled={submitBlockedReason !== null || isSubmitting}
|
||||
onClick={() => {
|
||||
void handleAutoRouterSubmit();
|
||||
}}
|
||||
>
|
||||
Add Auto Router
|
||||
</Button>
|
||||
</BlockedReasonTooltip>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</FieldGroup>
|
||||
</FieldGroup>
|
||||
</AutoRouterClassifierTabs>
|
||||
</form>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
|
|
|||
|
|
@ -48,6 +48,37 @@ const baseParams: BuildComplexityRouterConfigParams = {
|
|||
};
|
||||
|
||||
describe("buildComplexityRouterConfig", () => {
|
||||
it.each(["capability", "llm_v2", "heuristic"] as const)(
|
||||
"disables the removed overrides only for forecast creates: %s",
|
||||
(classifierType) => {
|
||||
const forecast = classifierType !== "heuristic";
|
||||
const params = {
|
||||
...baseParams,
|
||||
classifierType,
|
||||
adaptive: true,
|
||||
enableContextWindowEscalation: true,
|
||||
contextWindowEscalationBuffer: 0.9,
|
||||
};
|
||||
const config = buildComplexityRouterConfig(params);
|
||||
expect(config.adaptive).toBe(!forecast);
|
||||
expect(config.enable_context_window_escalation).toBe(!forecast);
|
||||
expect(config.escalation_keywords).toEqual(forecast ? [] : ["LITELLM ESCALATE"]);
|
||||
for (const key of [
|
||||
"adaptive_weights",
|
||||
"adaptive_eligible",
|
||||
"tier_distance_penalty",
|
||||
"context_window_escalation_buffer",
|
||||
]) {
|
||||
expect(Object.hasOwn(config, key)).toBe(!forecast);
|
||||
}
|
||||
if (forecast) {
|
||||
const untouched = buildComplexityRouterConfig({ ...baseParams, classifierType });
|
||||
expect(untouched.enable_context_window_escalation).toBe(false);
|
||||
expect(untouched.escalation_keywords).toEqual([]);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
it("carries Fast and reasoning overrides independently into a new router payload", () => {
|
||||
const params = { speed: "fast", reasoning_effort: "high", max_tokens: 1024 };
|
||||
const config = buildComplexityRouterConfig({
|
||||
|
|
|
|||
|
|
@ -1,3 +1,9 @@
|
|||
import {
|
||||
isForecastClassifier,
|
||||
withoutForecastPromptOverrides,
|
||||
type CapabilitySettings,
|
||||
type FuseSettings,
|
||||
} from "./forecast_classifier_config";
|
||||
import type { ModelGroup } from "../llm_calls/fetch_models";
|
||||
import { KeywordTierRule } from "./KeywordTierRules";
|
||||
import {
|
||||
|
|
@ -126,6 +132,48 @@ const scorerKnobPayload = ({
|
|||
};
|
||||
};
|
||||
|
||||
export interface StoredComplexityRouterConfig {
|
||||
tiers?: Partial<Record<keyof ComplexityTiers, unknown>>;
|
||||
enable_non_reasoning_tier?: boolean;
|
||||
tier_model_configs?: unknown;
|
||||
default_model?: string | null;
|
||||
plan_mode_min_tier?: unknown;
|
||||
classification_prompt?: unknown;
|
||||
classification_examples?: unknown;
|
||||
heuristic_first_max_tier?: unknown;
|
||||
hybrid_boundary_margin?: unknown;
|
||||
tier_labels?: unknown;
|
||||
classifier_type?: ClassifierType;
|
||||
capability_classifier_config?: unknown;
|
||||
llm_v2_config?: unknown;
|
||||
classifier_llm_config?: ClassifierLLMConfig;
|
||||
classifier_context_window_size?: unknown;
|
||||
classifier_context_budget_chars?: unknown;
|
||||
classifier_context_include_assistant_turns?: unknown;
|
||||
classifier_fallback?: unknown;
|
||||
classification_mode?: unknown;
|
||||
tier_boundaries?: unknown;
|
||||
token_thresholds?: unknown;
|
||||
dimension_weights?: unknown;
|
||||
custom_dimensions?: unknown;
|
||||
reasoning_override_min_score?: unknown;
|
||||
session_affinity?: unknown;
|
||||
session_affinity_ttl_seconds?: unknown;
|
||||
modality_routing?: unknown;
|
||||
modality_pin_override?: unknown;
|
||||
deployment_affinity?: unknown;
|
||||
adaptive?: boolean;
|
||||
adaptive_weights?: AdaptiveRouterWeights;
|
||||
tier_distance_penalty?: number;
|
||||
adaptive_eligible?: AdaptiveEligible;
|
||||
return_raw_model_name?: boolean;
|
||||
enable_context_window_escalation?: unknown;
|
||||
context_window_escalation_buffer?: unknown;
|
||||
stall_escalation_enabled?: unknown;
|
||||
stall_escalation_window?: unknown;
|
||||
stall_escalation_repeat_threshold?: unknown;
|
||||
}
|
||||
|
||||
export interface BuildComplexityRouterConfigParams {
|
||||
tiers: ComplexityTiers;
|
||||
enableNonReasoningTier?: boolean;
|
||||
|
|
@ -134,6 +182,8 @@ export interface BuildComplexityRouterConfigParams {
|
|||
planModeMinTier: string | undefined;
|
||||
tierLabels: ComplexityTierLabels | undefined;
|
||||
classifierType: ClassifierType;
|
||||
capabilityClassifierConfig?: CapabilitySettings;
|
||||
llmV2Config?: FuseSettings;
|
||||
classifierLlmConfig: ClassifierLLMConfigWire | undefined;
|
||||
classifierContextWindowSize: number | undefined;
|
||||
classifierContextBudgetChars: number | undefined;
|
||||
|
|
@ -198,6 +248,8 @@ export interface ComplexityRouterConfigPayload {
|
|||
plan_mode_min_tier?: string;
|
||||
tier_labels?: ComplexityTierLabels;
|
||||
classifier_type: ClassifierType;
|
||||
capability_classifier_config?: CapabilitySettings;
|
||||
llm_v2_config?: FuseSettings;
|
||||
classifier_llm_config?: ClassifierLLMConfig;
|
||||
classifier_context_window_size?: number;
|
||||
classifier_context_budget_chars?: number;
|
||||
|
|
@ -468,31 +520,34 @@ const classifierWireFields = (
|
|||
| "classifierContextBudgetChars"
|
||||
| "classifierContextIncludeAssistantTurns"
|
||||
>,
|
||||
): Partial<ComplexityRouterConfigPayload> => ({
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierLlmConfig && {
|
||||
classifier_llm_config:
|
||||
effectiveType === "capability" ? classifierLlmConfig : normalizeClassifierLlmConfig(classifierLlmConfig),
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
|
||||
...(effectiveType === "heuristic_first" &&
|
||||
heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
|
||||
...(effectiveType === "hybrid" &&
|
||||
hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextWindowSize !== undefined && {
|
||||
classifier_context_window_size: classifierContextWindowSize,
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextBudgetChars !== undefined && {
|
||||
classifier_context_budget_chars: classifierContextBudgetChars,
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextIncludeAssistantTurns !== undefined && {
|
||||
classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
|
||||
}),
|
||||
});
|
||||
): Partial<ComplexityRouterConfigPayload> => {
|
||||
const supportsFallback = usesLlmClassifier(effectiveType) && !isForecastClassifier(effectiveType);
|
||||
return {
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierLlmConfig && {
|
||||
classifier_llm_config: isForecastClassifier(effectiveType)
|
||||
? withoutForecastPromptOverrides(classifierLlmConfig)
|
||||
: normalizeClassifierLlmConfig(classifierLlmConfig),
|
||||
}),
|
||||
...(supportsFallback && classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
|
||||
...(effectiveType === "heuristic_first" &&
|
||||
heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
|
||||
...(effectiveType === "hybrid" &&
|
||||
hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextWindowSize !== undefined && {
|
||||
classifier_context_window_size: classifierContextWindowSize,
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextBudgetChars !== undefined && {
|
||||
classifier_context_budget_chars: classifierContextBudgetChars,
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierContextIncludeAssistantTurns !== undefined && {
|
||||
classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
|
||||
}),
|
||||
};
|
||||
};
|
||||
|
||||
export const buildComplexityRouterConfig = ({
|
||||
tiers,
|
||||
|
|
@ -502,6 +557,8 @@ export const buildComplexityRouterConfig = ({
|
|||
planModeMinTier,
|
||||
tierLabels,
|
||||
classifierType,
|
||||
capabilityClassifierConfig,
|
||||
llmV2Config,
|
||||
classifierLlmConfig,
|
||||
classifierContextWindowSize,
|
||||
classifierContextBudgetChars,
|
||||
|
|
@ -571,9 +628,11 @@ export const buildComplexityRouterConfig = ({
|
|||
// An edited tier set forces the LLM classifier, so llm-only inputs must survive a classifier_type
|
||||
// the form never rewrote. The UI gates the same controls on this, not on the raw value.
|
||||
const effectiveType: ClassifierType = customTierSet ? "llm" : classifierType;
|
||||
const forecast = isForecastClassifier(effectiveType);
|
||||
|
||||
const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType);
|
||||
const payload: ComplexityRouterConfigPayload = {
|
||||
tiers,
|
||||
tiers: forecast ? Object.fromEntries(Object.entries(tiers).filter(([, models]) => models.length > 0)) : tiers,
|
||||
// The backend rejects the flag beside a custom tier set.
|
||||
...(!customTierSet && enableNonReasoningTier && { enable_non_reasoning_tier: true }),
|
||||
...(serializedTierModelConfigs && { tier_model_configs: serializedTierModelConfigs }),
|
||||
|
|
@ -582,10 +641,13 @@ export const buildComplexityRouterConfig = ({
|
|||
...(cleanedTierLabels && { tier_labels: cleanedTierLabels }),
|
||||
classifier_type: classifierType,
|
||||
...classifierWireFields(effectiveType, classifierInputs),
|
||||
...(effectiveType === "capability" &&
|
||||
capabilityClassifierConfig && { capability_classifier_config: capabilityClassifierConfig }),
|
||||
...(effectiveType === "llm_v2" && { llm_v2_config: llmV2Config }),
|
||||
...(forecast && { adaptive: false }),
|
||||
// A built-in router's opening instructions. Suppressed beside a legacy whole-prompt override,
|
||||
// which the backend rejects as a second override of the same prompt.
|
||||
...(!customTierSet &&
|
||||
usesLlmClassifier(effectiveType) &&
|
||||
...(supportsOpeningPrompt &&
|
||||
!classifierLlmConfig?.system_prompt?.trim() && {
|
||||
...(classificationPrompt?.trim() && { classification_prompt: classificationPrompt.trim() }),
|
||||
...(classificationExamples?.trim() && { classification_examples: classificationExamples.trim() }),
|
||||
|
|
@ -597,7 +659,7 @@ export const buildComplexityRouterConfig = ({
|
|||
modality_pin_override: modalityPinOverride ?? false,
|
||||
...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
|
||||
...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
|
||||
escalation_keywords: cleanedEscalationKeywords,
|
||||
escalation_keywords: forecast ? [] : cleanedEscalationKeywords,
|
||||
// Only written when on: the backend rejects it alongside session_affinity, user_turn mode and
|
||||
// a custom tier set, so an off router must not carry the key into any of those saves.
|
||||
...(stallEscalationEnabled && {
|
||||
|
|
@ -612,19 +674,22 @@ export const buildComplexityRouterConfig = ({
|
|||
embedding_model: embeddingModel,
|
||||
match_threshold: matchThreshold,
|
||||
}),
|
||||
...(adaptive && {
|
||||
adaptive: true,
|
||||
adaptive_weights: adaptiveWeights,
|
||||
...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }),
|
||||
adaptive_eligible: adaptiveEligible,
|
||||
}),
|
||||
...(adaptive &&
|
||||
!forecast && {
|
||||
adaptive: true,
|
||||
adaptive_weights: adaptiveWeights,
|
||||
...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }),
|
||||
adaptive_eligible: adaptiveEligible,
|
||||
}),
|
||||
...(returnRawModelName && { return_raw_model_name: true }),
|
||||
...(enableContextWindowEscalation !== undefined && {
|
||||
enable_context_window_escalation: enableContextWindowEscalation,
|
||||
}),
|
||||
...(contextWindowEscalationBuffer !== undefined && {
|
||||
context_window_escalation_buffer: contextWindowEscalationBuffer,
|
||||
// Omission enables the backend default, so hidden forecast controls need an explicit opt-out.
|
||||
...((forecast || enableContextWindowEscalation !== undefined) && {
|
||||
enable_context_window_escalation: forecast ? false : enableContextWindowEscalation,
|
||||
}),
|
||||
...(!forecast &&
|
||||
contextWindowEscalationBuffer !== undefined && {
|
||||
context_window_escalation_buffer: contextWindowEscalationBuffer,
|
||||
}),
|
||||
...(sessionAffinityTtlSeconds !== undefined && {
|
||||
session_affinity_ttl_seconds: sessionAffinityTtlSeconds,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,80 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { transitionClassifierType } from "./classifier_type_transition";
|
||||
|
||||
const standard: ComplexityRouterConfigValue = {
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "judge", timeout_ms: 20000, classification_rubric: "business" },
|
||||
classifier_context_window_size: 8,
|
||||
classifier_context_budget_chars: 16000,
|
||||
classifier_context_include_assistant_turns: true,
|
||||
classifier_fallback: "default_model",
|
||||
tiers: { SIMPLE: ["efficient"], MEDIUM: ["middle"], COMPLEX: [], REASONING: ["capable"] },
|
||||
};
|
||||
|
||||
describe("transitionClassifierType", () => {
|
||||
it.each(["heuristic_first", "hybrid"] as const)("keeps existing LLM settings when switching to %s", (target) => {
|
||||
const result = transitionClassifierType(standard, target);
|
||||
const expectedSettings = {
|
||||
classifier_type: target,
|
||||
classifier_llm_config: standard.classifier_llm_config,
|
||||
classifier_context_window_size: 8,
|
||||
classifier_context_budget_chars: 16000,
|
||||
classifier_context_include_assistant_turns: true,
|
||||
classifier_fallback: "default_model",
|
||||
};
|
||||
expect(result).toMatchObject(expectedSettings);
|
||||
});
|
||||
|
||||
it.each(["capability", "llm_v2"] as const)("requires explicit policy input for a new %s classifier", (target) => {
|
||||
const result = transitionClassifierType(standard, target);
|
||||
expect(result.classifier_llm_config).toEqual({ model: "judge", timeout_ms: 20000 });
|
||||
expect(result.classifier_fallback).toBeUndefined();
|
||||
if (target === "capability") {
|
||||
expect(result.capability_classifier_config?.base_threshold).toBeNaN();
|
||||
} else {
|
||||
expect(result.llm_v2_config).toMatchObject({ efficient_profile: "", capable_profile: "", harness: "" });
|
||||
expect(result.llm_v2_config?.max_quality_gap).toBeNaN();
|
||||
}
|
||||
expect(standard.tiers.MEDIUM).toEqual(["middle"]);
|
||||
expect(standard.classifier_llm_config?.classification_rubric).toBe("business");
|
||||
});
|
||||
|
||||
it.each([
|
||||
["capability", "llm"],
|
||||
["capability", "heuristic_first"],
|
||||
["capability", "hybrid"],
|
||||
["llm_v2", "llm"],
|
||||
["llm_v2", "heuristic_first"],
|
||||
["llm_v2", "hybrid"],
|
||||
] as const)("restores the complexity rubric from %s to %s while preserving the judge", (source, target) => {
|
||||
const forecast = transitionClassifierType(standard, source);
|
||||
const result = transitionClassifierType(forecast, target);
|
||||
expect(result.classifier_llm_config).toEqual({
|
||||
model: "judge",
|
||||
timeout_ms: 20000,
|
||||
classification_rubric: "agentic",
|
||||
});
|
||||
expect(result.capability_classifier_config).toBeUndefined();
|
||||
expect(result.llm_v2_config).toBeUndefined();
|
||||
});
|
||||
|
||||
it("clears the inactive non-reasoning pool and plan floor when switching to local classification", () => {
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
...standard,
|
||||
tiers: { ...standard.tiers, NON_REASONING: ["chat"] },
|
||||
enable_non_reasoning_tier: true,
|
||||
plan_mode_min_tier: "NON_REASONING",
|
||||
};
|
||||
const result = transitionClassifierType(initial, "heuristic");
|
||||
expect(result.classifier_llm_config).toBeUndefined();
|
||||
expect(result.classifier_context_window_size).toBeUndefined();
|
||||
expect(result.classifier_context_budget_chars).toBeUndefined();
|
||||
expect(result.classifier_context_include_assistant_turns).toBeUndefined();
|
||||
expect(result.classifier_fallback).toBeUndefined();
|
||||
expect(result.tiers.NON_REASONING).toBeUndefined();
|
||||
expect(result.enable_non_reasoning_tier).toBeUndefined();
|
||||
expect(result.plan_mode_min_tier).toBeUndefined();
|
||||
expect(result.tiers.SIMPLE).toEqual(["efficient"]);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
import {
|
||||
type ClassifierType,
|
||||
type ComplexityRouterConfigValue,
|
||||
DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS,
|
||||
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
DEFAULT_HEURISTIC_FIRST_MAX_TIER,
|
||||
DEFAULT_HYBRID_BOUNDARY_MARGIN,
|
||||
NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
|
||||
usesLlmClassifier,
|
||||
} from "./ComplexityRouterConfig";
|
||||
import { isForecastClassifier, prepareForecastClassifier } from "./forecast_classifier_config";
|
||||
import { nonReasoningTierFields } from "./nonReasoningTierFields";
|
||||
|
||||
export const transitionClassifierType = (
|
||||
value: ComplexityRouterConfigValue,
|
||||
classifierType: ClassifierType,
|
||||
): ComplexityRouterConfigValue => {
|
||||
const startsLlmRubric =
|
||||
!value.classifier_llm_config ||
|
||||
(isForecastClassifier(value.classifier_type) && !isForecastClassifier(classifierType));
|
||||
const judgeConfig = value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS };
|
||||
const nextValue: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
classifier_llm_config: usesLlmClassifier(classifierType)
|
||||
? {
|
||||
...judgeConfig,
|
||||
...(startsLlmRubric && { classification_rubric: NEW_CLASSIFIER_CLASSIFICATION_RUBRIC }),
|
||||
}
|
||||
: undefined,
|
||||
classifier_context_window_size: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
|
||||
: undefined,
|
||||
classifier_context_budget_chars: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_include_assistant_turns
|
||||
: undefined,
|
||||
classifier_fallback: usesLlmClassifier(classifierType) ? value.classifier_fallback : undefined,
|
||||
heuristic_first_max_tier:
|
||||
classifierType === "heuristic_first"
|
||||
? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER
|
||||
: undefined,
|
||||
hybrid_boundary_margin:
|
||||
classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined,
|
||||
...nonReasoningTierFields(classifierType, value),
|
||||
};
|
||||
return prepareForecastClassifier(nextValue, classifierType);
|
||||
};
|
||||
|
|
@ -0,0 +1,284 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { getForecastConfigError, prepareForecastClassifier } from "./forecast_classifier_config";
|
||||
import { getKeywordTierRulesError } from "./build_complexity_router_config";
|
||||
import { activeTierRows } from "./tier_rows";
|
||||
import {
|
||||
buildUpdatedComplexityRouterConfig,
|
||||
hydrateComplexityRouterConfig,
|
||||
} from "../edit_auto_router/edit_auto_router_modal";
|
||||
|
||||
const capability: ComplexityRouterConfigValue = {
|
||||
classifier_type: "capability",
|
||||
classifier_llm_config: { model: "judge", timeout_ms: 20000 },
|
||||
tiers: { SIMPLE: ["efficient"], MEDIUM: [], COMPLEX: [], REASONING: ["capable"] },
|
||||
capability_classifier_config: {
|
||||
efficient_tier: "SIMPLE",
|
||||
capable_tier: "REASONING",
|
||||
base_threshold: 0.7,
|
||||
threshold_step: 0.1,
|
||||
},
|
||||
};
|
||||
const fuse: ComplexityRouterConfigValue = {
|
||||
...capability,
|
||||
classifier_type: "llm_v2",
|
||||
capability_classifier_config: undefined,
|
||||
adaptive: false,
|
||||
llm_v2_config: {
|
||||
efficient_profile: "A concise solver",
|
||||
capable_profile: "A solver with more reasoning budget",
|
||||
harness: "One attempt with shell and tests",
|
||||
max_quality_gap: 0.05,
|
||||
},
|
||||
};
|
||||
|
||||
describe("forecast classifier configuration", () => {
|
||||
it.each([
|
||||
{ version: "eval", slope: 21, intercept: 0 },
|
||||
{ version: "eval", slope: 1, intercept: -21 },
|
||||
{ version: " eval ", slope: 1, intercept: 0 },
|
||||
])("rejects capability calibration outside the server contract: %j", (calibration) => {
|
||||
const value = {
|
||||
...capability,
|
||||
capability_classifier_config: { ...capability.capability_classifier_config!, calibration },
|
||||
};
|
||||
expect(getForecastConfigError(value)).toContain("calibration");
|
||||
});
|
||||
|
||||
it.each([capability, fuse])("accepts a complete $classifier_type configuration", (value) => {
|
||||
expect(getForecastConfigError(value)).toBeNull();
|
||||
});
|
||||
|
||||
it("requires both solvers without requiring unused middle tiers", () => {
|
||||
expect(getForecastConfigError({ ...capability, tiers: { ...capability.tiers, REASONING: [] } })).toContain("both");
|
||||
expect(getForecastConfigError(capability)).toBeNull();
|
||||
});
|
||||
|
||||
it("rejects a stepped threshold that exceeds one", () => {
|
||||
expect(
|
||||
getForecastConfigError({
|
||||
...capability,
|
||||
capability_classifier_config: { ...capability.capability_classifier_config!, threshold_step: 0.2 },
|
||||
}),
|
||||
).toContain("twice");
|
||||
});
|
||||
|
||||
it.each([Number.NaN, -0.1, 1.1])("rejects an invalid probability %s", (base_threshold) => {
|
||||
expect(
|
||||
getForecastConfigError({
|
||||
...capability,
|
||||
capability_classifier_config: { ...capability.capability_classifier_config!, base_threshold },
|
||||
}),
|
||||
).not.toBeNull();
|
||||
});
|
||||
|
||||
it.each(["efficient_profile", "capable_profile", "harness"] as const)("requires %s for Fuse", (field) => {
|
||||
expect(getForecastConfigError({ ...fuse, llm_v2_config: { ...fuse.llm_v2_config!, [field]: " " } })).not.toBeNull();
|
||||
});
|
||||
|
||||
it("requires distinct single model groups and disables adaptive selection", () => {
|
||||
expect(getForecastConfigError({ ...fuse, tiers: { ...fuse.tiers, SIMPLE: ["capable"] } })).toContain("distinct");
|
||||
expect(getForecastConfigError({ ...fuse, tiers: { ...fuse.tiers, SIMPLE: ["efficient", "second"] } })).toContain(
|
||||
"distinct",
|
||||
);
|
||||
expect(getForecastConfigError({ ...fuse, adaptive: true })).toContain("adaptive");
|
||||
});
|
||||
|
||||
it("removes incompatible prompt and tier settings when switching to Fuse", () => {
|
||||
const previous: ComplexityRouterConfigValue = {
|
||||
...fuse,
|
||||
adaptive: true,
|
||||
classifier_fallback: "default_model",
|
||||
classifier_llm_config: {
|
||||
model: "judge",
|
||||
timeout_ms: 20000,
|
||||
system_prompt: "old rubric",
|
||||
classification_rubric: "business",
|
||||
},
|
||||
classification_prompt: "old prompt",
|
||||
classification_examples: "old example",
|
||||
tiers: { ...fuse.tiers, MEDIUM: ["extra"], SIMPLE: ["efficient", "second"] },
|
||||
};
|
||||
const next = prepareForecastClassifier(previous);
|
||||
const saved = buildUpdatedComplexityRouterConfig({}, next);
|
||||
expect(getForecastConfigError(next)).toBeNull();
|
||||
expect(saved.adaptive).toBe(false);
|
||||
expect(saved.tiers).toEqual({ SIMPLE: ["efficient"], REASONING: ["capable"] });
|
||||
expect(saved.classifier_llm_config).toEqual({ model: "judge", timeout_ms: 20000 });
|
||||
expect(saved).not.toHaveProperty("classification_prompt");
|
||||
expect(saved).not.toHaveProperty("classification_examples");
|
||||
expect(saved).not.toHaveProperty("classifier_fallback");
|
||||
});
|
||||
|
||||
it.each([capability, fuse])(
|
||||
"switching to $classifier_type removes hidden pools and their overrides while retaining the solver settings",
|
||||
(value) => {
|
||||
const efficientParams = { reasoning_effort: "low", speed: "fast", max_tokens: 1024 };
|
||||
const secondEfficientParams = { reasoning_effort: "medium", max_tokens: 2048 };
|
||||
const capableParams = { reasoning_effort: "high", max_tokens: 4096 };
|
||||
const secondCapableParams = { speed: "fast" };
|
||||
const previous: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
adaptive: true,
|
||||
plan_mode_min_tier: "MEDIUM",
|
||||
enable_non_reasoning_tier: true,
|
||||
tiers: {
|
||||
NON_REASONING: ["relay"],
|
||||
SIMPLE: ["efficient", "second-efficient"],
|
||||
MEDIUM: ["leftover-medium"],
|
||||
COMPLEX: ["leftover-complex"],
|
||||
REASONING: ["capable", "second-capable"],
|
||||
},
|
||||
tier_model_params: {
|
||||
NON_REASONING: { relay: { max_tokens: 128 } },
|
||||
SIMPLE: { efficient: efficientParams, "second-efficient": secondEfficientParams },
|
||||
MEDIUM: { "leftover-medium": { speed: "fast" } },
|
||||
COMPLEX: { "leftover-complex": { reasoning_effort: "high" } },
|
||||
REASONING: { capable: capableParams, "second-capable": secondCapableParams },
|
||||
LEGACY_CUSTOM: { "leftover-custom": { max_tokens: 512 } },
|
||||
},
|
||||
};
|
||||
const next = prepareForecastClassifier(previous);
|
||||
const preservesPools = value.classifier_type === "capability";
|
||||
const expectedTiers = {
|
||||
SIMPLE: preservesPools ? ["efficient", "second-efficient"] : ["efficient"],
|
||||
MEDIUM: [],
|
||||
COMPLEX: [],
|
||||
REASONING: preservesPools ? ["capable", "second-capable"] : ["capable"],
|
||||
};
|
||||
expect(next.tiers).toEqual(expectedTiers);
|
||||
expect(next.tier_model_params).toEqual({
|
||||
SIMPLE: {
|
||||
efficient: efficientParams,
|
||||
...(preservesPools && { "second-efficient": secondEfficientParams }),
|
||||
},
|
||||
REASONING: {
|
||||
capable: capableParams,
|
||||
...(preservesPools && { "second-capable": secondCapableParams }),
|
||||
},
|
||||
});
|
||||
expect(next.adaptive).toBe(preservesPools);
|
||||
expect(next.plan_mode_min_tier).toBeUndefined();
|
||||
expect(getForecastConfigError(next)).toBeNull();
|
||||
expect(activeTierRows(next).map((row) => row.name)).toEqual(["SIMPLE", "REASONING"]);
|
||||
expect(
|
||||
getKeywordTierRulesError([{ id: "old-middle", keywords: ["invoice"], tier: "MEDIUM" }], activeTierRows(next)),
|
||||
).toContain("no longer has");
|
||||
expect(
|
||||
getKeywordTierRulesError([{ id: "solver", keywords: ["audit"], tier: "REASONING" }], activeTierRows(next)),
|
||||
).toBeNull();
|
||||
const saved = buildUpdatedComplexityRouterConfig(previous, next);
|
||||
expect(saved.tiers).toEqual({ SIMPLE: expectedTiers.SIMPLE, REASONING: expectedTiers.REASONING });
|
||||
expect(saved.tier_model_configs).toEqual({
|
||||
SIMPLE: [
|
||||
{ model_name: "efficient", litellm_params: efficientParams },
|
||||
...(preservesPools ? [{ model_name: "second-efficient", litellm_params: secondEfficientParams }] : []),
|
||||
],
|
||||
REASONING: [
|
||||
{ model_name: "capable", litellm_params: capableParams },
|
||||
...(preservesPools ? [{ model_name: "second-capable", litellm_params: secondCapableParams }] : []),
|
||||
],
|
||||
});
|
||||
expect(saved).not.toHaveProperty("plan_mode_min_tier");
|
||||
},
|
||||
);
|
||||
|
||||
it("preserves fitted calibration through edits and removes it when disabled", () => {
|
||||
const stored = {
|
||||
...capability,
|
||||
capability_classifier_config: {
|
||||
...capability.capability_classifier_config!,
|
||||
calibration: { version: "eval-a", slope: 1.2, intercept: -0.3 },
|
||||
},
|
||||
};
|
||||
const hydrated = hydrateComplexityRouterConfig(stored, undefined);
|
||||
const edited = {
|
||||
...hydrated,
|
||||
capability_classifier_config: { ...hydrated.capability_classifier_config!, base_threshold: 0.6 },
|
||||
};
|
||||
expect(buildUpdatedComplexityRouterConfig(stored, edited).capability_classifier_config).toEqual({
|
||||
...stored.capability_classifier_config,
|
||||
base_threshold: 0.6,
|
||||
});
|
||||
const disabled = {
|
||||
...edited,
|
||||
capability_classifier_config: { ...edited.capability_classifier_config, calibration: undefined },
|
||||
};
|
||||
expect(buildUpdatedComplexityRouterConfig(stored, disabled).capability_classifier_config).toHaveProperty(
|
||||
"calibration",
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
it("preserves non-default tier assignments", () => {
|
||||
const value = {
|
||||
...capability,
|
||||
capability_classifier_config: {
|
||||
...capability.capability_classifier_config!,
|
||||
efficient_tier: "MEDIUM",
|
||||
capable_tier: "COMPLEX",
|
||||
},
|
||||
tiers: { SIMPLE: [], MEDIUM: ["efficient"], COMPLEX: ["capable"], REASONING: [] },
|
||||
};
|
||||
expect(getForecastConfigError(value)).toBeNull();
|
||||
expect(buildUpdatedComplexityRouterConfig({}, value).capability_classifier_config).toEqual(
|
||||
value.capability_classifier_config,
|
||||
);
|
||||
const previous = {
|
||||
...value,
|
||||
tiers: { ...value.tiers, SIMPLE: ["leftover-simple"], REASONING: ["leftover-reasoning"] },
|
||||
plan_mode_min_tier: "COMPLEX",
|
||||
tier_model_params: {
|
||||
MEDIUM: { efficient: { max_tokens: 1024 } },
|
||||
COMPLEX: { capable: { speed: "fast" } },
|
||||
SIMPLE: { "leftover-simple": { speed: "fast" } },
|
||||
},
|
||||
};
|
||||
const next = prepareForecastClassifier(previous);
|
||||
expect(next.tiers).toEqual(value.tiers);
|
||||
expect(next.plan_mode_min_tier).toBe("COMPLEX");
|
||||
expect(next.tier_model_params).toEqual({
|
||||
MEDIUM: { efficient: { max_tokens: 1024 } },
|
||||
COMPLEX: { capable: { speed: "fast" } },
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps configured extra Capability pools through hydration and an unrelated edit", () => {
|
||||
const stored = {
|
||||
...capability,
|
||||
adaptive: true,
|
||||
plan_mode_min_tier: "MEDIUM",
|
||||
tiers: { ...capability.tiers, MEDIUM: ["middle"] },
|
||||
tier_model_configs: { MEDIUM: [{ model_name: "middle", litellm_params: { speed: "fast", max_tokens: 1024 } }] },
|
||||
};
|
||||
const hydrated = hydrateComplexityRouterConfig(stored, undefined);
|
||||
expect(hydrated.tiers.MEDIUM).toEqual(["middle"]);
|
||||
expect(hydrated.plan_mode_min_tier).toBe("MEDIUM");
|
||||
expect(activeTierRows(hydrated).map((row) => row.name)).toEqual(["SIMPLE", "MEDIUM", "REASONING"]);
|
||||
expect(
|
||||
getKeywordTierRulesError([{ id: "kept", keywords: ["invoice"], tier: "MEDIUM" }], activeTierRows(hydrated)),
|
||||
).toBeNull();
|
||||
const saved = buildUpdatedComplexityRouterConfig(stored, {
|
||||
...hydrated,
|
||||
capability_classifier_config: { ...hydrated.capability_classifier_config!, base_threshold: 0.6 },
|
||||
});
|
||||
expect(saved.tiers).toEqual({ SIMPLE: ["efficient"], MEDIUM: ["middle"], REASONING: ["capable"] });
|
||||
expect(saved.tier_model_configs).toEqual(stored.tier_model_configs);
|
||||
expect(saved.adaptive).toBe(false);
|
||||
expect(saved.plan_mode_min_tier).toBe("MEDIUM");
|
||||
});
|
||||
|
||||
it("keeps empty built-in tiers available to standard classifiers", () => {
|
||||
const standard: ComplexityRouterConfigValue = { ...capability, classifier_type: "llm" };
|
||||
expect(activeTierRows(standard).map((row) => row.name)).toEqual(["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]);
|
||||
expect(
|
||||
getKeywordTierRulesError([{ id: "middle", keywords: ["invoice"], tier: "MEDIUM" }], activeTierRows(standard)),
|
||||
).toBeNull();
|
||||
});
|
||||
|
||||
it.each([capability, fuse])("drops $classifier_type settings when switching to the heuristic", (value) => {
|
||||
const saved = buildUpdatedComplexityRouterConfig(value, { ...value, classifier_type: "heuristic" });
|
||||
expect(saved).not.toHaveProperty("capability_classifier_config");
|
||||
expect(saved).not.toHaveProperty("llm_v2_config");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,188 @@
|
|||
import { z } from "zod";
|
||||
import type { ClassifierType, ComplexityRouterConfigValue, ComplexityTiers } from "./ComplexityRouterConfig";
|
||||
import { pruneTierModelParams } from "./complexity_router_tiers";
|
||||
import { tierOrderFor } from "./tier_rows";
|
||||
|
||||
const probability = z.number().finite().min(0).max(1);
|
||||
const version = z.string().trim().min(1).max(512);
|
||||
const profile = z.string().trim().min(1).max(4000);
|
||||
const transport = {
|
||||
max_output_tokens: z.number().int().positive().optional(),
|
||||
response_format: z.enum(["json_schema", "json_object"]).optional(),
|
||||
};
|
||||
const coefficients = z.object({ slope: z.number().finite().positive(), intercept: z.number().finite() });
|
||||
|
||||
const capabilityShape = {
|
||||
efficient_tier: z.string().min(1),
|
||||
capable_tier: z.string().min(1),
|
||||
base_threshold: probability,
|
||||
threshold_step: z.number().finite().nonnegative().optional(),
|
||||
...transport,
|
||||
calibration: z
|
||||
.object({
|
||||
version: z
|
||||
.string()
|
||||
.min(1)
|
||||
.max(128)
|
||||
.regex(/^\S(?:.*\S)?$/),
|
||||
slope: z.number().finite().min(0).max(20),
|
||||
intercept: z.number().finite().min(-20).max(20),
|
||||
})
|
||||
.nullable()
|
||||
.optional(),
|
||||
};
|
||||
export const capabilitySettingsSchema = z.object(capabilityShape);
|
||||
|
||||
const fuseCalibrationShape = {
|
||||
version,
|
||||
prompt_version: z.literal("llm-v2-1"),
|
||||
efficient: coefficients,
|
||||
capable: coefficients,
|
||||
};
|
||||
const fuseShape = {
|
||||
efficient_tier: z.string().min(1).optional(),
|
||||
capable_tier: z.string().min(1).optional(),
|
||||
efficient_profile: profile,
|
||||
capable_profile: profile,
|
||||
harness: profile,
|
||||
max_quality_gap: probability,
|
||||
...transport,
|
||||
calibration: z.object(fuseCalibrationShape).nullable().optional(),
|
||||
};
|
||||
export const fuseSettingsSchema = z.object(fuseShape);
|
||||
|
||||
export type CapabilitySettings = z.infer<typeof capabilitySettingsSchema>;
|
||||
export type FuseSettings = z.infer<typeof fuseSettingsSchema>;
|
||||
|
||||
export const isForecastClassifier = (type: ClassifierType): boolean => type === "capability" || type === "llm_v2";
|
||||
|
||||
export const newCapabilitySettings = (): CapabilitySettings => ({
|
||||
efficient_tier: "SIMPLE",
|
||||
capable_tier: "REASONING",
|
||||
base_threshold: Number.NaN,
|
||||
});
|
||||
|
||||
export const newFuseSettings = (): FuseSettings => ({
|
||||
efficient_profile: "",
|
||||
capable_profile: "",
|
||||
harness: "",
|
||||
max_quality_gap: Number.NaN,
|
||||
});
|
||||
|
||||
export const forecastTierNames = (
|
||||
value: Pick<ComplexityRouterConfigValue, "classifier_type" | "capability_classifier_config" | "llm_v2_config">,
|
||||
): readonly [string, string] => {
|
||||
const settings = value.classifier_type === "capability" ? value.capability_classifier_config : value.llm_v2_config;
|
||||
return [settings?.efficient_tier ?? "SIMPLE", settings?.capable_tier ?? "REASONING"];
|
||||
};
|
||||
|
||||
export const forecastModels = (tiers: ComplexityTiers, tier: string): string[] =>
|
||||
Object.entries(tiers).find(([name]) => name === tier)?.[1] ?? [];
|
||||
|
||||
export const withoutForecastPromptOverrides = <T extends { system_prompt?: string; classification_rubric?: string }>(
|
||||
config: T,
|
||||
): Omit<T, "system_prompt" | "classification_rubric"> => {
|
||||
const { system_prompt: _prompt, classification_rubric: _rubric, ...rest } = config;
|
||||
return rest;
|
||||
};
|
||||
|
||||
export const prepareForecastClassifier = (
|
||||
previous: ComplexityRouterConfigValue,
|
||||
classifierType: ClassifierType = previous.classifier_type,
|
||||
): ComplexityRouterConfigValue => {
|
||||
const value = { ...previous, classifier_type: classifierType };
|
||||
if (!isForecastClassifier(classifierType))
|
||||
return {
|
||||
...value,
|
||||
capability_classifier_config: undefined,
|
||||
llm_v2_config: undefined,
|
||||
};
|
||||
// Capture routing identity before replacing the source classifier's settings.
|
||||
const [sourceEfficient, sourceCapable] = forecastTierNames(previous);
|
||||
const transferredPair =
|
||||
isForecastClassifier(previous.classifier_type) && previous.classifier_type !== classifierType
|
||||
? { efficient_tier: sourceEfficient, capable_tier: sourceCapable }
|
||||
: {};
|
||||
const configured = {
|
||||
...value,
|
||||
capability_classifier_config:
|
||||
value.classifier_type === "capability"
|
||||
? { ...(value.capability_classifier_config ?? newCapabilitySettings()), ...transferredPair }
|
||||
: undefined,
|
||||
llm_v2_config:
|
||||
value.classifier_type === "llm_v2"
|
||||
? { ...(value.llm_v2_config ?? newFuseSettings()), ...transferredPair }
|
||||
: undefined,
|
||||
};
|
||||
const [efficient, capable] = forecastTierNames(configured);
|
||||
const solverModels = (tier: string) => {
|
||||
const models = forecastModels(value.tiers, tier);
|
||||
return value.classifier_type === "llm_v2" ? models.slice(0, 1) : models;
|
||||
};
|
||||
const tiers: ComplexityTiers = {
|
||||
SIMPLE: [],
|
||||
MEDIUM: [],
|
||||
COMPLEX: [],
|
||||
REASONING: [],
|
||||
[efficient]: solverModels(efficient),
|
||||
[capable]: solverModels(capable),
|
||||
};
|
||||
return {
|
||||
...configured,
|
||||
tiers,
|
||||
tier_model_params: Object.keys(value.tier_model_params ?? {}).reduce(
|
||||
(params, tier) => pruneTierModelParams(params, tier, forecastModels(tiers, tier)),
|
||||
value.tier_model_params,
|
||||
),
|
||||
plan_mode_min_tier:
|
||||
value.plan_mode_min_tier && forecastModels(tiers, value.plan_mode_min_tier).length > 0
|
||||
? value.plan_mode_min_tier
|
||||
: undefined,
|
||||
custom_tier_set: undefined,
|
||||
enable_non_reasoning_tier: false,
|
||||
classification_prompt: undefined,
|
||||
classification_examples: undefined,
|
||||
classifier_fallback: undefined,
|
||||
classifier_llm_config: value.classifier_llm_config && withoutForecastPromptOverrides(value.classifier_llm_config),
|
||||
...(value.classifier_type === "llm_v2" && { adaptive: false }),
|
||||
};
|
||||
};
|
||||
|
||||
export const getForecastConfigError = (value: ComplexityRouterConfigValue): string | null => {
|
||||
if (!isForecastClassifier(value.classifier_type) || value.custom_tier_set) return null;
|
||||
const timeout = value.classifier_llm_config?.timeout_ms;
|
||||
if (timeout !== undefined && (!Number.isInteger(timeout) || timeout <= 0))
|
||||
return "Enter a positive whole-number classifier timeout";
|
||||
const [efficient, capable] = forecastTierNames(value);
|
||||
const order: readonly string[] = tierOrderFor(value.enable_non_reasoning_tier);
|
||||
if (!order.includes(efficient) || !order.includes(capable) || order.indexOf(capable) <= order.indexOf(efficient))
|
||||
return "The capable tier must be higher than the efficient tier";
|
||||
const efficientModels = forecastModels(value.tiers, efficient);
|
||||
const capableModels = forecastModels(value.tiers, capable);
|
||||
if (!efficientModels.length || !capableModels.length)
|
||||
return "Select models for both the efficient and capable solvers";
|
||||
if (value.classifier_type === "capability") {
|
||||
const result = capabilitySettingsSchema.safeParse(value.capability_classifier_config);
|
||||
if (!result.success)
|
||||
return "Enter a solve threshold between 0 and 1 and valid capability settings, including any calibration coefficients";
|
||||
if (result.data.base_threshold + 2 * (result.data.threshold_step ?? 0) > 1)
|
||||
return "The solve threshold plus twice the boundary step must be at most 1";
|
||||
return null;
|
||||
}
|
||||
return getFuseConfigError(value, efficient, capable);
|
||||
};
|
||||
|
||||
const getFuseConfigError = (value: ComplexityRouterConfigValue, efficient: string, capable: string): string | null => {
|
||||
const efficientModels = forecastModels(value.tiers, efficient);
|
||||
const capableModels = forecastModels(value.tiers, capable);
|
||||
if (value.adaptive) return "Turn off adaptive routing for Fuse v2";
|
||||
if (value.enable_non_reasoning_tier) return "Fuse v2 does not support the non-reasoning tier";
|
||||
if (efficientModels.length !== 1 || capableModels.length !== 1 || efficientModels[0] === capableModels[0])
|
||||
return "Fuse v2 requires one distinct model group for each solver";
|
||||
if (Object.entries(value.tiers).some(([tier, models]) => models.length > 0 && ![efficient, capable].includes(tier)))
|
||||
return "Fuse v2 supports only its efficient and capable tiers";
|
||||
const result = fuseSettingsSchema.safeParse(value.llm_v2_config);
|
||||
if (!result.success)
|
||||
return "Complete both solver profiles, the harness, and a quality gap between 0 and 1; any calibration needs valid coefficients and a version";
|
||||
return null;
|
||||
};
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import type { ComplexityTiers } from "./ComplexityRouterConfig";
|
||||
import { isForecastClassifier } from "./forecast_classifier_config";
|
||||
import type { ClassifierType, ComplexityTiers } from "./ComplexityRouterConfig";
|
||||
import type { ComplexityTier } from "./KeywordTierRules";
|
||||
import type { TierModelParams, TierModelParamsByTier } from "./complexity_router_tiers";
|
||||
|
||||
|
|
@ -32,6 +33,7 @@ export const MAX_TIER_NAME_CHARS = 64;
|
|||
export const MAX_TIER_DEFINITION_CHARS = 500;
|
||||
|
||||
export interface ActiveTierSet {
|
||||
classifier_type?: ClassifierType;
|
||||
tiers: ComplexityTiers;
|
||||
enable_non_reasoning_tier?: boolean;
|
||||
custom_tier_set?: CustomTierSet;
|
||||
|
|
@ -62,7 +64,12 @@ export const activeTierRows = (value: ActiveTierSet): ActiveTierRow[] => {
|
|||
const rows =
|
||||
value.custom_tier_set?.tiers ??
|
||||
tierOrderFor(value.enable_non_reasoning_tier).map((tier) => builtInRow(tier, value.tiers));
|
||||
return rows.map((row) => ({ ...row, params: value.tier_model_params?.[row.id] ?? {} }));
|
||||
return rows
|
||||
.filter(
|
||||
(row) =>
|
||||
value.custom_tier_set || !isForecastClassifier(value.classifier_type ?? "heuristic") || row.models.length > 0,
|
||||
)
|
||||
.map((row) => ({ ...row, params: value.tier_model_params?.[row.id] ?? {} }));
|
||||
};
|
||||
|
||||
// The wire shape of an edited tier set, shared by the payload builder and the prompt preview so the
|
||||
|
|
|
|||
|
|
@ -23,6 +23,35 @@ const apply = (value: ComplexityRouterConfigValue, action: Parameters<typeof app
|
|||
applyTierSetAction(value, rules, action);
|
||||
|
||||
describe("applyTierSetAction", () => {
|
||||
it.each(["llm", "capability", "llm_v2"] as const)(
|
||||
"clears an emptied plan floor and can repopulate its pool for %s",
|
||||
(classifier_type) => {
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
...builtIn,
|
||||
classifier_type,
|
||||
plan_mode_min_tier: "SIMPLE",
|
||||
tier_model_params: { SIMPLE: { "gpt-3.5-turbo": { speed: "fast" } } },
|
||||
};
|
||||
const { value: cleared } = apply(initial, { kind: "models", id: "SIMPLE", models: [] });
|
||||
expect(cleared.plan_mode_min_tier).toBeUndefined();
|
||||
expect(cleared.tier_model_params).toBeUndefined();
|
||||
expect(cleared.tiers.SIMPLE).toEqual([]);
|
||||
const { value: restored } = apply(cleared, { kind: "models", id: "SIMPLE", models: ["gpt-4"] });
|
||||
expect(restored.tiers.SIMPLE).toEqual(["gpt-4"]);
|
||||
expect(restored.plan_mode_min_tier).toBeUndefined();
|
||||
expect(restored.tier_model_params).toBeUndefined();
|
||||
},
|
||||
);
|
||||
|
||||
it("preserves a populated custom plan floor after removing only one model", () => {
|
||||
const initial: ComplexityRouterConfigValue = { ...custom, plan_mode_min_tier: "sec" };
|
||||
const { value: changed } = apply(initial, { kind: "models", id: "sec", models: ["new-model"] });
|
||||
expect(changed.plan_mode_min_tier).toBe("sec");
|
||||
const { value: cleared } = apply(changed, { kind: "models", id: "sec", models: [] });
|
||||
expect(cleared.plan_mode_min_tier).toBeUndefined();
|
||||
expect(cleared.custom_tier_set?.tiers.find((row) => row.id === "sec")?.models).toEqual([]);
|
||||
});
|
||||
|
||||
it("adds a row and moves the form into an edited set, which the built-in record never leaves", () => {
|
||||
const { value } = apply(builtIn, { kind: "add" });
|
||||
expect(value.custom_tier_set?.tiers).toHaveLength(5);
|
||||
|
|
|
|||
|
|
@ -93,6 +93,30 @@ const exitToBuiltInTiers = (value: ComplexityRouterConfigValue, rows: readonly A
|
|||
return commitTierRows(activeTierRows(restored), "", restored);
|
||||
};
|
||||
|
||||
/** Model changes reconcile params and the plan floor against the resulting populated pools. */
|
||||
export const setTierModels = (
|
||||
value: ComplexityRouterConfigValue,
|
||||
id: string,
|
||||
models: string[],
|
||||
): ComplexityRouterConfigValue => {
|
||||
const next: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
...(value.custom_tier_set
|
||||
? {
|
||||
custom_tier_set: {
|
||||
...value.custom_tier_set,
|
||||
tiers: value.custom_tier_set.tiers.map((row) => (row.id === id ? { ...row, models } : row)),
|
||||
},
|
||||
}
|
||||
: { tiers: { ...value.tiers, [id]: models } }),
|
||||
tier_model_params: pruneTierModelParams(value.tier_model_params, id, models),
|
||||
};
|
||||
const floor = next.plan_mode_min_tier;
|
||||
return floor && !activeTierRows(next).some((row) => row.id === floor && row.models.length > 0)
|
||||
? { ...next, plan_mode_min_tier: undefined }
|
||||
: next;
|
||||
};
|
||||
|
||||
const nextTierSetValue = (
|
||||
value: ComplexityRouterConfigValue,
|
||||
rows: ActiveTierRow[],
|
||||
|
|
@ -102,11 +126,7 @@ const nextTierSetValue = (
|
|||
|
||||
switch (action.kind) {
|
||||
case "models":
|
||||
return commitTierRows(
|
||||
rows.map((row) => (row.id === action.id ? { ...row, models: action.models } : row)),
|
||||
fallbackId,
|
||||
{ ...value, tier_model_params: pruneTierModelParams(value.tier_model_params, action.id, action.models) },
|
||||
);
|
||||
return setTierModels(value, action.id, action.models);
|
||||
case "patch":
|
||||
return commitTierRows(
|
||||
rows.map((row) => (row.id === action.id ? { ...row, ...action.patch } : row)),
|
||||
|
|
|
|||
|
|
@ -46,6 +46,43 @@ const hydratedState: KeywordMatchingState = {
|
|||
};
|
||||
|
||||
describe("buildUpdatedComplexityRouterConfig keyword matching", () => {
|
||||
it.each(["capability", "llm_v2", "heuristic"] as const)(
|
||||
"handles enabled stored overrides when editing %s with or without keyword form state",
|
||||
(classifier_type) => {
|
||||
const stored = {
|
||||
...STORED,
|
||||
classifier_type,
|
||||
adaptive: classifier_type !== "llm_v2",
|
||||
adaptive_weights: { quality: 0.6, cost: 0.4 },
|
||||
adaptive_eligible: "all",
|
||||
tier_distance_penalty: 0.8,
|
||||
enable_context_window_escalation: true,
|
||||
context_window_escalation_buffer: 0.9,
|
||||
};
|
||||
const value = hydrateComplexityRouterConfig(stored, undefined);
|
||||
for (const keywordState of [undefined, hydratedState]) {
|
||||
const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, keywordState);
|
||||
const forecast = classifier_type !== "heuristic";
|
||||
expect(saved.adaptive).toBe(!forecast);
|
||||
expect(saved.enable_context_window_escalation).toBe(!forecast);
|
||||
expect(saved.escalation_keywords).toEqual(forecast ? [] : stored.escalation_keywords);
|
||||
for (const key of [
|
||||
"adaptive_weights",
|
||||
"adaptive_eligible",
|
||||
"tier_distance_penalty",
|
||||
"context_window_escalation_buffer",
|
||||
]) {
|
||||
expect(Object.hasOwn(saved, key)).toBe(!forecast);
|
||||
}
|
||||
expect(saved.keyword_tier_rules).toEqual(STORED.keyword_tier_rules);
|
||||
expect(saved.semantic_keyword_matching).toBe(true);
|
||||
expect(saved.some_future_backend_key).toEqual(STORED.some_future_backend_key);
|
||||
}
|
||||
expect(value.enable_context_window_escalation).toBe(true);
|
||||
expect(stored.escalation_keywords).toEqual(["urgent", "outage"]);
|
||||
},
|
||||
);
|
||||
|
||||
it("round-trips an untouched edit without changing any keyword-matching value", () => {
|
||||
// Opening the modal hydrates state from STORED; saving with nothing changed must be a
|
||||
// no-op. These keys are now MANAGED, so a hydration bug silently wipes them.
|
||||
|
|
@ -675,7 +712,11 @@ describe("managed keys survive an untouched open-and-save", () => {
|
|||
|
||||
// The opt-in fifth tier requires the LLM classifier, which this heuristic_first fixture is not,
|
||||
// so it gets its own round trip below.
|
||||
const KEYS_ANOTHER_TIER_LADDER_OWNS = new Set(["enable_non_reasoning_tier"]);
|
||||
const KEYS_ANOTHER_TIER_LADDER_OWNS = new Set([
|
||||
"capability_classifier_config",
|
||||
"llm_v2_config",
|
||||
"enable_non_reasoning_tier",
|
||||
]);
|
||||
|
||||
it("carries every managed key a built-in router can hold through hydrate then save", () => {
|
||||
const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined);
|
||||
|
|
@ -850,7 +891,12 @@ describe("LLM V2 configuration preservation", () => {
|
|||
harness: "Shell access, one attempt",
|
||||
max_quality_gap: 0.03,
|
||||
response_format: "json_object",
|
||||
calibration: { version: "pair-v1", prompt_version: "llm-v2-1" },
|
||||
calibration: {
|
||||
version: "pair-v1",
|
||||
prompt_version: "llm-v2-1",
|
||||
efficient: { slope: 1.2, intercept: -0.3 },
|
||||
capable: { slope: 0.9, intercept: 0.1 },
|
||||
},
|
||||
};
|
||||
const stored = {
|
||||
tiers: { SIMPLE: ["efficient"], REASONING: ["capable"] },
|
||||
|
|
|
|||
|
|
@ -1,3 +1,12 @@
|
|||
import AutoRouterClassifierTabs from "../add_model/AutoRouterClassifierTabs";
|
||||
import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config";
|
||||
export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config";
|
||||
import {
|
||||
getForecastConfigError,
|
||||
isForecastClassifier,
|
||||
capabilitySettingsSchema,
|
||||
fuseSettingsSchema,
|
||||
} from "../add_model/forecast_classifier_config";
|
||||
import React, { useEffect, useMemo, useState } from "react";
|
||||
import {
|
||||
complexityRouterSchema,
|
||||
|
|
@ -62,13 +71,8 @@ import {
|
|||
hydrateTokenThresholds,
|
||||
} from "../add_model/heuristic_scoring_knobs";
|
||||
import ComplexityRouterConfig, {
|
||||
AdaptiveEligible,
|
||||
AdaptiveRouterWeights,
|
||||
ClassifierLLMConfig,
|
||||
ClassifierType,
|
||||
effectiveClassifierType,
|
||||
ComplexityRouterConfigValue,
|
||||
ComplexityTiers,
|
||||
effectiveClassifierType,
|
||||
heuristicScoringRole,
|
||||
DEFAULT_ADAPTIVE_WEIGHTS,
|
||||
DEFAULT_SESSION_AFFINITY,
|
||||
|
|
@ -97,47 +101,6 @@ interface EditAutoRouterModalProps {
|
|||
// Keys this modal rewrites from its own form state on save. Anything absent from this set is
|
||||
// carried through untouched from the stored config, so a key only belongs here once the modal
|
||||
// actually renders a control that can set it.
|
||||
/** The complexity_router_config as it comes back from the proxy, before any hydration. Fields the
|
||||
* hydrators validate themselves stay `unknown`; the ones assigned straight through carry their type. */
|
||||
export interface StoredComplexityRouterConfig {
|
||||
tiers?: Partial<Record<keyof ComplexityTiers, unknown>>;
|
||||
enable_non_reasoning_tier?: boolean;
|
||||
tier_model_configs?: unknown;
|
||||
default_model?: string | null;
|
||||
plan_mode_min_tier?: unknown;
|
||||
classification_prompt?: unknown;
|
||||
classification_examples?: unknown;
|
||||
heuristic_first_max_tier?: unknown;
|
||||
hybrid_boundary_margin?: unknown;
|
||||
tier_labels?: unknown;
|
||||
classifier_type?: ClassifierType;
|
||||
classifier_llm_config?: ClassifierLLMConfig;
|
||||
classifier_context_window_size?: unknown;
|
||||
classifier_context_budget_chars?: unknown;
|
||||
classifier_context_include_assistant_turns?: unknown;
|
||||
classifier_fallback?: unknown;
|
||||
classification_mode?: unknown;
|
||||
tier_boundaries?: unknown;
|
||||
token_thresholds?: unknown;
|
||||
dimension_weights?: unknown;
|
||||
custom_dimensions?: unknown;
|
||||
reasoning_override_min_score?: unknown;
|
||||
session_affinity?: unknown;
|
||||
session_affinity_ttl_seconds?: unknown;
|
||||
modality_routing?: unknown;
|
||||
modality_pin_override?: unknown;
|
||||
deployment_affinity?: unknown;
|
||||
adaptive?: boolean;
|
||||
adaptive_weights?: AdaptiveRouterWeights;
|
||||
tier_distance_penalty?: number;
|
||||
adaptive_eligible?: AdaptiveEligible;
|
||||
return_raw_model_name?: boolean;
|
||||
enable_context_window_escalation?: unknown;
|
||||
context_window_escalation_buffer?: unknown;
|
||||
stall_escalation_enabled?: unknown;
|
||||
stall_escalation_window?: unknown;
|
||||
stall_escalation_repeat_threshold?: unknown;
|
||||
}
|
||||
|
||||
/**
|
||||
* The stored complexity_router_config as form state. Every key in MANAGED_COMPLEXITY_ROUTER_KEYS is
|
||||
|
|
@ -164,6 +127,8 @@ export const hydrateComplexityRouterConfig = (
|
|||
plan_mode_min_tier: hydratePlanModeMinTier(parsedConfig.plan_mode_min_tier, custom_tier_set),
|
||||
tier_labels: hydrateTierLabels(parsedConfig.tier_labels),
|
||||
classifier_type: parsedConfig.classifier_type || "heuristic",
|
||||
capability_classifier_config: capabilitySettingsSchema.safeParse(parsedConfig.capability_classifier_config).data,
|
||||
llm_v2_config: fuseSettingsSchema.safeParse(parsedConfig.llm_v2_config).data,
|
||||
classifier_llm_config: parsedConfig.classifier_llm_config,
|
||||
classifier_context_window_size:
|
||||
typeof parsedConfig.classifier_context_window_size === "number"
|
||||
|
|
@ -251,6 +216,8 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
|||
"plan_mode_min_tier",
|
||||
"tier_labels",
|
||||
"classifier_type",
|
||||
"capability_classifier_config",
|
||||
"llm_v2_config",
|
||||
"classifier_llm_config",
|
||||
"classifier_context_window_size",
|
||||
"classifier_context_budget_chars",
|
||||
|
|
@ -339,8 +306,8 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
keywordMatching?: KeywordMatchingState,
|
||||
): Record<string, unknown> => {
|
||||
const isManaged = (key: string): boolean => {
|
||||
if (key === "llm_v2_config" && effectiveClassifierType(value) !== "llm_v2") return true;
|
||||
if (MANAGED_COMPLEXITY_ROUTER_KEYS.has(key)) return true;
|
||||
if (key === "escalation_keywords" && isForecastClassifier(effectiveClassifierType(value))) return true;
|
||||
if (keywordMatching !== undefined && KEYWORD_MATCHING_KEYS.has(key)) return true;
|
||||
return customTechnicalKeywords !== undefined && key === "custom_technical_keywords";
|
||||
};
|
||||
|
|
@ -362,6 +329,8 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
classificationMode: value.classification_mode,
|
||||
tierLabels: value.tier_labels,
|
||||
classifierType: value.classifier_type,
|
||||
capabilityClassifierConfig: value.capability_classifier_config,
|
||||
llmV2Config: value.llm_v2_config,
|
||||
classifierLlmConfig: value.classifier_llm_config,
|
||||
classifierContextWindowSize: value.classifier_context_window_size,
|
||||
classifierContextBudgetChars: value.classifier_context_budget_chars,
|
||||
|
|
@ -399,7 +368,7 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
|
||||
// Keys this call does not own stay as the stored config left them.
|
||||
const unowned: readonly string[] = [
|
||||
...(keywordMatching === undefined ? KEYWORD_MATCHING_KEYS : []),
|
||||
...(keywordMatching === undefined ? [...KEYWORD_MATCHING_KEYS].filter((key) => !isManaged(key)) : []),
|
||||
...(customTechnicalKeywords === undefined ? ["custom_technical_keywords"] : []),
|
||||
];
|
||||
return {
|
||||
|
|
@ -458,6 +427,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
getPlanModeTierError(complexityRouterConfig.plan_mode_min_tier, activeTierRows(complexityRouterConfig)) ??
|
||||
getKeywordTierRulesError(keywordTierRules, activeTierRows(complexityRouterConfig)) ??
|
||||
getClassifierModelError(complexityRouterConfig) ??
|
||||
getForecastConfigError(complexityRouterConfig) ??
|
||||
(heuristicScoringRole(complexityRouterConfig) === "decides"
|
||||
? customDimensionsError(complexityRouterConfig.custom_dimensions)
|
||||
: null);
|
||||
|
|
@ -589,6 +559,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
}
|
||||
const classifierError =
|
||||
getClassifierModelError(complexityRouterConfig) ??
|
||||
getForecastConfigError(complexityRouterConfig) ??
|
||||
(heuristicScoringRole(complexityRouterConfig) === "decides"
|
||||
? customDimensionsError(complexityRouterConfig.custom_dimensions)
|
||||
: null);
|
||||
|
|
@ -742,6 +713,14 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
{ value: "custom", label: "Enter custom model name" },
|
||||
];
|
||||
|
||||
const routerNameField = (
|
||||
<FormField control={form.control} name="auto_router_name" label="Auto Router Name">
|
||||
{({ ref, ...field }) => (
|
||||
<Input {...field} ref={ref} readOnly={isMemberManaged} placeholder="e.g., auto_router_1, smart_routing" />
|
||||
)}
|
||||
</FormField>
|
||||
);
|
||||
|
||||
return (
|
||||
<Dialog open={isVisible} onOpenChange={(open) => !open && onCancel()}>
|
||||
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
|
||||
|
|
@ -755,48 +734,41 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
|
||||
<form onSubmit={(event) => event.preventDefault()} noValidate>
|
||||
<FieldGroup>
|
||||
<FormField control={form.control} name="auto_router_name" label="Auto Router Name">
|
||||
{({ ref, ...field }) => (
|
||||
<Input
|
||||
{...field}
|
||||
ref={ref}
|
||||
readOnly={isMemberManaged}
|
||||
placeholder="e.g., auto_router_1, smart_routing"
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
{routerNameField}
|
||||
|
||||
{isComplexityRouterModel ? (
|
||||
/* Complexity Router Configuration */
|
||||
<div className="w-full">
|
||||
<ComplexityRouterConfig
|
||||
editingTiers={editingTiers}
|
||||
onEditingTiersChange={setEditingTiers}
|
||||
showValidationErrors={showValidationErrors}
|
||||
modelInfo={modelInfo}
|
||||
value={complexityRouterConfig}
|
||||
onChange={(config) => {
|
||||
setComplexityRouterConfig(config);
|
||||
}}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
|
||||
keywordTierRules={keywordTierRules}
|
||||
onKeywordTierRulesChange={setKeywordTierRules}
|
||||
keywordRulesError={getKeywordTierRulesError(
|
||||
keywordTierRules,
|
||||
activeTierRows(complexityRouterConfig),
|
||||
)}
|
||||
semanticMatchingEnabled={semanticMatchingEnabled}
|
||||
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={setEmbeddingModel}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
|
||||
/>
|
||||
<AutoRouterClassifierTabs value={complexityRouterConfig} onChange={setComplexityRouterConfig}>
|
||||
<ComplexityRouterConfig
|
||||
editingTiers={editingTiers}
|
||||
onEditingTiersChange={setEditingTiers}
|
||||
showValidationErrors={showValidationErrors}
|
||||
modelInfo={modelInfo}
|
||||
value={complexityRouterConfig}
|
||||
onChange={(config) => {
|
||||
setComplexityRouterConfig(config);
|
||||
}}
|
||||
customTechnicalKeywords={customTechnicalKeywords}
|
||||
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
|
||||
keywordTierRules={keywordTierRules}
|
||||
onKeywordTierRulesChange={setKeywordTierRules}
|
||||
keywordRulesError={getKeywordTierRulesError(
|
||||
keywordTierRules,
|
||||
activeTierRows(complexityRouterConfig),
|
||||
)}
|
||||
semanticMatchingEnabled={semanticMatchingEnabled}
|
||||
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
|
||||
embeddingModel={embeddingModel}
|
||||
onEmbeddingModelChange={setEmbeddingModel}
|
||||
matchThreshold={matchThreshold}
|
||||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
|
||||
/>
|
||||
</AutoRouterClassifierTabs>
|
||||
</div>
|
||||
) : (
|
||||
<>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue