Merge branch 'litellm_internal_staging' into litellm_lit_7022_azure_ai_passthrough_config

This commit is contained in:
mateo-berri 2026-09-09 19:35:55 -07:00
commit a6681950e8
125 changed files with 8392 additions and 1110 deletions

View file

@ -0,0 +1,15 @@
-- DropForeignKey
DO $$
BEGIN
IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
END IF;
END $$;
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END $$;

View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

View file

@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
provider_specific_fields: Mapping[str, object]
def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
string payload in ``input`` rather than ``arguments``; both map to
@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Tool calls accumulate into the single trailing tool_calls choice
# like the typed branches above; a choice per call would hide every
# call after choices[0] from chat clients
accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
tool_call_index += 1
elif handle_raw_dict_callback is not None:
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
# New output item added
output_item = parsed_chunk.get("item", {})
if output_item.get("type") in ("function_call", "custom_tool_call"):
converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
provider_specific_fields: Final = converted.get("provider_specific_fields")
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
index=0,
delta=Delta(
tool_calls=(
_tool_call_dict_from_output_item(
tool_call_dict_from_output_item(
output_item, parsed_chunk.get("output_index", 0)
),
)

View file

@ -398,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
minimum=1,
maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
)
TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range(
"TOKEN_COUNTER_MAX_EXACT_CHARS",
default=4_000_000,
minimum=1,
maximum=1_000_000_000,
)
TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range(
"TOKEN_COUNTER_MAX_CONCURRENT_COUNTS",
default=4,
minimum=1,
maximum=256,
)
MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))

View file

@ -2,6 +2,7 @@
# On success, logs events to Langfuse
import inspect
import os
import re
import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, dict) else None
def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]:
"""Header pairs with the key type widened back to what a caller-supplied dict can actually hold."""
return mapping.items()
def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool:
if not isinstance(trace_id, str) or not isinstance(session_id, str):
return False
request: Final = _object_mapping(proxy_server_request)
raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None
if raw_headers is None:
return False
headers: Final = MappingProxyType(
{key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)}
)
if headers.get("x-litellm-trace-id"):
return False
if headers.get("langfuse_trace_id") is not None:
return False
if trace_id != session_id and headers.get("langfuse_session_id") != session_id:
return False
if headers.get("x-litellm-session-id") == trace_id:
return True
if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None:
return False
user_agent: Final = headers.get("user-agent")
codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None
return any(
value == trace_id
and (
key == "x-session-id"
or re.fullmatch(r"x-.+-session-id", key) is not None
or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id"))
)
for key, value in headers.items()
)
class _UsageObject(Protocol):
"""Token-count surface the Langfuse logger reads off a response usage payload."""
@ -609,6 +648,18 @@ class LangFuseLogger:
# This allows continuing an existing trace while still returning the correct trace_id
if existing_trace_id is not None:
trace_id = existing_trace_id
resolved_trace_id: Final = (
litellm_call_id or trace_id
if existing_trace_id is None
and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request"))
else trace_id
)
if resolved_trace_id != trace_id:
verbose_logger.debug(
"Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace",
trace_id,
resolved_trace_id,
)
requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
update_trace_keys: Final = (
requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else ()
@ -663,7 +714,7 @@ class LangFuseLogger:
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
else: # don't overwrite an existing trace
trace_params = {
"id": trace_id,
"id": resolved_trace_id,
"name": trace_name,
"session_id": session_id,
"input": masked_input if not mask_input else "redacted-by-litellm",
@ -845,13 +896,13 @@ class LangFuseLogger:
# Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value
# to match expected test behavior
if hasattr(generation_client, "trace_id") and generation_client.trace_id:
if generation_client.trace_id != trace_id:
if generation_client.trace_id != resolved_trace_id:
verbose_logger.warning(
"Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.",
trace_id,
resolved_trace_id,
generation_client.trace_id,
)
return trace_id, generation_id
return resolved_trace_id, generation_id
except Exception:
verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc())
return None, None

View file

@ -3,7 +3,7 @@ import json
import re
import time
import traceback
from collections.abc import Iterable, Sequence
from collections.abc import Mapping, Sequence
from typing import Final, Literal, cast
import litellm
@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
del choice.enhancements
def _invalid_choices_message(response_object: Mapping[str, object]) -> str:
raw_keys: Final = list(response_object.keys())
if "choices" not in response_object:
return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}"
return (
f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). "
f"Raw keys: {raw_keys}"
)
async def convert_to_streaming_response_async(
response_object: dict | None = None,
):
@ -179,14 +189,12 @@ async def convert_to_streaming_response_async(
choice_list: Final[list[StreamingChoices]] = []
if not response_object.get("choices"):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)
@ -287,14 +295,12 @@ def convert_to_streaming_response(
model_response_object: Final = ModelResponseStream()
choice_list: Final[list[StreamingChoices]] = []
if not response_object.get("choices"):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)
@ -623,15 +629,12 @@ def convert_to_model_response_object(
return convert_to_streaming_response(response_object=response_object)
choice_list: Final[list[Choices]] = []
if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
"LiteLLM: provider returned a response with no 'choices'. "
f"Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)

View file

@ -1473,17 +1473,14 @@ class CustomStreamWrapper:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "cached_response":
cached_chunk: Final = cast(ModelResponseStream, chunk)
chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason
cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None
chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None
response_obj = {
"text": cached_chunk.choices[0].delta.content,
"text": cached_choice.delta.content if cached_choice is not None else None,
"is_finished": chunk_finish_reason is not None,
"finish_reason": chunk_finish_reason,
"original_chunk": cached_chunk,
"tool_calls": (
cached_chunk.choices[0].delta.tool_calls
if hasattr(cached_chunk.choices[0].delta, "tool_calls")
else None
),
"tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None),
}
completion_obj["content"] = response_obj["text"]

View file

@ -3,11 +3,15 @@
import base64
import io
import struct
from collections.abc import Callable, Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from typing import Final, Literal, cast
import anyio
import anyio.lowlevel
import httpx
import tiktoken
from tokenizers import Tokenizer
from typing_extensions import ParamSpec, TypeVar
import litellm
from litellm import verbose_logger
@ -21,7 +25,10 @@ from litellm.constants import (
MAX_TILE_HEIGHT,
MAX_TILE_WIDTH,
TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
TOKEN_COUNTER_MAX_CONCURRENT_COUNTS,
TOKEN_COUNTER_MAX_EXACT_CHARS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@ -324,6 +331,32 @@ TokenCounterFunction = Callable[[str], int]
Type for a function that counts tokens in a string.
"""
EXTRAPOLATION_SAMPLES: Final = 16
T_ParamSpec: Final = ParamSpec("T_ParamSpec")
T_Retval = TypeVar("T_Retval")
_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter")
def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter:
existing: Final = _COUNT_OFFLOAD_LIMITER.get(None)
if existing is not None:
return existing
created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS)
_COUNT_OFFLOAD_LIMITER.set(created)
return created
def offload_token_count(
function: Callable[T_ParamSpec, T_Retval],
) -> Callable[T_ParamSpec, Awaitable[T_Retval]]:
async def offloaded(
*args: T_ParamSpec.args,
**kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract
) -> T_Retval:
return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs)
return offloaded
def _get_tiktoken_count_function(
encode_length: Callable[[str], int],
@ -545,9 +578,40 @@ def _count_extra(
return num_tokens
def _get_extrapolating_count_function(
count_exactly: TokenCounterFunction,
max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS,
) -> TokenCounterFunction:
def count_tokens(text: str) -> int:
if len(text) <= max_exact_chars:
return count_exactly(text)
samples: Final = _evenly_spaced_samples(text, max_exact_chars)
sampled_chars: Final = sum(len(sample) for sample in samples)
return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars)
return count_tokens
def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]:
sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars)
sample_chars: Final = total_chars // sample_count
last_start: Final = len(text) - sample_chars
return tuple(
text[start : start + sample_chars]
for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count))
)
def _get_count_function(
model: str | None,
custom_tokenizer: dict | SelectTokenizerResponse | None = None,
) -> TokenCounterFunction:
return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer))
def _get_exact_count_function(
model: str | None,
custom_tokenizer: dict | SelectTokenizerResponse | None = None,
) -> TokenCounterFunction:
"""
Get the function to count tokens based on the model and custom tokenizer."""
@ -556,10 +620,10 @@ def _get_count_function(
if model is not None or custom_tokenizer is not None:
tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
if tokenizer_json["type"] == "huggingface_tokenizer":
tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"]
def count_tokens(text: str) -> int:
enc: Final = tokenizer_json["tokenizer"].encode(text)
return len(enc.ids)
return len(tokenizer.encode_batch_fast([text])[0])
return count_tokens
elif tokenizer_json["type"] == "openai_tokenizer":

View file

@ -213,11 +213,17 @@ class AnthropicMessagesHandler(BaseTranslation):
"""
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
def post_call_hook_response(self, response: object) -> object:
if not isinstance(response, ModelResponse):
return response
return self.adapter.translate_openai_response_to_anthropic(response)
@staticmethod
def _build_streaming_usage_response(
responses_so_far: Sequence[object],

View file

@ -1487,8 +1487,9 @@ class LiteLLMAnthropicMessagesAdapter:
anthropic_content.insert(0, polyfill_result.compaction_block)
## extract finish reason
openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop"
translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic(
openai_finish_reason=response.choices[0].finish_reason
openai_finish_reason=openai_finish_reason
)
anthropic_finish_reason: Final = (
"refusal"

View file

@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import BaseOpenAILLM
from litellm.secret_managers.get_azure_ad_token_provider import (
@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM):
if scope is None:
scope = "https://cognitiveservices.azure.com/.default"
max_retries: Final = litellm_params.get("max_retries")
configured_max_retries: Final = litellm_params.get("max_retries")
max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries
timeout: Final = litellm_params.get("timeout")
if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret:
verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth")
@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM):
else:
azure_client_params["http_client"] = self._get_sync_http_client()
if max_retries is not None:
azure_client_params["max_retries"] = max_retries
azure_client_params["max_retries"] = max_retries
if timeout is not None:
azure_client_params["timeout"] = timeout

View file

@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
"""
Get the appropriate image edit config for an Azure AI model.
- MAI models use /mai/v1/images/edits with multipart form data and size
- MAI models use /mai/v1/images/edits with multipart form data
- FLUX 2 models use JSON with base64 image
- FLUX 1 models use multipart/form-data
"""

View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final
import httpx
from httpx._types import RequestFiles
@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import (
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.llms.openai import FileTypes
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
@ -26,65 +25,8 @@ if TYPE_CHECKING:
class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig):
"""Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5)."""
DEFAULT_SIZE = "1024x1024"
def get_supported_openai_params(self, model: str) -> list:
return ["prompt", "image", "model", "n", "size"]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
optional_params: Final[dict[str, Any]] = {}
supported_params: Final = self.get_supported_openai_params(model)
for key, value in dict(image_edit_optional_params).items():
if value is None or key in optional_params:
continue
if key in supported_params:
if key == "size" and value:
size_param = cast(str, value)
self._validate_size_param(size_param)
optional_params[key] = size_param
else:
optional_params[key] = value
elif not drop_params:
raise ValueError(
f"Parameter {key} is not supported for model {model}. "
f"Supported parameters are {supported_params}. "
f"Set drop_params=True to drop unsupported parameters."
)
if "size" not in optional_params:
optional_params["size"] = self.DEFAULT_SIZE
return optional_params
def _validate_size_param(self, size: str) -> None:
known_sizes: Final = {
"1024x1024",
"1792x1024",
"1024x1792",
"512x512",
"256x256",
}
if size in known_sizes:
return
if "x" in size:
try:
tuple(map(int, size.lower().split("x", 1)))
return
except ValueError:
raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
raise ValueError(
f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
)
return ["prompt", "image", "model", "n"]
def validate_environment(
self,

View file

@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.exceptions import UnsupportedParamsError
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
@ -21,6 +22,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
MAX_IMAGES_PER_REQUEST: Final = 1
MIN_DIMENSION_PX: Final = 768
MAX_TOTAL_PX: Final = 1_056_768
@staticmethod
def get_mai_image_generation_url(
api_base: str | None,
@ -145,16 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
if k in supported_params:
if k == "size" and v:
self._map_size_param(v, optional_params)
self._map_size_param(v, optional_params, model)
elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST:
if not drop_params:
raise self._unsupported(
model,
f"n={v} is not supported for model {model}. The Azure AI MAI image "
f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per "
"request and ignores any count, so a larger value would silently "
"return fewer images than requested. Send one request per image, or "
"set drop_params=True to drop n.",
)
else:
optional_params[k] = v
elif k in ("width", "height"):
optional_params[k] = v
elif not drop_params:
raise ValueError(
raise self._unsupported(
model,
f"Parameter {k} is not supported for model {model}. "
f"Supported parameters are {supported_params} and width/height. "
f"Set drop_params=True to drop unsupported parameters."
f"Set drop_params=True to drop unsupported parameters.",
)
if "width" not in optional_params:
@ -165,7 +181,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
optional_params.pop("size", None)
return optional_params
def _map_size_param(self, size: str, optional_params: dict) -> None:
@staticmethod
def _unsupported(model: str, message: str) -> UnsupportedParamsError:
return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model)
def _image_count(self, n: object, model: str) -> int:
if isinstance(n, int):
return n
try:
return int(str(n))
except ValueError:
raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.")
def _map_size_param(self, size: str, optional_params: dict, model: str) -> None:
size_mapping: Final = {
"1024x1024": (1024, 1024),
"1792x1024": (1792, 1024),
@ -176,19 +204,36 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
if size in size_mapping:
width, height = size_mapping[size]
optional_params["width"] = width
optional_params["height"] = height
elif "x" in size:
try:
width, height = map(int, size.lower().split("x"))
optional_params["width"] = width
optional_params["height"] = height
except ValueError:
raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
raise self._unsupported(
model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
)
else:
raise ValueError(
raise self._unsupported(
model,
f"Unsupported size value: '{size}'. "
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.",
)
self._validate_dimensions(model=model, size=size, width=width, height=height)
optional_params["width"] = width
optional_params["height"] = height
def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None:
if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX:
raise self._unsupported(
model,
f"Unsupported size value: '{size}'. Azure AI MAI image models require width and "
f"height of at least {self.MIN_DIMENSION_PX} pixels.",
)
if width * height > self.MAX_TOTAL_PX:
raise self._unsupported(
model,
f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most "
f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).",
)
def transform_image_generation_response(

View file

@ -61,6 +61,20 @@ class BaseTranslation(ABC):
on every other translation are undeliverable: the pipeline executor
discards them and releases the original chunks."""
assembles_streamed_response: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` stores the assembled response of an
ended stream under ``request_data["response"]`` before scanning it, the way the chat,
Responses, and Messages translations do. A streaming pipeline runs a guardrail that only
has the legacy post-call hook against that response, so on a translation without it such
a guardrail keeps running on its own."""
def post_call_hook_response(self, response: object) -> object:
"""The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
the object the translation stores under ``request_data["response"]`` while scanning an
ended stream. Chat and Responses scan that shape already; a translation that scans a
different one (Messages scans an OpenAI-shaped ModelResponse) overrides this."""
return response
@staticmethod
def transform_user_api_key_dict_to_metadata(
user_api_key_dict: Any | None,

View file

@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse`
import copy
import json
import re
import time
import types
from collections.abc import Mapping
@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig):
llm_provider="bedrock",
)
@staticmethod
def _is_openai_gpt_reasoning_model(model: str) -> bool:
return re.search(r"openai\.gpt-\d", model) is not None
def _is_nova_2_model(self, model: str) -> bool:
"""
Check if the model is a Nova 2 model that supports reasoningConfig.
@ -422,15 +427,15 @@ class AmazonConverseConfig(BaseConfig):
"""
Handle the reasoning_effort parameter based on the model type.
- GPT-OSS models: passed through unchanged via additionalModelRequestFields.
- OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
- GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields.
- OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
- Nova 2 models: transformed to reasoningConfig.
- Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on
adaptive Claude 4.6 / 4.7).
"""
if "gpt-oss" in model:
if "gpt-oss" in model or "deepseek" in model:
optional_params["reasoning_effort"] = reasoning_effort
elif "openai.gpt-5" in model:
elif self._is_openai_gpt_reasoning_model(model):
reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort}
optional_params["reasoning"] = reasoning
elif self._is_nova_2_model(model):
@ -509,6 +514,36 @@ class AmazonConverseConfig(BaseConfig):
)
thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS
def _is_deepseek_model(self, model: str, base_model: str) -> bool:
return "deepseek" in model or "deepseek" in base_model
def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool:
return "deepseek.r1" in model or "deepseek.r1" in base_model
def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool:
"""Whether the model accepts the Anthropic-shaped ``thinking`` request field.
Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons
natively: R1 returns a 400 when the field is sent and V3 silently ignores it.
"""
if self._is_deepseek_model(model=model, base_model=base_model):
return False
return (
"claude-3-7" in model
or "claude-sonnet-4" in model
or "claude-opus-4" in model
or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
)
def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool:
"""Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse.
DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw
``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts.
"""
return self._is_deepseek_r1_model(model=model, base_model=base_model)
def get_supported_openai_params(self, model: str) -> list[str]:
from litellm.utils import supports_function_calling
@ -564,23 +599,20 @@ class AmazonConverseConfig(BaseConfig):
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
supported_params.append("tool_choice")
if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model:
if (
"gpt-oss" in model
or self._is_openai_gpt_reasoning_model(model)
or self._is_openai_gpt_reasoning_model(base_model)
):
supported_params.append("reasoning_effort")
elif self._is_deepseek_model(model=model, base_model=base_model):
if not self._is_deepseek_r1_model(model=model, base_model=base_model):
supported_params.append("reasoning_effort")
elif self._is_nova_2_model(model):
# Nova 2 models support reasoning_effort (transformed to reasoningConfig)
# These models use a different reasoning structure than Anthropic's thinking parameter
supported_params.append("reasoning_effort")
elif (
"claude-3-7" in model
or "claude-sonnet-4" in model
or "claude-opus-4" in model
or "deepseek.r1" in model
or supports_reasoning(
model=model,
custom_llm_provider=self.custom_llm_provider,
)
or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
):
elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model):
supported_params.append("thinking")
supported_params.append("reasoning_effort")
supported_params.append("output_config")
@ -872,6 +904,11 @@ class AmazonConverseConfig(BaseConfig):
drop_params: bool,
) -> dict:
is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
base_model: Final = BedrockModelInfo.get_base_model(model)
drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model)
drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param(
model=model, base_model=base_model
)
for param, value in non_default_params.items():
if param == "response_format" and isinstance(value, dict):
@ -920,7 +957,12 @@ class AmazonConverseConfig(BaseConfig):
optional_params["_parallel_tool_use_config"] = {
"tool_choice": {"type": "auto", "disable_parallel_tool_use": not value}
}
if param == "thinking" and "openai.gpt-5" not in model:
if param == "thinking" and drop_thinking_param:
verbose_logger.debug(
"Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.",
model,
)
elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model):
if (
isinstance(value, dict)
and value.get("type") == "adaptive"
@ -946,6 +988,11 @@ class AmazonConverseConfig(BaseConfig):
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=model, optional_params=optional_params, custom_llm_provider="bedrock"
)
elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param:
verbose_logger.debug(
"Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.",
model,
)
elif param == "reasoning_effort" and isinstance(value, str):
self._handle_reasoning_effort_parameter(
model=model, reasoning_effort=value, optional_params=optional_params
@ -1805,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig):
data=request_data,
messages=messages,
encoding=encoding,
json_mode=json_mode,
)
def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str:
@ -2237,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig):
data: dict | str,
messages: list,
encoding,
json_mode: bool | None = None,
) -> ModelResponse:
## LOGGING
if logging_obj is not None:
@ -2247,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig):
additional_args={"complete_input_dict": data},
)
json_mode: Final[bool | None] = optional_params.get("json_mode", None)
resolved_json_mode: Final[bool | None] = (
json_mode if json_mode is not None else optional_params.get("json_mode", None)
)
## RESPONSE OBJECT
try:
completion_response: Final = ConverseResponseBlock(**response.json())
@ -2339,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks)
chat_completion_message["content"] = content_str
filtered_tools: Final = self._filter_json_mode_tools(
json_mode=json_mode,
json_mode=resolved_json_mode,
tools=tools,
chat_completion_message=chat_completion_message,
)
@ -2363,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig):
# When json_mode filtered out all synthetic tool calls the response
# is plain content, not a pending tool invocation. Fix finish_reason
# so callers (e.g. OpenAI SDK) don't misinterpret it.
if json_mode and not filtered_tools and tools:
if resolved_json_mode and not filtered_tools and tools:
initial_finish_reason = "stop"
(

View file

@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
json_mode=json_mode,
)
elif provider == "twelvelabs":
return litellm.AmazonTwelveLabsPegasusConfig().transform_response(

View file

@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
_should_convert_tool_call_to_json_mode,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation
strip_litellm_internal_message_fields,
strip_name_from_message,
)
@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.databricks import (
AllDatabricksContentValues,
DatabricksChoice,
DatabricksDelta,
DatabricksFunction,
DatabricksMessage,
DatabricksResponse,
DatabricksTool,
)
@ -247,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
litellm_params: dict,
stream: bool | None = None,
) -> str:
api_base = self._get_api_base(api_base)
complete_url: Final = f"{api_base}/chat/completions"
use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2
api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway)
url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base
complete_url: Final = f"{url_base}/chat/completions"
return complete_url
def get_supported_openai_params(self, model: str | None = None) -> list:
@ -534,6 +539,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
thinking_blocks.append(thinking_block)
return reasoning_content, thinking_blocks
@staticmethod
def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None:
return delta.get("reasoning_content")
@staticmethod
def resolve_reasoning_and_content(
message: DatabricksMessage, block_reasoning_content: str | None
) -> tuple[str | None, str | None]:
content_str: Final = DatabricksConfig.extract_content_str(message["content"])
if block_reasoning_content is not None:
return block_reasoning_content, content_str
return _extract_reasoning_content({**message, "content": content_str})
@staticmethod
def extract_citations(
content: AllDatabricksContentValues | None,
@ -577,14 +595,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
finish_reason = "stop"
if translated_message is None:
## get the content str
content_str = DatabricksConfig.extract_content_str(choice["message"]["content"])
## get the reasoning content
(
reasoning_content,
block_reasoning_content,
thinking_blocks,
) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content"))
reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content(
choice["message"], block_reasoning_content
)
citations = DatabricksConfig.extract_citations(choice["message"].get("content"))
@ -738,12 +755,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator):
# extract the reasoning content
(
reasoning_content,
block_reasoning_content,
thinking_blocks,
) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content"))
choice["delta"]["content"] = content_str
choice["delta"]["reasoning_content"] = reasoning_content
choice["delta"]["reasoning_content"] = (
block_reasoning_content
if block_reasoning_content is not None
else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"])
)
choice["delta"]["thinking_blocks"] = thinking_blocks
translated_choices.append(choice)
return ModelResponseStream(

View file

@ -177,19 +177,13 @@ class DatabricksBase:
# Default: just litellm
return f"litellm/{version}"
def _get_api_base(self, api_base: str | None) -> str:
"""
Get the Databricks API base URL.
If not provided, attempts to get it from the Databricks SDK.
"""
def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str:
if api_base is None:
try:
from databricks.sdk import WorkspaceClient
databricks_client: Final = WorkspaceClient()
api_base = f"{databricks_client.config.host}/serving-endpoints"
return api_base
except ImportError:
raise DatabricksException(
status_code=400,
@ -198,6 +192,18 @@ class DatabricksBase:
"or install the databricks-sdk Python library."
),
)
if not use_ai_gateway:
return api_base
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai-gateway/mlflow/v1"):
return normalized_api_base
if normalized_api_base.endswith("/serving-endpoints"):
return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1"
api_base_parts: Final = urlsplit(normalized_api_base)
if api_base_parts.path in ("", "/"):
return f"{normalized_api_base}/ai-gateway/mlflow/v1"
return api_base
def _get_oauth_m2m_token(

View file

@ -81,6 +81,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""

View file

@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
OpenAiResponsesToChatCompletionStreamIterator,
tool_call_dict_from_output_item,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
@ -84,7 +84,6 @@ from litellm.types.llms.openai import (
)
from litellm.types.responses.main import (
GenericResponseOutputItem,
OutputFunctionToolCall,
OutputText,
)
from litellm.types.utils import GenericGuardrailAPIInputs
@ -106,6 +105,19 @@ class _ToolCallShape(NamedTuple):
arguments: str
class _ToolCallFunctionFields(BaseModel):
model_config = ConfigDict(frozen=True)
name: str | None = None
arguments: str = ""
class _ToolCallFields(BaseModel):
model_config = ConfigDict(frozen=True)
function: _ToolCallFunctionFields
def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
return tuple(
_ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
@ -113,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl
)
def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
try:
fields: Final = _ToolCallFields.model_validate(payload)
except ValidationError:
return None
return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
def _post_guardrail_tool_call_shapes(
returned_tool_calls: Sequence[object] | None,
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str | None,
) -> tuple[_ToolCallShape, ...]:
if not pre_guardrail_tool_calls:
return pre_guardrail_tool_calls
if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
verbose_proxy_logger.warning(
"OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
"leaving the tool call output items unchanged",
guardrail_name,
"no" if returned_tool_calls is None else len(returned_tool_calls),
len(pre_guardrail_tool_calls),
)
return pre_guardrail_tool_calls
returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
if len(validated_shapes) != len(returned_shapes):
verbose_proxy_logger.warning(
"OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
"leaving the tool call output items unchanged",
guardrail_name,
)
return pre_guardrail_tool_calls
return validated_shapes
def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
class ResponseOutputEnvelope(TypedDict, total=False):
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
@ -140,8 +193,18 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
)
_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset(
{"response.function_call_arguments.delta", "response.function_call_arguments.done"}
_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call": "arguments", "custom_tool_call": "input"}
)
_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset(
{"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"}
)
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"}
)
_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset(
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
)
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
@ -180,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts
def _is_function_call_item(item: object) -> bool:
return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
def _is_tool_call_item(item: object) -> bool:
return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES
def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
return None
if isinstance(item, Mapping):
return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects
return item.model_dump() if isinstance(item, BaseModel) else None
def _is_tool_call_output_item(item: object) -> bool:
return _tool_call_output_item_mapping(item) is not None
def _last_message_role(messages: Sequence[object]) -> str | None:
@ -205,7 +280,7 @@ def _provenance_unit_bounds(
start_indexes: Final = tuple(
index
for index in range(len(raw_input))
if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
)
return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
@ -357,6 +432,7 @@ class OpenAIResponsesHandler(BaseTranslation):
"""
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
@ -603,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation):
- response.output is a list of output items
- Each output item can be:
* GenericResponseOutputItem with a content list of OutputText objects
* ResponseFunctionToolCall with tool call data
* ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data
- Each OutputText object has a text field
"""
@ -668,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
inputs["model"] = response_model
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
@ -676,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
@ -683,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
self._write_tool_call_rewrites_to_output(
tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=post_guardrail_tool_calls,
)
verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
@ -779,11 +866,10 @@ class OpenAIResponsesHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
post_guardrail_tool_calls: Final = _tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
else tool_calls_to_check
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Write guardrailed texts back into the output items in-place.
@ -933,11 +1019,12 @@ class OpenAIResponsesHandler(BaseTranslation):
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites into the completed
envelope's ``function_call`` items and sync the earlier stream events,
keyed by ``call_id``. The guardrail sees the envelope's function calls
in output order, which is how a rewritten call finds its ``call_id``;
the stream events find their call through the ``call_id`` on
``output_item`` events and the ``item_id`` on argument events, since an
envelope's ``function_call`` and ``custom_tool_call`` items and sync the
earlier stream events, keyed by ``call_id``. The guardrail sees the
envelope's tool calls in output order, which is how a rewritten call
finds its ``call_id``; the stream events find their call through the
``call_id`` on ``output_item`` events and the ``item_id`` on argument
and custom-input events, since an
event's ``output_index`` need not match the envelope's (the chat bridge
numbers tool calls from 1 while the envelope lists them after the
message). A rewrite whose calls do not line up with the envelope, or
@ -945,32 +1032,30 @@ class OpenAIResponsesHandler(BaseTranslation):
pipeline executor discards it and releases the original events."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
function_call_items: Final = tuple(
output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call"
)
tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
call_ids: Final = tuple(
call_id
for output_item in function_call_items
for output_item in tool_call_items
if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
)
stream_events: Final = responses_so_far[:-1]
call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events)
call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events)
event_call_ids: Final = tuple(
self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events
self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events
)
rewrites_by_call_id: Final = MappingProxyType(
{
call_id: after
call_id: _tool_call_rewrite(before, after)
for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
}
)
unresolved_argument_event: Final = any(
call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES
call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
for event, call_id in zip(stream_events, event_call_ids)
)
if (
len(call_ids) != len(function_call_items)
len(call_ids) != len(tool_call_items)
or len(frozenset(call_ids)) != len(call_ids)
or len(call_ids) != len(post_guardrail_tool_calls)
or unresolved_argument_event
@ -981,10 +1066,10 @@ class OpenAIResponsesHandler(BaseTranslation):
raise UndeliverableStreamRewrite(guardrail_name)
for output_item, rewrite in (
(output_item, rewrites_by_call_id[call_id])
for output_item, call_id in zip(function_call_items, call_ids)
for output_item, call_id in zip(tool_call_items, call_ids)
if call_id in rewrites_by_call_id
):
self._write_function_call_item(output_item, rewrite.name, rewrite.arguments)
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
delta_replacements: Final = MappingProxyType(
{call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
)
@ -992,16 +1077,18 @@ class OpenAIResponsesHandler(BaseTranslation):
if call_id not in rewrites_by_call_id:
continue
match stream_item_field(event, "type"):
case "response.function_call_arguments.delta":
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES:
self._write_event_field(event, "delta", next(delta_replacements[call_id]))
case "response.function_call_arguments.done":
self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments)
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS:
self._write_event_field(
event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments
)
case "response.output_item.added":
self._write_function_call_item(
self._write_tool_call_item(
stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
)
case "response.output_item.done":
self._write_function_call_item(
self._write_tool_call_item(
stream_item_field(event, "item"),
rewrites_by_call_id[call_id].name,
rewrites_by_call_id[call_id].arguments,
@ -1009,8 +1096,23 @@ class OpenAIResponsesHandler(BaseTranslation):
case _:
pass
def _write_tool_call_rewrites_to_output(
self,
tool_call_items: Sequence[object],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
) -> None:
if len(tool_call_items) != len(post_guardrail_tool_calls):
return
for output_item, rewrite in (
(output_item, _tool_call_rewrite(before, after))
for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
):
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
@staticmethod
def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
items: Final = tuple(
stream_item_field(event, "item")
for event in stream_events
@ -1020,32 +1122,35 @@ class OpenAIResponsesHandler(BaseTranslation):
{
item_id: call_id
for item in items
if stream_item_field(item, "type") == "function_call"
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
and isinstance(item_id := stream_item_field(item, "id"), str)
and isinstance(call_id := stream_item_field(item, "call_id"), str)
}
)
@staticmethod
def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
event_type: Final = stream_item_field(event, "type")
if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES:
if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES:
item_id: Final = stream_item_field(event, "item_id")
return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
return None
item: Final = stream_item_field(event, "item")
call_id: Final = stream_item_field(item, "call_id")
return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None
return (
call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
)
@staticmethod
def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None:
def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
if item is None:
return
if name is not None:
OpenAIResponsesHandler._write_event_field(item, "name", name)
if arguments is not None:
OpenAIResponsesHandler._write_event_field(item, "arguments", arguments)
item_type: Final = stream_item_field(item, "type")
if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
@ -1073,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation):
def _completed_response_scan_key(response: object) -> StreamingScanKey:
output_items: Final = stream_item_items(response, "output")
message_items: Final = tuple(
item for item in output_items if stream_item_field(item, "type") != "function_call"
item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
)
return StreamingScanKey(
texts=tuple(
@ -1085,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation):
tool_calls=tuple(
stream_item_fingerprint(item)
for item in output_items
if stream_item_field(item, "type") == "function_call"
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
),
stream_ended=True,
)
@ -1196,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation):
Override this method to customize text/image/tool extraction logic.
"""
# Check if this is a tool call (OutputFunctionToolCall)
if isinstance(output_item, OutputFunctionToolCall) or (
isinstance(output_item, BaseModel)
and hasattr(output_item, "type")
and getattr(output_item, "type") == "function_call"
):
tool_call_item: Final = _tool_call_output_item_mapping(output_item)
if tool_call_item is not None:
if tool_calls_to_check is not None:
tool_call_dict = (
LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=output_item,
index=output_idx,
)
)
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
return
elif isinstance(output_item, dict) and output_item.get("type") == "function_call":
# Handle dict representation of tool call
if tool_calls_to_check is not None:
# Convert dict to ResponseFunctionToolCall for processing
try:
tool_call_obj: Final = ResponseFunctionToolCall(**output_item)
tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=tool_call_obj,
index=output_idx,
)
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
except Exception:
pass
tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
return
# Handle both GenericResponseOutputItem and dict

View file

@ -5398,6 +5398,14 @@ def completion(
if dynamic_api_key is not None:
api_key = dynamic_api_key
# check if user passed in any of the OpenAI optional params
bridges_to_responses_api: Final = (
responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge
)
allowed_openai_params: Final[list[str] | None] = (
[*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"]
if bridges_to_responses_api
else kwargs.get("allowed_openai_params")
)
optional_param_args: Final = {
"functions": functions,
"function_call": function_call,
@ -5442,7 +5450,7 @@ def completion(
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"allowed_openai_params": kwargs.get("allowed_openai_params"),
"allowed_openai_params": allowed_openai_params,
"base_model": base_model,
}
optional_params = get_optional_params(**optional_param_args, **non_default_params)

File diff suppressed because it is too large Load diff

View file

@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
### Route Every Claude Code Session Through the Proxy
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
lite --base-url https://your-proxy.example.com login --config-claude
```
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
#### Configuring Claude Code Once, With a Virtual Key or Your Login
`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto
claude
```
With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (the ones whose id contains `claude` or `anthropic`) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control
Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt
What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request
### QA Complexity-Based Auto-Routing Against Your Real Proxy
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t
The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.

View file

@ -16,6 +16,7 @@ from .cmd_quoting import quote_for_cmd
from .pi import (
LITELLM_PROXY_API_KEY_ENV,
PI_PROVIDER_NAME,
ListingFailure,
PiSyncError,
fetch_model_ids,
fetch_model_limits,
@ -165,7 +166,9 @@ def prepare_pi(
"""
ids: Final = fetch_model_ids(base_url, api_key, get=get)
if isinstance(ids, PiSyncError):
raise AgentRunError(ids.message)
raise AgentRunError(
f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message
)
limits: Final = fetch_model_limits(base_url, api_key, get=get)
path: Final = models_json_path(base_env)
error: Final = sync_models_json(path, base_url, ids, limits)

View file

@ -41,9 +41,15 @@ from litellm.litellm_core_utils.cli_token_utils import (
from .claude_settings import (
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
STARTING_MODEL_ROLE,
ApiKeyHelper,
ClaudeSettingsError,
write_claude_settings,
KeepModel,
configure_claude_settings,
refuse_while_owned,
resolve_api_key_helper,
)
from .pkce_login import (
Http,
@ -778,13 +784,23 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
def _configure_claude_code(base_url: str) -> None:
"""Point Claude Code at base_url by patching ~/.claude/settings.json."""
"""Point Claude Code at base_url by patching ~/.claude/settings.json, undoable with `lite unconfigure claude`."""
try:
write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
configure_claude_settings(
base_url,
ApiKeyHelper(resolve_api_key_helper(base_url)),
KeepModel(),
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
)
except ClaudeSettingsError as e:
raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
click.echo(
"Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. "
f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}."
)
def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None:
@ -853,6 +869,11 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"]
if config_claude:
try:
refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
except ClaudeSettingsError as e:
raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}")
try:
if pkce:

View file

@ -14,11 +14,13 @@ from ..claude_settings import (
AUTOROUTE_BACKUP_PATH,
CLAUDE_SETTINGS_PATH,
ClaudeSettingsError,
StaticToken,
load_json_or_empty,
merge_claude_settings,
)
from ..up import BackupRecord as ClaudeBackupRecord
from ..up import restore_claude_settings, write_backup
from .config import master_key_from_config
from .config import AUTOROUTER_MODEL_NAME, master_key_from_config
from .process import (
CONFIG_PATH,
DEFAULT_AUTOROUTE_PORT,
@ -37,7 +39,6 @@ from .process import (
terminate,
write_pid_record,
)
from .settings import merge_claude_settings_static_token
from .wizard import run_configure_wizard
_GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
@ -156,7 +157,9 @@ def up(port: int) -> None:
ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
AUTOROUTE_BACKUP_PATH,
)
merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key)
merged: Final = merge_claude_settings(
original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME
)
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
with secure_create(CLAUDE_SETTINGS_PATH) as f:
json.dump(merged, f, indent=2)

View file

@ -1,51 +0,0 @@
from typing import Final
from pydantic import JsonValue
from .config import AUTOROUTER_MODEL_NAME
ENV_KEY: Final = "env"
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
# Force every one of Claude Code's own model tiers to request the auto-router by name.
# Router's auto-router registry is keyed by the literal requested model string
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
# model_name can never work as a catch-all -- these overrides are what actually makes
# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
"ANTHROPIC_DEFAULT_SONNET_MODEL",
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
"ANTHROPIC_DEFAULT_OPUS_MODEL",
)
def merge_claude_settings_static_token(
settings: dict[str, JsonValue], base_url: str, auth_token: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to a local ephemeral proxy with a static token.
Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the
locally persisted autoroute master key, so a plain env var is simpler and correct. Any
existing apiKeyHelper is cleared so it can't fight with the static token.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final[dict[str, JsonValue]] = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
**base_env,
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
**{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
}
env.pop(ANTHROPIC_API_KEY_KEY, None)
merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env}
merged.pop(API_KEY_HELPER_KEY, None)
return merged
__all__ = ["merge_claude_settings_static_token"]

View file

@ -1,37 +1,70 @@
"""Shared handling of Claude Code's ~/.claude/settings.json.
`lite up` patches this file temporarily and restores it on exit; `lite login
--config-claude` patches it persistently. Both need the same merge and the same
apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
live here rather than in either command module.
`lite up` and `lite autoroute up` patch this file temporarily and restore it on
exit; `lite login --config-claude` and `lite configure claude` patch it
persistently and record how to undo it. All of them need the same merge and the
same apiKeyHelper command, and `up` already imports from `auth`, so the shared
parts live here rather than in any one command module.
"""
import hashlib
import json
import shlex
import shutil
import sys
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import reduce
from itertools import chain
from pathlib import Path
from typing import Final
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.private_json import write_private_json
from litellm.litellm_core_utils.private_json import (
commit_staged_json,
discard_staged_json,
ensure_private_dir,
stage_private_json,
)
from .cmd_quoting import quote_for_cmd
ENV_KEY: Final = "env"
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
MODEL_KEY: Final = "model"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
"ANTHROPIC_DEFAULT_SONNET_MODEL",
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
"ANTHROPIC_DEFAULT_OPUS_MODEL",
"ANTHROPIC_DEFAULT_FABLE_MODEL",
)
OWNED_ENV_KEYS: Final = (
ENABLE_TOOL_SEARCH_KEY,
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY,
ANTHROPIC_BASE_URL_KEY,
ANTHROPIC_AUTH_TOKEN_KEY,
ANTHROPIC_API_KEY_KEY,
)
OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY)
OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS)
_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY))
_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY)
_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}"
STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on"
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json"
@dataclass(frozen=True, slots=True)
@ -55,6 +88,102 @@ class ClaudeSettingsError(Exception):
"""Raised for any user-actionable failure while reading or writing Claude Code settings."""
@dataclass(frozen=True, slots=True)
class StaticToken:
"""A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN."""
token: str
@dataclass(frozen=True, slots=True)
class ApiKeyHelper:
"""A `lite auth print-token` command Claude Code runs per request, so a login renews in place."""
command: str
ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper
@dataclass(frozen=True, slots=True)
class KeepModel:
"""Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login)."""
@dataclass(frozen=True, slots=True)
class UnpinModel:
"""Let go of a `model` an earlier configure pinned; one the user set themselves stays."""
@dataclass(frozen=True, slots=True)
class StartOn:
"""Pin the top-level `model`, the row Claude Code starts on."""
model: str
ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn
class OwnedValue(BaseModel):
"""What one key held at a moment in time; `present=False` is an absent key, not a null one."""
model_config = ConfigDict(frozen=True)
present: bool
value: JsonValue = None
class ConfigureReceipt(BaseModel):
"""What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key).
Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the
value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and
carries the earlier one otherwise, so a key the user edited in between stops matching and is left
alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier
snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the
repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot
was captured beside, so a credential is only ever put back next to the server it was issued for.
No fingerprint is a second copy of a token.
"""
model_config = ConfigDict(frozen=True)
file_existed: bool
env_present: bool
env_was_object: bool
previous: Mapping[str, OwnedValue]
written: Mapping[str, str]
endpoints: Mapping[str, OwnedValue]
@dataclass(frozen=True, slots=True)
class WithheldCredential:
"""A credential left removed: captured beside `endpoint`, while the restored file points elsewhere."""
key: str
endpoint: str
@dataclass(frozen=True, slots=True)
class UnconfigureOutcome:
"""Keys whose value unconfigure changed back, keys the user changed since and so were left as they
are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the
URL points back), and whether no settings file remains."""
restored: tuple[str, ...]
kept: tuple[str, ...]
withheld: tuple[WithheldCredential, ...] = ()
file_removed: bool = False
@dataclass(frozen=True, slots=True)
class _Claim:
previous: OwnedValue
written: str | None
endpoint: OwnedValue | None
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
try:
content: Final = path.read_bytes() if path.exists() else b""
@ -70,29 +199,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
)
def merge_claude_settings(
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to route Claude Code through the proxy.
def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]:
raw_env: Final = settings.get(ENV_KEY)
if raw_env is None:
return MappingProxyType({})
if not isinstance(raw_env, dict):
raise ClaudeSettingsError(
f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.'
)
return raw_env
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host, and
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker
is filled from the proxy's /v1/models; existing values of both are left
alone. Every other key is preserved untouched.
def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
"""Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a
purely local check, so commands run it before any login prompt or request."""
for owner in owners:
if owner.backup_path.exists():
raise ClaudeSettingsError(
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
f"{owner.backup_path}) and will restore it when it stops. "
f"Run `{owner.stop_command}` first, then retry."
)
def _write_target(settings_path: Path) -> Path:
"""Write through a symlinked settings.json rather than replacing the link, which would silently
detach a file symlinked into a dotfiles repo."""
try:
return settings_path.resolve() if settings_path.is_symlink() else settings_path
except OSError as e:
raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e
def _stage(path: Path, document: Mapping[str, object]) -> str:
try:
return stage_private_json(str(path), document)
except OSError as e:
raise ClaudeSettingsError(f"Could not write {path}: {e}") from e
def _land(
path: Path,
staged: str | None,
also_discard: Sequence[str | None] = (),
commit: Callable[[str, str], None] = commit_staged_json,
) -> None:
"""Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a
filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are
discarded, so no temp file holding a token is left behind."""
try:
if staged is None:
path.unlink(missing_ok=True)
else:
commit(staged, str(path))
except OSError as e:
for other in also_discard:
if other is not None:
discard_staged_json(other)
raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e
def merge_claude_settings(
settings: Mapping[str, JsonValue],
base_url: str,
credential: ClaudeCredential,
default_model: str | None = None,
tier_model: str | None = None,
) -> Mapping[str, JsonValue]:
"""Return a new settings mapping wired to route Claude Code through the proxy.
A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper;
the other credential slots are removed either way, since Claude Code given two credentials may
send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their
defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts
on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one
group. Apart from those tier keys, exactly OWNED_PATHS are touched.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE,
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
}
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
current_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(
(
(ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE),
(ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE),
),
((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS),
((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),),
((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (),
((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None),
)
)
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(
((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)),
((ENV_KEY, env),),
((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (),
((MODEL_KEY, default_model),) if default_model is not None else (),
)
)
def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
@ -121,56 +325,255 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
"""Persistently point Claude Code at base_url, preserving every unrelated setting.
def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue:
return OwnedValue(present=key in container, value=container.get(key))
Refuses while any owner holds a backup: each restores its backup when it
stops, which would silently undo this write.
"""
for owner in owners:
if owner.backup_path.exists():
raise ClaudeSettingsError(
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
f"{owner.backup_path}) and will restore it when it stops. "
f"Run `{owner.stop_command}` first, then retry."
)
normalized_base_url: Final = base_url.rstrip("/")
api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
existing: Final = load_json_or_empty(settings_path)
raw_env: Final = existing.get(ENV_KEY)
if raw_env is not None and not isinstance(raw_env, dict):
raise ClaudeSettingsError(
f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
"Fix or remove it, then retry."
)
merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
# os.replace() swaps the symlink itself for a regular file, silently detaching a
# settings.json that is symlinked into a dotfiles repo. There is no backup to undo
# that here, unlike `lite up`, so write through to the link's target instead.
target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
def _fingerprint(owned: OwnedValue) -> str:
return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest()
def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
raw_env: Final = settings.get(ENV_KEY)
return raw_env if isinstance(raw_env, dict) else MappingProxyType({})
def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue:
section, _, key = path.rpartition(".")
return _owned(_env(settings) if section else settings, key)
def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ())
)
def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
"""`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes."""
section, _, key = path.rpartition(".")
if not section:
return _with_key(settings, key, owned)
return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned)))
def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]:
return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings)
def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool:
"""Whether the key still holds what a configure wrote (a key no configure ever changed is never ours)."""
return receipt.written.get(path) == _fingerprint(_lookup(settings, path))
def _claim(
path: str,
current: Mapping[str, JsonValue],
merged: Mapping[str, JsonValue],
earlier: ConfigureReceipt | None,
url_now: OwnedValue,
) -> _Claim:
"""What this configure records for one key; see ConfigureReceipt for the rules."""
before, after = _lookup(current, path), _lookup(merged, path)
carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None
return _Claim(
previous=before if carried is None else carried.previous.get(path, before),
written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)),
endpoint=None
if path not in _CREDENTIAL_PATHS
else (url_now if carried is None else carried.endpoints.get(path, url_now)),
)
def _receipt(
current: Mapping[str, JsonValue],
merged: Mapping[str, JsonValue],
earlier: ConfigureReceipt | None,
file_exists: bool,
) -> ConfigureReceipt:
url_now: Final = _lookup(current, _BASE_URL_PATH)
claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS})
return ConfigureReceipt(
file_existed=file_exists if earlier is None else earlier.file_existed,
env_present=ENV_KEY in current if earlier is None else earlier.env_present,
env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object,
previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}),
written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}),
endpoints=MappingProxyType(
{path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None}
),
)
def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None:
if not state_path.exists():
return None
try:
write_private_json(str(target), merged)
return ConfigureReceipt.model_validate_json(state_path.read_bytes())
except (OSError, ValidationError) as e:
raise ClaudeSettingsError(
f"{state_path} is not a readable `lite configure claude` receipt ({e}). "
"Remove it and edit Claude Code's settings by hand if they still point at the proxy."
) from e
def configure_claude_settings(
base_url: str,
credential: ClaudeCredential,
model: ModelChoice,
settings_path: Path,
state_path: Path,
owners: Sequence[SettingsFileOwner],
commit: Callable[[str, str], None] = commit_staged_json,
) -> None:
"""Persistently route Claude Code through base_url, recording how to undo it.
Both files are staged before either is committed, so a full disk or a read-only directory fails
before anything changes. The two commits are still two renames: a receipt rename that fails
discards the staged settings, and a settings rename that fails after the receipt landed puts the
earlier receipt back (or removes the new one), so the receipt on disk never describes settings
that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an
earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login).
"""
refuse_while_owned(settings_path, owners)
current: Final = load_json_or_empty(settings_path)
_env_object(current, settings_path)
earlier: Final = read_configure_receipt(state_path)
existing: Final = (
_with(current, MODEL_KEY, earlier.previous[MODEL_KEY])
if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier)
else current
)
merged: Final = merge_claude_settings(
existing, base_url, credential, model.model if isinstance(model, StartOn) else None
)
receipt: Final = _receipt(current, merged, earlier, settings_path.exists())
target: Final = _write_target(settings_path)
try:
ensure_private_dir(state_path.parent)
except OSError as e:
raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e
staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json"))
try:
staged_settings: Final = _stage(target, merged)
except ClaudeSettingsError:
discard_staged_json(staged_receipt)
raise
_land(state_path, staged_receipt, (staged_settings,), commit)
try:
_land(target, staged_settings, commit=commit)
except ClaudeSettingsError as settings_error:
try:
_land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json")))
except ClaudeSettingsError as receipt_error:
raise ClaudeSettingsError(
f"{settings_error} The receipt at {state_path} now describes settings that were not written and "
f"could not be put back either ({receipt_error}); remove it before retrying."
) from settings_error
raise
def _endpoint_text(endpoint: OwnedValue) -> str:
if not endpoint.present:
return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)"
return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value)
def unconfigure_claude_settings(
settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner]
) -> UnconfigureOutcome:
"""Undo `lite configure claude`: put back every key still holding what configure wrote, leave the
rest alone, and withhold a credential the restored file would send to a different server than it
was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish)."""
refuse_while_owned(settings_path, owners)
receipt: Final = read_configure_receipt(state_path)
if receipt is None:
raise ClaudeSettingsError(
f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo."
)
current: Final = load_json_or_empty(settings_path)
_env_object(current, settings_path)
ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt))
kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present)
put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours}))
url_after: Final = _lookup(put_back, _BASE_URL_PATH)
withheld: Final = tuple(
WithheldCredential(path, _endpoint_text(receipt.endpoints[path]))
for path in _CREDENTIAL_PATHS
if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after
)
absent: Final = OwnedValue(present=False)
trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld}))
settings: Final = (
trimmed
if _env(trimmed) or receipt.env_was_object
else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None))
)
target: Final = _write_target(settings_path)
file_removed: Final = not settings and not (receipt.file_existed and target.exists())
kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy
receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}})
if withheld
else None
)
staged_settings: Final = None if file_removed else _stage(target, settings)
try:
staged_receipt: Final = (
None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json"))
)
except ClaudeSettingsError:
if staged_settings is not None:
discard_staged_json(staged_settings)
raise
_land(target, staged_settings, (staged_receipt,))
_land(state_path, staged_receipt)
return UnconfigureOutcome(
restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)),
kept=kept,
withheld=withheld,
file_removed=file_removed,
)
__all__ = (
"ANTHROPIC_API_KEY_KEY",
"ANTHROPIC_AUTH_TOKEN_KEY",
"ANTHROPIC_BASE_URL_KEY",
"ANTHROPIC_DEFAULT_MODEL_ENV_KEYS",
"API_KEY_HELPER_KEY",
"AUTOROUTE_BACKUP_PATH",
"BACKUP_PATH",
"CLAUDE_SETTINGS_PATH",
"CONFIGURE_STATE_PATH",
"ENABLE_GATEWAY_MODEL_DISCOVERY_KEY",
"ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE",
"ENABLE_TOOL_SEARCH_KEY",
"ENABLE_TOOL_SEARCH_VALUE",
"ENV_KEY",
"MODEL_KEY",
"OWNED_ENV_KEYS",
"OWNED_PATHS",
"OWNED_TOP_LEVEL_KEYS",
"SETTINGS_FILE_OWNERS",
"STARTING_MODEL_ROLE",
"ApiKeyHelper",
"ClaudeCredential",
"ClaudeSettingsError",
"ConfigureReceipt",
"KeepModel",
"ModelChoice",
"OwnedValue",
"SettingsFileOwner",
"StartOn",
"StaticToken",
"UnconfigureOutcome",
"UnpinModel",
"WithheldCredential",
"configure_claude_settings",
"load_json_or_empty",
"merge_claude_settings",
"read_configure_receipt",
"refuse_while_owned",
"resolve_api_key_helper",
"write_claude_settings",
"unconfigure_claude_settings",
)

View file

@ -0,0 +1,252 @@
"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable."""
import re
import sys
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Final
import click
from InquirerPy import inquirer
from InquirerPy.base.control import Choice
from .auth import CliContextObj, context_secret_vault, get_stored_api_key
from .claude_settings import (
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
STARTING_MODEL_ROLE,
ApiKeyHelper,
ClaudeCredential,
ClaudeSettingsError,
ModelChoice,
StartOn,
StaticToken,
UnconfigureOutcome,
UnpinModel,
configure_claude_settings,
refuse_while_owned,
resolve_api_key_helper,
unconfigure_claude_settings,
)
from .pi import ListingFailure, PiSyncError, fetch_model_ids
from .up import ensure_fresh_login
_LISTED_MODELS_SHOWN: Final = 20
_CLAUDE_TARGET: Final = "claude"
_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),)
_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default"
_CLAUDE_CODE_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE)
_MODEL_OPTION_HELP: Final = (
f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, "
"Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude "
"Code's sub-agent or background tiers; `lite autoroute up` is the mode that does."
)
def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]:
"""The credential to write and the key to check the proxy with.
An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes
into settings.json as a static token. Without one, the stored `lite login` credential is used
the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a
day and renews in place there; a missing or stale login is refreshed first, as `lite up` does.
"""
ctx_obj: Final[CliContextObj] = ctx.obj
explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key"))
if explicit:
return StaticToken(explicit), explicit
base_url: Final = ctx_obj["base_url"]
ensure_fresh_login(ctx)
stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
if not stored:
raise ClaudeSettingsError("Login did not produce a usable token.")
return ApiKeyHelper(resolve_api_key_helper(base_url)), stored
def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]:
"""Every configure path begins the same way: the local ownership check first, so a `lite up`
session is refused before any login prompt or request, then the credential, then the listing."""
try:
refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
credential, key = resolve_credential(ctx, api_key)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
return credential, _listed_models(ctx.obj["base_url"], key)
def _listing_error(base_url: str, error: PiSyncError) -> str:
"""The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question."""
if error.kind is ListingFailure.REJECTED:
return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key."
if error.kind is ListingFailure.UNREACHABLE:
return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?"
if error.kind is ListingFailure.EMPTY:
return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model."
return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy."
def _listed_models(base_url: str, key: str) -> tuple[str, ...]:
listed: Final = fetch_model_ids(base_url, key)
if isinstance(listed, PiSyncError):
raise click.ClickException(_listing_error(base_url, listed))
return listed
def _model_choice(model: str | None) -> ModelChoice:
return StartOn(model) if model is not None else UnpinModel()
def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequence[str], model: str | None) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"]
if model is not None and model not in listed:
shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN])
more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else ""
raise click.ClickException(
f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}."
)
try:
configure_claude_settings(
base_url, credential, _model_choice(model), CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS
)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
in_picker: Final = sum(1 for listed_model in listed if _CLAUDE_CODE_PICKER_FILTER.search(listed_model))
click.echo(f"Configured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url}.")
click.echo(
"Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN."
if isinstance(credential, StaticToken)
else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it."
)
click.echo(
f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model."
if model is not None
else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or "
"pass --model to start on a proxy model."
)
click.echo(
f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing "
"'claude' or 'anthropic')."
)
click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.")
if isinstance(credential, StaticToken) and CLAUDE_SETTINGS_PATH.is_symlink():
click.echo(
f"Note: {CLAUDE_SETTINGS_PATH} is a symlink to {CLAUDE_SETTINGS_PATH.resolve()}, so your key now lives in "
"that file; keep it out of version control.",
err=True,
)
def _pick_targets() -> tuple[str, ...]:
picked: Final = inquirer.checkbox(
message="Which agents should route through LiteLLM?",
choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS],
validate=lambda chosen: len(chosen) > 0,
invalid_message="Pick at least one.",
).execute()
return tuple(str(value) for value in picked)
def _pick_model(listed: Sequence[str]) -> str | None:
picked: Final = inquirer.fuzzy(
message="Model Claude Code starts on (type to filter; /model switches any time):",
choices=[_KEEP_DEFAULT_MODEL, *listed],
).execute()
return None if picked == _KEEP_DEFAULT_MODEL else str(picked)
def interactive_configure(
ctx: click.Context,
pick_targets: Callable[[], tuple[str, ...]] = _pick_targets,
pick_model: Callable[[Sequence[str]], str | None] = _pick_model,
) -> None:
"""`lite configure` with no agent named: ask which agents to wire and which model to pin."""
targets: Final = pick_targets()
if _CLAUDE_TARGET not in targets:
return
credential, listed = _start(ctx, None)
_apply_claude(ctx, credential, listed, pick_model(listed))
@click.group(name="configure", invoke_without_command=True)
@click.pass_context
def configure_group(ctx: click.Context) -> None:
"""Persistently route a coding agent through your LiteLLM proxy.
With no agent named, asks which agents to wire and which proxy model to pin.
"""
if ctx.invoked_subcommand is not None:
return
if not sys.stdin.isatty():
raise click.ClickException(
"`lite configure` asks questions, so it needs a terminal. Non-interactively, run "
"`lite configure claude --api-key <key> --model <model>`."
)
interactive_configure(ctx)
@click.group(name="unconfigure")
def unconfigure_group() -> None:
"""Undo `lite configure` for a coding agent."""
@configure_group.command(name="claude")
@click.option(
"--api-key",
"api_key",
default=None,
help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / "
"LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.",
)
@click.option("--model", default=None, help=_MODEL_OPTION_HELP)
@click.pass_context
def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None:
"""Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`.
Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a
static token, or your `lite login` through apiKeyHelper), and gateway model discovery so
/model lists the proxy's models; --model picks the one Claude Code starts on. Every other
setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back.
Assumes the proxy is already running.
"""
credential, listed = _start(ctx, api_key)
_apply_claude(ctx, credential, listed, model)
@unconfigure_group.command(name="claude")
def unconfigure_claude() -> None:
"""Return Claude Code's settings to what they were before `lite configure claude`.
Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are
put back; anything you changed since is left as it is and named in the output.
"""
try:
outcome: Final = unconfigure_claude_settings(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
_report_unconfigure(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, outcome)
def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None:
"""Say what unconfigure did, naming only keys whose value it changed."""
if outcome.file_removed:
click.echo(
f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys."
)
elif outcome.restored:
click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.")
else:
click.echo(f"Nothing in {settings_path} was still ours to restore.")
if outcome.kept:
click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.")
if outcome.withheld:
click.echo(
"Left removed, since the file now points at a different server than they were issued for: "
+ "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld)
+ f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` "
"again to put them back, or delete that file to drop them."
)
__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group")

View file

@ -10,6 +10,7 @@ import os
import tempfile
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -20,11 +21,28 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
PI_PROVIDER_NAME: Final = "litellm"
LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
_REJECTED_STATUSES: Final = frozenset((401, 403))
class ListingFailure(StrEnum):
"""Why a proxy could not be listed, decided once where the HTTP outcome is classified.
`unreachable` means no response at all; the other kinds prove the proxy answered, so callers
must not suggest checking whether it is running.
"""
UNREACHABLE = "unreachable"
REJECTED = "rejected"
BAD_BODY = "bad_body"
EMPTY = "empty"
OTHER = "other"
@dataclass(frozen=True, slots=True)
class PiSyncError:
message: str
status: int | None = None
kind: ListingFailure | None = None
@dataclass(frozen=True, slots=True)
@ -65,16 +83,20 @@ def fetch_model_ids(
timeout=10,
)
except requests.RequestException as e:
return PiSyncError(f"Could not list models from the proxy: {e}")
return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE)
if resp.status_code != 200:
return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
return PiSyncError(
f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.",
resp.status_code,
ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER,
)
try:
listing: Final = _ModelList.model_validate(resp.json())
except (ValueError, ValidationError) as e:
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY)
ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
if not ids:
return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY)
return ids
@ -200,6 +222,7 @@ __all__ = (
"LITELLM_PROXY_API_KEY_ENV",
"PI_CONFIG_DIR_ENV",
"PI_PROVIDER_NAME",
"ListingFailure",
"ModelLimits",
"PiSyncError",
"fetch_model_ids",

View file

@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_
from .claude_settings import (
BACKUP_PATH,
CLAUDE_SETTINGS_PATH,
ApiKeyHelper,
ClaudeSettingsError,
load_json_or_empty,
merge_claude_settings,
@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool:
return token_data is not None and token_data.get("refresh_token") is not None
def _ensure_fresh_login(ctx: click.Context) -> None:
def ensure_fresh_login(ctx: click.Context) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"].rstrip("/")
vault: Final = context_secret_vault(ctx)
@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None:
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
ctx.invoke(login, pkce=pkce)
if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault):
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
raise UpError("Login did not produce a usable token.")
def _restore_and_report() -> None:
@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None:
base_url: Final = ctx.obj["base_url"]
try:
_ensure_fresh_login(ctx)
ensure_fresh_login(ctx)
api_key: Final = resolve_api_key(ctx)
verify_proxy_key(base_url, api_key)
@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None:
)
CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper))
with open(CLAUDE_SETTINGS_PATH, "w") as f:
json.dump(merged, f, indent=2)
except (AgentRunError, ClaudeSettingsError) as e:

View file

@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key,
from .commands.autoroute.commands import autoroute_group
from .commands.chat import chat
from .commands.config import config_commands, get_config_value, hidden_command_names
from .commands.configure import configure_group, unconfigure_group
from .commands.credentials import credentials
from .commands.debug import debug
from .commands.encryption import encryption
@ -162,6 +163,9 @@ cli.add_command(model_groups)
# Add the autoroute command group (QA auto-routing against your real proxy)
cli.add_command(autoroute_group, name="autoroute")
cli.add_command(config_commands)
# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key)
cli.add_command(configure_group)
cli.add_command(unconfigure_group)
if __name__ == "__main__":

View file

@ -3300,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing:
has completed.
Guardrails routed through unified_guardrail are skipped, since they already ran
via its streaming iterator. Guardrails that override
async_post_call_success_hook directly run here, including those that implement
apply_guardrail but keep their native lifecycle hooks.
via its streaming iterator, and so are guardrails a post_call policy pipeline
manages, since the pipeline ran them against the buffered stream. Guardrails
that override async_post_call_success_hook directly run here, including those
that implement apply_guardrail but keep their native lifecycle hooks.
This is audit-only content has already been delivered to the client.
@ -3312,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing:
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
from litellm.proxy.utils import (
_check_and_merge_model_level_guardrails,
stream_gated_guardrail_names,
)
guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router)
stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if cb.guardrail_name in stream_gated:
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,

View file

@ -31,6 +31,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
ESTIMATED_OUTPUT_TOKENS_FIELD,
@ -3307,7 +3308,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
min_configured_tpm_limit=min_configured_otpm_limit,
call_type=call_type,
)
raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens(
raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)(
data=data, model=requested_model, call_type=call_type
)
estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1)

View file

@ -4559,6 +4559,23 @@ async def delete_verification_tokens(
litellm_changed_by=litellm_changed_by,
)
# Snapshot before the delete: the FK cascade drops the mapping rows, but their
# cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380).
jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple(
cache_key
for keys_for_token in await asyncio.gather(
*(
get_jwt_key_mapping_cache_keys_for_token(
hashed_token=key.token,
prisma_client=prisma_client,
)
for key in authorized_keys
if key.token is not None
)
)
for cache_key in keys_for_token
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
@ -4571,6 +4588,8 @@ async def delete_verification_tokens(
if len(deleted_tokens) != len(tokens):
failed_tokens = [token for token in tokens if token not in deleted_tokens]
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
else:
raise Exception("DB not connected. prisma_client is None")
except Exception as e:

View file

@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
return None if texts is None else tuple(texts)
def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]:
return tuple(texts or ())
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
@ -133,6 +137,94 @@ class _StreamRewriteObserver(CustomGuardrail):
return outputs
class _ScannedTextRecorder(CustomGuardrail):
def __init__(self, guardrail_name: str) -> None:
super().__init__(guardrail_name=guardrail_name)
self.inputs: GenericGuardrailAPIInputs | None = None
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
self.inputs = inputs
return inputs
class _LegacyHookStreamAdapter(CustomGuardrail):
"""Runs a guardrail that only implements the legacy post-call hook (no unified
``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The
endpoint translation hands it the texts it scanned plus the assembled response under
``request_data["response"]``; the hook gets that response in the shape its route gives
non-streaming hooks, an exception it raises ends the stream through the executor's
fail/error classification, and the response it hands back, or the one it changed in place
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
client through the translation's ended-stream write-back. A
replacement whose scanned texts do not line up with the originals, or whose tool calls
differ from them, is undeliverable, so the executor releases the original chunks. A stream
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
long as the hook left the tool calls alone."""
def __init__(
self,
inner: CustomGuardrail,
endpoint_translation: "BaseTranslation",
user_api_key_dict: "UserAPIKeyAuth",
) -> None:
super().__init__(guardrail_name=inner.guardrail_name)
self.inner: Final = inner
self.endpoint_translation: Final = endpoint_translation
self.user_api_key_dict: Final = user_api_key_dict
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response"))
replacement: Final = await self.inner.async_post_call_success_hook(
data=request_data,
user_api_key_dict=self.user_api_key_dict,
response=hooked,
)
rewrite: Final = hooked if replacement is None else replacement
if rewrite is None:
return inputs
rescanned: Final = await self._rescan(rewrite, logging_obj)
if rescanned is None:
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
rewritten: Final = rescanned.get("texts")
if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))):
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")):
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
if not rewritten:
return inputs
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
return rewritten_inputs
async def _rescan(
self, response: object, logging_obj: "LiteLLMLoggingObj | None"
) -> GenericGuardrailAPIInputs | None:
recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown")
await self.endpoint_translation.process_output_response(
response=response,
guardrail_to_apply=recorder,
litellm_logging_obj=logging_obj,
user_api_key_dict=self.user_api_key_dict,
)
return recorder.inputs
def _prepare_hook_input(
step: PipelineStep,
callback: CustomGuardrail,
@ -300,18 +392,29 @@ class PipelineExecutor:
endpoint_translation: "BaseTranslation",
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
user_api_key_dict: "UserAPIKeyAuth | None",
user_api_key_dict: "UserAPIKeyAuth",
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text and tool-call rewrites on translations that support ended-stream write-back. A
rewrite that cannot reach the client yet (one on a translation without write-back, or
one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the
buffered chunks go back to the originals and the step passes, so the client gets the
stream the merge base sent."""
observer: Final = _StreamRewriteObserver(callback)
guardrail without the unified interface runs its legacy post-call hook against the
assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the
client yet (one on a translation without write-back, one that drops or adds a tool call,
or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is
discarded: the buffered chunks go back to the originals and the step passes, so the
client gets the stream the merge base sent, and the guardrail stays out of the
applied-guardrails header since its output never reached the client. The response an
earlier step's translation stored under ``request_data["response"]`` is dropped first,
so this step's hook sees the stream as the steps before it left it."""
scanner: Final = (
callback
if PipelineExecutor.supports_unified_execution(callback)
else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict)
)
observer: Final = _StreamRewriteObserver(scanner)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
originals: Final = copy.deepcopy(streaming_chunks)
hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored
try:
if deliver_rewrites:
await endpoint_translation.process_output_streaming_response(
@ -332,11 +435,12 @@ class PipelineExecutor:
)
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
else:
if observer.changed_tool_call_count or (
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if observer.changed_tool_call_count or (
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if not callback.records_own_guardrail_information:
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
@ -396,11 +500,11 @@ class PipelineExecutor:
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call" and streaming_chunks is not None:
if not use_unified or endpoint_translation is None:
if endpoint_translation is None:
return (
"error",
None,
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation",
None,
)
await PipelineExecutor._run_streaming_step(
@ -456,10 +560,22 @@ class PipelineExecutor:
@staticmethod
def supports_unified_execution(callback: CustomGuardrail) -> bool:
"""Whether this guardrail runs through the unified apply_guardrail path,
the interface streaming pipeline execution requires."""
"""Whether this guardrail runs through the unified apply_guardrail path."""
return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
@staticmethod
def supports_streaming_execution(callback: CustomGuardrail) -> bool:
"""Whether a streaming pipeline step can run this guardrail against the buffered
stream: through the unified path, or through its post-call hook on the assembled
response when that hook is its only streaming path. A guardrail with its own
streaming iterator hook, or with neither hook, keeps running on its own."""
callback_type: Final = type(callback)
return PipelineExecutor.supports_unified_execution(callback) or (
callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook
and callback_type.async_post_call_streaming_iterator_hook
is CustomLogger.async_post_call_streaming_iterator_hook
)
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""

View file

@ -63,11 +63,13 @@ from litellm.constants import (
LITELLM_UI_SESSION_DURATION,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.proxy._types import (
UI_TEAM_ID,
CallbackDelete,
@ -272,7 +274,6 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.litellm_core_utils.agentic_loop_settings import (
validated_max_agentic_loops,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
@ -12816,7 +12817,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
CustomHuggingfaceTokenizer | None,
model_info.get("custom_tokenizer", None),
)
_tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer)
_tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)(
model=model_to_use, custom_tokenizer=custom_tokenizer
)
tokenizer_used: Final = str(_tokenizer_used["type"])
system_message: Final = _system_message(system)
@ -12829,7 +12832,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
)
total_tokens: Final = await asyncify(litellm.token_counter)(
total_tokens: Final = await offload_token_count(litellm.token_counter)(
model=model_to_use,
text=prompt,
messages=counted_messages,

View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

View file

@ -101,6 +101,7 @@ from litellm.litellm_core_utils.core_helpers import (
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms import load_guardrail_translation_mappings
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import (
@ -460,7 +461,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
def _pipeline_managed_guardrail_names(
def pipeline_managed_guardrail_names(
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
) -> frozenset[str]:
return _pipeline_step_guardrail_names(
@ -523,9 +524,17 @@ def _merge_pipeline_metadata_writes(
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
if callback is None:
return False
if PipelineExecutor.supports_unified_execution(callback):
return True
return (
translation is not None
and type(translation).assembles_streamed_response
and PipelineExecutor.supports_streaming_execution(callback)
)
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
@ -582,7 +591,7 @@ def _withdraw_deferred_claims(
outside_by_policy: Final = MappingProxyType(
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
)
running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(
running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
)
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
@ -657,37 +666,51 @@ def _body_selected_deferrals(
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
unsupported: Final = tuple(
def _pipeline_unsupported_streaming_guardrails(
pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
) -> tuple[str, ...]:
return tuple(
dict.fromkeys(
step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
step.guardrail
for step in pipeline.steps
if not _pipeline_step_supports_streaming(step.guardrail, translation)
)
)
def _pipeline_is_streamable(
policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
) -> bool:
unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation)
if not unsupported:
return True
verbose_proxy_logger.warning(
"Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
"which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
"Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they "
"need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a "
"route whose translation assembles the streamed response. The stream skips the pipeline and its "
"guardrails run on their own: %s",
policy_name,
", ".join(unsupported),
)
return False
def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
return resolve_endpoint_translation(user_api_key_dict, None) is not None
def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None":
resolved: Final = resolve_endpoint_translation(user_api_key_dict, None)
return None if resolved is None else resolved[1]
def _stream_gated_guardrail_names(
def stream_gated_guardrail_names(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> frozenset[str]:
if not _route_supports_streaming_pipelines(user_api_key_dict):
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if translation is None:
return frozenset()
return _pipeline_step_guardrail_names(
tuple(
(policy_name, pipeline)
for policy_name, pipeline in _post_call_pipelines(request_data)
if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
)
)
@ -699,16 +722,19 @@ def _streamable_post_call_pipelines(
The post_call pipelines a streaming response can be gated through.
Streaming pipelines scan the buffered stream through the endpoint guardrail
translation of the request route, so every step's guardrail needs the
unified apply_guardrail interface and the route needs a translation. A
pipeline that cannot be run that way yet is left out and its guardrails
run on the stream on their own, the way they did before pipelines ran on
streams at all, with a warning naming the pipeline.
translation of the request route, so every step's guardrail needs either the
unified apply_guardrail interface or, on a route whose translation assembles
the streamed response, a post-call hook that is its only streaming path, and
the route needs a translation. A pipeline that
cannot be run that way yet is left out and its guardrails run on the stream
on their own, the way they did before pipelines ran on streams at all, with
a warning naming the pipeline.
"""
post_call_pipelines: Final = _post_call_pipelines(request_data)
if not post_call_pipelines:
return ()
if not _route_supports_streaming_pipelines(user_api_key_dict):
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if translation is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
@ -720,7 +746,7 @@ def _streamable_post_call_pipelines(
return tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if _pipeline_is_streamable(policy_name, pipeline)
if _pipeline_is_streamable(policy_name, pipeline, translation)
)
@ -2110,7 +2136,7 @@ class ProxyLogging:
)
# Get pipeline-managed guardrails to skip in normal loop
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
caps: Final = ProxyLogging._callback_capabilities()
# Skip the per-request callback walk entirely when nothing in
@ -2875,7 +2901,7 @@ class ProxyLogging:
original_exception=original_exception,
)
request_data.update(_failure_fields_to_lift(request_data))
request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
# Remove before callbacks iterate — not serialisable
request_data.pop("litellm_logging_obj", None)
@ -3114,7 +3140,7 @@ class ProxyLogging:
if pipeline_response is not None:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
try:
# Merge model-level guardrails before checking which guardrails to run
@ -3430,7 +3456,7 @@ class ProxyLogging:
_cached_guardrail_data: dict | None = None
_guardrail_data_computed = False
pipeline_gated: Final = (
_stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
)
for callback in litellm.callbacks:

View file

@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
response_created_event_data["temperature"] = self.responses_api_request["temperature"]
if "text" in self.responses_api_request:
response_created_event_data["text"] = self.responses_api_request["text"]
if "tool_choice" in self.responses_api_request:
# Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format
response_created_event_data["tool_choice"] = (
LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"])
or "auto"
response_created_event_data["tool_choice"] = (
LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
self.responses_api_request.get("tool_choice")
)
else:
response_created_event_data["tool_choice"] = "auto"
)
if "tools" in self.responses_api_request:
response_created_event_data["tools"] = self.responses_api_request["tools"]
else:

View file

@ -27,8 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import (
)
from openai.types.responses import ResponseFunctionToolCall
from openai.types.responses.response_create_params import ResponseInputParam
from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
@ -68,6 +70,7 @@ from litellm.types.llms.openai import (
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStatus,
ToolChoice,
ValidChatCompletionMessageContentTypes,
ValidChatCompletionMessageContentTypesLiteral,
)
@ -126,6 +129,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
_TEXT_ADAPTER: Final = TypeAdapter(str)
_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice)
@runtime_checkable
@ -267,6 +271,27 @@ class LiteLLMCompletionResponsesConfig:
# Return as-is for unknown formats
return tool_choice
@staticmethod
def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice:
if tool_choice is None:
return "auto"
try:
return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice)
except ValidationError:
return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice)
@staticmethod
def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice:
match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice):
case {"type": "custom"}, {"function": {"name": str(custom_name)}}:
return ToolChoiceCustomParam(type="custom", name=custom_name)
case _, {"type": "function", "function": {"name": str(function_name)}}:
return ToolChoiceFunctionParam(type="function", name=function_name)
case _, "none" | "auto" | "required" as normalized:
return normalized
case _, _:
return "auto"
@staticmethod
def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool:
"""
@ -2263,7 +2288,9 @@ class LiteLLMCompletionResponsesConfig:
),
parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False),
temperature=getattr(chat_completion_response, "temperature", 0),
tool_choice=getattr(chat_completion_response, "tool_choice", "auto"),
tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
responses_api_request.get("tool_choice")
),
tools=getattr(chat_completion_response, "tools", []),
top_p=getattr(chat_completion_response, "top_p", None),
max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None),

View file

@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti
import httpx
from openai._streaming import SSEDecoder
from pydantic import BaseModel, ValidationError
from typing_extensions import TypeIs
import litellm
@ -438,18 +439,7 @@ class BaseResponsesAPIStreamingIterator:
if self._persist_completed_response_before_logging:
self._persist_completed_response_to_cache(is_async=is_async)
# Create a copy for logging to avoid modifying the response object that will be returned to the user
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
logging_response = self.completed_response
if self.completed_response is not None and hasattr(self.completed_response, "model_dump"):
try:
logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump())
except Exception:
# Fallback to original if serialization fails
pass
logging_response: Final[object] = _logging_copy(self.completed_response)
self._restore_provider_response_headers(logging_response)
end_time: Final = datetime.now()
@ -488,10 +478,10 @@ class BaseResponsesAPIStreamingIterator:
def _restore_provider_response_headers(self, logging_response: object) -> None:
"""Re-apply the provider's response headers to the copy handed to logging callbacks.
``model_validate(model_dump())`` above drops pydantic private attributes, so the
``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the
``_hidden_params`` the provider transform set on the nested response are lost. Returns early
when that copy fell back to the original event, so logging-only state never lands on the
object the caller is iterating.
when the event was not a pydantic model and logging got the original, so logging-only state
never lands on the object the caller is iterating.
"""
if logging_response is self.completed_response:
return
@ -544,7 +534,7 @@ class BaseResponsesAPIStreamingIterator:
def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
if response_obj is None or self.logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
try:
@ -1293,14 +1283,46 @@ def _add_text_like_part_events(
)
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
deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow
copies of the event and its nested response still keep the caller's ``usage`` attribute separate."""
if not isinstance(event, BaseModel):
return event
try:
return type(event).model_validate(event.model_dump())
except Exception:
return _detached_shallow_copy(event)
def _detached_shallow_copy(event: BaseModel) -> BaseModel:
nested: Final[object] = getattr(event, "response", None)
if isinstance(nested, BaseModel):
return event.model_copy(update={"response": nested.model_copy()})
return event.model_copy()
def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
if isinstance(usage, ResponseAPIUsage):
return usage
if not isinstance(usage, dict):
return None
try:
return ResponseAPIUsage.model_validate(usage)
except ValidationError:
return None
def _stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:
if response_obj is None or logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives
if isinstance(getattr(usage_obj, "cost", None), (int, float)):
return
try:

View file

@ -67,7 +67,7 @@ from litellm.constants import (
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
coerce_token_limit,
@ -98,6 +98,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
mask_credentials_in_payload,
mask_sensitive_structure,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.base_llm.passthrough.transformation import replace_path_segment
from litellm.llms.base_llm.vector_store.transformation import (
RouterVectorStoreEmbeddingExecutor,
@ -11155,28 +11156,31 @@ class Router:
def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]:
"""
Deployment ids that could serve ``model`` for ``team_id``, unioned across the paths
the router resolves a route through: ``model_group_alias``, a routing group, the
``model_name`` and team indexes, and wildcard pattern routes. Read-only and
side-effect-free, unlike ``_common_checks_available_deployment`` which also applies
fallbacks and can raise. Lets a pre-call check tell a genuine cross-group route from
same-group unavailability without re-deriving that precedence at the call site, and
without leaking deployment ids into request kwargs bound for the provider.
Deployment ids that could serve ``model`` for ``team_id``, following the same
precedence ``_common_checks_available_deployment`` uses to build a candidate pool:
``model_group_alias``, then a routing group, then the first matching early-resolve
path for a name that is not a ``model_name`` (team route, wildcard pattern via
``get_deployments_by_pattern``, team pattern router, default deployment), then the
``model_name`` and team indexes. Delegating to the router's own resolvers keeps this
aligned with how a route actually resolves rather than re-deriving it, and unlike
``_common_checks_available_deployment`` it is read-only: it does not apply request
fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call
check tell a genuine cross-group route from same-group unavailability without leaking
deployment ids into request kwargs bound for the provider.
"""
resolved: Final = self._get_model_from_alias(model=model) or model
routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id)
if routing_group_members is not None:
return self._deployment_ids(routing_group_members)
if resolved in self.model_names:
return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id))
team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None
return self._deployment_ids(
(
*self._get_all_deployments(model_name=resolved, team_id=team_id),
*(self.pattern_router.route(resolved) or ()),
*((team_router.route(resolved) or ()) if team_router is not None else ()),
)
early: Final = self._try_early_resolve_deployments_for_model_not_in_names(
model=resolved, request_team_id=team_id
)
if early is not None:
early_deployments: Final = early[1]
return self._deployment_ids(
(early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments
)
return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id))
@staticmethod
def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]:
@ -12095,7 +12099,7 @@ class Router:
try:
if not self._pre_call_checks_need_token_count(model, healthy_deployments):
return None
return await asyncify(self._count_pre_call_check_tokens)(
return await offload_token_count(self._count_pre_call_check_tokens)(
messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter
input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter
request_kwargs=request_kwargs,

View file

@ -2568,14 +2568,14 @@ class ComplexityRouter(CustomLogger):
"""Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the
event loop; None when counting fails, and the gate then leaves the placement alone."""
import litellm
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.token_counter import offload_token_count
out_of_band: Final = self._out_of_band_request_text(request_kwargs)
try:
counted: Final = await asyncify(litellm.token_counter)(
counted: Final = await offload_token_count(litellm.token_counter)(
messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence
)
return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0)
return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0)
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e)
return None

View file

@ -21,6 +21,7 @@ import litellm
from litellm import token_counter
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.types.router import RouterCacheEnum, RouterErrors
from litellm.utils import get_utc_datetime
@ -466,7 +467,7 @@ async def async_io_token_pre_call_check(
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
_model: Final = (deployment.get("litellm_params") or {}).get("model") or ""
estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
dt: Final = get_utc_datetime()

View file

@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
)
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes, StandardLoggingPayload
from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt
@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
if request_kwargs is not None and request_kwargs.get("_target_order") is not None:
return healthy_deployments
if messages is not None and is_prompt_caching_valid_prompt(
if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)(
messages=messages,
model=model,
min_token_count=_get_min_token_count_for_deployments(healthy_deployments),
@ -139,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
return
## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider
if is_prompt_caching_valid_prompt(
if await offload_token_count(is_prompt_caching_valid_prompt)(
model=model,
messages=cast(list[AllMessageValues], messages),
):

View file

@ -2,6 +2,7 @@ from typing import Any, Literal
from pydantic import BaseModel
from typing_extensions import (
ReadOnly,
Required,
TypedDict,
)
@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False):
role: Required[str]
content: Required[AllDatabricksContentValues]
tool_calls: list[DatabricksTool] | None
reasoning_content: ReadOnly[str | None]
reasoning: ReadOnly[str | None]
class DatabricksDelta(TypedDict, total=False):
role: ReadOnly[str]
content: ReadOnly[AllDatabricksContentValues | None]
reasoning_content: ReadOnly[str | None]
class DatabricksChoice(TypedDict, total=False):

View file

@ -525,6 +525,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
input_cost_per_second: float | None
output_cost_per_second: float | None
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_1080p: float | None
output_cost_per_second_4k: ReadOnly[float | None]
num_retries: int | None

View file

@ -318,6 +318,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
float | None
) # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_4k: ReadOnly[float | None]
ocr_cost_per_page: float | None # for OCR models
ocr_cost_per_credit: float | None # for OCR models priced by credit
@ -3522,6 +3523,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
output_cost_per_second: float | None = None
output_cost_per_second_1080p: float | None = None
output_cost_per_second_480p: float | None = None
output_cost_per_second_720p: float | None = None
output_cost_per_second_4k: float | None = None
input_cost_per_pixel: float | None = None
output_cost_per_pixel: float | None = None

View file

@ -2293,15 +2293,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st
dict: A dictionary with the tokenizer and its type.
"""
try:
tokenizer = Tokenizer.from_pretrained(
identifier,
revision=revision,
auth_token=auth_token,
)
except Exception as e:
verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e)
tokenizer = Tokenizer.from_pretrained(identifier, revision=revision)
tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token)
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
@ -3412,7 +3404,7 @@ def get_optional_params_image_gen(
non_default_params=non_default_params,
optional_params=optional_params,
model=model or "",
drop_params=drop_params if drop_params is not None else False,
drop_params=litellm.drop_params is True or drop_params is True,
)
elif (
custom_llm_provider == "openai"
@ -5913,6 +5905,7 @@ def _get_model_info_helper(
output_cost_per_second=_model_info.get("output_cost_per_second", None),
output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None),
output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None),
output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None),
output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None),
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
output_cost_per_image=_model_info.get("output_cost_per_image", None),

File diff suppressed because it is too large Load diff

View file

@ -478,6 +478,10 @@
"type": "number",
"minimum": 0
},
"output_cost_per_second_720p": {
"type": "number",
"minimum": 0
},
"output_cost_per_token": {
"type": "number",
"minimum": 0,

View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

View file

@ -181,13 +181,15 @@ quota_management.<behavior>.<variant>.<assertion>
| team_multi_window | fallback | spend_counter
<spend_tracking> chat_completions | stream | messages_bridge | embeddings
| cache_hit | key_rollup | concurrent_burst | tags | end_user
| per_model | failure | spend_calculate | pagination
| per_model | failure | spend_calculate | pagination | key_attribution
assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm
| blocks_then_resets | resets_windows_independently | alerts_without_blocking
| isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys
| routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost
| matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
| writes_failure_row | returns_cost | keeps_total
| writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email
| health_rows_keep_service_account | retrieve_batch_cost_joins_retrieving_key
| poller_batch_cost_joins_creating_key
e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages]
quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions]
```

View file

@ -58,3 +58,8 @@
- {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"}
- {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"}
- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"}
- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"}
- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"}
- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"}
- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"}
- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"}

View file

@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel):
class KeyGenerateResponse(BaseModel):
key: str
token: str | None = None
key_alias: str | None = None
models: list[str] = []
max_budget: float | None = None
@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel):
class SpendLogMetadata(BaseModel):
user_api_key_alias: str | None = None
applied_guardrails: list[str] | None = None
guardrail_information: list[GuardrailRunRecord] | None = None

View file

@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = (
("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"),
("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"),
("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"),
("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"),
)

View file

@ -15,9 +15,12 @@ import time
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
from e2e_config import unique_marker
from e2e_http import (
FileUploadForm,
Headers,
NoBody,
ProbeResult,
Result,
@ -35,6 +38,8 @@ from models import (
DateRangeParams,
EmbedBody,
EmbedResponse,
KeyGenerateBody,
KeyGenerateResponse,
OpenAPISchema,
SpendCalculateBody,
SpendCalculateResponse,
@ -43,13 +48,27 @@ from models import (
SpendLogsPageParams,
SpendTagsResponse,
TagSpend,
UserDeleteBody,
UserDeleteResponse,
UserNewBody,
UserNewResponse,
UserRole,
)
from proxy_client import ProxyClient
from proxy_client import Converged, ProxyClient, await_converged
from pydantic import BaseModel, Field
__all__ = [
"BatchCreateBody",
"CallbackLogMetadata",
"CallbackLogPayload",
"BatchObject",
"DailyActivityKeyBreakdown",
"FileObject",
"ProbeResult",
"ResponseIdentity",
"SpendClient",
"SpendLogRow",
"StreamingResponse",
"build_client",
"is_ok",
"unique_marker",
@ -57,6 +76,139 @@ __all__ = [
]
class GeminiApiKeyHeaders(Headers):
x_goog_api_key: str = Field(serialization_alias="x-goog-api-key")
content_type: str = Field(default="application/json", serialization_alias="Content-Type")
class GeminiPart(BaseModel):
text: str
class GeminiContent(BaseModel):
parts: list[GeminiPart]
class GeminiGenerationConfig(BaseModel):
maxOutputTokens: int
class GeminiGenerateBody(BaseModel):
contents: list[GeminiContent]
generationConfig: GeminiGenerationConfig
class ResponsesBody(BaseModel):
model: str
input: str
cache: dict[str, bool] | None = {"no-cache": True}
class QueuedChatBody(ChatBody):
priority: int = 0
class ResponseIdentity(BaseModel):
id: str | None = None
class HealthParams(BaseModel):
model: str
class ModelQuery(BaseModel):
model: str
class FileObject(BaseModel):
id: str
class BatchCreateBody(BaseModel):
input_file_id: str
endpoint: str = "/v1/chat/completions"
completion_window: str = "24h"
model: str
metadata: dict[str, str]
class BatchObject(BaseModel):
id: str
status: str
class ProviderQuery(BaseModel):
provider: str
class CallbackLogMetadata(BaseModel):
user_api_key_hash: str
user_api_key_alias: str
user_api_key_user_id: str
class CallbackLogPayload(BaseModel):
id: str
litellm_call_id: str
model: str
call_type: str = "acompletion"
start_time: float = Field(serialization_alias="startTime")
end_time: float = Field(serialization_alias="endTime")
response_cost: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
metadata: CallbackLogMetadata
class CallbackLogRecord(BaseModel):
status: str = "success"
standard_logging_payload: CallbackLogPayload
class CallbackLogsRequest(BaseModel):
records: list[CallbackLogRecord]
class CallbackLogsResponse(BaseModel):
processed: int
failed: int
class DailyActivityParams(BaseModel):
start_date: str
end_date: str
api_key: str
class DailyActivityKeyMetadata(BaseModel):
key_alias: str | None = None
team_id: str | None = None
user_email: str | None = None
class DailyActivityKeyMetrics(BaseModel):
api_requests: int = 0
class DailyActivityKeyBreakdown(BaseModel):
metrics: DailyActivityKeyMetrics
metadata: DailyActivityKeyMetadata
class DailyActivityBreakdown(BaseModel):
api_keys: dict[str, DailyActivityKeyBreakdown] = {}
class DailyActivityRow(BaseModel):
date: str
breakdown: DailyActivityBreakdown
class DailyActivityResponse(BaseModel):
results: list[DailyActivityRow] = []
def _chat_body(
model: str,
content: str,
@ -207,6 +359,166 @@ class SpendClient:
def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
return self.proxy.transport.probe(path, params=params)
def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
return unwrap(
self.proxy.transport.post(
"/user/new",
headers=self.proxy.transport.master,
json=UserNewBody(user_email=email, user_role=role, user_id=user_id),
response_type=UserNewResponse,
)
).user_id
def delete_user(self, user_id: str) -> None:
_ = unwrap(
self.proxy.transport.post(
"/user/delete",
headers=self.proxy.transport.master,
json=UserDeleteBody(user_ids=[user_id]),
response_type=UserDeleteResponse,
)
)
def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse:
return unwrap(
self.proxy.transport.post(
"/key/generate",
headers=self.proxy.transport.master,
json=body,
response_type=KeyGenerateResponse,
)
)
def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",
headers=self.proxy.transport.bearer(key),
json=_chat_body(model, content, max_tokens=max_tokens),
)
def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/queue/chat/completions",
headers=self.proxy.transport.bearer(key),
json=QueuedChatBody(
model=model,
messages=[ChatMessage(role="user", content=content)],
max_tokens=max_tokens,
),
)
def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/v1/messages",
headers=self.proxy.transport.bearer(key),
json=AnthropicMessagesBody(
model=model,
messages=[ChatMessage(role="user", content=content)],
max_tokens=max_tokens,
),
)
def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/v1/responses",
headers=self.proxy.transport.bearer(key),
json=ResponsesBody(model=model, input=content),
)
def send_embed(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/embeddings",
headers=self.proxy.transport.bearer(key),
json=EmbedBody(model=model, input=content),
)
def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
f"/gemini/v1beta/models/{model}:generateContent",
headers=GeminiApiKeyHeaders(x_goog_api_key=key),
json=GeminiGenerateBody(
contents=[GeminiContent(parts=[GeminiPart(text=content)])],
generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens),
),
)
def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject:
return unwrap(
self.proxy.transport.upload(
"/v1/files",
headers=self.proxy.transport.bearer(key),
form=FileUploadForm(purpose="batch"),
filename="key_attribution.jsonl",
content=content,
params=ModelQuery(model=model),
response_type=FileObject,
)
)
def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject:
return unwrap(
self.proxy.transport.post(
"/v1/batches",
headers=self.proxy.transport.bearer(key),
json=body,
response_type=BatchObject,
)
)
def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject:
return unwrap(
self.proxy.transport.get(
f"/v1/batches/{batch_id}",
headers=self.proxy.transport.bearer(key),
params=ProviderQuery(provider=provider),
response_type=BatchObject,
)
)
def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse:
return unwrap(
self.proxy.transport.post(
"/v1/rust_control_plane/logs",
headers=self.proxy.transport.bearer(key),
json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]),
response_type=CallbackLogsResponse,
)
)
def health(self, model: str) -> ProbeResult:
return self.proxy.transport.probe("/health", params=HealthParams(model=model))
def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
response: Final = unwrap(
self.proxy.transport.get(
"/user/daily/activity",
headers=self.proxy.transport.master,
params=DailyActivityParams(
start_date=start.strftime("%Y-%m-%d"),
end_date=end.strftime("%Y-%m-%d"),
api_key=token,
),
response_type=DailyActivityResponse,
)
)
return next(
(row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys),
None,
)
def poll_daily_activity_for_key(
self, token: str, *, start: datetime, end: datetime, min_requests: int
) -> DailyActivityKeyBreakdown | None:
outcome: Final = await_converged(
lambda: self.daily_activity_for_key(token, start=start, end=end),
converged=lambda found: found is not None and found.metrics.api_requests >= min_requests,
timeout=self.proxy.poll_timeout,
interval=self.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
def openapi(self) -> OpenAPISchema:
return unwrap(
self.proxy.transport.get(

View file

@ -0,0 +1,405 @@
"""Every spend row a live proxy writes joins its virtual key (MAT-180).
One virtual key with an alias, owned by a user with an email, drives every spend
write path a key can reach: /chat/completions, /queue/chat/completions,
/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch
input file upload, a batch create, and a replayed callback log (POST
/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry
`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash
/key/generate returns as `token`), which is the join /spend/logs?api_key= and
/user/daily/activity rely on to report key_alias and user_email. A row keyed by a
re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a
key-hash-* row with no alias and no email in the customer's usage exports.
The health-check service account writes rows too; those must stay keyed by the
literal service-account name, never by a hash of it. A batch's cost row is
written by the retrieve that first sees the batch in a terminal state, so the
batch the run creates is one OpenAI fails at validation within seconds (its one
line targets /v1/embeddings under a /v1/chat/completions batch), and the test
retrieves it by its raw provider id with the same key until it is failed. A raw
id is never owned by the CheckBatchCost poller, so that retrieve prices the batch
inline against the retrieving key and its {provider_batch_id}_batch_cost row
must join the key's token with its alias. A completed batch with a positive
cost is out of a single run's reach (OpenAI's completion window is 24h, and a
stack booted fresh per run lists no earlier run's batches), so the poller's own
row is not asserted here.
/spend/logs carries no email field, so the email assertion lives on
/user/daily/activity alone; /spend/logs is held to the alias in metadata.
"""
import base64
import time
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from models import KeyGenerateBody
from proxy_client import Converged, await_converged
from pydantic import BaseModel
from spend_e2e_client import (
BatchCreateBody,
BatchObject,
CallbackLogMetadata,
CallbackLogPayload,
DailyActivityKeyBreakdown,
ResponseIdentity,
SpendClient,
SpendLogRow,
StreamingResponse,
unique_marker,
)
pytestmark = pytest.mark.e2e
CHAT_MODEL: Final = "gemini-2.5-flash"
MESSAGES_MODEL: Final = "claude-haiku-4-5"
RESPONSES_MODEL: Final = "openai-responses-codex"
EMBED_MODEL: Final = "openai-text-embedding-3-small"
BATCH_MODEL: Final = "openai-gpt-4o-mini"
BATCH_BACKEND_MODEL: Final = "gpt-4o-mini"
BATCH_PROVIDER: Final = "openai"
HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check"
BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"})
FAILED_BATCH_POLL_SECONDS: Final = 120.0
FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0
MAX_TOKENS: Final = 8
REPLAY_RESPONSE_COST: Final = 0.0001
REPLAY_PROMPT_TOKENS: Final = 5
REPLAY_COMPLETION_TOKENS: Final = 1
WRITE_PATHS: Final = (
"chat_completions",
"queue_chat_completions",
"messages",
"responses",
"embeddings",
"gemini_passthrough",
"batch_file_upload",
"batch_create",
"callback_replay",
)
class EmbeddingLineBody(BaseModel):
model: str
input: str
class EmbeddingLine(BaseModel):
custom_id: str
method: str = "POST"
url: str = "/v1/embeddings"
body: EmbeddingLineBody
@dataclass(frozen=True, slots=True)
class AttributedKey:
key: str
token: str
alias: str
email: str
user_id: str
@dataclass(frozen=True, slots=True)
class WritePath:
name: str
request_id: str
@dataclass(frozen=True, slots=True)
class DrivenKey:
identity: AttributedKey
paths: tuple[WritePath, ...]
started_at: datetime
def _body_id(name: str, sent: StreamingResponse) -> WritePath:
assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
response_id: Final = ResponseIdentity.model_validate_json(sent.body).id
assert response_id, f"{name} answered without a response id: {sent.body[:300]}"
return WritePath(name=name, request_id=response_id)
def _call_id(name: str, sent: StreamingResponse) -> WritePath:
assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
assert sent.call_id, f"{name} answered without an x-litellm-call-id header"
return WritePath(name=name, request_id=sent.call_id)
def _endpoint_mismatched_jsonl(marker: str) -> bytes:
line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker))
return f"{line.model_dump_json()}\n".encode()
def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]:
uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker))
created: Final = client.create_batch(
identity.key,
BatchCreateBody(
input_file_id=uploaded.id,
model=BATCH_MODEL,
metadata={"run": marker},
),
)
return (
WritePath(name="batch_file_upload", request_id=uploaded.id),
WritePath(name="batch_create", request_id=created.id),
)
def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath:
request_id: Final = f"callback-replay-{marker}"
finished_at: Final = time.time()
replayed: Final = client.replay_callback_log(
identity.key,
CallbackLogPayload(
id=request_id,
litellm_call_id=request_id,
model=CHAT_MODEL,
start_time=finished_at - 1,
end_time=finished_at,
response_cost=REPLAY_RESPONSE_COST,
prompt_tokens=REPLAY_PROMPT_TOKENS,
completion_tokens=REPLAY_COMPLETION_TOKENS,
total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS,
metadata=CallbackLogMetadata(
user_api_key_hash=identity.token,
user_api_key_alias=identity.alias,
user_api_key_user_id=identity.user_id,
),
),
)
assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}"
return WritePath(name="callback_replay", request_id=request_id)
def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]:
marker: Final = unique_marker()
prompt: Final = f"Reply with the word ok. {marker}"
key: Final = identity.key
return (
_body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
_body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
_body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)),
_body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)),
_call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)),
_call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
*_drive_batch(client, identity, marker),
_drive_callback_replay(client, identity, marker),
)
def _provider_batch_id(unified_batch_id: str) -> str:
encoded: Final = unified_batch_id.removeprefix("batch_")
decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode()
return decoded.removeprefix("litellm:").split(";", 1)[0]
def _driven_batch_id(driven: DrivenKey) -> str:
return next(path.request_id for path in driven.paths if path.name == "batch_create")
def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject:
outcome: Final = await_converged(
lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER),
converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES,
timeout=FAILED_BATCH_POLL_SECONDS,
interval=FAILED_BATCH_POLL_INTERVAL_SECONDS,
now=time.monotonic,
sleep=time.sleep,
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
return [
row
for row in client.proxy.spend_logs_window(
start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1)
)
if HEALTH_SERVICE_ACCOUNT in (row.request_tags or [])
]
def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
outcome: Final = await_converged(
lambda: _health_rows_between(client, started_at),
converged=lambda rows: bool(rows),
timeout=client.proxy.poll_timeout,
interval=client.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
class TestKeyAttribution:
@pytest.fixture(scope="class")
def driven(self, client: SpendClient) -> Iterator[DrivenKey]:
marker: Final = unique_marker()
user_id: Final = client.create_user(
email=f"key-attribution-{marker}@example.com",
role="proxy_admin",
user_id=f"key-attribution-{marker}",
)
record: Final = client.generate_key_record(
KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}")
)
assert record.token, "/key/generate answered without the key's token hash"
assert record.key_alias, "/key/generate dropped the key alias"
identity: Final = AttributedKey(
key=record.key,
token=record.token,
alias=record.key_alias,
email=f"key-attribution-{marker}@example.com",
user_id=user_id,
)
started_at: Final = datetime.now(timezone.utc)
try:
yield DrivenKey(
identity=identity,
paths=_drive_every_write_path(client, identity),
started_at=started_at,
)
finally:
client.proxy.delete_key(identity.key)
client.delete_user(identity.user_id)
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.joins_key",
exercised_on=[
"chat_completions",
"messages",
"responses",
"embeddings",
"batches",
"files",
"google_native",
"rust_control_plane",
],
)
def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None:
assert tuple(path.name for path in driven.paths) == WRITE_PATHS
found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths)
unwritten: Final = [path.name for path, rows in found if not rows]
assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}"
unjoined: Final = [
(path.name, row.call_type, row.api_key)
for path, rows in found
for row in rows
if row.api_key != driven.identity.token
]
assert not unjoined, (
"spend rows whose api_key does not join LiteLLM_VerificationToken.token "
f"{driven.identity.token}: {unjoined}"
)
unaliased: Final = [
(path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None)
for path, rows in found
for row in rows
if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias
]
assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}"
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.reports_alias_and_email",
exercised_on=[
"chat_completions",
"messages",
"responses",
"embeddings",
"batches",
"files",
"google_native",
"rust_control_plane",
],
)
def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None:
expected_ids: Final = frozenset(path.request_id for path in driven.paths)
rows: Final = client.poll_logs_for_key(
driven.identity.key,
min_rows=len(driven.paths),
predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found),
)
missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows)
assert not missing, (
f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: "
f"{sorted(path.name for path in driven.paths if path.request_id in missing)}"
)
aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows)
assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}"
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.reports_alias_and_email",
exercised_on=[
"chat_completions",
"messages",
"responses",
"embeddings",
"batches",
"files",
"google_native",
"rust_control_plane",
],
)
def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None:
breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key(
driven.identity.token,
start=driven.started_at - timedelta(days=1),
end=datetime.now(timezone.utc) + timedelta(days=1),
min_requests=len(driven.paths),
)
assert breakdown is not None, (
f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: "
"the key's rows did not aggregate under its token"
)
assert breakdown.metrics.api_requests >= len(driven.paths), (
f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, "
f"expected at least {len(driven.paths)}"
)
assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}"
assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}"
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.health_rows_keep_service_account",
exercised_on=["chat_completions"],
)
def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None:
started_at: Final = datetime.now(timezone.utc)
probe: Final = client.health(CHAT_MODEL)
assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}"
rows: Final = _health_rows_since(client, started_at)
assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row"
rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT]
assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}"
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key",
exercised_on=["batches"],
)
def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None:
provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven))
fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id)
assert fetched.status == "failed", (
f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after "
f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted"
)
cost_request_id: Final = f"{provider_batch_id}_batch_cost"
rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id)
assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}"
call_types: Final = tuple(sorted({row.call_type or "" for row in rows}))
assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}"
unjoined: Final = [
(row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None)
for row in rows
if row.api_key != driven.identity.token
or row.metadata is None
or row.metadata.user_api_key_alias != driven.identity.alias
]
assert not unjoined, (
f"batch cost rows that do not join the retrieving key's token {driven.identity.token} "
f"with alias {driven.identity.alias!r}: {unjoined}"
)

View file

@ -2,6 +2,7 @@ import glob
import os
import re
import sys
from pathlib import Path
import pytest
@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting:
harness.run()
assert len(harness.deploy_calls) == 1
assert harness.resolved == []
class TestJWTKeyMappingCascade:
"""Regression tests for issue #33702.
A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted
because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so
deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign
key violation. The mapping must be removed automatically when its key is
deleted, which the FK now enforces via ON DELETE CASCADE.
"""
_FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey"
def _effective_on_delete(self):
"""Replay every migration in order and return the last ON DELETE action
declared for the JWT key mapping FK."""
action = None
for _migration_name, sql in _get_all_migrations():
for match in re.finditer(
rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?'
r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)",
sql,
re.IGNORECASE | re.DOTALL,
):
action = re.sub(r"\s+", " ", match.group(1).upper())
return action
def test_fk_effective_on_delete_is_cascade(self):
"""The final FK definition across all migrations must cascade deletes."""
assert self._effective_on_delete() == "CASCADE", (
f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a "
"virtual key removes its JWT key mapping (issue #33702)"
)
def test_schema_declares_cascade_on_relation(self):
"""schema.prisma must declare onDelete: Cascade on the mapping relation
so the generated client and DB agree."""
schema_paths = glob.glob(
os.path.abspath(
os.path.join(
os.path.dirname(__file__), "../../**/schema.prisma"
)
),
recursive=True,
)
declaring = tuple(
(path, schema)
for path, schema in ((p, Path(p).read_text()) for p in schema_paths)
if "model LiteLLM_JWTKeyMapping" in schema
)
assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found"
for path, schema in declaring:
match = re.search(
r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)",
schema,
)
assert match is not None, (
f"{path} declares LiteLLM_JWTKeyMapping but its verification token "
"relation could not be parsed, so this test cannot vouch for it "
"(issue #33702)"
)
assert "onDelete: Cascade" in match.group(1), (
f"{path} must declare onDelete: Cascade on the JWT key mapping "
"relation (issue #33702)"
)

View file

@ -1623,15 +1623,11 @@ class TestMissingChoicesGuard:
assert "no 'choices'" in exc_info.value.message
def test_convert_to_model_response_object_empty_choices_raises_api_error(self):
"""Empty choices list raises APIError, same as missing/null choices.
def test_convert_to_model_response_object_empty_choices_returns_empty_list(self):
"""An empty choices list is a real provider answer, so it converts to choices=[] instead of raising.
Provider-specific repair (e.g. github_copilot synthesizing choices for
Anthropic-native responses) happens before this guard, in the provider
config; the core utility keeps treating empty choices as an error.
See: https://github.com/BerriAI/litellm/issues/40276
"""
from litellm.exceptions import APIError
response_object = {
"id": "msg_123",
"model": "some-model",
@ -1639,16 +1635,17 @@ class TestMissingChoicesGuard:
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
}
with pytest.raises(APIError) as exc_info:
convert_to_model_response_object(
response_object=response_object,
model_response_object=ModelResponse(),
)
result = convert_to_model_response_object(
response_object=response_object,
model_response_object=ModelResponse(),
)
assert "no 'choices'" in exc_info.value.message
assert isinstance(result, ModelResponse)
assert result.choices == []
assert result.usage.prompt_tokens == 10
def test_convert_to_model_response_object_null_choices_raises_api_error(self):
"""choices=None raises APIError."""
"""choices=None raises APIError that names the type instead of claiming the key is missing."""
from litellm.exceptions import APIError
response_object = {
@ -1664,7 +1661,7 @@ class TestMissingChoicesGuard:
model_response_object=ModelResponse(),
)
assert "no 'choices'" in exc_info.value.message
assert "'choices' that is not a list (NoneType)" in exc_info.value.message
def test_convert_to_streaming_response_no_choices_raises_api_error(self):
"""Missing choices in streaming cache-hit path raises APIError."""

View file

@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter
def _fake_hf_tokenizer(num_tokens: int) -> MagicMock:
encoding = MagicMock()
encoding.ids = list(range(num_tokens))
encoding.__len__.return_value = num_tokens
tokenizer = MagicMock()
tokenizer.encode.return_value = encoding
tokenizer.encode_batch_fast.return_value = [encoding]
return tokenizer
@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch):
)
)
mock_tokenizer_cls.from_pretrained.assert_called_once_with(
"my-org/custom-tokenizer", revision="v2", auth_token=None
)
mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None)
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.request_model == "my-embedding-model"
assert response.model_used == "self-hosted-embedder"
assert response.total_tokens > 0
assert response.total_tokens >= 7
@pytest.mark.asyncio

View file

@ -1341,6 +1341,257 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None):
)
@pytest.mark.parametrize("level", ["DEFAULT", "ERROR"])
@pytest.mark.parametrize(
"headers,metadata,expected_id",
[
({"x-litellm-session-id": "session-7125"}, {}, "call"),
({"X-Claude-Code-Session-Id": "session-7125"}, {}, "call"),
({"x-session-id": "session-7125"}, {}, "call"),
({"session-id": "session-7125", "user-agent": "codex_cli_rs/1.0"}, {}, "call"),
({"thread-id": "session-7125", "user-agent": "codex-tui"}, {}, "call"),
({"session_id": "session-7125", "user-agent": "Codex 1.0"}, {}, "call"),
({"conversation_id": "session-7125", "user-agent": "codex_vscode/1.0"}, {}, "call"),
({"x-litellm-session-id": "short"}, {}, "call"),
({"x-litellm-trace-id": "session-7125"}, {}, "session-7125"),
(
{"X-LiteLLM-Trace-Id": "session-7125", "x-litellm-session-id": "session-7125"},
{},
"session-7125",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_trace_id": "session-7125"},
{},
"session-7125",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_trace_id": "explicit-trace"},
{},
"explicit-trace",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_existing_trace_id": "existing-trace"},
{},
"existing-trace",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_session_id": "custom-session"},
{},
"call",
),
(
{"x-litellm-session-id": "short", "langfuse_session_id": "custom-session"},
{},
"call",
),
(
{"X-Claude-Code-Session-Id": "session-7125", "langfuse_session_id": "custom-session"},
{},
"call",
),
(
{"x-session-id": "session-7125", "langfuse_session_id": "custom-session"},
{},
"call",
),
(
{
"session-id": "session-7125",
"user-agent": "codex_cli_rs/1.0",
"langfuse_session_id": "custom-session",
},
{},
"call",
),
(
{
"x-litellm-session-id": "session-7125",
"langfuse_session_id": "custom-session",
"x-litellm-trace-id": "explicit-trace",
},
{},
"explicit-trace",
),
(
{
"x-litellm-session-id": "session-7125",
"langfuse_session_id": "custom-session",
"langfuse_trace_id": "explicit-trace",
},
{},
"explicit-trace",
),
(
{
"x-litellm-session-id": "session-7125",
"langfuse_session_id": "custom-session",
"langfuse_existing_trace_id": "existing-trace",
},
{},
"existing-trace",
),
({}, {"trace_id": "session-7125", "session_id": "session-7125"}, "session-7125"),
({}, {"trace_id": "explicit-trace", "session_id": "session-7125"}, "explicit-trace"),
(
{"x-vendor-session-id": "short"},
{"trace_id": "short", "session_id": "short"},
"short",
),
(
{"x-session-id": "invalid value"},
{"trace_id": "invalid value", "session_id": "invalid value"},
"invalid value",
),
(
{"session-id": "session-7125", "user-agent": "codexfoo/1.0"},
{"trace_id": "session-7125", "session_id": "session-7125"},
"session-7125",
),
(
{"x-vendor-session-id": "short"},
{"trace_id": "session-7125", "session_id": "session-7125"},
"session-7125",
),
({}, {}, "call"),
],
)
def test_session_header_trace_provenance(headers, metadata, expected_id, level):
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
clean_headers,
redact_credential_headers,
)
logger: Final = _steering_logger()
for turn in range(2):
call_id = f"call-{turn}"
request_headers = Headers(headers)
data = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=request_headers, data={"metadata": dict(metadata)}, _metadata_variable_name="metadata"
)
original_metadata = dict(data["metadata"])
now = datetime.datetime.now()
result = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": call_id,
"litellm_trace_id": data.get("litellm_trace_id"),
"litellm_params": {
"metadata": data["metadata"],
"proxy_server_request": {"headers": redact_credential_headers(clean_headers(request_headers))},
},
"messages": [{"role": "user", "content": f"turn {turn}"}],
"optional_params": {},
},
response_obj=(
None
if level == "ERROR"
else litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}])
),
start_time=now,
end_time=now,
level=level,
status_message="provider error" if level == "ERROR" else None,
)
trace_params = logger.Langfuse.trace.call_args.kwargs
assert trace_params["id"] == (call_id if expected_id == "call" else expected_id)
assert result["trace_id"] == trace_params["id"]
if expected_id != "existing-trace":
assert trace_params["session_id"] == headers.get("langfuse_session_id", original_metadata.get("session_id"))
steering = {key[len("langfuse_") :]: value for key, value in headers.items() if key.startswith("langfuse_")}
assert data["metadata"] == {**original_metadata, **steering}
def test_session_header_trace_without_call_id_keeps_session_alias():
logger: Final = _steering_logger()
now: Final = datetime.datetime.now()
result: Final = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": "",
"litellm_params": {
"metadata": {"trace_id": "session-7125", "session_id": "session-7125"},
"proxy_server_request": {"headers": {"x-litellm-session-id": "session-7125"}},
},
"messages": [{"role": "user", "content": "no call id"}],
"optional_params": {},
},
response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]),
start_time=now,
end_time=now,
)
assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125"
assert result["trace_id"] == "session-7125"
def test_every_proxy_session_header_shape_is_classified_as_a_session_alias():
"""The classifier must cover every header shape the proxy turns into a chain id."""
from litellm.integrations.langfuse.langfuse import _is_session_header_trace
from litellm.proxy.litellm_pre_call_utils import (
_CODEX_SESSION_ID_HEADERS,
get_chain_id_from_headers,
)
session: Final = "session-7125-abcdef"
session_shapes: Final = (
{"x-litellm-session-id": session},
{"X-Claude-Code-Session-Id": session},
{"x-session-id": session},
*({header: session, "user-agent": "codex_cli_rs/1.0"} for header in _CODEX_SESSION_ID_HEADERS),
)
for headers in session_shapes:
assert get_chain_id_from_headers(dict(headers)) == session, headers
assert _is_session_header_trace(session, session, {"headers": headers}) is True, headers
explicit_trace: Final = {"x-litellm-trace-id": session, "x-litellm-session-id": session}
assert get_chain_id_from_headers(dict(explicit_trace)) == session
assert _is_session_header_trace(session, session, {"headers": explicit_trace}) is False
@pytest.mark.parametrize(
"proxy_server_request",
[None, {}, {"headers": None}],
ids=["no-proxy-request", "no-headers-key", "null-headers"],
)
def test_sdk_caller_without_request_headers_keeps_its_trace(proxy_server_request):
"""A direct SDK caller has no request headers, so a session-shaped trace id stays the caller's."""
logger: Final = _steering_logger()
now: Final = datetime.datetime.now()
result: Final = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": "call-0",
"litellm_params": {
"metadata": {"trace_id": "session-7125", "session_id": "session-7125"},
"proxy_server_request": proxy_server_request,
},
"messages": [{"role": "user", "content": "sdk turn"}],
"optional_params": {},
},
response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]),
start_time=now,
end_time=now,
)
assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125"
assert result["trace_id"] == "session-7125"
def test_session_header_classifier_survives_non_string_header_keys():
"""A non-string header key must not cost the caller its whole trace."""
from litellm.integrations.langfuse.langfuse import _is_session_header_trace
session: Final = "session-7125-abcdef"
headers: Final = {7: "numeric key", "x-litellm-session-id": session}
assert _is_session_header_trace(session, session, {"headers": headers}) is True
assert _is_session_header_trace(session, session, {"headers": {7: "numeric key"}}) is False
def test_mask_input_header_false_keeps_the_prompt():
logger = _steering_logger()

View file

@ -0,0 +1,40 @@
import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Final, TypeVar
import litellm
T = TypeVar("T")
def warm_tokenizer(model: str) -> None:
litellm.token_counter(model=model, text="load the tokenizer before anything is timed")
async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]:
async def wake_lag() -> float:
started: Final = time.perf_counter()
await asyncio.sleep(0.001)
return time.perf_counter() - started - 0.001
return tuple([await wake_lag() for _ in iter(until.is_set, True)])
async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]:
finished: Final = asyncio.Event()
async def timed() -> tuple[T, float]:
await asyncio.sleep(0)
started: Final = time.perf_counter()
try:
return await run(), time.perf_counter() - started
finally:
finished.set()
(result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished))
return result, took, lags
def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None:
assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count"

View file

@ -1,5 +1,6 @@
import os
import json
from collections.abc import Mapping, Sequence
from pathlib import Path
import pytest
@ -11,8 +12,6 @@ from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebS
from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams
def test_web_search_cost_low():
web_search_options = WebSearchOptions(search_context_size="low")
model_info = litellm.get_model_info("gpt-4o-search-preview")
@ -683,12 +682,13 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model(
def _openai_responses_with_web_search_calls(model, num_calls):
from litellm.types.llms.openai import ResponsesAPIResponse
from openai.types.responses.response_function_web_search import (
ActionSearch,
ResponseFunctionWebSearch,
)
from litellm.types.llms.openai import ResponsesAPIResponse
output = [
ResponseFunctionWebSearch(
id=f"ws_{i}",
@ -859,11 +859,62 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map)
custom_llm_provider="openai",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(0.035), (
f"dated search-preview id must bill the $0.035 search fee, got ${cost}"
assert cost == pytest.approx(0.025), (
f"dated search-preview id must bill the $0.025 search fee, got ${cost}"
)
@pytest.mark.parametrize(
"web_search_options",
[
None,
WebSearchOptions(search_context_size="low"),
WebSearchOptions(search_context_size="medium"),
WebSearchOptions(search_context_size="high"),
],
)
def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias(
web_search_options: WebSearchOptions | None, local_model_cost_map: None
) -> None:
alias_info = litellm.get_model_info("gpt-4o-mini")
snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18")
assert not snapshot_info["supports_web_search"]
assert not alias_info["supports_web_search"]
snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=web_search_options, model_info=snapshot_info
)
alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=web_search_options, model_info=alias_info
)
assert snapshot_cost == alias_cost == 0.025
def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps():
repo_root = Path(__file__).parents[4]
cost_maps = tuple(
json.loads((repo_root / path).read_text(encoding="utf-8"))
for path in (
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
)
)
canonical, backup = cost_maps
expected_search_price = {
"search_context_size_low": 0.025,
"search_context_size_medium": 0.025,
"search_context_size_high": 0.025,
}
for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"):
canonical_entry = canonical[model_name]
backup_entry = backup[model_name]
assert canonical_entry["search_context_cost_per_query"] == expected_search_price
assert backup_entry["search_context_cost_per_query"] == expected_search_price
assert canonical_entry == backup_entry
# Note: File search integration test removed due to complex annotation detection logic
# The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage

View file

@ -1,4 +1,6 @@
from typing import Final
import pytest
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -99,3 +101,97 @@ def test_handle_invalid_parallel_tool_calls_skips_custom_tool_calls():
)
result = _handle_invalid_parallel_tool_calls([custom_tool_call, function_tool_call])
assert result == [custom_tool_call, function_tool_call]
def test_convert_empty_choices_response() -> None:
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_streaming_response,
)
resp: Final = {
"id": "x",
"created": 1,
"model": "gemini-3.5-flash",
"object": "chat.completion",
"choices": [],
"usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10},
"vertex_ai_safety_results": ["blocked"],
}
result: Final = convert_to_model_response_object(
response_object=resp,
model_response_object=ModelResponse(),
response_type="completion",
)
assert result.choices == []
assert getattr(result, "vertex_ai_safety_results") == ["blocked"]
sync_stream: Final = list(convert_to_streaming_response(response_object=resp))
assert len(sync_stream) == 1
assert sync_stream[0].choices == []
@pytest.mark.asyncio
async def test_convert_empty_choices_response_async() -> None:
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_streaming_response_async,
)
resp: Final = {
"id": "x",
"created": 1,
"model": "gemini-3.5-flash",
"object": "chat.completion",
"choices": [],
"usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10},
}
async_chunks: Final = [chunk async for chunk in convert_to_streaming_response_async(response_object=resp)]
assert len(async_chunks) == 1
assert async_chunks[0].choices == []
def test_convert_missing_choices_raises_api_error() -> None:
from litellm.exceptions import APIError
resp: Final = {
"id": "x",
"created": 1,
"model": "gemini-3.5-flash",
"object": "chat.completion",
}
with pytest.raises(APIError) as exc_info:
convert_to_model_response_object(
response_object=resp,
model_response_object=ModelResponse(),
response_type="completion",
)
assert "no 'choices'" in str(exc_info.value)
@pytest.mark.parametrize(("choices", "type_name"), [({}, "dict"), ("", "str"), (None, "NoneType"), (0, "int")])
@pytest.mark.asyncio
async def test_convert_non_list_choices_raises_api_error(choices: object, type_name: str) -> None:
from litellm.exceptions import APIError
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_streaming_response,
convert_to_streaming_response_async,
)
resp: Final = {
"id": "x",
"created": 1,
"model": "gemini-3.5-flash",
"object": "chat.completion",
"choices": choices,
}
expected: Final = f"'choices' that is not a list \\({type_name}\\)"
with pytest.raises(APIError, match=expected):
convert_to_model_response_object(
response_object=resp,
model_response_object=ModelResponse(),
response_type="completion",
)
with pytest.raises(APIError, match=expected):
list(convert_to_streaming_response(response_object=resp))
with pytest.raises(APIError, match=expected):
async for _ in convert_to_streaming_response_async(response_object=resp):
pass

View file

@ -6,7 +6,7 @@ import pytest
import asyncio
import traceback
from typing import Optional
from typing import Final, Optional
import litellm
from litellm import verbose_logger
@ -2633,6 +2633,48 @@ def test_dispatch_cached_response_extracts_delta(
assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1"
def test_dispatch_cached_response_without_choices_is_an_empty_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
"""A cached completion with no choices replays as an empty, unfinished chunk
instead of raising IndexError on choices[0]."""
initialized_custom_stream_wrapper.custom_llm_provider = "cached_response"
chunk: Final = ModelResponseStream(id="chatcmpl-cache-empty", choices=[])
result, model_response, completion_obj = _run_dispatch(
initialized_custom_stream_wrapper, chunk
)
assert isinstance(result, _ProviderChunkParsed)
assert completion_obj["content"] is None
assert initialized_custom_stream_wrapper.received_finish_reason is None
assert model_response.id == "chatcmpl-cache-empty"
@pytest.mark.asyncio
async def test_cached_response_without_choices_streams_a_single_stop_chunk(
logging_obj: Logging,
):
"""A stream cache hit on a completion stored with choices == [] ends with one
finish_reason=stop chunk, the same shape the live empty stream produced."""
async def cached_chunks():
yield ModelResponseStream(id="chatcmpl-cache-empty", choices=[])
wrapper: Final = CustomStreamWrapper(
completion_stream=cached_chunks(),
model="test-model",
logging_obj=logging_obj,
custom_llm_provider="cached_response",
)
chunks: Final = tuple([chunk async for chunk in wrapper])
assert len(chunks) == 1
assert tuple(choice.finish_reason for chunk in chunks for choice in chunk.choices) == ("stop",)
assert all(choice.delta.content in (None, "") for chunk in chunks for choice in chunk.choices)
def test_dispatch_vertex_ai_legacy_text_and_finish_reason(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):

View file

@ -1,11 +1,16 @@
#### What this tests ####
# This tests litellm.token_counter.token_counter() function
import asyncio
import base64
import importlib
import threading
import time
import traceback
from concurrent.futures import Future, wait
from typing import Final
from unittest.mock import MagicMock
import anyio.to_thread
import pytest
import tiktoken
@ -15,13 +20,23 @@ import litellm
from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
from litellm import token_counter as token_counter_old
import litellm.constants
from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.token_counter import (
_get_exact_count_function,
_get_extrapolating_count_function,
_get_tiktoken_count_function,
calculate_img_tokens,
high_detail_image_token_upper_bound,
offload_token_count,
)
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
from tests.test_litellm.litellm_core_utils.messages_with_counts import (
MESSAGES_TEXT,
MESSAGES_WITH_IMAGES,
@ -125,6 +140,135 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch):
importlib.reload(litellm.constants)
async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free():
warm_tokenizer("claude-fable-5")
tokens, took, lags = await timed_with_loop_lags(
lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100)
)
assert tokens > 0
assert_loop_stayed_free(took, lags)
@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500])
def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int):
count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk))
front_heavy: Final = "a" * 1_000 + "b" * 4_000
exact: Final = 1_000 + len(front_heavy)
estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy)
assert abs(estimate - exact) <= exact // 100
assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars
def test_count_at_or_below_the_cap_is_exact():
count_exactly: Final = MagicMock(side_effect=len)
assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000
assert count_exactly.call_args_list == [(("a" * 5_000,),)]
class _SlowEncoder:
def __init__(self) -> None:
self._lock: Final = threading.Lock()
self.in_flight = 0
self.peak_in_flight = 0
def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
with self._lock:
self.in_flight += 1
self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
time.sleep(0.1)
with self._lock:
self.in_flight -= 1
return [[0] * len(text) for text in texts]
@pytest.mark.asyncio
async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool():
encoder: Final = _SlowEncoder()
count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder})
shared_pool: Final = anyio.to_thread.current_default_thread_limiter()
burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]:
if counting.done():
return ()
await asyncio.sleep(0.01)
return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting))
counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst))))
borrowed: Final = await shared_pool_borrowed_until_done(counting)
assert await counting == [3] * burst
assert len(borrowed) > 1 and max(borrowed) == 0
assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None:
def slow_count(counted: str) -> int:
time.sleep(0.1)
return len(counted)
result.set_result(asyncio.run(offload_token_count(slow_count)(text)))
def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process():
loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
results: Final = tuple(Future[int]() for _ in range(loops))
threads: Final = tuple(
threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True)
for size, result in enumerate(results, start=1)
)
for thread in threads:
thread.start()
_, pending = wait(results, timeout=5)
assert not pending
assert tuple(result.result() for result in results) == tuple(range(1, loops + 1))
@pytest.mark.parametrize(
("configured", "expected"),
[("8", 8), ("0", 4), ("not-an-int", 4)],
)
def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured)
try:
assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected
finally:
monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS")
importlib.reload(litellm.constants)
def test_token_counter_applies_the_default_cap():
max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS
prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]
over_the_cap: Final = prose + "a" * 200_000
exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap)
estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap)
assert estimate != exact
assert abs(estimate - exact) <= exact // 100
@pytest.mark.parametrize(
("configured", "expected"),
[("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)],
)
def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured)
try:
assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected
finally:
monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS")
importlib.reload(litellm.constants)
def test_token_counter_with_prefix():
messages = [
{"role": "user", "content": "Who won the world cup in 2022?"},

View file

@ -2270,3 +2270,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
assert open_key == StreamingScanKey(texts=("hi",))
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
class TestAnthropicMessagesHandlerPostCallHookResponse:
def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
from litellm.types.utils import Choices, Message, ModelResponse, Usage
assembled = ModelResponse(
id="msg_1",
model="claude",
choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")],
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
)
hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled)
assert hook_response["type"] == "message"
assert hook_response["role"] == "assistant"
assert hook_response["content"] == [{"type": "text", "text": "hello world"}]
assert hook_response["stop_reason"] == "end_turn"
assert hook_response["usage"]["input_tokens"] == 1
assert hook_response["usage"]["output_tokens"] == 2
def test_anything_else_reaches_the_hook_untouched(self):
native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]}
assert AnthropicMessagesHandler().post_call_hook_response(native) is native

View file

@ -41,6 +41,21 @@ from litellm.types.utils import (
)
def test_translate_openai_response_to_anthropic_empty_choices() -> None:
response: Final = ModelResponse(
id="chatcmpl-empty",
model="gemini-3.5-flash",
choices=[],
usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10),
)
result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response)
assert result["content"] == []
assert result["stop_reason"] == "end_turn"
assert result["usage"]["input_tokens"] == 10
def test_translate_chat_refusal_to_anthropic_response():
response = ModelResponse(
id="chatcmpl-refusal",

View file

@ -385,6 +385,58 @@ def test_select_azure_base_url_called(setup_mocks):
setup_mocks["select_url"].assert_called_once()
def test_initialize_defaults_max_retries_to_litellm_default(setup_mocks):
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={},
api_key="test-api-key",
api_base="https://test.openai.azure.com",
model_name="gpt-4",
api_version="2023-06-01",
is_async=False,
)
assert result["max_retries"] == litellm.constants.DEFAULT_MAX_RETRIES
@pytest.mark.parametrize(
"configured, expected",
[(0, 0), (5, 5), (None, litellm.constants.DEFAULT_MAX_RETRIES)],
)
def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expected):
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={"max_retries": configured},
api_key="test-api-key",
api_base="https://test.openai.azure.com",
model_name="gpt-4",
api_version="2023-06-01",
is_async=False,
)
assert result["max_retries"] == expected
def test_default_max_retries_env_var_reaches_azure_sdk_client():
import subprocess
import sys
code = (
"from litellm.llms.azure.common_utils import BaseAzureLLM\n"
"client = BaseAzureLLM().get_azure_openai_client("
"api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01',"
" client=None, _is_async=True, litellm_params={}, model='gpt-4')\n"
"print(client.max_retries)"
)
completed = subprocess.run(
[sys.executable, "-c", code],
env={**os.environ, "DEFAULT_MAX_RETRIES": "0"},
capture_output=True,
text=True,
check=True,
)
assert completed.stdout.strip() == "0"
@pytest.mark.parametrize(
"call_type",
[

View file

@ -6,6 +6,7 @@ import pytest
import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.azure_ai.image_edit import (
AzureFoundryMAIImageEditConfig,
get_azure_ai_image_edit_config,
@ -70,44 +71,48 @@ class TestAzureMAIImageEdit:
assert "/mai/v1/images/edits" in url
assert "api-version=preview" in url
def test_map_openai_params_keeps_size(self):
config = AzureFoundryMAIImageEditConfig()
optional_params = config.map_openai_params(
image_edit_optional_params={"size": "1792x1024", "n": 1},
def test_get_optional_params_image_edit_size_raises_400(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(litellm.UnsupportedParamsError, match="size") as exc_info:
ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={"size": "1024x1024", "n": 1},
)
assert exc_info.value.status_code == 400
def test_get_optional_params_image_edit_size_dropped_with_drop_params(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
optional_params = ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={"size": "1024x1024", "n": 1},
drop_params=True,
)
assert optional_params["size"] == "1792x1024"
assert "size" not in optional_params
assert optional_params["n"] == 1
assert "width" not in optional_params
assert "height" not in optional_params
def test_map_openai_params_defaults_size(self):
config = AzureFoundryMAIImageEditConfig()
optional_params = config.map_openai_params(
image_edit_optional_params={},
def test_get_optional_params_image_edit_without_size_forwards_nothing_extra(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
optional_params = ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
drop_params=True,
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={},
)
assert optional_params["size"] == "1024x1024"
assert optional_params == {}
def test_map_openai_params_unsupported_size_raises(self):
config = AzureFoundryMAIImageEditConfig()
with pytest.raises(ValueError, match="Unsupported size value: 'auto'"):
config.map_openai_params(
image_edit_optional_params={"size": "auto"},
model="MAI-Image-2.5",
drop_params=True,
)
def test_map_openai_params_invalid_size_format_raises(self):
config = AzureFoundryMAIImageEditConfig()
with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"):
config.map_openai_params(
image_edit_optional_params={"size": "1024xabc"},
model="MAI-Image-2.5",
drop_params=True,
def test_image_edit_size_surfaces_as_400(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.image_edit(
model="azure_ai/MAI-Image-2.5",
image=io.BytesIO(b"fake-image-bytes"),
prompt="Turn this into a studio product shot",
size="1024x1024",
api_key="test-key",
api_base="https://my-resource.services.ai.azure.com",
)
assert exc_info.value.status_code == 400
def test_transform_image_edit_request_uses_image_field(self):
config = AzureFoundryMAIImageEditConfig()
@ -117,14 +122,14 @@ class TestAzureMAIImageEdit:
model="MAI-Image-2.5",
prompt="Turn this into a studio product shot",
image=image_bytes,
image_edit_optional_request_params={"size": "1024x1024", "n": 1},
image_edit_optional_request_params={"n": 1},
litellm_params={},
headers={},
)
assert data["model"] == "MAI-Image-2.5"
assert data["prompt"] == "Turn this into a studio product shot"
assert data["size"] == "1024x1024"
assert "size" not in data
assert data["n"] == 1
assert len(files) == 1
assert files[0][0] == "image"

View file

@ -3,8 +3,8 @@ from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.exceptions import UnsupportedParamsError
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation import get_azure_image_generation_config
from litellm.llms.azure.image_generation.http_utils import (
@ -29,9 +29,7 @@ from litellm.utils import get_optional_params_image_gen
class TestAzureMAIImageGeneration:
def test_is_mai_model(self):
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5")
assert AzureFoundryMAIImageGenerationConfig.is_mai_model(
"azure_ai/MAI-Image-2.5"
)
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("azure_ai/MAI-Image-2.5")
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash")
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e")
assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro")
@ -42,16 +40,10 @@ class TestAzureMAIImageGeneration:
api_base="https://my-resource.services.ai.azure.com",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
)
assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
def test_get_mai_image_generation_url_preserves_full_path(self):
api = (
"https://my-resource.services.ai.azure.com/mai/v1/images/generations"
"?api-version=preview"
)
api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
api_base=api,
api_version="preview",
@ -63,10 +55,7 @@ class TestAzureMAIImageGeneration:
api_base="https://my-resource.services.ai.azure.com/mai/v1",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
)
assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
def test_get_azure_ai_image_generation_config_returns_mai(self):
config = get_azure_ai_image_generation_config("MAI-Image-2.5")
@ -104,13 +93,13 @@ class TestAzureMAIImageGeneration:
config = AzureFoundryMAIImageGenerationConfig()
optional_params = get_optional_params_image_gen(
model="MAI-Image-2.5",
size="1792x1024",
size="1024x1024",
n=1,
custom_llm_provider="azure_ai",
provider_config=config,
drop_params=True,
)
assert optional_params["width"] == 1792
assert optional_params["width"] == 1024
assert optional_params["height"] == 1024
assert "size" not in optional_params
@ -127,10 +116,7 @@ class TestAzureMAIImageGeneration:
assert "api-version=preview" in url
def test_mai_json_body_keeps_model(self):
api = (
"https://my-resource.services.ai.azure.com/mai/v1/images/generations"
"?api-version=preview"
)
api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
data = {
"model": "MAI-Image-2.5",
"prompt": "A photograph of a red fox",
@ -176,7 +162,7 @@ class TestAzureMAIImageGeneration:
def test_map_openai_params_unsupported_size_raises(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="Unsupported size value: 'auto'"):
with pytest.raises(UnsupportedParamsError, match="Unsupported size value: 'auto'"):
config.map_openai_params(
non_default_params={"size": "auto"},
optional_params={},
@ -186,7 +172,7 @@ class TestAzureMAIImageGeneration:
def test_map_openai_params_invalid_custom_size_raises(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"):
with pytest.raises(UnsupportedParamsError, match="Invalid size format: '1024xabc'"):
config.map_openai_params(
non_default_params={"size": "1024xabc"},
optional_params={},
@ -194,9 +180,138 @@ class TestAzureMAIImageGeneration:
drop_params=True,
)
@pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"])
def test_map_openai_params_size_below_minimum_dimension_raises(self, size):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="at least 768 pixels"):
config.map_openai_params(
non_default_params={"size": size},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
@pytest.mark.parametrize("size", ["1792x1024", "1024x1792"])
def test_map_openai_params_size_over_total_pixel_budget_raises(self, size):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"):
config.map_openai_params(
non_default_params={"size": size},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
@pytest.mark.parametrize("size", ["1032x1024", "1376x768"])
def test_map_openai_params_size_at_live_pixel_cap_passes_through(self, size):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"size": size},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
assert optional_params["width"] * optional_params["height"] == 1_056_768
def test_map_openai_params_size_one_pixel_over_live_cap_raises(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"):
config.map_openai_params(
non_default_params={"size": "1033x1024"},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
def test_map_openai_params_explicit_width_height_not_range_checked(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"width": 1792, "height": 1024},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["width"] == 1792
assert optional_params["height"] == 1024
@pytest.mark.parametrize("n", [2, 4, "2", 0, -1])
def test_map_openai_params_n_other_than_one_raises(self, n):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"):
config.map_openai_params(
non_default_params={"n": n},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
def test_map_openai_params_non_numeric_n_raises_400(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="not a whole number of images") as exc_info:
config.map_openai_params(
non_default_params={"n": "abc"},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
assert exc_info.value.status_code == 400
def test_get_optional_params_image_gen_global_drop_params_drops_multi_image_n(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", True)
optional_params = get_optional_params_image_gen(
model="MAI-Image-2.5",
n=4,
custom_llm_provider="azure_ai",
provider_config=AzureFoundryMAIImageGenerationConfig(),
)
assert "n" not in optional_params
assert optional_params["width"] == 1024
def test_get_optional_params_image_gen_without_any_drop_params_still_raises(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"):
get_optional_params_image_gen(
model="MAI-Image-2.5",
n=4,
custom_llm_provider="azure_ai",
provider_config=AzureFoundryMAIImageGenerationConfig(),
)
def test_map_openai_params_multi_image_n_dropped_with_drop_params(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"n": 4},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert "n" not in optional_params
def test_map_openai_params_single_image_n_still_passes_through(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"n": 1},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
assert optional_params["n"] == 1
@pytest.mark.parametrize("params", [{"n": 2}, {"n": "abc"}, {"size": "512x512"}, {"size": "1792x1024"}])
def test_image_generation_rejected_params_surface_as_400(self, params):
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.image_generation(
model="azure_ai/MAI-Image-2.5",
prompt="A photograph of a red fox",
api_key="test-key",
api_base="https://my-resource.services.ai.azure.com",
**params,
)
assert exc_info.value.status_code == 400
def test_map_openai_params_unsupported_param_raises(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="Parameter quality is not supported"):
with pytest.raises(UnsupportedParamsError, match="Parameter quality is not supported"):
config.map_openai_params(
non_default_params={"quality": "hd"},
optional_params={},
@ -343,16 +458,12 @@ class TestAzureMAIImageGeneration:
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "azure_ai/MAI-Image-2.5"
model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai")
image_response = ImageResponse(
data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
)
image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")])
cost = azure_ai_image_cost_calculator(
model=model,
image_response=image_response,
)
assert (
cost == len(image_response.data or []) * model_info["output_cost_per_image"]
)
assert cost == len(image_response.data or []) * model_info["output_cost_per_image"]
assert cost > 0

View file

@ -1,5 +1,7 @@
import json
from unittest.mock import MagicMock
import httpx
import pytest
@ -190,3 +192,45 @@ def test_get_error_class_preserves_provider_headers():
assert isinstance(error, BedrockError)
assert error.headers == {"x-amzn-RequestId": "req-invoke-500"}
assert error.response.headers["x-amzn-requestid"] == "req-invoke-500"
def test_transform_response_hands_json_mode_to_nova():
"""The invoke dispatcher forwards its json_mode argument to Nova instead of dropping it."""
from litellm.types.utils import ModelResponse
response_json = {
"output": {
"message": {
"role": "assistant",
"content": [
{
"toolUse": {
"toolUseId": "tooluse_nova_json",
"name": "json_tool_call",
"input": {"city": "Paris", "temperature": 21},
}
}
],
}
},
"stopReason": "tool_use",
"usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
}
raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock"))
result = AmazonInvokeConfig().transform_response(
model="invoke/amazon.nova-lite-v1:0",
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=MagicMock(),
request_data={},
messages=[{"role": "user", "content": "weather"}],
optional_params={},
litellm_params={},
encoding=None,
api_key=None,
json_mode=True,
)
assert result.choices[0].message.tool_calls is None
assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21}

View file

@ -382,6 +382,8 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto():
"us.openai.gpt-5.6-sol",
"global.openai.gpt-5.6-terra",
"bedrock/converse/us.openai.gpt-5.6-luna",
"us.openai.gpt-6-astra",
"bedrock/converse/global.openai.gpt-6-astra",
],
)
def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(model, local_model_cost_map):
@ -412,6 +414,7 @@ def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(mode
[
"us.openai.gpt-5.6-sol",
"bedrock/converse/global.openai.gpt-5.6-luna",
"us.openai.gpt-6-astra",
],
)
def test_openai_gpt5_converse_never_forwards_thinking(model, local_model_cost_map):
@ -863,6 +866,191 @@ def test_get_supported_openai_params():
assert "reasoning_effort" in supported_params
@pytest.mark.parametrize(
"model",
[
"bedrock/us.deepseek.r1-v1:0",
"bedrock/converse/us.deepseek.r1-v1:0",
"bedrock/deepseek.v3-v1:0",
"bedrock/deepseek.v3.2",
],
)
def test_bedrock_deepseek_does_not_advertise_thinking(model):
"""DeepSeek reasons natively on Bedrock and does not take the Anthropic-shaped `thinking`
field (R1 400s on it, V3 ignores it), so it must not be advertised as supported."""
config = AmazonConverseConfig()
supported_params = config.get_supported_openai_params(model=model)
assert "thinking" not in supported_params
assert "output_config" not in supported_params
@pytest.mark.parametrize("model", ["bedrock/us.deepseek.r1-v1:0", "bedrock/converse/us.deepseek.r1-v1:0"])
def test_bedrock_deepseek_r1_does_not_advertise_reasoning_effort(model):
"""DeepSeek R1 always reasons and returns a 400 for any reasoning_effort shape."""
config = AmazonConverseConfig()
assert "reasoning_effort" not in config.get_supported_openai_params(model=model)
@pytest.mark.parametrize("model", ["bedrock/deepseek.v3-v1:0", "bedrock/deepseek.v3.2", "bedrock/us.deepseek.v3.2"])
def test_bedrock_deepseek_v3_advertises_reasoning_effort(model):
"""DeepSeek V3 on Bedrock accepts a raw reasoning_effort in additionalModelRequestFields."""
config = AmazonConverseConfig()
assert "reasoning_effort" in config.get_supported_openai_params(model=model)
@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"])
def test_bedrock_deepseek_thinking_raises_without_drop_params(model):
"""Passing `thinking` to Bedrock DeepSeek must fail client-side with a clear
UnsupportedParamsError instead of leaking through to Bedrock."""
with pytest.raises(litellm.UnsupportedParamsError):
litellm.utils.get_optional_params(
model=model,
custom_llm_provider="bedrock",
thinking={"type": "enabled", "budget_tokens": 1024},
)
def test_bedrock_deepseek_r1_reasoning_effort_raises_without_drop_params():
with pytest.raises(litellm.UnsupportedParamsError):
litellm.utils.get_optional_params(
model="us.deepseek.r1-v1:0",
custom_llm_provider="bedrock",
reasoning_effort="high",
)
@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"])
def test_bedrock_deepseek_thinking_dropped_does_not_leak_into_request(model):
"""With drop_params, `thinking` is dropped rather than forwarded into
additionalModelRequestFields for Bedrock DeepSeek."""
optional_params = litellm.utils.get_optional_params(
model=model,
custom_llm_provider="bedrock",
thinking={"type": "enabled", "budget_tokens": 1024},
drop_params=True,
)
assert "thinking" not in optional_params
config = AmazonConverseConfig()
request = config._transform_request(
model=f"bedrock/converse/{model}",
messages=[{"role": "user", "content": "Say hi in one word."}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert "thinking" not in (request.get("additionalModelRequestFields") or {})
@pytest.mark.parametrize("param", ["thinking", "reasoning_effort"])
def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param):
"""Even when map_openai_params is called directly (bypassing the supported-params
gate), DeepSeek R1 must not forward thinking/reasoning_effort into
additionalModelRequestFields, since Bedrock rejects both with a 400."""
config = AmazonConverseConfig()
model = "bedrock/converse/us.deepseek.r1-v1:0"
value = {"type": "enabled", "budget_tokens": 1024} if param == "thinking" else "high"
optional_params = config.map_openai_params(
non_default_params={param: value, "max_tokens": 100},
optional_params={},
model=model,
drop_params=False,
)
assert "thinking" not in optional_params
assert "reasoning_effort" not in optional_params
request = config._transform_request(
model=model,
messages=[{"role": "user", "content": "Say hi in one word."}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert request.get("additionalModelRequestFields") is None
def test_bedrock_deepseek_v3_reasoning_effort_forwarded_raw():
"""DeepSeek V3 takes reasoning_effort verbatim in additionalModelRequestFields, never
converted into the Anthropic `thinking` block that Claude models get."""
config = AmazonConverseConfig()
model = "bedrock/deepseek.v3.2"
optional_params = config.map_openai_params(
non_default_params={"reasoning_effort": "high", "max_tokens": 100},
optional_params={},
model=model,
drop_params=False,
)
assert "thinking" not in optional_params
request = config._transform_request(
model=model,
messages=[{"role": "user", "content": "Say hi in one word."}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert request["additionalModelRequestFields"] == {"reasoning_effort": "high"}
def test_bedrock_deepseek_v3_thinking_dropped_by_map():
config = AmazonConverseConfig()
optional_params = config.map_openai_params(
non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100},
optional_params={},
model="bedrock/deepseek.v3.2",
drop_params=False,
)
assert "thinking" not in optional_params
assert "reasoning_effort" not in optional_params
@pytest.mark.parametrize(
"model, param, value, kept_key",
[
(
"bedrock/us.anthropic.claude-opus-4-20250514-v1:0",
"thinking",
{"type": "enabled", "budget_tokens": 1024},
"thinking",
),
(
"bedrock/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123",
"thinking",
{"type": "enabled", "budget_tokens": 1024},
"thinking",
),
(
"bedrock/openai.gpt-oss-safeguard-20b-1:0",
"reasoning_effort",
"high",
"reasoning_effort",
),
(
"bedrock/us.amazon.nova-2-lite-v1:0",
"reasoning_effort",
"high",
"reasoningConfig",
),
],
)
def test_bedrock_non_deepseek_reasoning_params_preserved(model, param, value, kept_key):
"""The DeepSeek leak fix must only drop reasoning request params for DeepSeek.
Claude behind an application-inference-profile ARN, gpt-oss-safeguard (absent from the
cost map so `supports_reasoning` is False), and Nova 2 all reason via a request param and
must keep it. Regression guard against gating the drop on a positive allowlist, which
silently degraded reasoning for anything the allowlist/ARN introspection missed."""
config = AmazonConverseConfig()
optional_params = config.map_openai_params(
non_default_params={param: value, "max_tokens": 100},
optional_params={},
model=model,
drop_params=False,
)
assert kept_key in optional_params
def test_get_supported_openai_params_bedrock_converse():
"""
Test that all documented bedrock converse models have the same set of supported openai params when using
@ -6727,3 +6915,41 @@ def test_forced_tool_choice_forwarded_on_converse_models_that_support_it(
)
assert result == {"any": {}}
def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it():
response_json = {
"metrics": {"latencyMs": 900},
"output": {
"message": {
"content": [
{
"toolUse": {
"input": {"city": "Paris", "population": 2100000},
"name": "json_tool_call",
"toolUseId": "tooluse_invoke_nova_json",
}
}
],
"role": "assistant",
}
},
"stopReason": "tool_use",
"usage": {"inputTokens": 40, "outputTokens": 20, "totalTokens": 60},
}
raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock.test"))
logging_obj = MagicMock()
result = AmazonConverseConfig().transform_response(
model="bedrock/invoke/us.amazon.nova-micro-v1:0",
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
request_data={},
messages=[],
optional_params={"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": {}}}]},
litellm_params={},
encoding=None,
json_mode=True,
)
assert result.choices[0].message.tool_calls is None
assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000}

View file

@ -10,18 +10,23 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
import litellm
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
from litellm.llms.openai.common_utils import OpenAIError
from litellm.main import responses_api_bridge_check
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
class TestChatGPTResponsesAPITransformation:
@pytest.mark.parametrize(
"model_name",
[
"chatgpt/gpt-5.5",
"chatgpt/gpt-5.6-luna",
"chatgpt/gpt-5.6-sol",
"chatgpt/gpt-5.6-terra",
"chatgpt/gpt-5.4",
"chatgpt/gpt-5.4-pro",
"chatgpt/gpt-5.3-chat-latest",
@ -40,6 +45,52 @@ class TestChatGPTResponsesAPITransformation:
assert isinstance(config, ChatGPTResponsesAPIConfig)
assert config.custom_llm_provider == LlmProviders.CHATGPT
@pytest.mark.parametrize(
"model_name",
[
"chatgpt/gpt-5.5",
"chatgpt/gpt-5.6-luna",
"chatgpt/gpt-5.6-sol",
"chatgpt/gpt-5.6-terra",
],
)
def test_chatgpt_responses_model_metadata(self, model_name: str, local_model_cost_map: None) -> None:
model_info = litellm.get_model_info(model_name)
assert model_info["litellm_provider"] == "chatgpt"
assert model_info["mode"] == "responses"
assert model_info["supported_endpoints"] == [
"/v1/chat/completions",
"/v1/responses",
]
assert model_info["max_input_tokens"] == 1050000
assert model_info["max_output_tokens"] == 128000
@pytest.mark.parametrize(
"model_name",
[
"gpt-5.5",
"gpt-5.6-luna",
"gpt-5.6-sol",
"gpt-5.6-terra",
],
)
def test_chatgpt_models_bridge_chat_completions_to_responses(
self, model_name: str, local_model_cost_map: None
) -> None:
"""A chat completions request for these models must take the Responses bridge.
`gpt-5.6-*` also exists as an openai chat model, so an unregistered
chatgpt model resolves to mode "chat" here and never reaches the bridge.
"""
model_info, resolved_model = responses_api_bridge_check(
model=model_name,
custom_llm_provider="chatgpt",
)
assert model_info["mode"] == "responses"
assert resolved_model == model_name
@patch("litellm.llms.chatgpt.responses.transformation.Authenticator")
def test_chatgpt_responses_endpoint_url(self, mock_authenticator_class):
mock_auth_instance = MagicMock()

View file

@ -255,6 +255,19 @@ def test_transform_messages_sanitizes_empty_content():
assert result[1]["content"] == "Hi"
def test_transform_request_preserves_unity_model_service_name():
config = DatabricksConfig()
result = config.transform_request(
model="system.ai.kimi-k3",
messages=[{"role": "user", "content": "hello"}],
optional_params={},
litellm_params={},
headers={},
)
assert result["model"] == "system.ai.kimi-k3"
def test_transform_request_strips_thinking_blocks_and_reasoning_content():
"""Regression for LIT-6762: replaying an assistant turn that litellm decorated with
`thinking_blocks` / `reasoning_content` made Databricks 400 with
@ -590,3 +603,87 @@ def test_chunk_parser_without_usage_still_parses_content():
assert result.id == "chatcmpl-test"
assert result.model == "databricks-claude-sonnet-5"
assert result.choices[0]["delta"]["content"] == "hi"
@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"])
def test_transform_choices_surfaces_top_level_reasoning_content(reasoning_key: str) -> None:
config = DatabricksConfig()
databricks_choices = [
{
"message": {
"role": "assistant",
"content": "391",
reasoning_key: "We need answer just number. 17*23=391.",
},
"index": 0,
"finish_reason": "stop",
}
]
choices = config._transform_dbrx_choices(choices=databricks_choices)
assert choices[0].message.content == "391"
assert choices[0].message.reasoning_content == "We need answer just number. 17*23=391."
assert getattr(choices[0].message, "thinking_blocks", None) is None
def test_transform_choices_parses_think_tags_in_string_content():
config = DatabricksConfig()
databricks_choices = [
{
"message": {"role": "assistant", "content": "<think>17 times 23</think>391"},
"index": 0,
"finish_reason": "stop",
}
]
choices = config._transform_dbrx_choices(choices=databricks_choices)
assert choices[0].message.content == "391"
assert choices[0].message.reasoning_content == "17 times 23"
def test_transform_choices_prefers_reasoning_blocks_over_top_level_field():
config = DatabricksConfig()
databricks_choices = [
{
"message": {
"role": "assistant",
"content": [
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "from block"}]},
{"type": "text", "text": "391"},
],
"reasoning_content": "from field",
},
"index": 0,
"finish_reason": "stop",
}
]
choices = config._transform_dbrx_choices(choices=databricks_choices)
assert choices[0].message.reasoning_content == "from block"
assert choices[0].message.content == "391"
@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"])
def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> None:
iterator = DatabricksChatResponseIterator(None, sync_stream=True)
chunk = {
"id": "1",
"object": "chat.completion.chunk",
"created": 0,
"model": "lit-qa-deepseek-v4-flash",
"choices": [
{
"delta": {"role": "assistant", "content": None, reasoning_key: "We need answer"},
"index": 0,
"finish_reason": None,
}
],
}
parsed = iterator.chunk_parser(chunk)
assert parsed.choices[0].delta.reasoning_content == "We need answer"
assert parsed.choices[0].delta.content is None

View file

@ -657,6 +657,77 @@ class TestEndpointURLConstruction:
assert api_base.endswith("/chat/completions")
def test_chat_gateway_endpoint_for_unity_model_on_legacy_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/serving-endpoints",
api_key="test-key",
model="system.ai.kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
def test_chat_gateway_endpoint_preserves_explicit_gateway_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/ai-gateway/mlflow/v1/",
api_key="test-key",
model="system.ai.kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
def test_chat_gateway_preserves_unity_model_service_name_with_explicit_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
config = DatabricksConfig()
request = config.transform_request(
model="catalog.schema.kimi-k3",
messages=[{"role": "user", "content": "hello"}],
optional_params={},
litellm_params={},
headers={},
)
assert config.get_complete_url(
api_base="https://test.net/ai-gateway/mlflow/v1",
api_key="test-key",
model="catalog.schema.kimi-k3",
optional_params={},
litellm_params={},
) == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
assert request["model"] == "catalog.schema.kimi-k3"
def test_chat_legacy_endpoint_remains_default(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/serving-endpoints",
api_key="test-key",
model="databricks-kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/serving-endpoints/chat/completions"
def test_embeddings_endpoint(self, monkeypatch):
"""Embeddings endpoint is correctly appended."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)

View file

@ -10,11 +10,19 @@ from collections.abc import Callable
from typing import Any, List, Literal, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
import logging
import pytest
from fastapi import HTTPException
from openai.types.responses import ResponseFunctionToolCall
from pydantic import BaseModel
from openai.types.responses import (
ResponseCustomToolCall,
ResponseCustomToolCallInputDeltaEvent,
ResponseCustomToolCallInputDoneEvent,
ResponseFunctionToolCall,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -23,11 +31,12 @@ from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.types.llms.openai import ChatCompletionToolCallChunk
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
@ -57,6 +66,60 @@ class MockGuardrail(CustomGuardrail):
return inputs
class PersimmonMaskingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
tool_calls = [
{
**tool_call,
"function": {
**tool_call["function"],
"arguments": tool_call["function"]["arguments"].replace("persimmon", "[MASKED]"),
},
}
for tool_call in inputs.get("tool_calls", [])
]
return {**inputs, "tool_calls": tool_calls}
class FlatShapeGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])]
return {**inputs, "tool_calls": flat_tool_calls}
class DroppingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
return {**inputs, "tool_calls": []}
CUSTOM_TOOL_CALL_ITEM = {
"type": "custom_tool_call",
"id": "ctc_1",
"call_id": "call_exec_1",
"name": "exec",
"input": "echo persimmon",
"status": "completed",
}
class TestOpenAIResponsesHandlerDiscovery:
"""Test that the handler is properly discovered by the guardrail system"""
@ -557,7 +620,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
texts_to_check: List[str] = []
images_to_check: List[str] = []
tool_calls_to_check: List[Any] = []
tool_calls_to_check: List[ChatCompletionToolCallChunk] = []
task_mappings: List[Tuple[int, int]] = []
# Extract tool calls
@ -628,6 +691,123 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
== '{"location":"Boston, MA","unit":"celsius"}'
)
@pytest.mark.parametrize(
"output_item",
[
dict(CUSTOM_TOOL_CALL_ITEM),
CustomToolCallOutputItem(**CUSTOM_TOOL_CALL_ITEM),
ResponseCustomToolCall(**{key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "status"}),
],
ids=["dict", "litellm_typed", "openai_typed"],
)
def test_extract_custom_tool_call_input_as_arguments(self, output_item):
handler = OpenAIResponsesHandler()
texts_to_check: List[str] = []
tool_calls_to_check: List[Any] = []
handler._extract_output_text_and_images(
output_item=output_item,
output_idx=2,
texts_to_check=texts_to_check,
images_to_check=[],
task_mappings=[],
tool_calls_to_check=tool_calls_to_check,
)
assert texts_to_check == []
assert tool_calls_to_check == [
{
"id": "call_exec_1",
"type": "function",
"function": {"name": "exec", "arguments": "echo persimmon"},
"index": 2,
}
]
@pytest.mark.asyncio
@pytest.mark.parametrize("typed", [False, True], ids=["dict", "typed"])
async def test_process_output_response_writes_tool_call_rewrites_back(self, typed):
handler = OpenAIResponsesHandler()
function_call = {
"type": "function_call",
"id": "fc_1",
"call_id": "call_fn_1",
"name": "lookup_fruit",
"arguments": '{"fruit": "persimmon"}',
"status": "completed",
}
message = {
"type": "message",
"id": "msg_1",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "running persimmon", "annotations": []}],
}
payload = {
"id": "resp_1",
"created_at": 1,
"model": "gpt-5.6",
"object": "response",
"status": "completed",
"output": [message, function_call, dict(CUSTOM_TOOL_CALL_ITEM)],
}
response = ResponsesAPIResponse.model_validate(payload) if typed else payload
result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask"))
output = result.output if typed else result["output"]
function_item, custom_item = output[1], output[2]
assert (function_item.arguments if typed else function_item["arguments"]) == '{"fruit": "[MASKED]"}'
assert (custom_item.input if typed else custom_item["input"]) == "echo [MASKED]"
assert (custom_item.name if typed else custom_item["name"]) == "exec"
assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon"
@staticmethod
def _custom_tool_call_response(item: dict) -> dict:
return {
"id": "resp_1",
"created_at": 1,
"model": "gpt-5.6",
"object": "response",
"status": "completed",
"output": [item],
}
@pytest.mark.asyncio
async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self):
handler = OpenAIResponsesHandler()
response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat"))
assert result["output"][0]["input"] == "echo persimmon"
assert result["output"][0]["name"] == "exec"
@pytest.mark.asyncio
async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog):
handler = OpenAIResponsesHandler()
response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper"))
assert result["output"][0]["input"] == "echo persimmon"
assert any(
"dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage()
for record in caplog.records
)
@pytest.mark.asyncio
async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self):
handler = OpenAIResponsesHandler()
nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"}
response = self._custom_tool_call_response(nameless_item)
result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask"))
assert result["output"][0]["input"] == "echo [MASKED]"
assert "name" not in result["output"][0]
@pytest.mark.asyncio
async def test_process_output_response_with_tool_calls(self):
"""Test processing output response containing function tool calls"""
@ -1315,6 +1495,128 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}'
assert completed_event.response.output[0].name == "lookup_fruit"
@staticmethod
def _ended_custom_tool_call_stream_events() -> List[dict]:
def item(input_text: str, status: str) -> dict:
return {**CUSTOM_TOOL_CALL_ITEM, "input": input_text, "status": status}
return [
{"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
{"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "},
{"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"},
{"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"},
{"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"created_at": 1,
"model": "gpt-5.6",
"output": [item("echo persimmon", "completed")],
"status": "completed",
},
},
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_custom_tool_call_events(self):
handler = OpenAIResponsesHandler()
events = self._ended_custom_tool_call_stream_events()
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["item"]["input"] == ""
assert events[1]["delta"] == "echo [MASKED]"
assert events[2]["delta"] == ""
assert events[3]["input"] == "echo [MASKED]"
assert events[4]["item"]["input"] == "echo [MASKED]"
assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]"
assert events[5]["response"]["output"][0]["name"] == "exec"
assert "arguments" not in events[5]["response"]["output"][0]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self):
handler = OpenAIResponsesHandler()
events = self._ended_custom_tool_call_stream_events()
items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]]
for item in items:
del item["name"]
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert events[3]["input"] == "echo [MASKED]"
assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]"
assert all("name" not in item for item in items)
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self):
from litellm.types.llms.openai import (
OutputItemAddedEvent,
OutputItemDoneEvent,
ResponseCompletedEvent,
)
handler = OpenAIResponsesHandler()
typed_events: List[BaseModel] = [
model.model_validate({**event, "sequence_number": sequence_number})
for sequence_number, (model, event) in enumerate(
zip(
(
OutputItemAddedEvent,
ResponseCustomToolCallInputDeltaEvent,
ResponseCustomToolCallInputDeltaEvent,
ResponseCustomToolCallInputDoneEvent,
OutputItemDoneEvent,
ResponseCompletedEvent,
),
self._ended_custom_tool_call_stream_events(),
)
)
]
completed_event = typed_events[5]
assert isinstance(completed_event.response.output[0], CustomToolCallOutputItem)
await handler.process_output_streaming_response(
responses_so_far=typed_events,
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert typed_events[1].delta == "echo [MASKED]"
assert typed_events[2].delta == ""
assert typed_events[3].input == "echo [MASKED]"
assert typed_events[4].item.input == "echo [MASKED]"
assert completed_event.response.output[0].input == "echo [MASKED]"
assert completed_event.response.output[0].name == "exec"
@pytest.mark.asyncio
async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = self._ended_custom_tool_call_stream_events()
events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@staticmethod
def _bridged_function_call_stream_events() -> List[dict]:
reasoning = {"type": "reasoning", "id": "rs_1", "summary": []}
@ -2747,8 +3049,21 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
def test_completed_event_with_a_custom_tool_call_changes_the_key(self):
handler = OpenAIResponsesHandler()
message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
ended_key = handler.get_streaming_scan_key(
[self._delta(0, "hi"), self._completed(1, [message, dict(CUSTOM_TOOL_CALL_ITEM)])]
)
rewritten_key = handler.get_streaming_scan_key(
[self._delta(0, "hi"), self._completed(1, [message, {**CUSTOM_TOOL_CALL_ITEM, "input": "echo kumquat"}])]
)
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "echo persimmon" in ended_key.tool_calls[0]
assert rewritten_key != ended_key
def test_completed_event_reads_every_output_text_part(self):
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText
item = GenericResponseOutputItem(
type="message",

View file

@ -322,8 +322,8 @@ class TestModelCostEntry:
entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"]
assert entry["mode"] == "audio_transcription"
assert entry["litellm_provider"] == "vertex_ai"
assert entry["input_cost_per_audio_token"] == pytest.approx(2.5e-06)
assert entry["input_cost_per_token"] == pytest.approx(2.5e-06)
assert entry["input_cost_per_audio_token"] == pytest.approx(2e-06)
assert entry["input_cost_per_token"] == pytest.approx(2e-06)
assert entry["output_cost_per_token"] == pytest.approx(1.2e-05)
assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]

View file

@ -159,6 +159,10 @@ class TestUpCommand:
assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key"
assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true"
assert "apiKeyHelper" not in captured["settings"]
# The ephemeral proxy serves only the autorouter, so a starting model left by
# `lite configure claude --model` or a user pin would 400 on the first message.
assert captured["settings"]["model"] == "autorouter"
assert captured["settings"]["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter"
assert captured["settings_mode"] == 0o600
assert terminate_calls == [99999]

View file

@ -1,63 +0,0 @@
from litellm.proxy.client.cli.commands.autoroute.settings import (
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS,
merge_claude_settings_static_token,
)
def test_preserves_unrelated_top_level_keys():
merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc")
assert merged["theme"] == "dark"
def test_preserves_unrelated_env_keys():
settings = {"env": {"SOME_OTHER_VAR": "value"}}
merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert merged["env"]["SOME_OTHER_VAR"] == "value"
def test_sets_base_url_and_auth_token():
merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc")
assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000"
assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc"
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true"
def test_preserves_existing_tool_search():
settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}}
merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false"
def test_drops_stray_api_key():
settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}}
merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert "ANTHROPIC_API_KEY" not in merged["env"]
def test_removes_existing_api_key_helper():
settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"}
merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert "apiKeyHelper" not in merged
def test_does_not_mutate_input():
settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"}
merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"}
def test_forces_all_claude_code_default_model_tiers_to_the_autorouter():
# A bare "*" model_name deployment looks like the obvious way to catch every request
# regardless of which model Claude Code thinks it's using, but Router's auto-router
# registry is keyed by the literal requested model string with no wildcard resolution
# (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude
# Code's own tiers hit the auto-router is to override the env vars it reads per tier.
merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc")
for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS:
assert merged["env"][key] == "autorouter"
def test_overrides_a_preexisting_default_model_env_var():
settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}}
merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc")
assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter"

View file

@ -27,6 +27,7 @@ from litellm.proxy.client.cli.commands.auth import (
print_token,
whoami,
)
from litellm.proxy.client.cli.commands import auth as auth_module
from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner
@ -1398,8 +1399,10 @@ class TestLoginConfigClaude:
def setup_method(self):
self.runner = CliRunner()
def _run_login(self, tmp_path, args, base_url="https://test.example.com"):
def _run_login(self, tmp_path, monkeypatch, args, base_url="https://test.example.com"):
settings_path = tmp_path / "claude" / "settings.json"
monkeypatch.setattr(auth_module, "CLAUDE_SETTINGS_PATH", settings_path)
monkeypatch.setattr(auth_module, "CONFIGURE_STATE_PATH", tmp_path / "claude_configure_state.json")
backup_path = tmp_path / "claude_settings_backup.json"
poll_response = Mock()
poll_response.status_code = 200
@ -1416,7 +1419,6 @@ class TestLoginConfigClaude:
patch("requests.get", return_value=poll_response),
patch("litellm.proxy.client.cli.commands.auth.save_cli_token"),
patch("litellm.proxy.client.cli.interface.show_commands"),
patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path),
patch(
"litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS",
(SettingsFileOwner(backup_path, "lite up", "lite down"),),
@ -1429,16 +1431,16 @@ class TestLoginConfigClaude:
result = self.runner.invoke(login, args, obj={"base_url": base_url})
return result, settings_path, backup_path
def test_default_login_does_not_touch_claude_settings(self, tmp_path):
result, settings_path, _backup_path = self._run_login(tmp_path, [])
def test_default_login_does_not_touch_claude_settings(self, tmp_path, monkeypatch):
result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, [])
assert result.exit_code == 0
assert "Login successful!" in result.output
assert not settings_path.exists()
assert "Configured Claude Code" not in result.output
def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path):
result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path, monkeypatch):
result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"])
assert result.exit_code == 0
written = json.loads(settings_path.read_text())
@ -1446,25 +1448,43 @@ class TestLoginConfigClaude:
assert written["env"]["ENABLE_TOOL_SEARCH"] == "true"
assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token"
assert "Configured Claude Code" in result.output
assert "pins a proxy model for every tier" not in result.output
assert "the model Claude Code starts on" in result.output
def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path):
def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path, monkeypatch):
settings_path = tmp_path / "claude" / "settings.json"
settings_path.parent.mkdir(parents=True)
settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}}))
result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"])
assert result.exit_code == 0
written = json.loads(settings_path.read_text())
assert written["theme"] == "dark"
assert written["env"]["KEEP"] == "me"
def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path):
def test_refuses_before_logging_in_while_lite_up_holds_the_settings(self, tmp_path, monkeypatch):
# The local precondition comes first: no browser, no token stored, no "Login successful!".
backup_path = tmp_path / "claude_settings_backup.json"
backup_path.write_text("{}")
monkeypatch.setattr(auth_module, "CLAUDE_SETTINGS_PATH", tmp_path / "claude" / "settings.json")
monkeypatch.setattr(
auth_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),)
)
with patch("requests.post") as post, patch("webbrowser.open") as browser:
result = self.runner.invoke(login, ["--config-claude"], obj={"base_url": "https://test.example.com"})
assert result.exit_code != 0
assert "not logging in" in result.output and "lite down" in result.output
assert "Login successful!" not in result.output
post.assert_not_called()
browser.assert_not_called()
def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path, monkeypatch):
settings_path = tmp_path / "claude" / "settings.json"
settings_path.parent.mkdir(parents=True)
settings_path.write_text("not json at all {{{")
result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"])
assert result.exit_code != 0
assert "Login successful!" in result.output

View file

@ -1,4 +1,5 @@
import json
import os
import shlex
import stat
import time
@ -9,14 +10,25 @@ from click.testing import CliRunner
from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord
from litellm.proxy.client.cli import cli
from litellm.litellm_core_utils.private_json import commit_staged_json
from litellm.proxy.client.cli.commands.claude_settings import (
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS,
AUTOROUTE_BACKUP_PATH,
BACKUP_PATH,
OWNED_ENV_KEYS,
OWNED_TOP_LEVEL_KEYS,
SETTINGS_FILE_OWNERS,
ApiKeyHelper,
ClaudeSettingsError,
KeepModel,
SettingsFileOwner,
StartOn,
StaticToken,
UnpinModel,
configure_claude_settings,
merge_claude_settings,
resolve_api_key_helper,
write_claude_settings,
unconfigure_claude_settings,
)
@ -24,6 +36,7 @@ def _owners(*backup_paths):
"""Stand-in owners for the real `lite up` / `lite autoroute up` registry."""
return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths)
CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings"
AUTH_MODULE = "litellm.proxy.client.cli.commands.auth"
WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE"
@ -97,17 +110,28 @@ def lite_on_path():
yield
class TestWriteClaudeSettings:
def _helper_configure(base_url, settings_path, owners, state_path=None):
"""`lite login --config-claude`'s shape: the login credential behind apiKeyHelper, no pinned model."""
state = state_path if state_path is not None else settings_path.parent.parent / "state.json"
root = base_url.rstrip("/")
configure_claude_settings(
root, ApiKeyHelper(resolve_api_key_helper(root)), KeepModel(), settings_path, state, owners
)
class TestConfigureWithTheLoginHelper:
def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path):
settings_path, backup_path = paths
assert not settings_path.parent.exists()
write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com/", settings_path, _owners(backup_path))
written = json.loads(settings_path.read_text())
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com"
assert written["env"]["ENABLE_TOOL_SEARCH"] == "true"
assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token"
assert "model" not in written
def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path):
settings_path, backup_path = paths
@ -123,7 +147,7 @@ class TestWriteClaudeSettings:
)
)
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
written = json.loads(settings_path.read_text())
assert written["theme"] == "dark"
@ -135,26 +159,31 @@ class TestWriteClaudeSettings:
def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path):
settings_path, backup_path = paths
write_claude_settings("https://first.example.com", settings_path, _owners(backup_path))
write_claude_settings("https://second.example.com", settings_path, _owners(backup_path))
_helper_configure("https://first.example.com", settings_path, _owners(backup_path))
_helper_configure("https://second.example.com", settings_path, _owners(backup_path))
written = json.loads(settings_path.read_text())
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com"
assert "second.example.com" in written["apiKeyHelper"]
assert "first.example.com" not in written["apiKeyHelper"]
def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path):
def test_drops_stray_static_credentials_so_the_helper_token_wins(self, paths, lite_on_path):
# Claude Code prefers ANTHROPIC_AUTH_TOKEN over apiKeyHelper, so a virtual key left behind
# by an earlier `lite configure claude --api-key` would silently keep winning.
settings_path, backup_path = paths
settings_path.parent.mkdir(parents=True)
settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}}))
settings_path.write_text(
json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked", "ANTHROPIC_AUTH_TOKEN": "sk-old"}})
)
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"]
env = json.loads(settings_path.read_text())["env"]
assert "ANTHROPIC_API_KEY" not in env and "ANTHROPIC_AUTH_TOKEN" not in env
def test_written_file_is_owner_only(self, paths, lite_on_path):
settings_path, backup_path = paths
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600
def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path):
@ -162,7 +191,7 @@ class TestWriteClaudeSettings:
backup_path.write_text("{}")
with pytest.raises(ClaudeSettingsError, match="lite down"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert not settings_path.exists()
@ -172,7 +201,7 @@ class TestWriteClaudeSettings:
settings_path.write_text("not json at all {{{")
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert settings_path.read_text() == "not json at all {{{"
@ -180,7 +209,7 @@ class TestWriteClaudeSettings:
settings_path, backup_path = paths
with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None):
with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert not settings_path.exists()
@ -196,7 +225,7 @@ class TestWriteClaudeSettings:
settings_path.write_bytes(b'{"theme": "\xff\xfe"}')
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path):
"""An unreadable settings file must not surface as "Authentication failed".
@ -210,16 +239,18 @@ class TestWriteClaudeSettings:
settings_path.mkdir()
with pytest.raises(ClaudeSettingsError, match="Could not read"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path):
settings_path, backup_path = paths
with patch(
f"{CLAUDE_SETTINGS_MODULE}.write_private_json",
side_effect=OSError("Read-only file system"),
):
with pytest.raises(ClaudeSettingsError, match="Read-only file system"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
settings_path.parent.mkdir(parents=True)
settings_path.parent.chmod(0o500)
try:
with pytest.raises(ClaudeSettingsError, match="Could not write"):
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
finally:
settings_path.parent.chmod(0o700)
assert not settings_path.exists()
class TestApiKeyHelperIsActuallyInvocable:
@ -303,7 +334,7 @@ class TestConflictingOwnersOfTheSettingsFile:
backup.write_text("{}")
stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command)
with pytest.raises(ClaudeSettingsError, match="currently managing"):
write_claude_settings("https://proxy.example.com", settings_path, (stand_in,))
_helper_configure("https://proxy.example.com", settings_path, (stand_in,))
backup.unlink()
assert not settings_path.exists()
@ -314,9 +345,9 @@ class TestConflictingOwnersOfTheSettingsFile:
autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down")
with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"):
write_claude_settings("https://proxy.example.com", settings_path, (autoroute,))
_helper_configure("https://proxy.example.com", settings_path, (autoroute,))
with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"):
write_claude_settings("https://proxy.example.com", settings_path, (autoroute,))
_helper_configure("https://proxy.example.com", settings_path, (autoroute,))
def test_the_registry_matches_the_paths_the_commands_actually_use(self):
"""A second definition of the autoroute dir must not drift from this one."""
@ -341,7 +372,7 @@ class TestDoesNotDestroyUserOwnedStructure:
link.parent.mkdir()
link.symlink_to(real)
write_claude_settings("https://proxy.example.com", link, ())
_helper_configure("https://proxy.example.com", link, ())
assert link.is_symlink()
assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com"
@ -354,6 +385,456 @@ class TestDoesNotDestroyUserOwnedStructure:
settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"}))
with pytest.raises(ClaudeSettingsError, match="non-object"):
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
_helper_configure("https://proxy.example.com", settings_path, _owners(backup_path))
assert json.loads(settings_path.read_text())["env"] == "not-an-object"
class TestMergeClaudeSettings:
"""One merge for every way Claude Code gets wired: `lite up`, `lite login --config-claude`,
`lite configure claude` and `lite autoroute up`."""
def test_a_static_token_lands_in_env_and_the_helper_slot_is_cleared(self):
settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "leaked"}}
merged = merge_claude_settings(settings, "http://127.0.0.1:4000/", StaticToken("token-abc"))
assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000"
assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc"
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true"
assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
assert "ANTHROPIC_API_KEY" not in merged["env"]
assert "apiKeyHelper" not in merged
assert "model" not in merged
assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS)
def test_a_helper_lands_top_level_and_the_static_slots_are_cleared(self):
settings = {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-old", "ANTHROPIC_API_KEY": "leaked"}}
merged = merge_claude_settings(settings, "http://127.0.0.1:4000", ApiKeyHelper("lite auth print-token"))
assert merged["apiKeyHelper"] == "lite auth print-token"
assert "ANTHROPIC_AUTH_TOKEN" not in merged["env"] and "ANTHROPIC_API_KEY" not in merged["env"]
def test_keeps_existing_switch_values_and_unrelated_keys_without_mutating_the_input(self):
settings = {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}}
merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc"))
assert merged["theme"] == "dark"
assert merged["env"]["SOME_OTHER_VAR"] == "value"
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false"
assert settings == {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}}
def test_a_default_model_sets_only_the_row_claude_code_starts_on(self):
merged = merge_claude_settings(
{}, "http://127.0.0.1:4000", StaticToken("token-abc"), default_model="claude-auto"
)
assert merged["model"] == "claude-auto"
assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS)
def test_a_tier_model_forces_every_claude_code_tier_as_autoroute_needs(self):
# Router's auto-router registry is keyed by the literal requested model string with no
# wildcard resolution, so `lite autoroute up` overrides the env var each tier reads.
settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}}
merged = merge_claude_settings(
settings, "http://127.0.0.1:4000", StaticToken("token-abc"), tier_model="autorouter"
)
assert {merged["env"][key] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS} == {"autorouter"}
assert "model" not in merged
def test_touches_exactly_the_declared_owned_keys(self):
# The receipt and unconfigure restore exactly OWNED_*_KEYS, so a key the merge writes outside
# that table would be written by configure and never undone.
settings = {
"theme": "dark",
"permissions": {"allow": ["Bash"]},
"env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "old", "ENABLE_TOOL_SEARCH": "false"},
"apiKeyHelper": "old-helper",
"model": "old-model",
}
for credential in (StaticToken("token-abc"), ApiKeyHelper("helper")):
merged = merge_claude_settings(settings, "http://127.0.0.1:4000", credential, default_model="claude-auto")
changed_top_level = {key for key in set(settings) | set(merged) if settings.get(key) != merged.get(key)}
assert changed_top_level - {"env"} <= set(OWNED_TOP_LEVEL_KEYS)
changed_env = {
key
for key in set(settings["env"]) | set(merged["env"])
if settings["env"].get(key) != merged["env"].get(key)
}
assert changed_env <= set(OWNED_ENV_KEYS)
assert merged["permissions"] == {"allow": ["Bash"]}
assert merged["env"]["KEEP_ME"] == "1"
PROXY = "http://127.0.0.1:4000"
ANTHROPIC = "https://api.anthropic.com"
HELPER = ApiKeyHelper("lite auth print-token")
ORIGINAL = {
"theme": "dark",
"permissions": {"allow": ["Bash"]},
"env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC},
"apiKeyHelper": "/usr/local/bin/lite auth print-token",
"model": "claude-opus-5",
}
def _set(path, value):
"""A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file."""
def edit(settings):
section, _, key = path.rpartition(".")
container = settings.setdefault(section, {}) if section else settings
if value is _ABSENT:
container.pop(key, None)
else:
container[key] = value
return settings
return edit
_ABSENT = object()
class _Rig:
"""One settings file plus receipt under tmp_path, driven through the public functions only."""
def __init__(self, tmp_path, initial):
self.settings = tmp_path / "claude" / "settings.json"
self.state = tmp_path / "state" / "claude_configure_state.json"
if initial is not None:
self.settings.parent.mkdir(parents=True)
self.settings.write_text(json.dumps(initial))
def read(self):
return json.loads(self.settings.read_text()) if self.settings.exists() else None
def configure(self, credential=StaticToken("sk-virtual-key"), model=StartOn("claude-auto"), **kwargs):
configure_claude_settings(PROXY, credential, model, self.settings, self.state, (), **kwargs)
def edit(self, *edits):
settings = self.read()
for apply in edits:
settings = apply(settings)
self.settings.write_text(json.dumps(settings))
def unconfigure(self):
return unconfigure_claude_settings(self.settings, self.state, ())
# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure
# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used
# to be one test each; the receipt's rules are what make them all come out right.
UNDO_SCENARIOS = {
"plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}),
"no file before": (None, [], None, {"file_removed": True}),
"no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}),
"null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}),
"empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}),
"user edits stay and are named": (
ORIGINAL,
[_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()},
),
"user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}),
"user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}),
"user removed our key: neither restored nor kept": (
ORIGINAL,
[_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)],
ORIGINAL,
{"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()},
),
"restored names only what changed": (
{"model": "claude-opus-5"},
[],
{"model": "claude-opus-5"},
{
"restored": {
"env.ANTHROPIC_BASE_URL",
"env.ENABLE_TOOL_SEARCH",
"env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",
"apiKeyHelper",
},
"kept": (),
},
{"credential": HELPER, "model": KeepModel()},
),
"repeat across credential kinds keeps the first snapshot": (
ORIGINAL,
[
{"credential": HELPER, "model": UnpinModel()},
{"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")},
],
ORIGINAL,
{},
),
"repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}),
"repeat without a model lets go of our pin, user had one": (
{"model": "claude-opus-5"},
[{"model": UnpinModel()}],
{"model": "claude-opus-5"},
{},
),
"re-login keeps our pin": (None, [{"credential": HELPER, "model": KeepModel()}], None, {"file_removed": True}),
"edit between configures survives an unpin repeat": (
ORIGINAL,
[
_set("model", "my-favourite"),
_set("env.ENABLE_TOOL_SEARCH", "false"),
{"credential": HELPER, "model": UnpinModel()},
],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}},
),
"edit between configures survives a re-login": (
ORIGINAL,
[
_set("model", "my-favourite"),
_set("env.ENABLE_TOOL_SEARCH", "false"),
{"credential": HELPER, "model": KeepModel()},
],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}},
),
"edit between configures: a same-model repeat displaces it, so it is what comes back": (
ORIGINAL,
[_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": HELPER}],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}},
),
"base URL changed since: credentials withheld, receipt kept": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")],
{**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT},
{
"withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)},
"kept": {"env.ANTHROPIC_BASE_URL"},
"receipt_kept": True,
},
),
"base URL changed and back: judged against the URL the restored file holds": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)],
ORIGINAL,
{"withheld": ()},
),
"credential captured beside no URL goes back only beside no URL": (
{"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}},
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")],
{"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}},
{
"withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")},
"receipt_kept": True,
},
),
"restored document empty while a credential is withheld: file goes, receipt stays": (
None,
[
_set("env.ANTHROPIC_API_KEY", "sk-user"),
{"credential": HELPER, "model": KeepModel()},
_set("env.ANTHROPIC_BASE_URL", _ABSENT),
],
None,
{"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True},
{"credential": HELPER, "model": KeepModel()},
),
"a credential the user changed is kept, never also withheld": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")],
{
**ORIGINAL,
"env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"},
"apiKeyHelper": "/opt/mine/helper",
},
{
"withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)},
"kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"},
"receipt_kept": True,
},
),
}
def _expected_file(expected):
if expected is None:
return None
return {k: v for k, v in expected.items() if v is not _ABSENT}
class TestConfigureAndUnconfigure:
"""`configure_claude_settings` records how to undo itself; `unconfigure_claude_settings` undoes only that."""
@pytest.mark.parametrize("scenario", UNDO_SCENARIOS.values(), ids=UNDO_SCENARIOS.keys())
def test_undo_matrix(self, tmp_path, scenario):
initial, steps, expected, outcome_expectations, *first = scenario
rig = _Rig(tmp_path, initial)
rig.configure(**(first[0] if first else {}))
for step in steps:
if isinstance(step, dict):
rig.configure(**step)
elif rig.settings.exists() and step(json.loads(rig.settings.read_text())) is None:
rig.settings.unlink()
else:
rig.edit(step)
outcome = rig.unconfigure()
assert rig.read() == _expected_file(expected)
assert rig.state.exists() == outcome_expectations.get("receipt_kept", False)
for field, want in outcome_expectations.items():
if field == "withheld":
assert {(item.key, item.endpoint) for item in outcome.withheld} == set(want)
elif field == "not_restored":
assert not set(want) & set(outcome.restored) and not set(want) & set(outcome.kept)
elif field == "restored_includes":
assert set(want) <= set(outcome.restored)
elif field in ("restored", "kept"):
assert set(getattr(outcome, field)) == set(want)
elif field != "receipt_kept":
assert getattr(outcome, field) == want
assert not {item.key for item in outcome.withheld} & set(outcome.kept)
def test_configure_writes_owner_only_and_the_receipt_never_holds_the_key(self, tmp_path):
rig = _Rig(tmp_path, ORIGINAL)
rig.configure(credential=StaticToken("sk-virtual-key-never-on-disk-twice"))
configured = rig.read()
assert configured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key-never-on-disk-twice"
assert configured["env"]["ANTHROPIC_BASE_URL"] == PROXY and configured["model"] == "claude-auto"
assert "ANTHROPIC_API_KEY" not in configured["env"] and "apiKeyHelper" not in configured
assert stat.S_IMODE(rig.settings.stat().st_mode) == 0o600 == stat.S_IMODE(rig.state.stat().st_mode)
assert "sk-virtual-key-never-on-disk-twice" not in rig.state.read_text()
def test_withheld_credentials_come_back_once_the_url_points_at_their_server_again(self, tmp_path):
# The kept receipt owns only the withheld slots: the second unconfigure restores exactly those.
rig = _Rig(tmp_path, ORIGINAL)
rig.configure()
rig.edit(_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"))
rig.unconfigure()
rig.edit(_set("env.ANTHROPIC_BASE_URL", ANTHROPIC), _set("theme", "light"))
outcome = rig.unconfigure()
assert rig.read() == {**ORIGINAL, "theme": "light"}
assert set(outcome.restored) == {"env.ANTHROPIC_API_KEY", "apiKeyHelper"}
assert outcome.kept == () and outcome.withheld == () and not rig.state.exists()
@pytest.mark.parametrize(
("path", "value", "repeat_credential"),
[
("env.ANTHROPIC_API_KEY", "sk-user-added-later", HELPER),
("env.ANTHROPIC_AUTH_TOKEN", "sk-users-own-token", HELPER),
("apiKeyHelper", "/opt/mine/helper", StaticToken("sk-rotated")),
],
ids=["user-adds-api-key", "user-replaces-our-token", "user-sets-own-helper"],
)
def test_a_credential_the_user_set_between_two_configures_is_what_comes_back(
self, tmp_path, path, value, repeat_credential
):
# The repeat's merge clears the slot, so the displaced value is snapshotted and is what returns;
# it was set while the file pointed at the proxy, so it returns once the file points there again.
rig = _Rig(tmp_path, {"theme": "dark"})
rig.configure(credential=HELPER, model=KeepModel())
rig.edit(_set(path, value))
rig.configure(credential=repeat_credential, model=KeepModel())
assert not _lookup(rig.read(), path)
outcome = rig.unconfigure()
assert [(item.key, item.endpoint) for item in outcome.withheld] == [(path, PROXY)]
assert rig.read() == {"theme": "dark"} and rig.state.exists()
rig.settings.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_BASE_URL": PROXY}}))
outcome = rig.unconfigure()
assert _lookup(rig.read(), path) == value
assert outcome.restored == (path,) and outcome.withheld == () and not rig.state.exists()
def test_a_receipt_commit_that_fails_leaves_no_staged_token_behind(self, tmp_path):
rig = _Rig(tmp_path, {})
def commit_receipt_fails(staged, path):
if path == str(rig.state):
os.unlink(staged)
raise OSError("receipt rename failed")
commit_staged_json(staged, path)
with pytest.raises(ClaudeSettingsError, match=r"Could not write .*receipt rename failed"):
rig.configure(credential=StaticToken("sk-never-left-in-a-temp-file"), commit=commit_receipt_fails)
assert not list(rig.settings.parent.glob(".tmp-*")) and not list(rig.state.parent.glob(".tmp-*"))
assert rig.read() == {} and not rig.state.exists()
@pytest.mark.parametrize("configured_before", [False, True], ids=["first-configure", "repeat-configure"])
def test_a_settings_commit_that_fails_after_the_receipt_landed_puts_the_receipt_back(
self, tmp_path, configured_before
):
# The two renames are not atomic: a settings rename that fails after the receipt landed must
# not leave a receipt describing settings that were never written.
rig = _Rig(tmp_path, ORIGINAL)
if configured_before:
rig.configure()
receipt_before = rig.state.read_text() if configured_before else None
settings_before = rig.settings.read_text()
def commit_settings_fails(staged, path):
if path == str(rig.settings):
os.unlink(staged)
raise OSError("rename failed")
commit_staged_json(staged, path)
with pytest.raises(ClaudeSettingsError, match="rename failed"):
rig.configure(credential=StaticToken("sk-rotated"), commit=commit_settings_fails)
assert rig.settings.read_text() == settings_before
assert (rig.state.read_text() if rig.state.exists() else None) == receipt_before
if configured_before:
rig.unconfigure()
assert rig.read() == ORIGINAL
def test_a_failed_repeat_configure_leaves_the_earlier_undo_intact(self, tmp_path):
rig = _Rig(tmp_path, ORIGINAL)
rig.configure()
receipt_before = rig.state.read_text()
rig.settings.parent.chmod(0o500)
try:
with pytest.raises(ClaudeSettingsError, match="Could not write"):
rig.configure(credential=StaticToken("sk-rotated"))
finally:
rig.settings.parent.chmod(0o700)
assert rig.state.read_text() == receipt_before and not list(rig.state.parent.glob(".tmp-*"))
assert rig.read()["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key"
rig.unconfigure()
assert rig.read() == ORIGINAL
def test_unconfigure_reports_a_receipt_it_cannot_remove_as_a_settings_error(self, tmp_path):
rig = _Rig(tmp_path, ORIGINAL)
rig.configure()
rig.state.parent.chmod(0o500)
try:
with pytest.raises(ClaudeSettingsError, match="Could not remove"):
rig.unconfigure()
finally:
rig.state.parent.chmod(0o700)
def test_configure_writes_through_a_symlinked_settings_file(self, tmp_path):
target = tmp_path / "dotfiles" / "settings.json"
target.parent.mkdir()
target.write_text(json.dumps({"theme": "dark"}))
link = tmp_path / "settings.json"
link.symlink_to(target)
configure_claude_settings(PROXY, StaticToken("sk-virtual-key"), UnpinModel(), link, tmp_path / "state.json", ())
assert link.is_symlink()
assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key"
@pytest.mark.parametrize("operation", ["configure", "unconfigure"])
def test_refuses_while_a_temporary_owner_holds_a_backup(self, paths, tmp_path, operation):
settings_path, backup_path = paths
backup_path.write_text("{}")
owners = _owners(backup_path)
state = tmp_path / "state.json"
attempt = (
(lambda: configure_claude_settings(PROXY, StaticToken("k"), UnpinModel(), settings_path, state, owners))
if operation == "configure"
else (lambda: unconfigure_claude_settings(settings_path, state, owners))
)
with pytest.raises(ClaudeSettingsError, match="lite down"):
attempt()
assert not settings_path.exists()
def test_unconfigure_without_a_receipt_is_an_error_not_a_silent_no_op(self, tmp_path):
with pytest.raises(ClaudeSettingsError, match="nothing to undo"):
_Rig(tmp_path, None).unconfigure()
def _lookup(settings, path):
section, _, key = path.rpartition(".")
return (settings.get(section) or {}).get(key) if section else settings.get(key)

View file

@ -0,0 +1,331 @@
import json
import os
import stat
import click
import pytest
import requests
import responses
from click.testing import CliRunner
from litellm.proxy.client.cli import cli
from litellm.proxy.client.cli.commands import configure as configure_module
from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner
from litellm.proxy.client.cli.commands.configure import configure_claude, configure_group, interactive_configure
PROXY = "http://proxy.test:4000"
VALID_KEY = "sk-virtual-key"
LISTED_MODELS = ("claude-auto", "gpt-5.6-luna")
def _mock_models():
responses.get(
f"{PROXY}/v1/models",
json={"data": [{"id": model, "object": "model"} for model in LISTED_MODELS]},
match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}"})],
)
responses.get(f"{PROXY}/v1/models", status=401)
@pytest.fixture
def paths(monkeypatch, tmp_path):
settings_path = tmp_path / "claude" / "settings.json"
state_path = tmp_path / "litellm" / "claude_configure_state.json"
monkeypatch.setattr(configure_module, "CLAUDE_SETTINGS_PATH", settings_path)
monkeypatch.setattr(configure_module, "CONFIGURE_STATE_PATH", state_path)
return settings_path, state_path
@pytest.fixture
def lite_on_path(monkeypatch, tmp_path):
"""A real `lite` executable on PATH, so the apiKeyHelper command resolves without patching."""
bin_dir = tmp_path / "bin"
bin_dir.mkdir()
lite = bin_dir / "lite"
lite.write_text("#!/bin/sh\nexit 0\n")
lite.chmod(lite.stat().st_mode | stat.S_IXUSR)
monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ.get('PATH', '')}")
return str(lite)
@pytest.fixture
def runner():
return CliRunner()
@pytest.fixture
def lite_up_backup(monkeypatch, tmp_path):
"""A `lite up` session holding its backup, the local precondition every settings write refuses on."""
backup = tmp_path / "claude_settings_backup.json"
backup.write_text("{}")
monkeypatch.setattr(configure_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup, "lite up", "lite down"),))
return backup
def _configure(runner, *args):
return runner.invoke(cli, ["--base-url", PROXY, "configure", "claude", *args])
class TestConfigureClaudeWithAVirtualKey:
@responses.activate
def test_writes_settings_and_reports_without_echoing_the_key(self, runner, paths):
_mock_models()
settings_path, state_path = paths
result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto")
assert result.exit_code == 0, result.output
written = json.loads(settings_path.read_text())
assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY
assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY
assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
assert written["model"] == "claude-auto"
assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in written["env"]
assert state_path.exists()
assert VALID_KEY not in result.output
assert "Starting model: claude-auto" in result.output
assert "1 of the proxy's 2 models" in result.output
assert "lite unconfigure claude" in result.output
assert len(responses.calls) == 1
@responses.activate
def test_takes_the_key_from_the_global_option_and_keeps_claude_codes_default(self, runner, paths):
_mock_models()
settings_path, _ = paths
result = runner.invoke(cli, ["--base-url", PROXY, "--api-key", VALID_KEY, "configure", "claude"])
assert result.exit_code == 0, result.output
written = json.loads(settings_path.read_text())
assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY
assert "model" not in written
assert "Starting model: not pinned" in result.output
@responses.activate
def test_refuses_a_model_the_proxy_does_not_list(self, runner, paths):
_mock_models()
settings_path, _ = paths
result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-nope")
assert result.exit_code != 0
assert "'claude-nope' is not served" in result.output
assert "claude-auto, gpt-5.6-luna" in result.output
assert not settings_path.exists()
@responses.activate
def test_refuses_a_key_the_proxy_rejects(self, runner, paths):
_mock_models()
settings_path, _ = paths
result = _configure(runner, "--api-key", "sk-wrong")
assert result.exit_code != 0
assert "rejected your key (HTTP 401)" in result.output
assert not settings_path.exists()
@responses.activate
@pytest.mark.parametrize(
("mock", "expected", "unexpected"),
[
(
lambda: responses.get(f"{PROXY}/v1/models", body=requests.ConnectionError("refused")),
"Is the proxy at",
"answered",
),
(
lambda: responses.get(f"{PROXY}/v1/models", status=500),
"The proxy at http://proxy.test:4000 answered",
"Is the proxy at",
),
(
lambda: responses.get(f"{PROXY}/v1/models", body="<html>not json</html>"),
"answered, so check that it is a LiteLLM proxy",
"Is the proxy at",
),
(
lambda: responses.get(f"{PROXY}/v1/models", json={"data": []}),
"Claude Code would have nothing to run",
"Is the proxy at",
),
],
ids=["unreachable", "http-500", "non-json-body", "empty-list"],
)
def test_the_listing_hint_matches_how_the_listing_failed(self, runner, paths, mock, expected, unexpected):
# Only a proxy that never answered gets the "is it running" question; a 500, a non-JSON body or an
# empty list prove it is up, and the hint says so instead.
mock()
settings_path, _ = paths
result = _configure(runner, "--api-key", VALID_KEY)
assert result.exit_code != 0
assert expected in result.output and unexpected not in result.output
assert not settings_path.exists()
@responses.activate
@pytest.mark.parametrize("entry", ["virtual-key", "login", "interactive"])
def test_refuses_while_lite_up_holds_a_backup_before_any_login_or_request(
self, runner, paths, monkeypatch, lite_up_backup, entry
):
_mock_models()
def login_must_not_run(ctx):
raise AssertionError("the local precondition must be checked before a login is attempted")
monkeypatch.setattr(configure_module, "ensure_fresh_login", login_must_not_run)
if entry == "interactive":
ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": None})
with pytest.raises(click.ClickException, match="lite down"):
interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None)
else:
args = ["--api-key", VALID_KEY] if entry == "virtual-key" else []
result = runner.invoke(configure_claude, args, obj={"base_url": PROXY, "api_key": None})
assert result.exit_code != 0 and "lite down" in result.output
assert len(responses.calls) == 0
assert not paths[0].exists()
@responses.activate
def test_says_so_when_the_key_is_written_through_a_symlink(self, runner, paths, tmp_path):
_mock_models()
settings_path, _ = paths
target = tmp_path / "dotfiles" / "settings.json"
target.parent.mkdir()
target.write_text("{}")
settings_path.parent.mkdir(parents=True)
settings_path.symlink_to(target)
result = _configure(runner, "--api-key", VALID_KEY)
assert result.exit_code == 0, result.output
assert "keep it out of version control" in result.output
assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY
class TestConfigureClaudeWithTheLogin:
def _stored_login(self, monkeypatch):
monkeypatch.setattr(configure_module, "ensure_fresh_login", lambda ctx: None)
monkeypatch.setattr(configure_module, "get_stored_api_key", lambda expected_base_url, vault: VALID_KEY)
@responses.activate
def test_uses_the_login_through_the_helper_and_writes_no_secret(self, runner, paths, monkeypatch, lite_on_path):
_mock_models()
self._stored_login(monkeypatch)
settings_path, _ = paths
result = runner.invoke(
configure_claude,
["--model", "claude-auto"],
obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": True},
)
assert result.exit_code == 0, result.output
written = json.loads(settings_path.read_text())
assert written["apiKeyHelper"] == f"{lite_on_path} --base-url {PROXY} auth print-token"
assert "ANTHROPIC_AUTH_TOKEN" not in written["env"]
assert written["model"] == "claude-auto"
assert VALID_KEY not in settings_path.read_text()
assert "read through apiKeyHelper" in result.output
@responses.activate
def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths, monkeypatch, lite_on_path):
_mock_models()
self._stored_login(monkeypatch)
settings_path, _ = paths
result = runner.invoke(
configure_claude,
["--api-key", VALID_KEY],
obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True},
)
assert result.exit_code == 0, result.output
written = json.loads(settings_path.read_text())
assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY and "apiKeyHelper" not in written
class TestInteractiveConfigure:
@responses.activate
def test_asks_for_targets_and_a_starting_model_then_configures(self, paths):
_mock_models()
settings_path, _ = paths
asked = {}
def pick_model(listed):
asked["listed"] = tuple(listed)
return "claude-auto"
ctx = click.Context(
configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False}
)
interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model)
assert asked["listed"] == LISTED_MODELS
assert json.loads(settings_path.read_text())["model"] == "claude-auto"
def test_does_nothing_when_claude_code_is_not_picked(self, paths):
settings_path, _ = paths
ctx = click.Context(
configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False}
)
interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda listed: None)
assert not settings_path.exists()
def test_bare_configure_without_a_terminal_names_the_non_interactive_command(self, runner, paths):
result = runner.invoke(cli, ["--base-url", PROXY, "configure"])
assert result.exit_code != 0
assert "lite configure claude --api-key" in result.output
class TestUnconfigureClaude:
@responses.activate
def test_restores_the_original_file_and_removes_the_receipt(self, runner, paths):
_mock_models()
settings_path, state_path = paths
settings_path.parent.mkdir(parents=True)
original = {"theme": "dark", "model": "claude-opus-5"}
settings_path.write_text(json.dumps(original))
assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code == 0, result.output
assert json.loads(settings_path.read_text()) == original
assert not state_path.exists()
assert "Restored in" in result.output and "model" in result.output
assert "ANTHROPIC_API_KEY" not in result.output, "a key that never existed was not restored"
@responses.activate
def test_a_file_only_configure_created_is_reported_removed_not_restored(self, runner, paths):
_mock_models()
settings_path, _ = paths
assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code == 0, result.output
assert not settings_path.exists()
assert "No settings file remains" in result.output and "Restored" not in result.output
@responses.activate
def test_says_when_nothing_was_still_ours_and_names_what_it_kept(self, runner, paths):
_mock_models()
settings_path, _ = paths
settings_path.parent.mkdir(parents=True)
settings_path.write_text(json.dumps({"theme": "dark"}))
assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0
edited = json.loads(settings_path.read_text())
edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()}
edited["model"] = "mine"
settings_path.write_text(json.dumps(edited))
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code == 0, result.output
assert "Nothing in" in result.output and "was still ours to restore" in result.output
assert "Left as you changed them since:" in result.output and "model" in result.output
@responses.activate
def test_names_the_server_a_withheld_credential_was_captured_with_and_keeps_the_receipt(self, runner, paths):
_mock_models()
settings_path, state_path = paths
settings_path.parent.mkdir(parents=True)
settings_path.write_text(
json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://api.anthropic.com", "ANTHROPIC_API_KEY": "sk-ant"}})
)
assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0
edited = json.loads(settings_path.read_text())
edited["env"]["ANTHROPIC_BASE_URL"] = "http://other-proxy:4000"
settings_path.write_text(json.dumps(edited))
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code == 0, result.output
assert "env.ANTHROPIC_API_KEY (captured with https://api.anthropic.com)" in result.output
assert str(state_path) in result.output and state_path.exists()
assert "sk-ant" not in result.output
def test_refuses_while_lite_up_holds_a_backup(self, runner, paths, lite_up_backup):
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code != 0 and "lite down" in result.output
def test_without_a_receipt_it_fails_loudly(self, runner, paths):
result = runner.invoke(cli, ["unconfigure", "claude"])
assert result.exit_code != 0
assert "nothing to undo" in result.output

View file

@ -4,9 +4,11 @@ import stat
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import pytest
import requests
from litellm.proxy.client.cli.commands.pi import (
ListingFailure,
ModelLimits,
PiSyncError,
fetch_model_ids,
@ -28,6 +30,10 @@ class _FakeResponse:
return self._payload
def _refused(*args, **kwargs):
raise requests.ConnectionError("refused")
class TestFetchModelIds:
def test_returns_ids_in_proxy_order_deduped(self):
captured = {}
@ -53,9 +59,7 @@ class TestFetchModelIds:
assert "Could not list models" in result.message
def test_non_200_is_a_value(self):
result = fetch_model_ids(
"http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
)
result = fetch_model_ids("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500))
assert isinstance(result, PiSyncError)
assert "HTTP 500" in result.message
@ -75,6 +79,22 @@ class TestFetchModelIds:
)
assert isinstance(result, PiSyncError)
assert "no models" in result.message
assert result.kind is ListingFailure.EMPTY
@pytest.mark.parametrize(
("get", "kind"),
[
(_refused, ListingFailure.UNREACHABLE),
(lambda *a, **k: _FakeResponse(401), ListingFailure.REJECTED),
(lambda *a, **k: _FakeResponse(403), ListingFailure.REJECTED),
(lambda *a, **k: _FakeResponse(500), ListingFailure.OTHER),
(lambda *a, **k: _FakeResponse(200), ListingFailure.BAD_BODY),
],
ids=["unreachable", "401", "403", "500", "bad-body"],
)
def test_the_failure_kind_is_decided_where_the_response_is_classified(self, get, kind):
result = fetch_model_ids("http://localhost:4000", "sk-key", get=get)
assert isinstance(result, PiSyncError) and result.kind is kind
class TestFetchModelLimits:

View file

@ -11,11 +11,11 @@ from click.testing import CliRunner
from litellm.proxy.client.cli.commands import up as up_module
from litellm.proxy.client.cli.commands.agents import AgentRunError
from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError
from litellm.proxy.client.cli.commands.claude_settings import ApiKeyHelper, ClaudeSettingsError
from litellm.proxy.client.cli.commands.up import (
BackupRecord,
UpError,
_ensure_fresh_login,
ensure_fresh_login,
down,
load_json_or_empty,
merge_claude_settings,
@ -40,12 +40,12 @@ def _patch_paths(monkeypatch, tmp_path):
class TestMergeClaudeSettings:
def test_preserves_unrelated_top_level_keys(self):
merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper")
merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", ApiKeyHelper("helper"))
assert merged["theme"] == "dark"
def test_preserves_unrelated_env_keys(self):
settings = {"env": {"SOME_OTHER_VAR": "value"}}
merged = merge_claude_settings(settings, "http://localhost:4000", "helper")
merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper"))
assert merged["env"]["SOME_OTHER_VAR"] == "value"
def test_overrides_base_url_and_helper(self):
@ -53,7 +53,7 @@ class TestMergeClaudeSettings:
"env": {"ANTHROPIC_BASE_URL": "https://old.example.com"},
"apiKeyHelper": "old-helper",
}
merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper")
merged = merge_claude_settings(settings, "http://localhost:4000/", ApiKeyHelper("new-helper"))
assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true"
assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
@ -61,21 +61,21 @@ class TestMergeClaudeSettings:
def test_preserves_existing_gateway_model_discovery(self):
settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}}
merged = merge_claude_settings(settings, "http://localhost:4000", "helper")
merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper"))
assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0"
def test_preserves_existing_tool_search(self):
settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}}
merged = merge_claude_settings(settings, "http://localhost:4000", "helper")
merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper"))
assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false"
def test_drops_stray_api_key(self):
settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}}
merged = merge_claude_settings(settings, "http://localhost:4000", "helper")
merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper"))
assert "ANTHROPIC_API_KEY" not in merged["env"]
def test_works_from_empty_settings(self):
merged = merge_claude_settings({}, "http://localhost:4000", "helper")
merged = merge_claude_settings({}, "http://localhost:4000", ApiKeyHelper("helper"))
assert merged["env"] == {
"ANTHROPIC_BASE_URL": "http://localhost:4000",
"ENABLE_TOOL_SEARCH": "true",
@ -85,7 +85,7 @@ class TestMergeClaudeSettings:
def test_does_not_mutate_input(self):
settings = {"env": {"FOO": "bar"}}
merge_claude_settings(settings, "http://localhost:4000", "helper")
merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper"))
assert settings == {"env": {"FOO": "bar"}}
@ -327,7 +327,7 @@ class TestEnsureFreshLogin:
monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True)
login_calls = _capture_login(monkeypatch)
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
assert login_calls == []
@ -339,7 +339,7 @@ class TestEnsureFreshLogin:
monkeypatch, on_login=lambda: store.log_in({"key": "sk-b", "base_url": "http://proxy-b:4000"}, "sk-b")
)
_ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
assert login_calls == [("http://proxy-b:4000", False)]
assert store.key_requests == ["http://proxy-b:4000", "http://proxy-b:4000"]
@ -353,7 +353,7 @@ class TestEnsureFreshLogin:
on_login=lambda: store.log_in({"key": "sk-a", "base_url": "http://proxy-a:4000"}, "sk-a"),
)
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
assert login_calls == [("http://proxy-a:4000", False)]
@ -363,7 +363,7 @@ class TestEnsureFreshLogin:
monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True)
with pytest.raises(UpError, match="Run `lite login` first"):
_ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
def test_trusts_a_pkce_credential_that_was_renewed_on_the_way_in(self, monkeypatch):
"""A --pkce key inside its freshness buffer is renewed by `get_stored_api_key`, so `lite up`
@ -377,7 +377,7 @@ class TestEnsureFreshLogin:
)
login_calls = _capture_login(monkeypatch)
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
assert login_calls == []
assert store.key_requests == ["http://proxy-a:4000"]
@ -390,7 +390,7 @@ class TestEnsureFreshLogin:
on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"),
)
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
assert login_calls == [("http://proxy-a:4000", True)]
@ -399,7 +399,7 @@ class TestEnsureFreshLogin:
_FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=-10), {})
with pytest.raises(UpError, match="Run `lite login --pkce` first"):
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
def test_trusts_the_key_the_cli_group_already_resolved_instead_of_reading_the_token_file_again(
self, monkeypatch
@ -409,7 +409,7 @@ class TestEnsureFreshLogin:
store = _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=86_400), {})
login_calls = _capture_login(monkeypatch)
_ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group"))
ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group"))
assert login_calls == []
assert store.key_requests == []
@ -423,7 +423,7 @@ class TestEnsureFreshLogin:
)
with pytest.raises(UpError, match="Run `lite login --pkce` first"):
_ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None))
ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None))
assert store.key_requests == []
@ -435,7 +435,7 @@ class TestEnsureFreshLogin:
on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"),
)
_ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None))
ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None))
assert login_calls == [("http://proxy-a:4000", True)]
assert store.key_requests == ["http://proxy-a:4000"]

View file

@ -3670,5 +3670,42 @@ async def test_post_call_success_hook_contains_header_merge_failures(
)
@pytest.mark.asyncio
async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter):
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
handler, _cache = rate_limiter
stash = get_or_create_request_stash()
warm_tokenizer("claude-fable-5")
data: dict[str, object] = {
"model": "claude-fable-5",
"messages": [{"role": "user", "content": text * 100}],
}
itpm_descriptor = {
"key": PROJECT_ITPM_DESCRIPTOR_KEY,
"value": "proj-loop:claude-fable-5",
"rate_limit": {"tokens_per_unit": 10_000_000, "window_size": 60},
}
_, took, lags = await timed_with_loop_lags(
lambda: handler._reserve_project_io_tokens_or_raise(
descriptors=[itpm_descriptor],
data=data,
requested_model="claude-fable-5",
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-itpm-loop"), project_id="proj-loop"),
tpm_reservation_scopes=[],
tpm_reservation_amount=0,
)
)
assert stash.rate_limit_response is not None
assert_loop_stayed_free(took, lags)
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])

View file

@ -265,7 +265,7 @@ class TestSuggesterRejectsModelsWithoutToolCalling:
def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map):
supported_params = litellm.get_supported_openai_params(
model="amazon.nova-pro-v1:0",
model="meta.llama4-scout-17b-instruct-v1:0",
custom_llm_provider="bedrock",
)

View file

@ -5085,6 +5085,104 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
assert len(deleted_keys) == 2
class _JWTMappingRow:
def __init__(self, token, jwt_claim_name, jwt_claim_value):
self.token = token
self.jwt_claim_name = jwt_claim_name
self.jwt_claim_value = jwt_claim_value
class _CascadingJWTMappingTable:
"""Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key is deleted."""
def __init__(self, rows):
self.rows = rows
async def find_many(self, where, **kwargs):
return [row for row in self.rows if row.token == where["token"]]
def cascade(self, deleted_tokens):
self.rows = [row for row in self.rows if row.token not in deleted_tokens]
class _RecordingEvict:
def __init__(self):
self.cache_keys = ()
async def __call__(self, cache_keys, user_api_key_cache):
self.cache_keys = tuple(cache_keys)
@pytest.mark.asyncio
async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypatch):
"""Deleting a key must evict its jwt_key_mapping cache entries (LIT-5380).
The FK cascade removes the mapping rows, so a surviving cache entry would keep
resolving the deleted token hash and 401 every JWT call from that identity until
virtual_key_mapping_cache_ttl expires, instead of auto-registering again.
"""
jwt_table = _CascadingJWTMappingTable(
[_JWTMappingRow("hashed-token-1", "email", "user@example.com")]
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id=None,
key_alias="jwt-mapped-key",
spend=0.0,
max_budget=None,
models=[],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[key1]
)
mock_prisma_client.db.litellm_jwtkeymapping = jwt_table
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
async def cascading_delete_data(tokens):
jwt_table.cascade(tokens)
return list(tokens)
mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data)
recording_evict = _RecordingEvict()
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.evict_and_broadcast",
recording_evict,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
lambda token: token,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
await delete_verification_tokens(
tokens=["hashed-token-1"],
user_api_key_cache=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
),
)
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
@pytest.mark.asyncio
async def test_delete_key_fn_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import KeyRequest

View file

@ -675,8 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch):
],
)
monkeypatch.setattr(litellm, "callbacks", [])
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
@ -1149,6 +1147,7 @@ def _assert_passed_with_discard_warning(result, caplog):
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", [])
@pytest.mark.asyncio
@ -1329,3 +1328,350 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
class _LegacyHookGuardrail(CustomGuardrail):
"""A guardrail with only the legacy post-call hook: it never defines apply_guardrail."""
def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None):
super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True)
self.replacement = replacement
self.raises = raises
self.rewrite_in_place = rewrite_in_place
self.calls = []
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response})
if self.raises is not None:
raise self.raises
if self.rewrite_in_place is not None:
response["text"] = self.rewrite_in_place
return self.replacement
class _NativeHooksGuardrail(_LegacyHookGuardrail):
use_native_lifecycle_hooks = True
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail")
class _LegacyScanningTranslation:
"""Stores the assembled response under request_data["response"] before scanning, like the
chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one
text per entry of a replacement's "texts"."""
delivers_ended_stream_rewrites = True
def post_call_hook_response(self, response):
return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]}
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
request_data.setdefault(
"response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}
)
outputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]},
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
responses_so_far[0]["text"] = outputs["texts"][0]
return responses_so_far
async def process_output_response(
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
):
inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])}
if response.get("tool_calls"):
inputs["tool_calls"] = list(response["tool_calls"])
await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data={"response": response},
input_type="response",
logging_obj=litellm_logging_obj,
)
return response
class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation):
"""Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from
the inputs, while the non-streaming scan of the same response sends an empty list."""
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]})
await guardrail_to_apply.apply_guardrail(
inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]},
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
async def process_output_response(
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
):
await guardrail_to_apply.apply_guardrail(
inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])},
request_data={"response": response},
input_type="response",
logging_obj=litellm_logging_obj,
)
return response
def _tool_only_chunk():
return {"text": "", "tool_call": _chunk()["tool_call"]}
def _native(text):
return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]}
def _legacy_replacement(*texts, tool_calls=None):
return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls}
async def _run_legacy_streaming_step(
monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None
):
return await _run_legacy_streaming_steps(
monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation
)
async def _run_legacy_streaming_steps(
monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None
):
monkeypatch.setattr(litellm, "callbacks", list(guardrails))
return await PipelineExecutor.execute_steps(
steps=[
PipelineStep(
guardrail=guardrail.guardrail_name,
on_pass="next" if position + 1 < len(guardrails) else "allow",
on_fail=on_fail,
on_error=on_error,
)
for position, guardrail in enumerate(guardrails)
],
mode="post_call",
data={"model": "m"},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="p",
streaming_chunks=chunks,
endpoint_translation=_LegacyScanningTranslation() if translation is None else translation,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail])
async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class):
guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world"))
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert chunks[0]["text"] == "[REWRITTEN] hello world"
assert [call["response"] for call in guardrail.calls] == [_native("hello world")]
assert guardrail.calls[0]["data"]["model"] == "m"
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world")
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert chunks[0]["text"] == "[REWRITTEN] hello world"
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=None)
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
assert result.terminal_action == "allow"
assert len(guardrail.calls) == 1
assert chunks == [_chunk()]
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch):
exc = HTTPException(status_code=400, detail={"error": "output blocked"})
chunks = [_chunk()]
result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks)
assert result.terminal_action == "block"
assert [step.outcome for step in result.step_results] == ["fail"]
assert result.original_exception is exc
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch):
chunks = [_chunk()]
result = await _run_legacy_streaming_step(
monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block"
)
assert result.terminal_action == "block"
assert [step.outcome for step in result.step_results] == ["error"]
assert result.step_results[0].error_detail == "boom"
@pytest.mark.asyncio
async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two"))
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog):
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
guardrail = _LegacyHookGuardrail(
replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call])
)
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[]))
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=None)
chunks = [_tool_only_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
)
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
assert chunks == [_tool_only_chunk()]
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog):
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call]))
chunks = [_tool_only_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_tool_only_chunk()]
class _NoHooksGuardrail(CustomGuardrail):
pass
class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
async for item in response:
yield item
class _UnscannableRewriteTranslation(_LegacyScanningTranslation):
"""Like the chat handler on a response whose choices are plain dicts: the non-streaming scan
never hands anything to the guardrail."""
async def process_output_response(
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
):
return response
def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path():
assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True
assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True
assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False
assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False
@pytest.mark.asyncio
async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]"))
chunks = [_chunk()]
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation())
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch):
masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world"))
auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor")
chunks = [_chunk()]
result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next")
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass", "pass"]
assert chunks[0]["text"] == "[REWRITTEN] hello world"
assert [call["response"] for call in masker.calls] == [_native("hello world")]
assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")]

View file

@ -671,6 +671,95 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp
assert routed.native_hooks_ran == []
@pytest.mark.asyncio
async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch):
"""A post_call pipeline step already ran the opted-out guardrail's own hook against
the buffered stream, so the deferred audit must not run it a second time."""
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
from litellm.types.utils import Choices, Message, ModelResponse
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
captured_logging_obj=_streaming_logging_obj(),
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
cache_hit=False,
)
assert pipeline_managed.native_hooks_ran == []
@pytest.mark.asyncio
async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch):
"""A pipeline step with neither streaming interface keeps the whole pipeline off the
stream, so the deferred audit is the only place the opted-out guardrail's own hook
still runs, the way it did before pipelines ran on streams."""
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
from litellm.types.utils import Choices, Message, ModelResponse
class NeitherHookGuardrail(CustomGuardrail):
pass
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call)
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither])
pipeline = GuardrailPipeline(
mode="post_call",
steps=[
PipelineStep(guardrail="keeps_native", on_fail="next"),
PipelineStep(guardrail="gr-neither", on_fail="block"),
],
)
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
captured_logging_obj=_streaming_logging_obj(),
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
cache_hit=False,
)
assert pipeline_managed.native_hooks_ran == ["post_call"]
@pytest.mark.asyncio
async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch):
"""A route with no endpoint guardrail translation cannot gate the stream through its
pipelines, so the deferred audit still owes the opted-out guardrail its own hook."""
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
from litellm.types.utils import Choices, Message, ModelResponse
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"),
captured_logging_obj=_streaming_logging_obj(),
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
cache_hit=False,
)
assert pipeline_managed.native_hooks_ran == ["post_call"]
@pytest.mark.asyncio
async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch):
"""The realtime path calls apply_guardrail directly, so the opt-out has to be

View file

@ -6,6 +6,7 @@ import os
import re
import socket
import subprocess
import time
import types
from datetime import datetime, timedelta, timezone
from pathlib import Path
@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import app, initialize
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -12954,3 +12955,59 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough()
import litellm.proxy.proxy_server as ps
assert ps.general_settings["enable_openai_websocket_passthrough"] is False
async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch):
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
warm_tokenizer("claude-fable-5")
response, took, lags = await timed_with_loop_lags(
lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100))
)
assert response.total_tokens > 0
assert_loop_stayed_free(took, lags)
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch):
from tokenizers import Tokenizer
from litellm import Router
from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"]
class SlowHubTokenizer:
@staticmethod
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
time.sleep(0.3)
return claude_tokenizer
monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer)
monkeypatch.setattr(
"litellm.proxy.proxy_server.llm_router",
Router(
model_list=[
{
"model_name": "self-hosted",
"litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
"model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}},
}
]
),
)
response, took, lags = await timed_with_loop_lags(
lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop"))
)
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens > 0
assert_loop_stayed_free(took, lags)

View file

@ -1823,6 +1823,44 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs():
assert lifted["standard_logging_object"] == {"id": "log-1"}
@pytest.mark.asyncio
async def test_a_dispatched_failure_is_counted_off_the_event_loop():
from unittest.mock import AsyncMock, patch
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
warm_tokenizer("claude-fable-5")
request_data = {
"litellm_logging_obj": _LoggingObj(
{
"first_api_call_start_time": 1700000000.0,
"call_type": "acompletion",
"model": "claude-fable-5",
"messages": [{"role": "user", "content": text * 100}],
}
),
"metadata": {},
}
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.alert_types = []
with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
_, took, lags = await timed_with_loop_lags(
lambda: proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=Exception("boom"),
user_api_key_dict=UserAPIKeyAuth(),
)
)
assert request_data["combined_usage_object"].prompt_tokens > 0
assert_loop_stayed_free(took, lags)
@pytest.mark.asyncio
async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch):
"""Regression for LIT-6043: an expected 4xx must not format a traceback for

View file

@ -11,6 +11,7 @@ from __future__ import annotations
import asyncio
import json
from copy import deepcopy
import logging
from collections.abc import Iterator
from typing import Any, Callable, Dict, List
@ -29,7 +30,7 @@ from litellm.integrations.prometheus import PrometheusLogger
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
from litellm.types.llms.openai import ResponsesAPIResponse
@ -1723,21 +1724,85 @@ async def _async_chunk_iter(chunks: List[Any]):
yield chunk
def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
def _legacy_hook_stream_guardrail(
seen: Dict[str, Any],
rewrite: Callable[[Any], Any] | None = None,
raises: Exception | None = None,
native_lifecycle: bool = False,
) -> CustomGuardrail:
class LegacyHookGuardrail(CustomGuardrail):
use_native_lifecycle_hooks = native_lifecycle
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["count"] = seen.get("count", 0) + 1
seen["data"] = data
seen["user_api_key_dict"] = user_api_key_dict
seen["response"] = deepcopy(response)
if raises is not None:
raise raises
return None if rewrite is None else rewrite(response)
if native_lifecycle:
class NativeLifecycleGuardrail(LegacyHookGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail")
return NativeLifecycleGuardrail(
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
)
return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail:
class IteratorHookGuardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
seen["count"] = seen.get("count", 0) + 1
async for item in response:
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
yield item
return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True)
def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail:
class IteratorAndLegacyHookGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1
return None
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1
async for item in response:
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
yield item
return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True)
def _rewritten_model_response(response: Any) -> litellm.ModelResponse:
payload = response.model_dump()
payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"]
return litellm.ModelResponse(**payload)
def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only(
make_user_api_key_auth, monkeypatch, caplog
):
class NativeOnlyGuardrail(CustomGuardrail):
pass
supported = _unified_stream_guardrail({})
native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call)
monkeypatch.setattr(litellm, "callbacks", [supported, native_only])
governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
legacy = _legacy_hook_stream_guardrail({})
legacy.guardrail_name = "gr-legacy"
iterator_only = _iterator_hook_only_guardrail("gr-iterator", {})
monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only])
governed = GuardrailPipeline(
mode="post_call",
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")],
)
ungoverned = GuardrailPipeline(
mode="post_call",
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")],
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")],
)
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")])
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")])
data = {
"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}
}
@ -1746,8 +1811,44 @@ def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions"))
assert streamable == (("governed", governed),)
assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog))
assert not any("'governed'" in message for message in _warnings(caplog))
assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog))
assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog))
@pytest.mark.parametrize(
"request_route",
["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
)
def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response(
make_user_api_key_auth, monkeypatch, caplog, request_route
):
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})])
legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}}
auth = make_user_api_key_auth(request_route=request_route)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
streamable = _streamable_post_call_pipelines(data, auth)
assert streamable == ()
assert stream_gated_guardrail_names(data, auth) == frozenset()
assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog))
def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path(
make_user_api_key_auth, monkeypatch, caplog
):
monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})])
both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}}
auth = make_user_api_key_auth(request_route="/v1/chat/completions")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
streamable = _streamable_post_call_pipelines(data, auth)
assert streamable == ()
assert stream_gated_guardrail_names(data, auth) == frozenset()
assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog))
def test_streamable_post_call_pipelines_is_empty_on_route_without_translation(
@ -1797,55 +1898,123 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u
@pytest.mark.asyncio
@pytest.mark.parametrize("native_lifecycle", [False, True])
async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support(
async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite(
proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog
):
seen: Dict[str, Any] = {}
if native_lifecycle:
class NativeOnlyGuardrail(CustomGuardrail):
use_native_lifecycle_hooks = True
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
seen["count"] = seen.get("count", 0) + 1
return inputs
else:
class NativeOnlyGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["count"] = seen.get("count", 0) + 1
return response
monkeypatch.setattr(
litellm,
"callbacks",
[NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
)
guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = _stream_chunks()
delivered: List[Any] = []
auth = make_user_api_key_auth(request_route="/v1/chat/completions")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(),
data=data,
call_type="completion",
guardrails_only=True,
user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True
)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
assert out is not None and out.get("stream") is True
assert seen["count"] == 1
assert isinstance(seen["response"], litellm.ModelResponse)
assert seen["response"].choices[0].message.content == "hello world"
assert seen["data"]["messages"] == data["messages"]
assert seen["user_api_key_dict"] is auth
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world"
assert delivered[1].choices[0].delta.content in (None, "")
assert delivered[1].choices[0].finish_reason == "stop"
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
assert _warnings(caplog) == []
@pytest.mark.asyncio
async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = _stream_chunks()
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(chunks),
request_data=data,
)
]
assert seen["count"] == 1
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {}
blocked = HTTPException(status_code=400, detail={"error": "output blocked"})
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
delivered: List[Any] = []
async def _drain() -> None:
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(_stream_chunks()),
request_data=data,
):
delivered.append(item)
assert out is not None
assert out.get("stream") is True
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
assert len(delivered) == 2
assert seen.get("count") is None
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
with pytest.raises(HTTPException) as info:
await _drain()
assert seen["count"] == 1
assert delivered == []
assert info.value is blocked
@pytest.mark.asyncio
async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {}
def rewrite(response: Any) -> Dict[str, Any]:
return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]}
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
response=_async_chunk_iter(_anthropic_sse_chunks()),
request_data=data,
)
]
assert seen["count"] == 1
assert seen["response"]["content"][0]["text"] == "hello world"
assert seen["response"]["role"] == "assistant"
raw = b"".join(delivered).decode()
assert "[REWRITTEN] hello world" in raw
assert raw.count("event: content_block_delta") == 1
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
assert f"event: {expected_event}" in raw
@pytest.mark.asyncio
@ -1853,19 +2022,7 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
class IteratorHookGuardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
seen["count"] = seen.get("count", 0) + 1
async for item in response:
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
yield item
monkeypatch.setattr(
litellm,
"callbacks",
[IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
)
monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
@ -1884,6 +2041,30 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
@pytest.mark.asyncio
async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(_stream_chunks()),
request_data=data,
)
]
assert seen == {"iterator_hook_calls": 1}
assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"]
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
@pytest.mark.asyncio
@pytest.mark.parametrize(
"rewrite_attribute, value",

View file

@ -1,4 +1,5 @@
import json
from typing import Final
import pytest
@ -1421,6 +1422,88 @@ class TestToolChoiceTransformation:
)
assert result == "required"
@pytest.mark.parametrize(
"request_tool_choice,expected",
[
({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}),
({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}),
({"type": "custom", "name": "ApplyPatch"}, {"type": "custom", "name": "ApplyPatch"}),
({"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}),
({"type": "function"}, "required"),
({"type": "tool"}, "required"),
({"type": "auto"}, "auto"),
("required", "required"),
("none", "none"),
(None, "auto"),
("any", "auto"),
("run_command", "auto"),
({"name": "run_command"}, "auto"),
],
)
def test_transform_tool_choice_for_responses_api_response(
self, request_tool_choice: object, expected: str | dict[str, str]
) -> None:
result: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
request_tool_choice
)
assert result == expected
def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self) -> None:
chat_completion_response: Final = ModelResponse(
id="chatcmpl-named-tool-choice",
created=1748575031,
model="claude-haiku-4-5",
object="chat.completion",
choices=[
Choices(
index=0,
finish_reason="tool_calls",
message=Message(
role="assistant",
content=None,
tool_calls=[
ChatCompletionMessageToolCall(
id="call_pwd",
type="function",
function=Function(name="run_command", arguments='{"command":"pwd"}'),
)
],
),
)
],
)
responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
request_input="Run the command pwd.",
responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}},
chat_completion_response=chat_completion_response,
)
assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"}
def test_non_streamed_response_with_unrecognized_tool_choice_echoes_auto(self) -> None:
chat_completion_response: Final = ModelResponse(
id="chatcmpl-unrecognized-tool-choice",
created=1748575031,
model="claude-haiku-4-5",
object="chat.completion",
choices=[
Choices(
index=0,
finish_reason="stop",
message=Message(role="assistant", content="/Users/dev"),
)
],
)
responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
request_input="Run the command pwd.",
responses_api_request={"tool_choice": "any"},
chat_completion_response=chat_completion_response,
)
assert responses_api_response.tool_choice == "auto"
class TestContentTypeTransformation:
"""Test content type transformation from Responses API to Chat Completion format"""

View file

@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve
"""
import json
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -628,3 +629,79 @@ def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id():
assert item_dones[0].item.call_id == "toolu_01AbCdEf"
for evt in deltas + dones:
assert evt.item_id == added[0].item.id
def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id=CHAT_COMPLETION_ID,
created=1748575031,
model="claude-haiku-4-5",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(
role="assistant",
content=None,
tool_calls=[
{
"id": "call_pwd",
"type": "function",
"function": {"name": "run_command", "arguments": '{"command":"pwd"}'},
"index": 0,
}
],
),
finish_reason=finish_reason,
)
],
)
def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None:
iterator: Final = LiteLLMCompletionStreamingIterator(
model="claude-haiku-4-5",
litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]),
request_input="Run the command pwd.",
responses_api_request={
"tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}],
"tool_choice": {"type": "function", "name": "run_command"},
},
custom_llm_provider="anthropic",
litellm_metadata={},
)
events: Final = list(iterator)
response_events: Final = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES]
assert [event.type for event in response_events] == [
"response.created",
"response.in_progress",
"response.completed",
]
assert [event.response.tool_choice for event in response_events] == [
{"type": "function", "name": "run_command"},
{"type": "function", "name": "run_command"},
{"type": "function", "name": "run_command"},
]
assert any(getattr(event, "type", None) == "response.output_item.done" for event in events)
def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None:
iterator: Final = LiteLLMCompletionStreamingIterator(
model="claude-haiku-4-5",
litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]),
request_input="Run the command pwd.",
responses_api_request={
"tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}],
"tool_choice": "any",
},
custom_llm_provider="anthropic",
litellm_metadata={},
)
response_events: Final = [
event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES
]
assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"]

View file

@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import (
SyncResponsesAPIStreamingIterator,
)
from litellm.types.llms.openai import (
ResponseAPIUsage,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
from litellm.types.llms.openai import ResponseAPIUsage
return ResponsesAPIResponse(
id="resp_lit6427",
created_at=int(datetime(2025, 1, 1).timestamp()),
@ -368,6 +367,53 @@ def test_stamp_responses_usage_cost_keeps_provider_reported_cost():
logging_obj._response_cost_calculator.assert_not_called()
def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse:
return ResponsesAPIResponse.model_construct(
id="resp_lit7391",
created_at=int(datetime(2025, 1, 1).timestamp()),
status="completed",
model="perplexity/deepseek-v4-flash-0731",
object="response",
output=[],
truncation="",
usage=usage,
)
def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
response = _unvalidated_response_with_dict_usage(
{
"input_tokens": 29,
"output_tokens": 120,
"output_tokens_details": {"reasoning_tokens": 117},
"total_tokens": 149,
"cost": {"currency": "USD", "input_cost": 0, "output_cost": 3e-05, "total_cost": 3e-05},
}
)
logging_obj = Mock(spec=LiteLLMLoggingObj)
_stamp_responses_usage_cost(response, logging_obj)
assert isinstance(response.usage, ResponseAPIUsage)
assert response.usage.cost == pytest.approx(3e-05)
assert response.usage.output_tokens_details.reasoning_tokens == 117
logging_obj._response_cost_calculator.assert_not_called()
def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149})
logging_obj = Mock(spec=LiteLLMLoggingObj)
logging_obj._response_cost_calculator.return_value = 0.000704
_stamp_responses_usage_cost(response, logging_obj)
assert isinstance(response.usage, ResponseAPIUsage)
assert response.usage.cost == pytest.approx(0.000704)
logging_obj._response_cost_calculator.assert_called_once_with(result=response)
def test_stamp_responses_usage_cost_survives_calculator_failure():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
@ -535,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")):
iterator._log_completed_response(is_async=True)
assert logged == [iterator.completed_response]
assert len(logged) == 1
assert logged[0] is not iterator.completed_response
assert logged[0].response is not iterator.completed_response.response
assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1"
assert iterator.completed_response.response._hidden_params == {}
def _unvalidated_completed_config() -> Mock:
"""Config whose completed event carries a Perplexity-style response that fails validation
(``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``."""
mock_config = Mock(spec=BaseResponsesAPIConfig)
def _transform(model, parsed_chunk, logging_obj):
response = _unvalidated_response_with_dict_usage(
ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001})
)
return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response)
mock_config.transform_streaming_response.side_effect = _transform
return mock_config
@pytest.mark.asyncio
async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation():
"""LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging
rewrites the assembled response's usage to chat shape in place, so the event handed to logging
must never be the one the caller receives."""
logging_obj = _logging_obj_stub()
logging_obj.stream = True
logged: list[object] = []
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
logging_obj._on_deferred_stream_complete = None
iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj)
events = [event async for event in iterator]
assert len(logged) == 1
now = datetime.now()
LiteLLMLoggingObj._get_assembled_streaming_response(
logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[]
)
assert logged[0].response.usage["prompt_tokens"] == 29
client_usage = events[-1].response.usage
assert isinstance(client_usage, ResponseAPIUsage)
assert client_usage.input_tokens == 29
assert client_usage.cost == pytest.approx(0.0001)

View file

@ -477,3 +477,65 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost
assert deployments[0]["litellm_params"]["model"] == "anthropic/claude-opus-4-6"
assert _get_min_token_count_for_deployments(deployments) == 4096
@pytest.mark.asyncio
async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop():
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
warm_tokenizer("anthropic/claude-fable-5")
check = PromptCachingDeploymentCheck(cache=DualCache())
deployments = _deployments("anthropic/claude-fable-5")
messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}])
result, took, lags = await timed_with_loop_lags(
lambda: check.async_filter_deployments(
model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages
)
)
assert result == deployments
assert_loop_stayed_free(took, lags)
@pytest.mark.asyncio
async def test_async_log_success_event_counts_the_prompt_off_the_event_loop():
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
warm_tokenizer("anthropic/claude-fable-5")
cache = DualCache()
check = PromptCachingDeploymentCheck(cache=cache)
messages = cast(
List[AllMessageValues],
[{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}],
)
standard_logging_object = {
"call_type": "acompletion",
"model": "anthropic/claude-fable-5",
"messages": messages,
"model_id": "dep-1",
}
_, took, lags = await timed_with_loop_lags(
lambda: check.async_log_success_event(
kwargs={"standard_logging_object": standard_logging_object},
response_obj=None,
start_time=None,
end_time=None,
)
)
assert await PromptCachingCache(cache=cache).async_get_model_id(messages=messages, tools=None) == {
"model_id": "dep-1"
}
assert_loop_stayed_free(took, lags)

Some files were not shown because too many files have changed in this diff Show more