mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
commit
0647d5bafd
56 changed files with 2678 additions and 535 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
],
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -2441,7 +2441,7 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
|
|||
|
||||
class TranscriptionUsageDurationObject(BaseModel):
|
||||
type: Literal["duration"]
|
||||
seconds: int
|
||||
seconds: float
|
||||
|
||||
|
||||
class TranscriptionUsageInputTokenDetailsObject(BaseModel):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
0
tests/test_litellm/llms/deepseek/chat/__init__.py
Normal file
0
tests/test_litellm/llms/deepseek/chat/__init__.py
Normal 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"
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
},
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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={() => {}}
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue