fix: resolve UP045 lint violations (Optional[X] -> X | None)

Convert Optional[X] type annotations to X | None syntax to satisfy ruff strict gate.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Sameer Kankute 2026-06-24 17:22:35 +05:30
commit 0647d5bafd
No known key found for this signature in database
56 changed files with 2678 additions and 535 deletions

View file

@ -3,6 +3,7 @@
Flow:
1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL
2. PUT <signed_url> with CSV content
3. PATCH /metrics/agent/ai/{connection_id} → advance metricsMarker
"""
from __future__ import annotations
@ -127,8 +128,6 @@ class FocusMavvrikDestination(FocusDestination):
timeout=30.0,
)
if resp.status_code == 410:
# Connector has been disconnected in Mavvrik — reset flag so next
# delivery attempt re-registers after it becomes active again.
self._registered = False
raise RuntimeError(
"Mavvrik FOCUS destination: connector is disconnected (410). "
@ -273,14 +272,35 @@ class FocusMavvrikDestination(FocusDestination):
pass
raise
async def _update_metrics_marker(self, date_epoch: int) -> None:
"""PATCH agent endpoint to advance metricsMarker after a successful upload."""
resp = await self._http.client.request(
method="PATCH",
url=self._agent_url,
headers=self._auth_headers,
json={"metricsMarker": date_epoch},
timeout=30.0,
)
if resp.status_code == 410:
self._registered = False
raise RuntimeError(
"Mavvrik FOCUS destination: connector is disconnected (410). "
"Re-enable the connection in the Mavvrik dashboard."
)
if resp.status_code >= 400:
verbose_logger.warning(
"Mavvrik FOCUS destination: failed to update metricsMarker (%s): %s",
resp.status_code,
resp.text[:200],
)
return
verbose_logger.debug(
"Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch
)
async def get_metrics_marker(self) -> Optional[int]:
"""Register with Mavvrik and return the current metricsMarker.
The metricsMarker is a Unix timestamp (seconds) representing the last
date Mavvrik has successfully ingested. Called on every scheduled run
so the logger can detect and catch up any dates missed due to previous
export failures.
Always calls the Mavvrik register API — unlike deliver() which skips
registration once _registered is True, catch-up requires a fresh
marker value on every run.
@ -328,6 +348,7 @@ class FocusMavvrikDestination(FocusDestination):
return
date_str = time_window.start_time.strftime("%Y-%m-%d")
date_epoch = int(time_window.start_time.timestamp())
verbose_logger.debug(
"Mavvrik FOCUS destination: uploading %d bytes for date=%s (%s)",
@ -339,6 +360,7 @@ class FocusMavvrikDestination(FocusDestination):
await self._ensure_registered()
signed_url = await self._get_signed_url(date_str)
await self._upload_to_gcs(signed_url, content)
await self._update_metrics_marker(date_epoch)
verbose_logger.debug(
"Mavvrik FOCUS destination: upload complete for date=%s", date_str

View file

@ -149,8 +149,8 @@ class MavvrikFocusLogger(FocusLogger):
On each run:
1. Register with Mavvrik → get metricsMarker (last successfully ingested date)
2. If metricsMarker is behind yesterday, catch up missed dates (capped at
_MAX_CATCHUP_DAYS to avoid runaway loops on long outages)
2. If metricsMarker is behind yesterday (or 0/None for a fresh connector),
catch up missed dates (capped at _MAX_CATCHUP_DAYS)
3. Export yesterday (today's daily window)
This ensures a failed export on day N is automatically retried on day N+1
@ -177,13 +177,21 @@ class MavvrikFocusLogger(FocusLogger):
last_ingested = _parse_metrics_marker(marker)
# Catch up missed dates, capped at _MAX_CATCHUP_DAYS
if last_ingested and last_ingested < yesterday:
# Never go further back than _MAX_CATCHUP_DAYS from yesterday
earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1)
catch_up_date = max(last_ingested + timedelta(days=1), earliest_catchup)
# Catch up missed dates, capped at _MAX_CATCHUP_DAYS.
# last_ingested=None means metricsMarker=0 (fresh connector, never ingested) —
# treat the same as being _MAX_CATCHUP_DAYS behind so we export all available history.
earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1)
if last_ingested is None or last_ingested < yesterday:
catch_up_date = (
earliest_catchup
if last_ingested is None
else max(last_ingested + timedelta(days=1), earliest_catchup)
)
if last_ingested + timedelta(days=1) < earliest_catchup:
if (
last_ingested is not None
and last_ingested + timedelta(days=1) < earliest_catchup
):
verbose_proxy_logger.warning(
"Mavvrik FOCUS export: metricsMarker is more than %d days behind "
"(%s). Catching up from %s only; earlier data will not be re-exported.",
@ -197,18 +205,24 @@ class MavvrikFocusLogger(FocusLogger):
"Mavvrik FOCUS export: catching up missed date %s",
catch_up_date.date(),
)
# Use now as end_time for catch-up windows too — rows for old dates
# may have been flushed to DB well after their calendar day ended.
catch_up_end = min(catch_up_date + timedelta(days=1), now)
window = FocusTimeWindow(
start_time=catch_up_date,
end_time=catch_up_date + timedelta(days=1),
end_time=catch_up_end,
frequency="daily",
)
await self._export_window(window=window, limit=None)
catch_up_date += timedelta(days=1)
# Export yesterday's window (the normal daily run)
# Export yesterday's window (the normal daily run).
# Use `now` as end_time so spend rows flushed after midnight are included.
# LiteLLM's DailyUserSpend rows for a given date keep getting updated_at
# bumped as the flush job runs; capping at midnight would miss those updates.
window = FocusTimeWindow(
start_time=yesterday,
end_time=yesterday + timedelta(days=1),
end_time=now,
frequency="daily",
)
await self._export_window(window=window, limit=None)
@ -253,6 +267,21 @@ class MavvrikFocusLogger(FocusLogger):
)
if type(cb) is MavvrikFocusLogger
]
if not loggers and "mavvrik" in litellm.callbacks:
# The logger is registered as the string "mavvrik" but hasn't been
# instantiated yet (lazy init happens on first LLM call). Force it now
# so the scheduler can register the daily export job at startup.
from litellm.litellm_core_utils.litellm_logging import ( # noqa: PLC0415
_init_custom_logger_compatible_class,
)
instance = _init_custom_logger_compatible_class(
logging_integration="mavvrik",
internal_usage_cache=None,
llm_router=None,
)
if isinstance(instance, MavvrikFocusLogger):
loggers = [instance]
if not loggers:
verbose_proxy_logger.debug(
"No MavvrikFocusLogger registered; skipping scheduler"

View file

@ -629,6 +629,7 @@ class ModelResponseIterator:
ChatCompletionToolCallChunk | None,
List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]],
Dict[str, Any],
str | None,
]:
"""
Helper function to handle the content block delta
@ -636,6 +637,7 @@ class ModelResponseIterator:
text = ""
tool_use: ChatCompletionToolCallChunk | None = None
provider_specific_fields = {}
reasoning_content: str | None = None
content_block = ContentBlockDelta(**chunk) # type: ignore
thinking_blocks: List[
Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
@ -670,14 +672,24 @@ class ModelResponseIterator:
thinking_content = content_block["delta"].get("thinking")
if isinstance(thinking_content, str) and thinking_content:
self.reasoning_content_chunks.append(thinking_content)
thinking_blocks = [
ChatCompletionThinkingBlock(
type="thinking",
thinking=thinking_content or "",
signature=str(content_block["delta"].get("signature") or ""),
)
]
provider_specific_fields["thinking_blocks"] = thinking_blocks
reasoning_content = thinking_content
signature = content_block["delta"].get("signature")
if isinstance(signature, str) and signature:
thinking_blocks = [
ChatCompletionThinkingBlock(
type="thinking",
thinking="".join(
cast(str, block["delta"].get("thinking"))
for block in self.content_blocks
if isinstance(block["delta"].get("thinking"), str)
),
signature=signature,
)
]
provider_specific_fields["thinking_blocks"] = thinking_blocks
if reasoning_content is None:
reasoning_content = ""
elif (
"content" in content_block["delta"]
and content_block["delta"].get("type") == "compaction_delta"
@ -688,25 +700,13 @@ class ModelResponseIterator:
"content": content_block["delta"]["content"],
}
return text, tool_use, thinking_blocks, provider_specific_fields
def _handle_reasoning_content(
self,
thinking_blocks: List[
Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
],
) -> str | None:
"""
Handle the reasoning content
"""
reasoning_content = None
for block in thinking_blocks:
thinking_content = cast(str | None, block.get("thinking"))
if reasoning_content is None:
reasoning_content = ""
if thinking_content is not None:
reasoning_content += thinking_content
return reasoning_content
return (
text,
tool_use,
thinking_blocks,
provider_specific_fields,
reasoning_content,
)
def _handle_redacted_thinking_content(
self,
@ -796,11 +796,8 @@ class ModelResponseIterator:
tool_use,
thinking_blocks,
provider_specific_fields,
reasoning_content,
) = self._content_block_delta_helper(chunk=chunk)
if thinking_blocks:
reasoning_content = self._handle_reasoning_content(
thinking_blocks=thinking_blocks
)
elif type_chunk == "content_block_start":
"""
event: content_block_start

View file

@ -377,6 +377,17 @@ class BaseConfig(ABC):
) -> "ModelResponse":
pass
def transform_parsed_response_dict(self, parsed_response: dict) -> dict:
"""
Repair a parsed OpenAI-format response dict before generic conversion.
Providers routed through the OpenAI SDK handler bypass transform_response,
which calls convert_to_model_response_object directly on the SDK's parsed
output. Override this to normalize a malformed response (e.g. github_copilot
returning empty choices for Anthropic-native Claude responses).
"""
return parsed_response
@abstractmethod
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]

View file

@ -85,6 +85,7 @@ class BaseRerankConfig(ABC):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
pass

View file

@ -601,6 +601,15 @@ def extract_model_name_from_bedrock_arn(model: str) -> str:
return model
def is_bedrock_application_inference_profile_arn(model: str) -> bool:
"""
An application inference profile ARN ends in an opaque id with no provider
substring, so the invoke path cannot resolve a provider from it. Such ARNs
must use the converse route, which needs no provider.
"""
return ":application-inference-profile/" in model
def strip_bedrock_routing_prefix(model: str) -> str:
"""Strip LiteLLM routing prefixes from model name."""
for prefix in ["bedrock/", "converse/", "invoke/", "openai/", "nova-2/", "nova/"]:
@ -916,6 +925,9 @@ class BedrockModelInfo(BaseLLMModelInfo):
) or _model_after_bedrock.startswith("nova/"):
return "converse"
if is_bedrock_application_inference_profile_arn(model):
return "converse"
base_model = BedrockModelInfo.get_base_model(model)
alt_model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model)
if (

View file

@ -224,6 +224,10 @@ class BedrockEmbedding(BaseAWSLLM):
returned_response = AmazonTitanV2Config()._transform_response(
response_list=response_list, model=model
)
elif model == "amazon.titan-embed-g1-text-02":
returned_response = AmazonTitanG1Config()._transform_response(
response_list=response_list, model=model
)
elif provider == "twelvelabs":
returned_response = (
TwelveLabsMarengoEmbeddingConfig()._transform_response(
@ -447,6 +451,7 @@ class BedrockEmbedding(BaseAWSLLM):
"amazon.titan-embed-image-v1",
"amazon.titan-embed-text-v1",
"amazon.titan-embed-text-v2:0",
"amazon.titan-embed-g1-text-02",
]:
batch_data = []
for i in input:
@ -464,6 +469,10 @@ class BedrockEmbedding(BaseAWSLLM):
transformed_request = AmazonTitanV2Config()._transform_request(
input=i, inference_params=inference_params
)
elif model == "amazon.titan-embed-g1-text-02":
transformed_request = AmazonTitanG1Config()._transform_request(
input=i, inference_params=inference_params
)
else:
raise Exception(
"Unmapped model. Received={}. Expected={}".format(
@ -472,6 +481,7 @@ class BedrockEmbedding(BaseAWSLLM):
"amazon.titan-embed-image-v1",
"amazon.titan-embed-text-v1",
"amazon.titan-embed-text-v2:0",
"amazon.titan-embed-g1-text-02",
],
)
)

View file

@ -26,11 +26,18 @@ class CohereRerankHandler(BaseTranslation):
The handler specifically processes:
- The 'query' parameter (string)
- The 'instruction' parameter (string), when present
Note: Documents are not processed by guardrails as they are the corpus
being searched, not user input.
"""
# User-controlled free-text fields that reach the model and must be
# scanned. 'instruction' is folded into the prompt by instruction-aware
# rerankers (e.g. hosted vLLM / Qwen3-Reranker), so it is as sensitive as
# 'query'; omitting it would let a caller smuggle content past guardrails.
_SCANNED_FIELDS = ("query", "instruction")
async def process_input_messages(
self,
data: dict,
@ -38,42 +45,55 @@ class CohereRerankHandler(BaseTranslation):
litellm_logging_obj: Optional[Any] = None,
) -> Any:
"""
Process input query by applying guardrails.
Process input text fields ('query' and 'instruction') by applying
guardrails and writing the sanitized values back.
Args:
data: Request data dictionary containing 'query'
data: Request data dictionary containing 'query' and optionally
'instruction'
guardrail_to_apply: The guardrail instance to apply
Returns:
Modified data with guardrails applied to query only
Modified data with guardrails applied to query/instruction only
"""
# Process query only
query = data.get("query")
if query is not None and isinstance(query, str):
inputs = GenericGuardrailAPIInputs(texts=[query])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
# Collect every scannable text field in a stable order so the
# guardrailed results can be written back to the right key by index.
fields_to_scan = [
(key, data[key])
for key in self._SCANNED_FIELDS
if isinstance(data.get(key), str)
]
if not fields_to_scan:
verbose_proxy_logger.debug(
"Rerank: No query/instruction to process or not strings"
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["query"] = guardrailed_texts[0] if guardrailed_texts else query
return data
verbose_proxy_logger.debug(
"Rerank: Applied guardrail to query. "
"Original length: %d, New length: %d",
len(query),
len(data["query"]),
)
else:
verbose_proxy_logger.debug(
"Rerank: No query to process or query is not a string"
)
inputs = GenericGuardrailAPIInputs(texts=[value for _, value in fields_to_scan])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
for idx, (key, original) in enumerate(fields_to_scan):
# Defensive: only write back when the guardrail returned a value for
# this index; otherwise keep the original (never forward unscanned).
if idx < len(guardrailed_texts):
data[key] = guardrailed_texts[idx]
verbose_proxy_logger.debug(
"Rerank: Applied guardrail to %s. "
"Original length: %d, New length: %d",
key,
len(original),
len(data[key]),
)
return data

View file

@ -57,6 +57,7 @@ class CohereRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map Cohere rerank params

View file

@ -49,6 +49,7 @@ class CohereRerankV2Config(CohereRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map Cohere rerank params

View file

@ -116,6 +116,7 @@ class DashScopeRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
# qwen3-rerank accepts query/documents/top_n/return_documents. The
# rest (rank_fields, max_*_per_doc) are silently dropped.

View file

@ -104,6 +104,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
# Start with the basic parameters
optional_rerank_params = {}

View file

@ -146,6 +146,56 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
and (optional_params.get("thinking") or {}).get("type") == "enabled"
)
@staticmethod
def _drop_unsupported_tools(optional_params: dict) -> dict:
"""
DeepSeek's /chat/completions only accepts tools of type "function".
Requests bridged from /v1/responses can carry responses-API-native tool
types (e.g. a Codex CLI tool typed "namespace"); DeepSeek rejects the
whole request with `unknown variant '<type>', expected 'function'` (issue
#30722). Drop the unsupported entries so the function tools still go
through, and drop the now-dangling tool_choice/parallel_tool_calls when
nothing callable survives.
Only non-`function` tools are ever dropped, so a `tool_choice` that names
a specific function still points at a surviving tool and is left intact;
`tool_choice`/`parallel_tool_calls` are cleared only when no function
tool remains.
"""
tools = optional_params.get("tools")
if not isinstance(tools, list) or not tools:
return optional_params
def _is_function_tool(tool: object) -> bool:
return isinstance(tool, dict) and tool.get("type") == "function"
function_tools = [tool for tool in tools if _is_function_tool(tool)]
if len(function_tools) == len(tools):
return optional_params
dropped_types = sorted(
{
str(tool.get("type")) if isinstance(tool, dict) else type(tool).__name__
for tool in tools
if not _is_function_tool(tool)
}
)
litellm.verbose_logger.warning(
"DeepSeek chat completions only supports function tools; dropping "
"unsupported tool type(s) %s before sending the request",
dropped_types,
)
cleaned = {k: v for k, v in optional_params.items() if k != "tools"}
if function_tools:
return {**cleaned, "tools": function_tools}
return {
k: v
for k, v in cleaned.items()
if k not in ("tool_choice", "parallel_tool_calls")
}
def transform_request(
self,
model: str,
@ -163,6 +213,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
(user explicitly enabled it), preventing spurious injection on models
like deepseek-v3.2 that support thinking as opt-in but not always-on.
"""
optional_params = self._drop_unsupported_tools(optional_params)
if self._thinking_mode_active(model=model, optional_params=optional_params):
messages = self._fill_reasoning_content(messages)
return super().transform_request(
@ -185,6 +236,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
Async equivalent of transform_request — applies the same reasoning_content
fix for multi-turn thinking-mode conversations.
"""
optional_params = self._drop_unsupported_tools(optional_params)
if self._thinking_mode_active(model=model, optional_params=optional_params):
messages = self._fill_reasoning_content(messages)
return await super().async_transform_request(

View file

@ -67,6 +67,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict[str, Any]:
"""
Map Cohere rerank params to Fireworks AI rerank params

View file

@ -194,6 +194,88 @@ class GithubCopilotConfig(OpenAIConfig):
)
return text_content, tool_calls, thinking_blocks
@staticmethod
def _normalize_anthropic_usage(usage: dict) -> dict:
normalized = dict(usage)
if "input_tokens" in usage and "prompt_tokens" not in usage:
normalized["prompt_tokens"] = usage["input_tokens"]
if "output_tokens" in usage and "completion_tokens" not in usage:
normalized["completion_tokens"] = usage["output_tokens"]
if "total_tokens" not in normalized:
normalized["total_tokens"] = normalized.get(
"prompt_tokens", 0
) + normalized.get("completion_tokens", 0)
return normalized
@classmethod
def _synthesize_choices_for_anthropic_native(cls, response_json: dict) -> dict:
"""
Synthesize a `choices` array from an Anthropic-native Copilot response.
Newer Copilot Claude models (e.g. opus-4.7, opus-4.8) return content
blocks and `stop_reason` without an OpenAI-style `choices` array, and the
max_tokens=1 probe returns no content at all. Returns the response
unchanged when it already carries choices.
See: https://github.com/BerriAI/litellm/issues/29391
"""
if response_json.get("choices"):
return response_json
content = ""
tool_calls: List[ChatCompletionToolCallChunk] = []
thinking_blocks: List[Any] | None = None
raw_content = response_json.get("content")
if isinstance(raw_content, list):
content, tool_calls, thinking_blocks = cls._parse_anthropic_native_content(
raw_content
)
elif isinstance(raw_content, str):
content = raw_content
stop_reason = response_json.get("stop_reason")
finish_reason_map = {
"end_turn": "stop",
"max_tokens": "length",
"stop_sequence": "stop",
"tool_use": "tool_calls",
}
if tool_calls:
finish_reason = "tool_calls"
elif stop_reason in finish_reason_map:
finish_reason = finish_reason_map[stop_reason]
elif content:
finish_reason = "stop"
else:
finish_reason = "length"
message: dict = {
"role": "assistant",
"content": content if content or not tool_calls else None,
}
if tool_calls:
message["tool_calls"] = tool_calls
if thinking_blocks:
message["thinking_blocks"] = thinking_blocks
synthesized = {
**response_json,
"choices": [
{"index": 0, "message": message, "finish_reason": finish_reason}
],
}
usage = response_json.get("usage")
if isinstance(usage, dict):
synthesized["usage"] = cls._normalize_anthropic_usage(usage)
return synthesized
def transform_parsed_response_dict(self, parsed_response: dict) -> dict:
"""
Repair the OpenAI-SDK-parsed response on the handler path that bypasses
transform_response. See: https://github.com/BerriAI/litellm/issues/30927
"""
return self._synthesize_choices_for_anthropic_native(parsed_response)
def transform_response(
self,
model: str,
@ -208,15 +290,6 @@ class GithubCopilotConfig(OpenAIConfig):
api_key: str | None = None,
json_mode: bool | None = None,
) -> "ModelResponse":
"""
Handle newer Copilot models (e.g. claude-opus-4.7, claude-opus-4.8) that
return Anthropic-native format responses without a `choices` array.
Synthesizes the missing `choices` from Anthropic-native fields, then
delegates to the parent so all standard post-processing applies.
See: https://github.com/BerriAI/litellm/issues/29391
"""
try:
response_json = raw_response.json()
except Exception:
@ -235,70 +308,12 @@ class GithubCopilotConfig(OpenAIConfig):
)
if not response_json.get("choices"):
content = ""
tool_calls: List[ChatCompletionToolCallChunk] = []
thinking_blocks: List[Any] | None = None
if "content" in response_json and isinstance(
response_json["content"], list
):
content, tool_calls, thinking_blocks = (
self._parse_anthropic_native_content(response_json["content"])
)
elif isinstance(response_json.get("content"), str):
content = response_json["content"]
stop_reason = response_json.get("stop_reason")
finish_reason_map = {
"end_turn": "stop",
"max_tokens": "length",
"stop_sequence": "stop",
"tool_use": "tool_calls",
}
# Prefer tool_calls when blocks were extracted; otherwise map stop_reason.
if tool_calls:
finish_reason = "tool_calls"
elif stop_reason in finish_reason_map:
finish_reason = finish_reason_map[stop_reason]
elif content:
finish_reason = "stop"
else:
finish_reason = "length"
message: dict = {
"role": "assistant",
"content": content if content or not tool_calls else None,
}
if tool_calls:
message["tool_calls"] = tool_calls
if thinking_blocks:
message["thinking_blocks"] = thinking_blocks
response_json["choices"] = [
{
"index": 0,
"message": message,
"finish_reason": finish_reason,
}
]
if "usage" in response_json:
usage = response_json["usage"]
if "input_tokens" in usage and "prompt_tokens" not in usage:
usage["prompt_tokens"] = usage["input_tokens"]
if "output_tokens" in usage and "completion_tokens" not in usage:
usage["completion_tokens"] = usage["output_tokens"]
if "total_tokens" not in usage:
usage["total_tokens"] = usage.get("prompt_tokens", 0) + usage.get(
"completion_tokens", 0
)
# Build a patched response so super() sees valid JSON with choices
patched = httpx.Response(
response_json = self._synthesize_choices_for_anthropic_native(response_json)
raw_response = httpx.Response(
status_code=raw_response.status_code,
headers=raw_response.headers,
content=json.dumps(response_json).encode(),
)
raw_response = patched
return super().transform_response(
model=model,

View file

@ -61,6 +61,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
"top_n",
"rank_fields",
"return_documents",
"instruction",
]
def map_cohere_rerank_params(
@ -76,6 +77,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map parameters for Hosted VLLM rerank
@ -83,16 +85,22 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
if max_chunks_per_doc is not None:
raise ValueError("Hosted VLLM does not support max_chunks_per_doc")
return dict(
OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
)
mapped_params = OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
)
# `instruction` is a vLLM-supported passthrough (folded into the model's
# chat_template_kwargs). Only forward it when explicitly set so omitting
# it leaves the request unchanged.
if instruction is not None:
mapped_params["instruction"] = instruction
return dict(mapped_params)
def validate_environment(
self,
headers: dict,
@ -135,6 +143,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
top_n=optional_rerank_params.get("top_n", None),
rank_fields=optional_rerank_params.get("rank_fields", None),
return_documents=optional_rerank_params.get("return_documents", None),
instruction=optional_rerank_params.get("instruction", None),
)
return rerank_request.model_dump(exclude_none=True)

View file

@ -100,6 +100,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
optional_rerank_params = {}
if non_default_params is not None:

View file

@ -45,6 +45,7 @@ class JinaAIRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
optional_params = {}
supported_params = self.get_supported_cohere_rerank_params(model)

View file

@ -238,11 +238,11 @@ class MoonshotChatConfig(OpenAIGPTConfig):
https://platform.moonshot.ai/docs/guide/migrating-from-openai-to-kimi#about-tool_choice
"""
messages.append(
optional_params.pop("tool_choice")
return [
*messages,
{
"role": "user",
"content": "Please select a tool to handle the current issue.", # Usually, the Kimi large language model understands the intention to invoke a tool and selects one for invocation
}
)
optional_params.pop("tool_choice")
return messages
},
]

View file

@ -117,6 +117,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map Cohere/OpenAI rerank params to Nvidia NIM format.

View file

@ -7,7 +7,10 @@ These models use token-based pricing instead of pixel-based pricing like DALL-E.
from typing import Optional
from litellm import verbose_logger
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_image_response_cost_from_usage,
generic_cost_per_token,
)
from litellm.types.utils import ImageResponse, Usage
@ -16,54 +19,40 @@ def cost_calculator(
image_response: ImageResponse,
custom_llm_provider: Optional[str] = None,
) -> float:
"""
Calculate cost for OpenAI gpt-image models.
Uses the same usage format as Responses API, so we reuse the helper
to transform to chat completion format and use generic_cost_per_token.
Args:
model: The model name (e.g., "gpt-image-1", "gpt-image-2")
image_response: The ImageResponse containing usage data
custom_llm_provider: Optional provider name
Returns:
float: Total cost in USD
"""
"""Calculate cost for OpenAI gpt-image models (token-based pricing)."""
usage = getattr(image_response, "usage", None)
if usage is None:
verbose_logger.debug(
f"No usage data available for {model}, cannot calculate token-based cost"
)
return 0.0
# If usage is already a Usage object with completion_tokens_details set,
# use it directly (it was already transformed in convert_to_image_response)
provider = custom_llm_provider or "openai"
# A chat Usage with an explicit output breakdown: cost via generic_cost_per_token.
if isinstance(usage, Usage) and usage.completion_tokens_details is not None:
chat_usage = usage
else:
# Transform ImageUsage to Usage using the existing helper
# ImageUsage has the same format as ResponseAPIUsage
from litellm.responses.utils import ResponseAPILoggingUtils
chat_usage = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
prompt_cost, completion_cost = generic_cost_per_token(
model=model, usage=usage, custom_llm_provider=provider
)
return prompt_cost + completion_cost
# Use generic_cost_per_token for cost calculation
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=chat_usage,
custom_llm_provider=custom_llm_provider or "openai",
)
# ImageUsage / ResponseAPIUsage: reuse the shared helper (same path as
# azure_ai/gemini/vertex_ai). It prices generated output tokens at
# output_cost_per_image_token, classifying them as image tokens when the provider
# does not itemize output and splitting text/image when it does.
if getattr(usage, "input_tokens", None) is not None:
token_based_cost = calculate_image_response_cost_from_usage(
model=model, image_response=image_response, custom_llm_provider=provider
)
if token_based_cost is not None:
return token_based_cost
total_cost = prompt_cost + completion_cost
# Fallback: a Usage with no output breakdown that the image helper can't read —
# cost via generic_cost_per_token (text rate) instead of returning 0.0.
if isinstance(usage, Usage):
prompt_cost, completion_cost = generic_cost_per_token(
model=model, usage=usage, custom_llm_provider=provider
)
return prompt_cost + completion_cost
verbose_logger.debug(
f"OpenAI gpt-image cost calculation for {model}: "
f"prompt_cost=${prompt_cost:.6f}, completion_cost=${completion_cost:.6f}, "
f"total=${total_cost:.6f}"
)
return total_cost
return 0.0

View file

@ -785,7 +785,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
)
logging_obj.model_call_details["response_headers"] = headers
stringified_response = response.model_dump()
stringified_response = (
provider_config.transform_parsed_response_dict(
response.model_dump()
)
)
logging_obj.post_call(
input=messages,
api_key=api_key,
@ -933,7 +937,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
timeout=timeout,
logging_obj=logging_obj,
)
stringified_response = response.model_dump()
stringified_response = provider_config.transform_parsed_response_dict(
response.model_dump()
)
logging_obj.post_call(
input=data["messages"],
api_key=api_key,

View file

@ -242,6 +242,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map Cohere rerank params to Vertex AI format

View file

@ -39,6 +39,7 @@ class VoyageRerankConfig(BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
# Voyage AI uses 'top_k' instead of 'top_n'
optional_params: Dict[str, Any] = {"query": query, "documents": documents}

View file

@ -104,6 +104,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> Dict:
"""
Map Cohere rerank params to IBM watsonx.ai rerank params

File diff suppressed because it is too large Load diff

View file

@ -16,8 +16,14 @@ class ColdStorageHandler:
This class is responsible for handling Getting/Setting the proxy server request from cold storage.
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
The cold storage logger can be injected for testing; when omitted it is
resolved from the configured ``litellm.cold_storage_custom_logger``.
"""
def __init__(self, cold_storage_logger: CustomLogger | None = None):
self._injected_cold_storage_logger = cold_storage_logger
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,
object_key: str,
@ -31,33 +37,26 @@ class ColdStorageHandler:
Returns:
Optional[dict]: The proxy server request dict or None if not found
"""
# select the custom logger to use for cold storage
custom_logger_name: _custom_logger_compatible_callbacks_literal | None = (
self._select_custom_logger_for_cold_storage()
custom_logger = (
self._injected_cold_storage_logger or self._resolve_cold_storage_logger()
)
# if no custom logger name is configured, return None
if custom_logger_name is None:
if custom_logger is None:
return None
# get the active/initialized custom logger
custom_logger: CustomLogger | None = (
return await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key,
)
def _resolve_cold_storage_logger(self) -> CustomLogger | None:
custom_logger_name = self._select_custom_logger_for_cold_storage()
if custom_logger_name is None:
return None
return (
litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(
custom_logger_name
)
)
# if no custom logger is found, return None
if custom_logger is None:
return None
proxy_server_request = await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key,
)
return proxy_server_request
def _select_custom_logger_for_cold_storage(
self,
) -> _custom_logger_compatible_callbacks_literal | None:

View file

@ -3,7 +3,17 @@ import collections
import json
import os
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Mapping,
NamedTuple,
Optional,
Union,
)
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, status
@ -29,6 +39,7 @@ from litellm.repositories.verification_token_repository import (
if TYPE_CHECKING:
from litellm.proxy.proxy_server import PrismaClient
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
else:
PrismaClient = Any
@ -2175,6 +2186,89 @@ async def ui_view_spend_logs(
raise handle_exception_on_proxy(e)
class RequestResponsePayload(NamedTuple):
messages: Union[str, list, dict] | None
response: Union[str, list, dict] | None
proxy_server_request: Union[str, dict] | None
_EMPTY_SPEND_LOG_VALUES = frozenset({"", "{}", "[]", "null"})
def _spend_log_field_has_content(value: Union[str, list, dict] | None) -> bool:
if value is None:
return False
if isinstance(value, str):
return value.strip() not in _EMPTY_SPEND_LOG_VALUES
if isinstance(value, (list, dict)):
return len(value) > 0
return True
def _cold_storage_object_key_from_metadata(
metadata: Union[str, dict] | None,
) -> str | None:
if isinstance(metadata, str):
try:
metadata = json.loads(metadata)
except (json.JSONDecodeError, TypeError):
return None
if not isinstance(metadata, dict):
return None
object_key = metadata.get("cold_storage_object_key")
return object_key if isinstance(object_key, str) and object_key else None
async def _resolve_request_response_payload(
row: Mapping[str, Any],
cold_storage_handler: "ColdStorageHandler",
) -> RequestResponsePayload:
"""
Decide where the prompt/response come from for a single spend-log row.
PG holds the content when ``store_prompts_in_spend_logs`` is on; otherwise it
holds ``"{}"`` placeholders and the real payload lives in cold storage keyed
by ``metadata.cold_storage_object_key``. The choice is made on actual row
content, not config flags, so historical and mixed-storage rows both resolve
correctly.
"""
messages = row.get("messages")
response = row.get("response")
proxy_server_request = row.get("proxy_server_request")
pg_payload = RequestResponsePayload(messages, response, proxy_server_request)
if (
_spend_log_field_has_content(messages)
or _spend_log_field_has_content(response)
or _spend_log_field_has_content(proxy_server_request)
):
return pg_payload
object_key = _cold_storage_object_key_from_metadata(row.get("metadata"))
if object_key is None:
return pg_payload
try:
payload = await cold_storage_handler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key
)
except Exception:
verbose_proxy_logger.warning(
"Failed to fetch cold storage payload for key %s; falling back to DB values",
object_key,
exc_info=True,
)
return pg_payload
if payload is None:
return pg_payload
return RequestResponsePayload(
messages=payload.get("messages"),
response=payload.get("response"),
proxy_server_request=payload.get("proxy_server_request"),
)
@router.get(
"/spend/logs/ui/{request_id}",
tags=["Budget & Spend Tracking"],
@ -2241,26 +2335,27 @@ async def ui_view_request_response_for_request_id(
if payload is not None:
return payload
# Fallback: fetch heavy columns directly from the database.
# The list endpoint (/spend/logs/ui) intentionally excludes messages,
# response, and proxy_server_request for performance. When no custom
# logger (S3, GCS, etc.) is configured, we still need to serve these
# fields from the DB for the detail/drawer view.
# Fallback: the list endpoint omits the heavy columns for performance, so
# serve them here. When prompts were offloaded to cold storage the DB holds
# only placeholders, so _resolve_request_response_payload fetches the real
# payload from the configured cold storage backend by object key.
if prisma_client is not None:
from litellm.proxy.spend_tracking.cold_storage_handler import (
ColdStorageHandler,
)
sql_query = """
SELECT messages, response, proxy_server_request
SELECT messages, response, proxy_server_request, metadata
FROM "LiteLLM_SpendLogs"
WHERE request_id = $1
LIMIT 1
"""
db_result = await prisma_client.db.query_raw(sql_query, request_id)
if db_result and len(db_result) > 0:
row = db_result[0]
return {
"messages": row.get("messages"),
"response": row.get("response"),
"proxy_server_request": row.get("proxy_server_request"),
}
resolved = await _resolve_request_response_payload(
db_result[0], cold_storage_handler=ColdStorageHandler()
)
return resolved._asdict()
return None

View file

@ -86,6 +86,11 @@ def rerank(
"""
Reranks a list of documents based on their relevance to the query
"""
# `instruction` is read from kwargs rather than declared as a named param.
# The router forwards rerank calls via an untyped `**kwargs` unpack, and a
# typed named param there would trip the basedpyright budget gate without
# adding real safety; it stays typed downstream via get_optional_rerank_params.
instruction: str | None = kwargs.get("instruction", None)
headers: dict | None = kwargs.get("headers") # type: ignore
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
@ -138,6 +143,7 @@ def rerank(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
max_tokens_per_doc=max_tokens_per_doc,
instruction=instruction,
non_default_params=kwargs,
)
verbose_logger.info(f"optional_rerank_params: {optional_rerank_params}")

View file

@ -15,6 +15,7 @@ def get_optional_rerank_params(
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
non_default_params: dict | None = None,
) -> Dict:
all_non_default_params = non_default_params or {}
@ -30,6 +31,11 @@ def get_optional_rerank_params(
all_non_default_params["max_chunks_per_doc"] = max_chunks_per_doc
if max_tokens_per_doc is not None:
all_non_default_params["max_tokens_per_doc"] = max_tokens_per_doc
if instruction is not None:
# Also surfaced in non_default_params so providers that read it from
# there (e.g. DeepInfra) keep working now that `rerank()` consumes
# `instruction` as a named param instead of leaving it in **kwargs.
all_non_default_params["instruction"] = instruction
return rerank_provider_config.map_cohere_rerank_params(
model=model,
drop_params=drop_params,
@ -41,5 +47,6 @@ def get_optional_rerank_params(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
max_tokens_per_doc=max_tokens_per_doc,
instruction=instruction,
non_default_params=all_non_default_params,
)

View file

@ -72,7 +72,7 @@ def get_fallback_model_group(
elif list(item.keys())[0] == "*": # check generic fallback
generic_fallback_idx = idx
elif isinstance(item, str):
fallback_model_group = [fallbacks.pop(idx)] # returns single-item list
fallback_model_group = [item]
## if none, check for generic fallback
if fallback_model_group is None:
if stripped_model_fallback is not None:

View file

@ -19,6 +19,10 @@ class RerankRequest(BaseModel):
return_documents: Optional[bool] = None
max_chunks_per_doc: Optional[int] = None
max_tokens_per_doc: Optional[int] = None
# Optional task/query instruction passed through to providers that support it
# (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing
# request when None, so this is fully backward-compatible.
instruction: Optional[str] = None
class OptionalRerankParams(TypedDict, total=False):
@ -29,6 +33,7 @@ class OptionalRerankParams(TypedDict, total=False):
return_documents: Optional[bool]
max_chunks_per_doc: Optional[int]
max_tokens_per_doc: Optional[int]
instruction: Optional[str]
class RerankBilledUnits(TypedDict, total=False):

View file

@ -2441,7 +2441,7 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
class TranscriptionUsageDurationObject(BaseModel):
type: Literal["duration"]
seconds: int
seconds: float
class TranscriptionUsageInputTokenDetailsObject(BaseModel):

View file

@ -5844,6 +5844,9 @@ def _is_potential_model_name_in_model_cost(
)
_ABOVE_THRESHOLD_COST_KEY = re.compile(r"_above_\d+k?_tokens$")
def _get_model_info_helper(
model: str,
custom_llm_provider: Optional[str] = None,
@ -6021,7 +6024,7 @@ def _get_model_info_helper(
)
_output_cost_per_token = 0
return ModelInfoBase(
returned_model_info = ModelInfoBase(
key=key,
max_tokens=_model_info.get("max_tokens", None),
max_input_tokens=_model_info.get("max_input_tokens", None),
@ -6238,6 +6241,13 @@ def _get_model_info_helper(
uses_embed_content=_model_info.get("uses_embed_content", None),
supports_image_size=_model_info.get("supports_image_size", None),
)
for cost_key, cost_value in _model_info.items():
if (
cost_key not in returned_model_info
and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None
):
returned_model_info[cost_key] = cost_value # type: ignore[literal-required]
return returned_model_info
except Exception as e:
verbose_logger.debug(f"Error getting model info: {e}")
raise Exception(

View file

@ -570,6 +570,15 @@
"output_cost_per_token": 0.0,
"output_vector_size": 1536
},
"amazon.titan-embed-g1-text-02": {
"input_cost_per_token": 1e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 8192,
"max_tokens": 8192,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 1536
},
"amazon.titan-embed-text-v2:0": {
"input_cost_per_token": 2e-08,
"litellm_provider": "bedrock",
@ -22393,7 +22402,7 @@
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_priority": 1.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -22441,7 +22450,7 @@
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_priority": 1.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -22487,7 +22496,7 @@
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_batches": 1e-07,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -22532,7 +22541,7 @@
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_batches": 1e-07,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -22573,9 +22582,9 @@
"input_cost_per_token": 1.5e-05,
"input_cost_per_token_batches": 7.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 272000,
"max_tokens": 272000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00012,
"output_cost_per_token_batches": 6e-05,
@ -22609,9 +22618,9 @@
"input_cost_per_token": 1.5e-05,
"input_cost_per_token_batches": 7.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 272000,
"max_tokens": 272000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00012,
"output_cost_per_token_batches": 6e-05,
@ -29866,6 +29875,22 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/z-ai/glm-5.1": {
"input_cost_per_token": 1.05e-06,
"output_cost_per_token": 3.5e-06,
"cache_read_input_token_cost": 5.25e-07,
"cache_creation_input_token_cost": 0.0,
"litellm_provider": "openrouter",
"max_input_tokens": 202752,
"max_output_tokens": 65535,
"max_tokens": 65535,
"mode": "chat",
"source": "https://openrouter.ai/z-ai/glm-5.1",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/minimax/minimax-m2.1": {
"input_cost_per_token": 2.7e-07,
"output_cost_per_token": 1.2e-06,
@ -31365,13 +31390,13 @@
"output_cost_per_token": 0.0
},
"sambanova/MiniMax-M2.7": {
"input_cost_per_token": 3e-07,
"input_cost_per_token": 6e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 204800,
"max_input_tokens": 196608,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"output_cost_per_token": 2.4e-06,
"source": "https://cloud.sambanova.ai/plans/pricing",
"supports_function_calling": true,
"supports_reasoning": true,
@ -31388,6 +31413,7 @@
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/DeepSeek-R1-Distill-Llama-70B": {
"deprecation_date": "2026-03-20",
"input_cost_per_token": 7e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 131072,
@ -31398,6 +31424,7 @@
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/DeepSeek-V3-0324": {
"deprecation_date": "2026-04-14",
"input_cost_per_token": 3e-06,
"litellm_provider": "sambanova",
"max_input_tokens": 32768,
@ -31428,6 +31455,7 @@
"supports_vision": true
},
"sambanova/Llama-4-Scout-17B-16E-Instruct": {
"deprecation_date": "2025-06-19",
"input_cost_per_token": 4e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 8192,
@ -31444,6 +31472,7 @@
"supports_tool_choice": true
},
"sambanova/Meta-Llama-3.1-405B-Instruct": {
"deprecation_date": "2025-06-25",
"input_cost_per_token": 5e-06,
"litellm_provider": "sambanova",
"max_input_tokens": 16384,
@ -31457,6 +31486,7 @@
"supports_tool_choice": true
},
"sambanova/Meta-Llama-3.1-8B-Instruct": {
"deprecation_date": "2026-04-14",
"input_cost_per_token": 1e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 16384,
@ -31470,6 +31500,7 @@
"supports_tool_choice": true
},
"sambanova/Meta-Llama-3.2-1B-Instruct": {
"deprecation_date": "2025-06-25",
"input_cost_per_token": 4e-08,
"litellm_provider": "sambanova",
"max_input_tokens": 16384,
@ -31480,6 +31511,7 @@
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/Meta-Llama-3.2-3B-Instruct": {
"deprecation_date": "2025-06-25",
"input_cost_per_token": 8e-08,
"litellm_provider": "sambanova",
"max_input_tokens": 4096,
@ -31503,6 +31535,7 @@
"supports_tool_choice": true
},
"sambanova/Meta-Llama-Guard-3-8B": {
"deprecation_date": "2025-06-25",
"input_cost_per_token": 3e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 16384,
@ -31513,6 +31546,7 @@
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/QwQ-32B": {
"deprecation_date": "2025-06-25",
"input_cost_per_token": 5e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 16384,
@ -31523,6 +31557,7 @@
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/Qwen2-Audio-7B-Instruct": {
"deprecation_date": "2025-06-19",
"input_cost_per_token": 5e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 4096,
@ -31534,6 +31569,7 @@
"supports_audio_input": true
},
"sambanova/Qwen3-32B": {
"deprecation_date": "2026-04-06",
"input_cost_per_token": 4e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 8192,
@ -31547,9 +31583,9 @@
"supports_tool_choice": true
},
"sambanova/DeepSeek-V3.1": {
"max_tokens": 32768,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3e-06,
"output_cost_per_token": 4.5e-06,
"litellm_provider": "sambanova",
@ -31563,13 +31599,36 @@
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 2.2e-07,
"output_cost_per_token": 5.9e-07,
"litellm_provider": "sambanova",
"mode": "chat",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/DeepSeek-V3.2": {
"max_tokens": 32768,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
"input_cost_per_token": 3e-06,
"output_cost_per_token": 4.5e-06,
"litellm_provider": "sambanova",
"mode": "chat",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"sambanova/gemma-4-31B-it": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3.8e-07,
"output_cost_per_token": 1.15e-06,
"litellm_provider": "sambanova",
"mode": "chat",
"supports_vision": true,
"source": "https://cloud.sambanova.ai/plans/pricing"
},
"snowflake/claude-3-5-sonnet": {
@ -37872,6 +37931,21 @@
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-5.1": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 2.6e-07,
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 4.4e-06,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-5-code": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 3e-07,
@ -37902,6 +37976,21 @@
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-4.7-flash": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 0,
"input_cost_per_token": 0,
"output_cost_per_token": 0,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-4.6": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 1.1e-07,

View file

@ -34,6 +34,11 @@ img_base_64 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIh
"text",
titan_embedding_response,
), # V2 text model
(
"bedrock/amazon.titan-embed-g1-text-02",
"text",
titan_embedding_response,
), # G1 text model
(
"bedrock/amazon.titan-embed-image-v1",
"image",
@ -459,3 +464,13 @@ def test_bedrock_embedding_region_bug_reproduction():
os.environ["AWS_REGION_NAME"] = original_region_name
else:
os.environ.pop("AWS_REGION_NAME", None)
def test_bedrock_titan_g1_text_02_model_info():
"""Test that amazon.titan-embed-g1-text-02 has correct pricing metadata"""
model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02")
assert model_info is not None, "Model info should not be None"
assert model_info["litellm_provider"] == "bedrock"
assert model_info["mode"] == "embedding"
assert model_info["input_cost_per_token"] == 1e-07
assert model_info["max_input_tokens"] == 8192

View file

@ -1627,7 +1627,12 @@ 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."""
"""Empty choices list raises APIError, same as missing/null choices.
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.
"""
from litellm.exceptions import APIError
response_object = {
@ -1683,7 +1688,9 @@ class TestMissingChoicesGuard:
assert "no 'choices'" in exc_info.value.message
def test_convert_to_model_response_object_stream_true_no_choices_raises_api_error(self):
def test_convert_to_model_response_object_stream_true_no_choices_raises_api_error(
self,
):
"""Missing choices via stream=True path raises APIError when generator is consumed."""
from litellm.exceptions import APIError
@ -2471,6 +2478,13 @@ class TestConvertToModelResponseObjectCompletion:
def test_model_response_none_raises(self):
with pytest.raises(Exception):
convert_to_model_response_object(
response_object={"choices": [{"message": {"content": "hi", "role": "assistant"}, "finish_reason": "stop"}]},
response_object={
"choices": [
{
"message": {"content": "hi", "role": "assistant"},
"finish_reason": "stop",
}
]
},
model_response_object=None,
)

View file

@ -34,6 +34,12 @@ def _dest(**overrides) -> FocusMavvrikDestination:
return FocusMavvrikDestination(prefix="mavvrik_focus_exports", config=config)
def _patch_resp(status: int = 204) -> MagicMock:
r = MagicMock()
r.status_code = status
return r
def test_missing_api_key_raises():
with pytest.raises(ValueError, match="MAVVRIK_API_KEY"):
FocusMavvrikDestination(
@ -127,6 +133,8 @@ async def test_large_content_uploads_in_multiple_chunks():
chunk2_resp = MagicMock()
chunk2_resp.status_code = 200
patch_resp = _patch_resp(204)
mock_http = MagicMock()
mock_http.client = MagicMock()
mock_http.client.request = AsyncMock(
@ -136,6 +144,7 @@ async def test_large_content_uploads_in_multiple_chunks():
init_resp,
chunk1_resp,
chunk2_resp,
patch_resp,
]
)
dest._http = mock_http
@ -152,15 +161,20 @@ async def test_large_content_uploads_in_multiple_chunks():
filename="usage.csv",
)
# register + get_signed_url + init + 2 chunk PUTs = 5 calls
assert mock_http.client.request.call_count == 5
# register + get_signed_url + init + 2 chunk PUTs + PATCH = 6 calls
assert mock_http.client.request.call_count == 6
# Check Content-Range headers
put_calls = mock_http.client.request.call_args_list[3:]
# Check Content-Range headers on the chunk PUTs (calls 3 and 4)
put_calls = mock_http.client.request.call_args_list[3:5]
assert "bytes" in put_calls[0].kwargs["headers"]["Content-Range"]
assert "/*" in put_calls[0].kwargs["headers"]["Content-Range"] # intermediate
assert "/*" not in put_calls[1].kwargs["headers"]["Content-Range"] # final
# Verify the PATCH call advanced metricsMarker
patch_call = mock_http.client.request.call_args_list[5]
assert patch_call.kwargs["method"] == "PATCH"
assert "metricsMarker" in patch_call.kwargs["json"]
@pytest.mark.asyncio
async def test_deliver_calls_register_get_url_and_upload():
@ -180,12 +194,14 @@ async def test_deliver_calls_register_get_url_and_upload():
upload_resp = MagicMock()
upload_resp.status_code = 200
patch_resp = _patch_resp(204)
mock_http = MagicMock()
mock_http.client = MagicMock()
# All 4 calls go through self._http.client.request:
# 1. register, 2. get_signed_url, 3. GCS session init POST, 4. GCS PUT
# All 5 calls go through self._http.client.request:
# 1. register, 2. get_signed_url, 3. GCS session init POST, 4. GCS PUT, 5. PATCH marker
mock_http.client.request = AsyncMock(
side_effect=[register_resp, signed_url_resp, init_resp, upload_resp]
side_effect=[register_resp, signed_url_resp, init_resp, upload_resp, patch_resp]
)
dest._http = mock_http
@ -196,10 +212,14 @@ async def test_deliver_calls_register_get_url_and_upload():
)
assert dest._registered is True
assert mock_http.client.request.call_count == 4
assert mock_http.client.request.call_count == 5
# Verify Content-Range header was set on the PUT
put_call = mock_http.client.request.call_args_list[3]
assert "Content-Range" in put_call.kwargs["headers"]
# Verify PATCH was called last with metricsMarker
patch_call = mock_http.client.request.call_args_list[4]
assert patch_call.kwargs["method"] == "PATCH"
assert "metricsMarker" in patch_call.kwargs["json"]
@pytest.mark.asyncio
@ -224,17 +244,19 @@ async def test_register_called_only_once_across_multiple_deliveries():
mock_http = MagicMock()
mock_http.client = MagicMock()
# First delivery: register, get_signed_url, GCS init, GCS PUT
# Second delivery: get_signed_url, GCS init, GCS PUT (register skipped)
# First delivery: register, get_signed_url, GCS init, GCS PUT, PATCH
# Second delivery: get_signed_url, GCS init, GCS PUT, PATCH (register skipped)
mock_http.client.request = AsyncMock(
side_effect=[
register_resp,
_signed_url_resp(),
init_resp,
upload_resp,
_patch_resp(204),
_signed_url_resp(),
init_resp,
upload_resp,
_patch_resp(204),
]
)
dest._http = mock_http
@ -243,8 +265,8 @@ async def test_register_called_only_once_across_multiple_deliveries():
await dest.deliver(content=b"header\nrow1\n", time_window=window, filename="1.csv")
await dest.deliver(content=b"header\nrow2\n", time_window=window, filename="2.csv")
# 7 total: register(1) + [get_url+init+put](2) × 2 deliveries
assert mock_http.client.request.call_count == 7
# 9 total: register(1) + [get_url+init+put+patch](4) × 2 deliveries
assert mock_http.client.request.call_count == 9
# First call was register
first_call = mock_http.client.request.call_args_list[0]
assert first_call.kwargs["method"] == "POST"
@ -749,3 +771,43 @@ async def test_gcs_session_cancelled_on_chunk_failure():
delete_call = calls[4]
assert delete_call.kwargs["method"] == "DELETE"
assert "storage.googleapis.com/session" in delete_call.kwargs["url"]
@pytest.mark.asyncio
async def test_update_metrics_marker_warns_on_non_410_error():
"""_update_metrics_marker must log a warning on any >=400 (non-410) status but not raise."""
dest = _dest()
fail_resp = MagicMock()
fail_resp.status_code = 500
fail_resp.text = "Internal Server Error"
mock_http = MagicMock()
mock_http.client = MagicMock()
mock_http.client.request = AsyncMock(return_value=fail_resp)
dest._http = mock_http
# Must not raise — warning only
await dest._update_metrics_marker(1234567890)
assert mock_http.client.request.call_count == 1
@pytest.mark.asyncio
async def test_update_metrics_marker_raises_on_410():
"""_update_metrics_marker must raise RuntimeError and reset _registered on 410."""
dest = _dest()
dest._registered = True
resp_410 = MagicMock()
resp_410.status_code = 410
resp_410.text = "Gone"
mock_http = MagicMock()
mock_http.client = MagicMock()
mock_http.client.request = AsyncMock(return_value=resp_410)
dest._http = mock_http
with pytest.raises(RuntimeError, match="disconnected"):
await dest._update_metrics_marker(1234567890)
assert dest._registered is False

View file

@ -384,6 +384,44 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens():
assert round(completion_cost, 10) == round(expected_completion, 10)
def test_generic_cost_per_token_honors_non_standard_above_threshold():
"""Regression for #30344: get_model_info must keep arbitrary
input/output_cost_per_token_above_<N>_tokens thresholds, not only the hard-coded
128k/200k/272k/512k set, so a custom tier boundary is applied past its limit."""
model = "litellm-test-non-standard-tier"
custom_llm_provider = "openai"
litellm.register_model(
{
model: {
"litellm_provider": custom_llm_provider,
"mode": "chat",
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
"input_cost_per_token_above_500k_tokens": 9e-6,
"output_cost_per_token_above_500k_tokens": 18e-6,
}
}
)
try:
prompt_tokens = 600000
completion_tokens = 1000
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
assert round(prompt_cost, 10) == round(9e-6 * prompt_tokens, 10)
assert round(completion_cost, 10) == round(18e-6 * completion_tokens, 10)
finally:
litellm.model_cost.pop(model, None)
def test_generic_cost_per_token_gpt55():
"""gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input."""
model = "gpt-5.5"

View file

@ -74,6 +74,137 @@ def test_redacted_thinking_content_block_delta():
assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields
def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
chunks = [
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 1. "},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 2."},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "signature_delta", "signature": "sig-final"},
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
expected_thinking_block = {
"type": "thinking",
"thinking": "Step 1. Step 2.",
"signature": "sig-final",
}
assert reasoning_content == "Step 1. Step 2."
assert thinking_blocks == (expected_thinking_block,)
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {
"thinking_blocks": [expected_thinking_block]
}
def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
chunks = [
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 1. "},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 2."},
},
{"type": "content_block_stop", "index": 0},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
assert thinking_blocks == ()
def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
chunks = [
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 1. "},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "Step 2."},
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
assert thinking_blocks == ()
def test_handle_json_mode_chunk_response_format_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True

View file

@ -158,6 +158,68 @@ def test_deepseek_cris():
assert bedrock_route == "converse"
def test_application_inference_profile_arn_routes_to_converse():
"""
Regression for #18258: a bare application-inference-profile ARN passed as
`bedrock/arn:...` must route to converse. The ARN ends in an opaque id with
no provider substring, so the invoke path cannot build a provider-native
body and raises "Unknown provider=None". Converse needs no provider, so it
is the correct route.
"""
route = BedrockModelInfo.get_bedrock_route(
model="bedrock/arn:aws:bedrock:us-west-2:123412341234:application-inference-profile/a1b2c3"
)
assert route == "converse"
def test_explicit_invoke_prefix_wins_over_application_inference_profile_arn():
"""
An explicit invoke/ prefix is respected even for an application-inference-profile
ARN; only the bare `bedrock/arn:...` form is auto-routed to converse. The
explicit invoke path remains a dead end for these ARNs (no provider can be
derived, so completion raises "Unknown provider=None") by design: a caller
that explicitly asks for invoke gets invoke. The auto-route only rescues the
documented bare form.
"""
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
model = "bedrock/invoke/arn:aws:bedrock:us-west-2:123412341234:application-inference-profile/a1b2c3"
assert BedrockModelInfo.get_bedrock_route(model) == "invoke"
assert BaseAWSLLM.get_bedrock_invoke_provider(model) is None
def test_system_defined_inference_profile_arn_still_routes_to_converse():
"""
A system-defined cross-region inference-profile ARN embeds a known model, so
get_base_model resolves it and it already routes to converse. Guards that the
application-inference-profile fix does not change this working case.
"""
route = BedrockModelInfo.get_bedrock_route(
model="bedrock/arn:aws:bedrock:us-east-1:123:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
)
assert route == "converse"
def test_other_opaque_arn_types_still_route_to_invoke():
"""
Only application-inference-profile ARNs are auto-routed to converse. Other
opaque ARNs (provisioned-model, imported-model, custom-model-deployment)
also yield no invoke provider, but they are frequently invoke-only with
provider-specific body formats, so routing them to converse could break
them. Guards the deliberate scope against an over-broad "any opaque ARN ->
converse" generalization.
"""
for arn_segment in (
"provisioned-model/abcdefgh1234",
"imported-model/abcdefgh1234",
"custom-model-deployment/abcdefgh1234",
):
route = BedrockModelInfo.get_bedrock_route(
model=f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}"
)
assert route == "invoke", f"{arn_segment} should stay on invoke route"
def test_govcloud_cross_region_inference_prefix():
"""
Test that GovCloud models with cross-region inference prefix (us-gov.) are parsed correctly

View file

@ -2,10 +2,8 @@
Unit tests for Cohere Rerank Guardrail Translation Handler
"""
import asyncio
import os
import sys
from typing import List, Optional, Tuple
import pytest
@ -94,6 +92,74 @@ class TestInputProcessing:
"id": "doc2",
}
@pytest.mark.asyncio
async def test_process_query_and_instruction(self):
"""Both query and instruction are guardrailed; documents untouched"""
handler = CohereRerankHandler()
guardrail = MockGuardrail(guardrail_name="test")
data = {
"model": "qwen3-reranker",
"query": "What is machine learning?",
"instruction": "Rank by relevance to ML research",
"documents": ["Doc 1", "Doc 2"],
}
result = await handler.process_input_messages(data, guardrail)
# Both user-controlled text fields are scanned and written back
assert result["query"] == "What is machine learning? [GUARDRAILED]"
assert result["instruction"] == "Rank by relevance to ML research [GUARDRAILED]"
# Documents unchanged
assert result["documents"] == ["Doc 1", "Doc 2"]
@pytest.mark.asyncio
async def test_instruction_masked_with_pii(self):
"""A masking guardrail rewrites instruction, not just query"""
class PIIMaskingGuardrail(CustomGuardrail):
async def apply_guardrail(
self, inputs: dict, request_data: dict, input_type: str, **kwargs
) -> dict:
texts = inputs.get("texts", [])
return {"texts": [t.replace("John Doe", "[NAME_REDACTED]") for t in texts]}
handler = CohereRerankHandler()
guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii")
data = {
"model": "qwen3-reranker",
"query": "find records",
"instruction": "prioritize anything authored by John Doe",
"documents": ["Doc 1"],
}
result = await handler.process_input_messages(data, guardrail)
# The sensitive value in instruction is sanitized before forwarding
assert "John Doe" not in result["instruction"]
assert "[NAME_REDACTED]" in result["instruction"]
assert result["documents"] == ["Doc 1"]
@pytest.mark.asyncio
async def test_non_string_instruction_not_scanned(self):
"""A non-string instruction is left as-is (only strings are scanned)"""
handler = CohereRerankHandler()
guardrail = MockGuardrail(guardrail_name="test")
data = {
"model": "qwen3-reranker",
"query": "hello",
"instruction": 12345, # invalid type; backend will reject it
"documents": ["Doc 1"],
}
result = await handler.process_input_messages(data, guardrail)
# Query still guardrailed; non-string instruction untouched
assert result["query"] == "hello [GUARDRAILED]"
assert result["instruction"] == 12345
@pytest.mark.asyncio
async def test_process_no_query(self):
"""Test processing when query is missing"""

View file

@ -0,0 +1,103 @@
from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig
def _function_tool(name: str) -> dict:
return {
"type": "function",
"function": {"name": name, "parameters": {"type": "object"}},
}
def test_drop_unsupported_tools_keeps_function_tools_only():
optional_params = {
"tools": [
_function_tool("shell"),
{"type": "namespace", "name": "container.exec"},
_function_tool("apply_patch"),
],
"tool_choice": "auto",
}
result = DeepSeekChatConfig._drop_unsupported_tools(optional_params)
assert [tool["function"]["name"] for tool in result["tools"]] == [
"shell",
"apply_patch",
]
assert all(tool["type"] == "function" for tool in result["tools"])
assert result["tool_choice"] == "auto"
def test_drop_unsupported_tools_drops_dangling_tool_choice_when_none_survive():
optional_params = {
"tools": [{"type": "namespace", "name": "container.exec"}],
"tool_choice": "required",
"parallel_tool_calls": True,
"temperature": 0.2,
}
result = DeepSeekChatConfig._drop_unsupported_tools(optional_params)
assert "tools" not in result
assert "tool_choice" not in result
assert "parallel_tool_calls" not in result
assert result["temperature"] == 0.2
def test_drop_unsupported_tools_is_noop_for_function_only():
optional_params = {
"tools": [_function_tool("shell")],
"tool_choice": "auto",
}
result = DeepSeekChatConfig._drop_unsupported_tools(optional_params)
assert result is optional_params
def test_drop_unsupported_tools_is_noop_without_tools():
optional_params = {"temperature": 0.7}
result = DeepSeekChatConfig._drop_unsupported_tools(optional_params)
assert result is optional_params
def test_transform_request_strips_unsupported_tools_from_body():
config = DeepSeekChatConfig()
body = config.transform_request(
model="deepseek-chat",
messages=[{"role": "user", "content": "hi"}],
optional_params={
"tools": [
_function_tool("shell"),
{"type": "namespace", "name": "container.exec"},
],
"tool_choice": "auto",
},
litellm_params={},
headers={},
)
assert [tool["type"] for tool in body["tools"]] == ["function"]
assert body["tools"][0]["function"]["name"] == "shell"
async def test_async_transform_request_strips_unsupported_tools_from_body():
config = DeepSeekChatConfig()
body = await config.async_transform_request(
model="deepseek-chat",
messages=[{"role": "user", "content": "hi"}],
optional_params={
"tools": [
_function_tool("shell"),
{"type": "namespace", "name": "container.exec"},
],
"tool_choice": "auto",
},
litellm_params={},
headers={},
)
assert [tool["type"] for tool in body["tools"]] == ["function"]
assert body["tools"][0]["function"]["name"] == "shell"

View file

@ -878,3 +878,107 @@ class TestGithubCopilotTransformResponse:
litellm_params={},
encoding=None,
)
class TestGithubCopilotTransformParsedResponseDict:
"""
Tests for GithubCopilotConfig.transform_parsed_response_dict, the hook the
OpenAI SDK handler calls on its parsed response. That handler bypasses
transform_response, so this is the seam that repairs empty-choices responses
from newer Copilot Claude models on the live completion path.
See: https://github.com/BerriAI/litellm/issues/30927
"""
def test_synthesizes_choices_from_anthropic_content(self):
config = GithubCopilotConfig()
parsed = {
"id": "msg_vrtx_01",
"model": "claude-opus-4.8",
"object": "chat.completion",
"choices": [],
"content": [{"type": "text", "text": "Hello!"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
}
repaired = config.transform_parsed_response_dict(parsed)
assert len(repaired["choices"]) == 1
choice = repaired["choices"][0]
assert choice["message"]["content"] == "Hello!"
assert choice["finish_reason"] == "stop"
assert repaired["usage"]["prompt_tokens"] == 10
assert repaired["usage"]["completion_tokens"] == 5
assert repaired["usage"]["total_tokens"] == 15
def test_passthrough_when_choices_present(self):
config = GithubCopilotConfig()
parsed = {
"id": "chatcmpl-1",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
}
assert config.transform_parsed_response_dict(parsed) is parsed
@patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client")
@patch(
"litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request"
)
def test_openai_handler_repairs_github_copilot_empty_choices(
mock_request, mock_get_client
):
"""
The OpenAI SDK handler calls convert_to_model_response_object directly on the
SDK's parsed output, bypassing transform_response. convert raises APIError on
empty choices, so the handler must route github_copilot responses through
transform_parsed_response_dict first. Removing that wiring (or resolving a
config without the override) fails this test with APIError.
See: https://github.com/BerriAI/litellm/issues/30927
"""
from litellm.llms.openai.openai import OpenAIChatCompletion
mock_get_client.return_value = MagicMock()
class _FakeSDKResponse:
def model_dump(self):
return {
"id": "msg_vrtx_01",
"model": "claude-opus-4.8",
"object": "chat.completion",
"choices": [],
"content": [{"type": "text", "text": "Hi there"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 12, "output_tokens": 3},
}
mock_request.return_value = ({}, _FakeSDKResponse())
result = OpenAIChatCompletion().completion(
model="claude-opus-4.8",
messages=[{"role": "user", "content": "Hi"}],
model_response=ModelResponse(),
timeout=60.0,
optional_params={},
litellm_params={},
logging_obj=MagicMock(),
custom_llm_provider="github_copilot",
client=MagicMock(),
api_key="gh.test-key-123456789",
acompletion=False,
)
assert isinstance(result, ModelResponse)
assert result.choices[0].message.content == "Hi there"
assert result.choices[0].finish_reason == "stop"
mock_request.assert_called_once()

View file

@ -4,6 +4,7 @@ import sys
import pytest
from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
from litellm.rerank_api.rerank_utils import get_optional_rerank_params
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
@ -37,6 +38,54 @@ class TestHostedVLLMRerankTransform:
assert params["rank_fields"] == ["field1"]
assert params["return_documents"] is True
def test_map_cohere_rerank_params_omits_instruction_when_absent(self):
# Backward-compat: when no instruction is supplied, it must not appear
# in the mapped params (and therefore not in the outgoing request body).
params = self.config.map_cohere_rerank_params(
non_default_params=None,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
)
assert "instruction" not in params
def test_map_cohere_rerank_params_passes_instruction_when_set(self):
params = self.config.map_cohere_rerank_params(
non_default_params=None,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
instruction="Rank by relevance to genomics",
)
assert params["instruction"] == "Rank by relevance to genomics"
def test_transform_request_includes_instruction_when_set(self):
body = self.config.transform_rerank_request(
model=self.model,
optional_rerank_params={
"query": "test query",
"documents": ["doc1", "doc2"],
"instruction": "Rank by relevance to genomics",
},
headers={},
)
assert body["instruction"] == "Rank by relevance to genomics"
def test_transform_request_omits_instruction_when_absent(self):
# exclude_none must drop the field entirely so the body matches the
# pre-existing (instruction-less) shape exactly.
body = self.config.transform_rerank_request(
model=self.model,
optional_rerank_params={
"query": "test query",
"documents": ["doc1", "doc2"],
},
headers={},
)
assert "instruction" not in body
def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self):
with pytest.raises(
ValueError, match="Hosted VLLM does not support max_chunks_per_doc"
@ -74,6 +123,7 @@ class TestHostedVLLMRerankTransform:
}
result = self.config._transform_response(response_dict)
assert result.id == "abc123"
assert result.results is not None
assert len(result.results) == 2
assert result.results[0]["index"] == 0
assert result.results[0]["relevance_score"] == 0.9
@ -94,3 +144,32 @@ class TestHostedVLLMRerankTransform:
}
with pytest.raises(ValueError, match="Missing required fields in the result="):
self.config._transform_response(response_dict)
class TestGetOptionalRerankParamsInstruction:
"""`instruction` is threaded through get_optional_rerank_params only when set."""
def setup_method(self):
self.config = HostedVLLMRerankConfig()
self.model = "hosted-vllm-model"
def test_instruction_threaded_when_set(self):
params = get_optional_rerank_params(
rerank_provider_config=self.config,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
instruction="Rank by relevance to genomics",
)
assert params["instruction"] == "Rank by relevance to genomics"
def test_instruction_absent_when_not_set(self):
params = get_optional_rerank_params(
rerank_provider_config=self.config,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
)
assert "instruction" not in params

View file

@ -305,6 +305,34 @@ class TestMoonshotConfig:
assert len(result["messages"]) == 2
assert result["messages"][1]["content"] == "Please select a tool to handle the current issue."
def test_tool_choice_required_does_not_mutate_input_messages(self):
"""tool_choice='required' must not mutate the caller's messages list.
The handling appends a "select a tool" user message; building it in
place corrupts the caller's conversation history and makes
transform_request non-idempotent across retries.
"""
config = MoonshotChatConfig()
messages = [{"role": "user", "content": "What's the weather like?"}]
for _ in range(2):
optional_params = {
"tool_choice": "required",
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
}
result = config.transform_request(
model="moonshot-v1-8k",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={},
)
# The returned request carries the extra message.
assert len(result["messages"]) == 2
# The caller's list is untouched, so repeated calls stay idempotent.
assert messages == [{"role": "user", "content": "What's the weather like?"}]
def test_tool_choice_non_required_preserved(self):
"""Test that non-'required' tool_choice values are preserved"""
config = MoonshotChatConfig()

View file

@ -13,7 +13,52 @@ from litellm.cost_calculator import completion_cost
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
)
from litellm.types.utils import TranscriptionResponse
from litellm.types.utils import (
TranscriptionResponse,
TranscriptionUsageDurationObject,
)
class TestDiarizedJsonUsageParsing:
"""gpt-4o-transcribe / diarized_json returns a fractional `usage.seconds`."""
def test_fractional_duration_seconds_does_not_raise(self):
"""
A diarized_json response carries usage={"type": "duration", "seconds": <float>}.
OpenAI specs `seconds` as a float, so a fractional value must parse cleanly
instead of raising and getting retried until the upstream rate-limits.
"""
response_object = {
"text": "speaker_1: Olá",
"task": "transcribe",
"duration": 295.8,
"segments": [
{
"id": "seg_001",
"speaker": "speaker_1",
"start": 0.0,
"end": 1.0,
"text": "Olá",
"type": "transcript.text.segment",
}
],
"usage": {"type": "duration", "seconds": 295.8},
}
result = convert_to_model_response_object(
response_object=response_object,
model_response_object=TranscriptionResponse(),
response_type="audio_transcription",
)
assert isinstance(result.usage, TranscriptionUsageDurationObject)
assert result.usage.seconds == 295.8
def test_usage_duration_object_accepts_float_seconds(self):
assert (
TranscriptionUsageDurationObject(type="duration", seconds=295.8).seconds
== 295.8
)
class TestTranscriptionDurationNotInResponseBody:

View file

@ -3786,3 +3786,298 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict
assert body["data"][0]["metadata"] == {}
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
class _FakeColdStorageLogger:
"""Injectable cold storage logger that records the object key it was asked for."""
def __init__(self, payload):
self._payload = payload
self.requested_object_keys = []
async def get_proxy_server_request_from_cold_storage_with_object_key(
self, object_key
):
self.requested_object_keys.append(object_key)
return self._payload
def _cold_storage_handler(payload):
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
logger = _FakeColdStorageLogger(payload)
return ColdStorageHandler(cold_storage_logger=logger), logger
@pytest.mark.parametrize(
"value, expected",
[
(None, False),
("", False),
(" ", False),
("{}", False),
("[]", False),
("null", False),
('{"a": 1}', True),
({}, False),
({"a": 1}, True),
([], False),
([1], True),
(5, True),
],
)
def test_spend_log_field_has_content(value, expected):
assert spend_management_endpoints._spend_log_field_has_content(value) is expected
@pytest.mark.parametrize(
"metadata, expected",
[
(None, None),
("{}", None),
("not-json", None),
({"cold_storage_object_key": ""}, None),
({"cold_storage_object_key": "k/req-1.json"}, "k/req-1.json"),
('{"cold_storage_object_key": "k/req-2.json"}', "k/req-2.json"),
],
)
def test_cold_storage_object_key_from_metadata(metadata, expected):
assert (
spend_management_endpoints._cold_storage_object_key_from_metadata(metadata)
== expected
)
@pytest.mark.asyncio
async def test_resolve_payload_prefers_pg_and_skips_cold_storage():
handler, logger = _cold_storage_handler({"messages": "X", "response": "Y"})
row = {
"messages": "{}",
"response": '{"choices": [{"message": {"content": "hi"}}]}',
"proxy_server_request": "{}",
"metadata": {"cold_storage_object_key": "k/req.json"},
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert resolved.response == '{"choices": [{"message": {"content": "hi"}}]}'
assert logger.requested_object_keys == []
@pytest.mark.asyncio
async def test_resolve_payload_fetches_from_cold_storage_when_pg_empty():
cold_payload = {
"messages": [{"role": "user", "content": "what is 2+2"}],
"response": {"choices": [{"message": {"content": "4"}}]},
"proxy_server_request": {"body": {"model": "gpt-4o-mini"}},
}
handler, logger = _cold_storage_handler(cold_payload)
row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": {"cold_storage_object_key": "llm-gateway/prod/req-42.json"},
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert logger.requested_object_keys == ["llm-gateway/prod/req-42.json"]
assert resolved.messages == cold_payload["messages"]
assert resolved.response == cold_payload["response"]
assert resolved.proxy_server_request == cold_payload["proxy_server_request"]
@pytest.mark.asyncio
async def test_resolve_payload_metadata_as_json_string():
cold_payload = {"messages": "in", "response": "out", "proxy_server_request": None}
handler, logger = _cold_storage_handler(cold_payload)
row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": json.dumps({"cold_storage_object_key": "k/str-meta.json"}),
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert logger.requested_object_keys == ["k/str-meta.json"]
assert resolved.response == "out"
@pytest.mark.asyncio
async def test_resolve_payload_no_object_key_returns_empty_without_fetch():
handler, logger = _cold_storage_handler({"messages": "should-not-be-used"})
row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": {},
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert logger.requested_object_keys == []
assert resolved == spend_management_endpoints.RequestResponsePayload(
"{}", "{}", "{}"
)
@pytest.mark.asyncio
async def test_resolve_payload_cold_storage_miss_falls_back_to_pg_values():
handler, logger = _cold_storage_handler(None)
row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": {"cold_storage_object_key": "k/missing.json"},
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert logger.requested_object_keys == ["k/missing.json"]
assert resolved == spend_management_endpoints.RequestResponsePayload(
"{}", "{}", "{}"
)
@pytest.mark.asyncio
async def test_resolve_payload_cold_storage_exception_falls_back_to_pg_values():
"""A backend error during fetch degrades to PG values instead of bubbling a 500."""
class _RaisingLogger:
async def get_proxy_server_request_from_cold_storage_with_object_key(
self, object_key
):
raise RuntimeError("cold storage backend unavailable")
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
handler = ColdStorageHandler(cold_storage_logger=_RaisingLogger())
row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": {"cold_storage_object_key": "k/boom.json"},
}
resolved = await spend_management_endpoints._resolve_request_response_payload(
row, cold_storage_handler=handler
)
assert resolved == spend_management_endpoints.RequestResponsePayload(
"{}", "{}", "{}"
)
@pytest.mark.asyncio
async def test_cold_storage_handler_uses_injected_logger():
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
logger = _FakeColdStorageLogger({"messages": "in", "response": "out"})
handler = ColdStorageHandler(cold_storage_logger=logger)
result = await handler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key="k/req.json"
)
assert result == {"messages": "in", "response": "out"}
assert logger.requested_object_keys == ["k/req.json"]
@pytest.mark.asyncio
async def test_cold_storage_handler_returns_none_when_no_logger_configured(monkeypatch):
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
monkeypatch.setattr(litellm, "cold_storage_custom_logger", None, raising=False)
handler = ColdStorageHandler()
result = await handler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key="k/req.json"
)
assert result is None
@pytest.mark.asyncio
async def test_cold_storage_handler_resolves_configured_logger_from_registry(monkeypatch):
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
logger = _FakeColdStorageLogger({"messages": "from-registry"})
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2", raising=False)
monkeypatch.setattr(
litellm.logging_callback_manager,
"get_active_custom_logger_for_callback_name",
lambda name: logger if name == "s3_v2" else None,
)
handler = ColdStorageHandler()
result = await handler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key="k/req.json"
)
assert result == {"messages": "from-registry"}
assert logger.requested_object_keys == ["k/req.json"]
def test_ui_view_request_response_reads_from_cold_storage(client, monkeypatch):
"""End-to-end: a placeholder row with a cold_storage_object_key is served from
cold storage through the detail endpoint."""
from types import SimpleNamespace
placeholder_row = {
"messages": "{}",
"response": "{}",
"proxy_server_request": "{}",
"metadata": {"cold_storage_object_key": "k/cold.json"},
}
async def _query_raw(_sql, *_args):
return [placeholder_row]
fake_prisma = SimpleNamespace(db=SimpleNamespace(query_raw=_query_raw))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
cold_logger = _FakeColdStorageLogger(
{
"messages": [{"role": "user", "content": "hi"}],
"response": {"choices": [{"message": {"content": "hello"}}]},
"proxy_server_request": None,
}
)
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2", raising=False)
monkeypatch.setattr(
litellm.logging_callback_manager,
"get_active_additional_logging_utils_from_custom_logger",
lambda: [],
)
monkeypatch.setattr(
litellm.logging_callback_manager,
"get_active_custom_logger_for_callback_name",
lambda name: cold_logger if name == "s3_v2" else None,
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_1"
)
try:
response = client.get(
"/spend/logs/ui/req-cold",
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
body = response.json()
assert body["messages"] == [{"role": "user", "content": "hi"}]
assert body["response"] == {"choices": [{"message": {"content": "hello"}}]}
assert cold_logger.requested_object_keys == ["k/cold.json"]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)

View file

@ -2,7 +2,10 @@ import json
import pytest
from litellm.router_utils.fallback_event_handlers import run_async_fallback
from litellm.router_utils.fallback_event_handlers import (
get_fallback_model_group,
run_async_fallback,
)
class StreamingWrapper:
@ -137,3 +140,16 @@ async def test_run_async_fallback_skips_original_model_group():
)
assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
def test_get_fallback_model_group_does_not_mutate_fallbacks():
"""A string fallback must be resolved without mutating the caller's
fallbacks list, which is the live router config shared across requests."""
fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}, "gpt-4o-mini"]
fallback_model_group, _ = get_fallback_model_group(
fallbacks=fallbacks, model_group="unmatched-model"
)
assert fallback_model_group == ["gpt-4o-mini"]
assert fallbacks == [{"gpt-3.5-turbo": ["claude-3-haiku"]}, "gpt-4o-mini"]

View file

@ -381,5 +381,93 @@ class TestCompletionCostIntegration:
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
class TestGPTImage2OutputImageTokensNoBreakdown:
"""
Regression test: the OpenAI Images endpoints (/v1/images/generations and
/v1/images/edits) return usage with NO output token breakdown — litellm's
ImageUsage has no ``output_tokens_details`` field. Before the fix, the
generated-image OUTPUT tokens were priced at the text rate
(``output_cost_per_token`` = $10/1M for gpt-image-2) instead of the image rate
(``output_cost_per_image_token`` = $30/1M), a ~3x undercount on the dominant
cost component.
"""
def test_gpt_image_2_output_priced_as_image_when_no_breakdown(self):
from litellm.llms.openai.image_generation.cost_calculator import (
cost_calculator,
)
# Mirrors a real gpt-image-2 /v1/images/edits response: input breakdown is
# present, but there is no usable output token breakdown.
usage = ImageUsage(
input_tokens=3987,
output_tokens=5488,
total_tokens=9475,
input_tokens_details=ImageUsageInputTokensDetails(
text_tokens=943,
image_tokens=3044,
),
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(b64_json="test")],
)
image_response.usage = usage
image_response._hidden_params = {"custom_llm_provider": "openai"}
cost = cost_calculator(
model="gpt-image-2",
image_response=image_response,
custom_llm_provider="openai",
)
# gpt-image-2 pricing:
# text input: 943 * $5/1M = 0.004715
# image input: 3044 * $8/1M = 0.024352
# image output: 5488 * $30/1M = 0.164640 (NOT text output $10/1M = 0.054880)
expected_cost = 943 * 5e-6 + 3044 * 8e-6 + 5488 * 3e-5
assert abs(cost - expected_cost) < 1e-6, (
f"Expected {expected_cost}, got {cost}. Generated image output tokens "
f"are likely being priced at the text output_cost_per_token rate."
)
def test_gpt_image_2_chat_usage_without_breakdown_is_costed_not_zero(self):
"""A chat ``Usage`` with ``completion_tokens_details=None`` must still be
costed via ``generic_cost_per_token`` (output at the text rate) rather than
erroring or silently returning 0.0."""
from litellm.llms.openai.image_generation.cost_calculator import (
cost_calculator,
)
usage = Usage(
prompt_tokens=600,
completion_tokens=5000,
total_tokens=5600,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=100,
image_tokens=500,
),
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(b64_json="test")],
)
image_response.usage = usage
image_response._hidden_params = {"custom_llm_provider": "openai"}
cost = cost_calculator(
model="gpt-image-2",
image_response=image_response,
custom_llm_provider="openai",
)
# No output breakdown -> output priced at the text rate (output_cost_per_token):
# text in 100*$5/1M + image in 500*$8/1M + output 5000*$10/1M
expected_cost = 100 * 5e-6 + 500 * 8e-6 + 5000 * 1e-5
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
if __name__ == "__main__":
pytest.main([__file__, "-v"])

View file

@ -467,10 +467,15 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
enableSorting: true,
cell: (info) => {
const maxBudget = info.getValue() as number | null;
if (maxBudget === null) {
return "Unlimited";
if (maxBudget !== null) {
return `$${formatNumberWithCommas(maxBudget)}`;
}
return `$${formatNumberWithCommas(maxBudget)}`;
const teamId = info.row.original.team_id;
const team = teams?.find((t) => t.team_id === teamId);
if (team?.max_budget != null) {
return `$${formatNumberWithCommas(team.max_budget)} (Team)`;
}
return "Unlimited";
},
},
{

View file

@ -1,7 +1,7 @@
import { renderWithProviders } from "../../../tests/test-utils";
import { screen, waitFor } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { KeyResponse } from "../key_team_helpers/key_list";
import { KeyResponse, Team } from "../key_team_helpers/key_list";
import KeyInfoView from "./key_info_view";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import useTeams from "@/app/(dashboard)/hooks/useTeams";
@ -103,8 +103,26 @@ const baseAuthorized = {
userEmail: null,
disabledPersonalKeyCreation: null,
showSSOBanner: false,
isLoading: false,
isAuthorized: true,
};
const makeTeam = (overrides: Partial<Team>): Team => ({
team_id: "team-default",
team_alias: "Default Team",
models: [],
max_budget: null,
budget_duration: null,
tpm_limit: null,
rpm_limit: null,
organization_id: "",
created_at: "2026-01-01T00:00:00Z",
keys: [],
members_with_roles: [],
spend: 0,
...overrides,
});
describe("KeyInfoView overview budget display (LIT-2845)", () => {
beforeEach(() => {
vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() });
@ -151,7 +169,64 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => {
it("renders 'Unlimited' when max_budget is null", async () => {
renderWithProviders(
<KeyInfoView
keyData={{ ...MOCK_KEY_DATA, max_budget: null }}
keyData={{ ...MOCK_KEY_DATA, max_budget: null } as unknown as KeyResponse}
onClose={() => {}}
keyId={"test-key-id"}
onKeyDataUpdate={() => {}}
teams={[]}
/>,
);
await waitFor(() => {
expect(screen.getByText(/of Unlimited/)).toBeInTheDocument();
});
});
it("renders team budget with alias and duration when key has no own budget but team has one", async () => {
vi.mocked(useTeams).mockReturnValue({
teams: [makeTeam({ team_id: "team-123", team_alias: "Test Budget", max_budget: 1200, budget_duration: "30d" })],
setTeams: vi.fn(),
});
renderWithProviders(
<KeyInfoView
keyData={{ ...MOCK_KEY_DATA, max_budget: null, team_id: "team-123" } as unknown as KeyResponse}
onClose={() => {}}
keyId={"test-key-id"}
onKeyDataUpdate={() => {}}
teams={[]}
/>,
);
await waitFor(() => {
expect(screen.getByText(/of \$1,200\.00 \(Team: Test Budget \/ 30d\)/)).toBeInTheDocument();
});
});
it("renders team budget without duration when team has no budget_duration", async () => {
vi.mocked(useTeams).mockReturnValue({
teams: [makeTeam({ team_id: "team-456", team_alias: "No Duration Team", max_budget: 500 })],
setTeams: vi.fn(),
});
renderWithProviders(
<KeyInfoView
keyData={{ ...MOCK_KEY_DATA, max_budget: null, team_id: "team-456" } as unknown as KeyResponse}
onClose={() => {}}
keyId={"test-key-id"}
onKeyDataUpdate={() => {}}
teams={[]}
/>,
);
await waitFor(() => {
expect(screen.getByText(/of \$500\.00 \(Team: No Duration Team\)/)).toBeInTheDocument();
});
});
it("renders 'Unlimited' when key has no budget and team also has no budget", async () => {
vi.mocked(useTeams).mockReturnValue({
teams: [makeTeam({ team_id: "team-789", team_alias: "Free Team" })],
setTeams: vi.fn(),
});
renderWithProviders(
<KeyInfoView
keyData={{ ...MOCK_KEY_DATA, max_budget: null, team_id: "team-789" } as unknown as KeyResponse}
onClose={() => {}}
keyId={"test-key-id"}
onKeyDataUpdate={() => {}}

View file

@ -411,6 +411,15 @@ export default function KeyInfoView({
});
};
const parentTeam = currentKeyData.team_id ? teamsData?.find((team) => team.team_id === currentKeyData.team_id) : null;
const budgetDisplay =
currentKeyData.max_budget !== null
? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}`
: parentTeam?.max_budget != null
? `$${formatNumberWithCommas(parentTeam.max_budget, 2)} (Team: ${parentTeam.team_alias || parentTeam.team_id}${parentTeam.budget_duration ? ` / ${parentTeam.budget_duration}` : ""})`
: "Unlimited";
return (
<div className="w-full h-full overflow-y-auto p-4">
<KeyInfoHeader
@ -520,12 +529,7 @@ export default function KeyInfoView({
<Text>Spend</Text>
<div className="mt-2">
<Title>${formatNumberWithCommas(currentKeyData.spend, 4)}</Title>
<Text>
of{" "}
{currentKeyData.max_budget !== null
? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}`
: "Unlimited"}
</Text>
<Text>of {budgetDisplay}</Text>
</div>
</Card>

View file

@ -27,6 +27,17 @@ describe("PrettyMessagesView", () => {
expect(screen.getByText("Hi there!")).toBeInTheDocument();
});
it("renders input when request is a bare messages array (cold storage payload)", () => {
const request = [{ role: "user", content: "Write me a poem" }];
const response = {
choices: [{ message: { role: "assistant", content: "A quiet moment." } }],
};
render(<PrettyMessagesView request={request} response={response} />);
expect(screen.getByText("Write me a poem")).toBeInTheDocument();
expect(screen.getByText("A quiet moment.")).toBeInTheDocument();
});
it("should render the realtime pretty view for realtime API responses", () => {
const request = {};
const response = {

View file

@ -39,18 +39,24 @@ export const ROLE_STYLES: Record<string, RoleStyle> = {
* Parse request messages and response message from log data
*/
export const parseMessages = (request: any, response: any): ParsedMessages => {
// Parse request messages
// Parse request messages. `request` is either the raw request body
// ({ messages: [...] }) or, when prompts come from cold storage, the bare
// messages array itself.
const requestMessages: ParsedMessage[] = [];
if (request?.messages && Array.isArray(request.messages)) {
request.messages.forEach((msg: any) => {
requestMessages.push({
role: msg.role || "user",
content: parseMessageContent(msg.content),
toolCallId: msg.tool_call_id,
});
const requestMessageList = Array.isArray(request)
? request
: Array.isArray(request?.messages)
? request.messages
: [];
requestMessageList.forEach((msg: any) => {
requestMessages.push({
role: msg.role || "user",
content: parseMessageContent(msg.content),
toolCallId: msg.tool_call_id,
});
}
});
// Parse response message
let responseMessage: ParsedMessage | null = null;