diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 49a36aa23f0..2d4e85630dc 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -218,6 +218,7 @@ jobs: tests/proxy_unit_tests/test_gemini_agents_endpoints.py tests/proxy_unit_tests/test_get_favicon.py tests/proxy_unit_tests/test_get_image.py + tests/proxy_unit_tests/test_reducto_ocr_route.py tests/proxy_unit_tests/test_ui_path_detection.py tests/proxy_unit_tests/test_prompt_test_endpoint.py tests/proxy_unit_tests/test_check_batch_cost.py diff --git a/litellm/__init__.py b/litellm/__init__.py index d8d48b5865f..3365abe3256 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -413,6 +413,12 @@ internal_user_budget_duration: Optional[str] = None tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None max_end_user_budget: Optional[float] = None max_end_user_budget_id: Optional[str] = None +# When True, end-user IDs extracted from requests are validated against +# LiteLLM_EndUserTable / LiteLLM_UserTable. Values that do not resolve to a +# known row are dropped before reaching spend logs. Defaults to False for +# backwards compatibility — arbitrary client-supplied identifiers still +# pass through unchanged. +validate_end_user_id_in_db: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None @@ -636,6 +642,7 @@ minimax_models: Set = set() aws_polly_models: Set = set() gigachat_models: Set = set() llamagate_models: Set = set() +reducto_models: Set = set() bedrock_mantle_models: Set = set() @@ -903,6 +910,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): gigachat_models.add(key) elif value.get("litellm_provider") == "llamagate": llamagate_models.add(key) + elif value.get("litellm_provider") == "reducto": + reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": bedrock_mantle_models.add(key) @@ -1014,6 +1023,7 @@ model_list = list( | ovhcloud_models | lemonade_models | docker_model_runner_models + | reducto_models | bedrock_mantle_models | set(clarifai_models) ) @@ -1120,6 +1130,7 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, + "reducto": reducto_models, "bedrock_mantle": bedrock_mantle_models, } diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index e3cbf422e5d..51abbbf729b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -30,6 +30,11 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) +from litellm.responses.sse_output_recovery import ( + parse_sse_json_chunk, + record_output_item_chunk, + record_output_text_chunk, +) from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, @@ -97,7 +102,7 @@ def _build_reasoning_item( def _reasoning_item_to_response_input( - r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]] + r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]], ) -> Dict[str, Any]: """Convert a stored ChatCompletionReasoningItem back to a Responses API input item.""" r_input: Dict[str, Any] = { @@ -601,6 +606,79 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return choices + @classmethod + def _extract_output_from_completed_event( + cls, parsed_chunk: Dict[str, Any] + ) -> Optional[List[Dict[str, Any]]]: + response_payload = parsed_chunk.get("response") + if not isinstance(response_payload, dict): + return None + response_output = response_payload.get("output") + if not isinstance(response_output, list) or len(response_output) == 0: + return None + return cast(List[Dict[str, Any]], response_output) + + @classmethod + def _recover_output_items_from_raw_sse( + cls, raw_sse: Optional[str] + ) -> List[Dict[str, Any]]: + if not raw_sse or not isinstance(raw_sse, str): + return [] + + recovered_output_items: Dict[int, Dict[str, Any]] = {} + recovered_text_only_items: Dict[int, Dict[str, Any]] = {} + + for chunk in raw_sse.splitlines(): + parsed_chunk = parse_sse_json_chunk(chunk) + if parsed_chunk is None: + continue + + event_type = parsed_chunk.get("type") + + if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + recovered_output = cls._extract_output_from_completed_event( + parsed_chunk + ) + if recovered_output is not None: + return recovered_output + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + record_output_item_chunk( + parsed_chunk=parsed_chunk, + output_items=recovered_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE: + record_output_text_chunk( + parsed_chunk=parsed_chunk, + output_items=recovered_output_items, + text_only_items=recovered_text_only_items, + ) + continue + + # Merge text-only items into the recovered output items. Real + # OUTPUT_ITEM_DONE events take precedence at any given output_index, + # but text-only items at indices without a matching OUTPUT_ITEM_DONE + # must still be preserved (e.g. multi-output responses where some + # indices only emitted OUTPUT_TEXT_DONE). + merged_items: Dict[int, Dict[str, Any]] = {**recovered_text_only_items} + merged_items.update(recovered_output_items) + + if merged_items: + return [item for _, item in sorted(merged_items.items())] + + return [] + + @classmethod + def _recover_output_items_from_logging( + cls, logging_obj: "LiteLLMLoggingObj" + ) -> List[Dict[str, Any]]: + model_call_details = getattr(logging_obj, "model_call_details", {}) or {} + original_response = model_call_details.get("original_response") + return cls._recover_output_items_from_raw_sse(original_response) + def transform_response( # noqa: PLR0915 self, model: str, @@ -625,9 +703,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if raw_response.error is not None: raise ValueError(f"Error in response: {raw_response.error}") + output_items = raw_response.output + if len(output_items) == 0: + recovered_output_items = self._recover_output_items_from_logging( + logging_obj + ) + if recovered_output_items: + output_items = cast(Any, recovered_output_items) + raw_response.output = cast(Any, recovered_output_items) + verbose_logger.warning( + "Recovered empty Responses API output from raw SSE for model=%s", + model, + ) + # Convert response output to choices using the static helper choices = self._convert_response_output_to_choices( - output_items=raw_response.output, + output_items=output_items, handle_raw_dict_callback=self._handle_raw_dict_response_item, ) @@ -641,7 +732,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) else: raise ValueError( - f"Unknown items in responses API response: {raw_response.output}" + f"Unknown items in responses API response: {output_items}" ) setattr(model_response, "choices", choices) @@ -1237,7 +1328,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): raise ValueError( f"Chat provider: Invalid function argument delta {parsed_chunk}" ) - elif event_type == "response.output_item.done": + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") == "function_call": diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 2257861aff6..98e00cf5788 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1879,10 +1879,6 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - pages_processed = response.usage_info.pages_processed - if pages_processed is None: - raise ValueError("OCR response pages_processed is None") - try: model_info: Optional[ModelInfo] = litellm.get_model_info( model=model, custom_llm_provider=custom_llm_provider @@ -1890,9 +1886,49 @@ def ocr_cost( except Exception: model_info = None - ocr_cost_per_page: float = 0.0 + credits = getattr(response.usage_info, "credits", None) + cost_per_credit = None if model_info is not None: - ocr_cost_per_page = model_info.get("ocr_cost_per_page") or 0.0 + cost_per_credit = model_info.get("ocr_cost_per_credit") + if credits is not None and cost_per_credit is not None: + return cost_per_credit * credits, 0.0 + + ocr_cost_per_page: Optional[float] = None + if model_info is not None: + ocr_cost_per_page = model_info.get("ocr_cost_per_page") + + pages_processed = response.usage_info.pages_processed + if pages_processed is None: + if cost_per_credit is not None or ocr_cost_per_page is None: + # Surface missing usage data instead of silently under-reporting + # cost. The previous behavior raised ValueError; we now return 0.0 + # for credit-priced or unpriced models, so log a warning to keep + # the regression visible to operators. + verbose_logger.warning( + "OCR cost: model=%s custom_llm_provider=%s response.usage_info." + "pages_processed is None and credits=%s; returning 0.0 cost.", + model, + custom_llm_provider, + credits, + ) + return 0.0, 0.0 + raise ValueError("OCR response pages_processed is None") + + if ocr_cost_per_page is None: + # No per-page pricing configured. Either the model is on credit-based + # pricing (and credits weren't returned, so the credit branch above did + # not match) or the model has no OCR pricing entry at all. Surface a + # warning so that missing pricing entries are visible rather than + # silently producing zero cost for billable usage. + verbose_logger.warning( + "OCR cost: model=%s custom_llm_provider=%s reported " + "pages_processed=%s but no ocr_cost_per_page is configured; " + "returning 0.0 cost.", + model, + custom_llm_provider, + pages_processed, + ) + return 0.0, 0.0 total_ocr_processing_cost: float = ocr_cost_per_page * pages_processed return total_ocr_processing_cost, 0.0 diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index f9d4496c21f..86eae0e7954 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -14,22 +14,38 @@ from litellm.integrations.custom_logger import CustomLogger class CustomBatchLogger(CustomLogger): + preserve_events_added_during_flush = False + + # Default cap on the in-memory log queue. Prevents unbounded memory growth + # if ``async_send_batch`` consistently fails (e.g. the destination is + # unreachable) and events are preserved across flush attempts. Subclasses + # may override by passing ``max_queue_size`` or by setting the attribute + # directly (see ``RubrikLogger`` for an example). + DEFAULT_MAX_QUEUE_SIZE = 50_000 + def __init__( self, flush_lock: Optional[asyncio.Lock] = None, batch_size: Optional[int] = None, flush_interval: Optional[int] = None, + max_queue_size: Optional[int] = None, **kwargs, ) -> None: """ Args: flush_lock (Optional[asyncio.Lock], optional): Lock to use when flushing the queue. Defaults to None. Only used for custom loggers that do batching + max_queue_size (Optional[int], optional): Maximum number of events to retain in ``log_queue``. When the limit is exceeded (e.g. because the send destination is unreachable and events are preserved for retry), the oldest events are dropped. Defaults to ``DEFAULT_MAX_QUEUE_SIZE``. """ self.log_queue: List = [] self.flush_interval = flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE self.last_flush_time = time.time() self.flush_lock = flush_lock + self.max_queue_size: int = ( + max_queue_size + if max_queue_size is not None + else self.DEFAULT_MAX_QUEUE_SIZE + ) super().__init__(**kwargs) @@ -47,11 +63,40 @@ class CustomBatchLogger(CustomLogger): async with self.flush_lock: if self.log_queue: + log_queue_length = len(self.log_queue) verbose_logger.debug( "CustomLogger: Flushing batch of %s events", len(self.log_queue) ) - await self.async_send_batch() - self.log_queue.clear() + try: + await self.async_send_batch() + except Exception: + # If the underlying batch send raised, do NOT drop the + # in-flight events. They will be retried on the next flush. + # Most existing async_send_batch implementations swallow + # their own errors, so this only affects loggers that opt + # in to surfacing failures (e.g. Rubrik). + verbose_logger.exception( + "CustomLogger: async_send_batch raised; preserving " + "%s events in queue for retry", + log_queue_length, + ) + # Guard against unbounded queue growth if the destination + # is persistently unreachable. Drop the oldest events + # beyond ``max_queue_size``. + overflow = len(self.log_queue) - self.max_queue_size + if overflow > 0: + del self.log_queue[:overflow] + verbose_logger.warning( + "CustomLogger: log queue exceeded max_queue_size=%s; " + "dropped %s oldest events.", + self.max_queue_size, + overflow, + ) + return + if self.preserve_events_added_during_flush: + del self.log_queue[:log_queue_length] + else: + self.log_queue.clear() self.last_flush_time = time.time() async def async_send_batch(self, *args, **kwargs): diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py new file mode 100644 index 00000000000..af396ecdc73 --- /dev/null +++ b/litellm/integrations/rubrik.py @@ -0,0 +1,605 @@ +"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" + +import asyncio +import os +import random +import time +import urllib.parse +import uuid +from collections import Counter +from typing import TYPE_CHECKING, Any, Literal, Optional + +import httpx +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + GenericGuardrailAPIInputs, + StandardLoggingPayload, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + +_ENDPOINT_ANTHROPIC_MESSAGES = "/v1/messages" +_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch" +_MAX_QUEUE_SIZE = 10_000 +_DROP_WARNING_INTERVAL_SECONDS = 60.0 + + +class _MalformedToolBlockingResponseError(Exception): + """Raised when the tool blocking service returns a structurally invalid + response (e.g. empty ``choices``). + + Distinct from transient network/HTTP errors so callers can surface a + louder, misconfiguration-style log instead of treating it as a routine + fail-open. + """ + + +class RubrikLogger(CustomGuardrail, CustomBatchLogger): + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + **kwargs, + ): + self.flush_lock = asyncio.Lock() + kwargs.setdefault("guardrail_name", "rubrik") + # `initialize_guardrail` always passes these kwargs explicitly, with + # value `None` when the user omits `mode` / `default_on` from the + # guardrail config. Coerce None (omitted) to the desired default + # while preserving any explicit value the caller did set -- + # in particular `default_on=False` if the user wants the guardrail + # off by default. + kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call + if kwargs.get("default_on") is None: + kwargs["default_on"] = True + super().__init__( + flush_lock=self.flush_lock, + **kwargs, + ) + + verbose_logger.debug("initializing rubrik logger") + + self.sampling_rate = 1.0 + rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") + if rbrk_sampling_rate is not None: + try: + parsed_rate = float(rbrk_sampling_rate.strip()) + self.sampling_rate = max(0.0, min(1.0, parsed_rate)) + if parsed_rate != self.sampling_rate: + verbose_logger.warning( + f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to " + f"{self.sampling_rate}" + ) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0" + ) + + self.key = api_key or os.getenv("RUBRIK_API_KEY") + if not self.key: + verbose_logger.warning( + "Rubrik: No API key configured. Requests will be unauthenticated." + ) + _batch_size = os.getenv("RUBRIK_BATCH_SIZE") + + if _batch_size: + try: + self.batch_size = int(_batch_size) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default" + ) + + # Cap the in-memory retry queue so a Rubrik webhook outage cannot let + # authenticated traffic accumulate prompt/response payloads until the + # proxy runs out of memory. Once the cap is reached, oldest events are + # dropped to make room for fresh ones (drop-oldest backpressure). + self.max_queue_size = _MAX_QUEUE_SIZE + self._dropped_since_warning = 0 + self._last_drop_warning_time = 0.0 + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + + if _webhook_url is None: + raise ValueError( + "Rubrik webhook URL not configured. " + "Set RUBRIK_WEBHOOK_URL or pass api_base." + ) + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self.tool_blocking_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_TOOL_BLOCKING}" + self.logging_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" + + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + self.tool_blocking_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback, + params={"timeout": httpx.Timeout(5.0, connect=2.0)}, + ) + + self._headers: dict[str, str] = {"Content-Type": "application/json"} + if self.key: + self._headers["Authorization"] = f"Bearer {self.key}" + + # Periodic flush is started lazily on the first log event so that + # low-traffic deployments still get their batches drained even when the + # logger is instantiated outside a running event loop (sync init). + self._flush_task: Optional[asyncio.Task[Any]] = ( + self._start_periodic_flush_task() + ) + + def _start_periodic_flush_task(self) -> Optional[asyncio.Task[Any]]: + """Start the periodic flush task only when an event loop is already running.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + verbose_logger.debug( + "Rubrik logger init: no running event loop, " + "periodic flush will start on first log event." + ) + return None + return loop.create_task(self.periodic_flush()) + + def _ensure_periodic_flush_task(self) -> None: + # Synchronous helper: in asyncio's cooperative model there is no await + # between the check and assignment, so two callers cannot race here. + if self._flush_task is None or self._flush_task.done(): + self._flush_task = self._start_periodic_flush_task() + + async def aclose(self): + """Close the dedicated HTTP clients used by this logger.""" + # Cancel the periodic flush task before closing the HTTP clients so + # the loop doesn't wake up and try to POST via a closed client. + if self._flush_task is not None and not self._flush_task.done(): + self._flush_task.cancel() + try: + await self._flush_task + except (asyncio.CancelledError, Exception): + pass + self._flush_task = None + await self.tool_blocking_client.close() + await self.async_httpx_client.close() + + # -- Guardrail hook -------------------------------------------------------- + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Validate tool calls against the blocking service (fail-open).""" + if input_type != "response": + return inputs + + tool_calls = inputs.get("tool_calls") + if not tool_calls: + return inputs + + try: + return await self._check_tool_calls( + inputs, tool_calls, request_data, logging_obj + ) + except ModifyResponseException: + raise + except _MalformedToolBlockingResponseError as e: + # Distinct from transient errors: the service responded but the + # payload was structurally invalid, which usually indicates a + # misconfigured webhook or a breaking change in its response + # format. Log loudly so operators notice their tool-blocking + # policy is not actually being enforced. + verbose_logger.critical( + "Tool blocking service returned a malformed response: %s. " + "Tool calls are NOT being checked -- verify the webhook " + "configuration. Returning original response unchanged.", + e, + exc_info=True, + ) + return inputs + except Exception as e: + verbose_logger.error( + f"Tool blocking hook failed: {e}. " + "Returning original response unchanged.", + exc_info=True, + ) + return inputs + + async def _check_tool_calls( + self, + inputs: GenericGuardrailAPIInputs, + tool_calls: Any, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send tool calls to blocking service, raise if any are blocked.""" + message_tool_calls = self._normalize_tool_calls(tool_calls) + + call_details = ( + getattr(logging_obj, "model_call_details", {}) if logging_obj else {} + ) + response = request_data.get("response") + request_id = getattr(response, "id", None) if response else None + if logging_obj and not call_details: + verbose_logger.warning( + "Rubrik: logging_obj present but model_call_details is empty " + "-- request context will be missing" + ) + + response_data = self._build_tool_call_payload(message_tool_calls, request_id) + req_data = self._extract_request_data(call_details) + + service_response = await self._post_to_tool_blocking_service( + response_data, req_data + ) + blocked_explanation = self._extract_blocked_tools( + service_response, message_tool_calls + ) + + if blocked_explanation is not None: + model = self._resolve_model(request_data, call_details) + raise ModifyResponseException( + message=blocked_explanation, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + return inputs + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> list[ChatCompletionMessageToolCall]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + result = [] + for tc in tool_calls: + if isinstance(tc, ChatCompletionMessageToolCall): + result.append(tc) + elif isinstance(tc, dict): + func = tc.get("function", {}) + result.append( + ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + ) + elif hasattr(tc, "id") and hasattr(tc, "function"): + result.append( + ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + ) + else: + raise TypeError( + f"Cannot normalize tool_call of type {type(tc).__name__}" + ) + return result + + @staticmethod + def _build_tool_call_payload( + tool_calls: list[ChatCompletionMessageToolCall], + request_id: str | None, + ) -> dict[str, Any]: + """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + return { + "id": request_id or f"chatcmpl-{uuid.uuid4()}", + "object": "chat.completion", + "created": int(time.time()), + "model": "", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + tc.model_dump(exclude_none=True) for tc in tool_calls + ], + }, + "finish_reason": "tool_calls", + } + ], + } + + @staticmethod + def _extract_request_data(call_details: dict[str, Any]) -> dict[str, Any]: + """Extract original request data from model_call_details.""" + if not call_details: + return {} + litellm_params = call_details.get("litellm_params", {}) or {} + return { + "messages": call_details.get("messages"), + "model": call_details.get("model"), + "proxy_server_request": RubrikLogger._sanitize_proxy_server_request( + litellm_params.get("proxy_server_request") + ), + } + + @staticmethod + def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + """Allowlist only routing fields (``url``, ``method``) when forwarding + ``proxy_server_request`` to the external Rubrik webhook, dropping + inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw + request ``body`` so proxy credentials are not exfiltrated.""" + if not isinstance(proxy_server_request, dict): + return proxy_server_request + return { + key: proxy_server_request[key] + for key in ("url", "method") + if key in proxy_server_request + } + + @staticmethod + def _resolve_model( + request_data: dict[str, Any], call_details: dict[str, Any] + ) -> str: + """Get the model name for the ModifyResponseException.""" + response = request_data.get("response") + if response and hasattr(response, "model"): + return response.model or "unknown" + return call_details.get("model", "unknown") + + # -- Logging hooks --------------------------------------------------------- + + async def _prepare_log_payload( + self, kwargs: dict, event_type: str + ) -> StandardLoggingPayload | None: + """Shared logic for success and failure logging.""" + if random.random() > self.sampling_rate: + verbose_logger.debug( + f"Skipping Rubrik {event_type} logging " + f"(sampling_rate={self.sampling_rate})" + ) + return None + + # Deep-copy so mutations don't affect other callbacks sharing this object + standard_logging_payload: StandardLoggingPayload = safe_deep_copy( + kwargs["standard_logging_object"] + ) + + # For Anthropic /v1/messages requests, LiteLLM creates a separate + # ModelResponse (with a generated chatcmpl-* id) for logging, which + # differs from the original Anthropic msg-* id on the response dict. + # Normalize to litellm_call_id so that the logging and tool-blocking + # endpoints see the same request identifier. + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_request = litellm_params.get("proxy_server_request", {}) or {} + url_path = urllib.parse.urlparse(proxy_request.get("url", "")).path + if url_path.endswith(_ENDPOINT_ANTHROPIC_MESSAGES): + _litellm_call_id = kwargs.get("litellm_call_id") + if _litellm_call_id: + standard_logging_payload["id"] = _litellm_call_id # type: ignore[literal-required] + + if "system" in kwargs: + system_prompt_msg_list = kwargs["system"] + try: + if system_prompt_msg_list: + system_scaffold = { + "role": "system", + "content": system_prompt_msg_list, + } + if isinstance(standard_logging_payload["messages"], list): + standard_logging_payload["messages"].insert(0, system_scaffold) + elif isinstance(standard_logging_payload["messages"], (dict, str)): + standard_logging_payload["messages"] = [ + system_scaffold, + standard_logging_payload["messages"], + ] + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + return standard_logging_payload + + async def _enqueue_log_event(self, kwargs: dict, event_type: str): + try: + self._ensure_periodic_flush_task() + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + + self.log_queue.append(payload) + self._enforce_max_queue_size() + + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() + except Exception as e: + verbose_logger.error( + f"Rubrik {event_type} logging hook failed: {e}. " + "Skipping logging for this event.", + exc_info=True, + ) + + def _enforce_max_queue_size(self) -> None: + overflow = len(self.log_queue) - self.max_queue_size + if overflow <= 0: + return + del self.log_queue[:overflow] + self._dropped_since_warning += overflow + now = time.time() + if now - self._last_drop_warning_time >= _DROP_WARNING_INTERVAL_SECONDS: + verbose_logger.warning( + "Rubrik: log queue exceeded max_queue_size=%s; dropped %s " + "oldest events since the last warning. The Rubrik webhook may " + "be unhealthy or undersized for current traffic.", + self.max_queue_size, + self._dropped_since_warning, + ) + self._dropped_since_warning = 0 + self._last_drop_warning_time = now + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._enqueue_log_event(kwargs, "success") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await self._enqueue_log_event(kwargs, "failure") + + # -- Batch logging --------------------------------------------------------- + + async def _log_batch_to_rubrik(self, data): + # NOTE: this method intentionally re-raises on failure so the parent + # CustomBatchLogger.flush_queue keeps the unsent events in the queue + # for the next flush attempt instead of silently dropping them. + try: + response = await self.async_httpx_client.post( + url=self.logging_endpoint, + json=data, + headers=self._headers, + ) + response.raise_for_status() + except httpx.HTTPStatusError as e: + verbose_logger.exception( + f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}" + ) + raise + except Exception: + verbose_logger.exception("Rubrik Layer Error") + raise + + async def async_send_batch(self): + """Handles sending batches of responses to Rubrik. + + Note: the canonical flush path is :meth:`flush_queue`, which takes a + single snapshot used for both sending and queue draining. This method + is kept for direct callers / tests; it intentionally does NOT remove + events from the queue. + """ + if not self.log_queue: + return + + log_queue_snapshot = list(self.log_queue) + verbose_logger.debug( + "Rubrik: Flushing batch of %s events", len(log_queue_snapshot) + ) + await self._log_batch_to_rubrik( + data=log_queue_snapshot, + ) + + async def flush_queue(self): + """Snapshot, send, and drain in one consistent step. + + Overrides the base implementation so the same snapshot drives both + the HTTP send and the queue truncation. This avoids the subtle + coupling where the base class captures `len(self.log_queue)` + separately from the snapshot taken inside `async_send_batch`, + which could otherwise drift in a future refactor and cause + duplicate deliveries to Rubrik. + """ + if self.flush_lock is None: + return + + async with self.flush_lock: + if not self.log_queue: + return + snapshot = list(self.log_queue) + verbose_logger.debug("Rubrik: Flushing batch of %s events", len(snapshot)) + try: + await self._log_batch_to_rubrik(data=snapshot) + except Exception: + # Already logged with traceback inside _log_batch_to_rubrik. + # Preserve the in-flight events for retry on the next flush. + return + del self.log_queue[: len(snapshot)] + self.last_flush_time = time.time() + + # -- Tool blocking service ------------------------------------------------- + + async def _post_to_tool_blocking_service( + self, + response_data: dict[str, Any], + request_data: dict[str, Any], + ) -> dict[str, Any]: + """Post a payload to the tool blocking service and return the response. + + Args: + response_data: The OpenAI-formatted response payload to send. + request_data: Original LLM request data to include alongside + the response for additional context. Empty dict if unavailable. + + Raises: + Exception: If the service is unavailable or returns an error. + """ + envelope = { + "request": request_data, + "response": response_data, + } + verbose_logger.debug( + f"Sending request to tool blocking service: " + f"{self.tool_blocking_endpoint}" + ) + http_response = await self.tool_blocking_client.post( + self.tool_blocking_endpoint, + json=envelope, + headers=self._headers, + ) + http_response.raise_for_status() + result: dict[str, Any] = http_response.json() + return result + + @staticmethod + def _extract_blocked_tools( + service_response: dict[str, Any], + all_tool_calls: list[ChatCompletionMessageToolCall], + ) -> Optional[str]: + """Return the blocking explanation if any tool calls were blocked. + + Compares the service response (which contains only allowed tools) against + the full set of tool calls. Returns ``None`` if all tools are allowed, or + the explanation string (prefixed with newlines) otherwise. + + Expects service_response in OpenAI chat completion format: + {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} + """ + choices = service_response.get("choices", []) + if not choices: + raise _MalformedToolBlockingResponseError( + "Tool blocking service returned empty response" + ) + + message = choices[0].get("message", {}) + returned_tool_calls = message.get("tool_calls") or [] + blocking_explanation = message.get("content", "") + + allowed_id_counts: Counter = Counter( + tc["id"] + for tc in returned_tool_calls + if isinstance(tc, dict) and tc.get("id") + ) + required_id_counts: Counter = Counter(tc.id for tc in all_tool_calls if tc.id) + + all_allowed = len(returned_tool_calls) >= len(all_tool_calls) and all( + allowed_id_counts.get(tc_id, 0) >= count + for tc_id, count in required_id_counts.items() + ) + + if all_allowed: + return None + + explanation = blocking_explanation or "Tool call blocked by policy." + return f"\n\n{explanation}" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 876f1b167db..af0460956a6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1769,9 +1769,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = 0.0 elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] - elif self.model_call_details.get("response_cost") is not None: + elif ( + existing_cost := self.model_call_details.get("response_cost") + ) is not None and existing_cost != 0: # Preserve response_cost if already calculated (e.g., by pass-through - # handlers like Gemini/Vertex which call completion_cost directly) + # handlers like Gemini/Vertex which call completion_cost directly). + # Do not preserve 0 from failure_handler on intermediate router retries. pass else: self.model_call_details["response_cost"] = self._response_cost_calculator( diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index b7f4d8e3b2d..263e0c094ce 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -54,6 +54,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" pages_processed: Optional[int] = None + credits: Optional[float] = None doc_size_bytes: Optional[int] = None model_config = {"extra": "allow"} diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index c883ab68dff..d9599b8b9c4 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, List, Optional import httpx from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.litellm_core_utils.litellm_logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -22,6 +23,7 @@ from litellm.llms.bedrock.common_utils import ( from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -169,6 +171,24 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("output_format", None) + if not ( + _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_output_config", + ) + or AnthropicConfig._model_supports_effort_param(model) + ): + if anthropic_request.pop("output_config", None) is not None: + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` for " + "model=%s — neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 151e0e404a0..69b61298d33 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -45,6 +45,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import ModelResponseStream +from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -557,7 +558,29 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request=anthropic_messages_request, ) - # 5a. Remove `custom` field from tools (Bedrock doesn't support it) + # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, + # but older models do not — strip it to avoid request rejection. + # Ref: https://github.com/BerriAI/litellm/issues/22797 + if not ( + _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_output_config", + ) + or AnthropicConfig._model_supports_effort_param(model) + ): + if anthropic_messages_request.pop("output_config", None) is not None: + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` for " + "model=%s — neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + + # 5b. Remove `custom` field from tools (Bedrock doesn't support it) # Claude Code sends `custom: {defer_loading: true}` on tool definitions, # which causes Bedrock to reject the request with "Extra inputs are not permitted" # Ref: https://github.com/BerriAI/litellm/issues/22847 diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 66acd933416..56b61b66c84 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -1,7 +1,5 @@ -import json -from typing import Any, Optional +from typing import Any, Dict, Optional -from litellm.constants import STREAM_SSE_DONE_STRING from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -9,13 +7,17 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses.sse_output_recovery import ( + parse_sse_json_chunk, + record_output_item_chunk, + record_output_text_chunk, +) from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders -from litellm.utils import CustomStreamWrapper from ..authenticator import Authenticator from ..common_utils import ( @@ -111,86 +113,139 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response: Any, logging_obj: Any, ): - content_type = (raw_response.headers or {}).get("content-type", "") body_text = raw_response.text or "" - if "text/event-stream" not in content_type.lower(): - trimmed_body = body_text.lstrip() - if not ( - trimmed_body.startswith("event:") - or trimmed_body.startswith("data:") - or "\nevent:" in body_text - or "\ndata:" in body_text - ): - return super().transform_response_api_response( - model=model, - raw_response=raw_response, - logging_obj=logging_obj, - ) + if not self._should_parse_as_sse( + raw_response=raw_response, body_text=body_text + ): + return super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) logging_obj.post_call( original_response=raw_response.text, additional_args={"complete_input_dict": {}}, ) - completed_response = None - error_message = None - for chunk in body_text.splitlines(): - stripped_chunk = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) - if not stripped_chunk: - continue - stripped_chunk = stripped_chunk.strip() - if not stripped_chunk: - continue - if stripped_chunk == STREAM_SSE_DONE_STRING: - break - try: - parsed_chunk = json.loads(stripped_chunk) - except json.JSONDecodeError: - continue - if not isinstance(parsed_chunk, dict): - continue - event_type = parsed_chunk.get("type") - if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: - response_payload = parsed_chunk.get("response") - if isinstance(response_payload, dict): - response_payload = dict(response_payload) - if "created_at" in response_payload: - response_payload["created_at"] = _safe_convert_created_field( - response_payload["created_at"] - ) - try: - completed_response = ResponsesAPIResponse(**response_payload) - except Exception: - completed_response = ResponsesAPIResponse.model_construct( - **response_payload - ) - break - if event_type in ( - ResponsesAPIStreamEvents.RESPONSE_FAILED, - ResponsesAPIStreamEvents.ERROR, - ): - error_obj = parsed_chunk.get("error") or ( - parsed_chunk.get("response") or {} - ).get("error") - if error_obj is not None: - if isinstance(error_obj, dict): - error_message = error_obj.get("message") or str(error_obj) - else: - error_message = str(error_obj) - + completed_response, error_message = self._extract_completed_response_from_sse( + body_text=body_text + ) if completed_response is None: raise OpenAIError( message=error_message or raw_response.text, status_code=raw_response.status_code, ) + self._attach_response_headers( + completed_response=completed_response, raw_response=raw_response + ) + return completed_response + + def _should_parse_as_sse(self, raw_response: Any, body_text: str) -> bool: + content_type = (raw_response.headers or {}).get("content-type", "") + if "text/event-stream" in content_type.lower(): + return True + trimmed_body = body_text.lstrip() + return bool( + trimmed_body.startswith("event:") + or trimmed_body.startswith("data:") + or "\nevent:" in body_text + or "\ndata:" in body_text + ) + + def _extract_completed_response_from_sse( + self, body_text: str + ) -> tuple[Optional[ResponsesAPIResponse], Optional[str]]: + completed_response = None + error_message = None + streamed_output_items: Dict[int, dict] = {} + text_only_output_items: Dict[int, dict] = {} + for chunk in body_text.splitlines(): + parsed_chunk = parse_sse_json_chunk(chunk) + if parsed_chunk is None: + continue + + event_type = parsed_chunk.get("type") + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + record_output_item_chunk( + parsed_chunk=parsed_chunk, + output_items=streamed_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE: + record_output_text_chunk( + parsed_chunk=parsed_chunk, + output_items=streamed_output_items, + text_only_items=text_only_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + # Real OUTPUT_ITEM_DONE events take precedence at any given + # output_index, but text-only items at indices without a + # matching OUTPUT_ITEM_DONE must still be preserved (e.g. + # providers that emit only OUTPUT_TEXT_DONE for some indices). + merged_items: Dict[int, dict] = {**text_only_output_items} + merged_items.update(streamed_output_items) + completed_response = self._build_completed_response_from_chunk( + parsed_chunk=parsed_chunk, + streamed_output_items=merged_items, + ) + break + + if event_type in ( + ResponsesAPIStreamEvents.RESPONSE_FAILED, + ResponsesAPIStreamEvents.ERROR, + ): + extracted_error = self._extract_error_message(parsed_chunk) + if extracted_error is not None: + error_message = extracted_error + + return completed_response, error_message + + def _build_completed_response_from_chunk( + self, parsed_chunk: Dict[str, Any], streamed_output_items: Dict[int, dict] + ) -> Optional[ResponsesAPIResponse]: + response_payload = parsed_chunk.get("response") + if not isinstance(response_payload, dict): + return None + response_payload = dict(response_payload) + if not response_payload.get("output") and streamed_output_items: + response_payload["output"] = [ + item for _, item in sorted(streamed_output_items.items()) + ] + if "created_at" in response_payload: + response_payload["created_at"] = _safe_convert_created_field( + response_payload["created_at"] + ) + try: + return ResponsesAPIResponse(**response_payload) + except Exception: + return ResponsesAPIResponse.model_construct(**response_payload) + + def _extract_error_message(self, parsed_chunk: Dict[str, Any]) -> Optional[str]: + error_obj = parsed_chunk.get("error") or ( + parsed_chunk.get("response") or {} + ).get("error") + if error_obj is None: + return None + if isinstance(error_obj, dict): + return error_obj.get("message") or str(error_obj) + return str(error_obj) + + def _attach_response_headers( + self, + completed_response: ResponsesAPIResponse, + raw_response: Any, + ) -> None: raw_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_headers) if not hasattr(completed_response, "_hidden_params"): setattr(completed_response, "_hidden_params", {}) completed_response._hidden_params["additional_headers"] = processed_headers completed_response._hidden_params["headers"] = raw_headers - return completed_response def get_complete_url( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d2af0a3dd52..96fdf4494f9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1409,6 +1409,8 @@ class BaseLLMHTTPHandler: document=document, optional_params=optional_params, headers=headers, + api_key=api_key, + api_base=api_base, ) # All providers return OCRRequestData @@ -1477,6 +1479,8 @@ class BaseLLMHTTPHandler: document=document, optional_params=optional_params, headers=headers, + api_key=api_key, + api_base=api_base, ) # All providers return OCRRequestData diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index eaf01c5fe18..d39adf0b6f4 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -4,6 +4,7 @@ from typing import Any, List, Literal, Optional, Tuple, Union, cast import httpx import litellm +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -26,6 +27,7 @@ from litellm.types.utils import ( ProviderSpecificModelInfo, ) from litellm.utils import ( + get_model_cost_mutation_generation, supports_function_calling, supports_reasoning, supports_tool_choice, @@ -112,6 +114,19 @@ class FireworksAIConfig(OpenAIGPTConfig): # Only add tools for models that support function calling if supports_function_calling(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("tools") + supported_params.append("parallel_tool_calls") + else: + # Historically every Fireworks model advertised tool support, so a + # JSON entry that flips `supports_function_calling` to false will + # silently drop `tools` from requests. Surface this so users can + # tell why their tool calls suddenly stop working. + verbose_logger.debug( + "fireworks_ai model %r is marked as not supporting " + "function calling in model_prices_and_context_window.json; " + "`tools` and `parallel_tool_calls` will be dropped from the " + "request.", + model, + ) # Only add tool_choice for models that explicitly support it if supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"): @@ -251,34 +266,100 @@ class FireworksAIConfig(OpenAIGPTConfig): return messages - def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: - # Models that support reasoning_effort - reasoning_supported_models = [ - "qwen3-8b", - "qwen3-32b", - "qwen3-coder-480b-a35b-instruct", - "deepseek-v3p1", - "deepseek-v3p2", - "glm-4p5", - "glm-4p5-air", - "glm-4p6", - "gpt-oss-120b", - "gpt-oss-20b", + # Cached index of fireworks_ai/* entries from litellm.model_cost. Building + # this index requires a full scan of model_cost (tens of thousands of + # entries), so we memoize it. The cache key is (id(model_cost), + # mutation_generation): the generation counter is bumped on every + # register_model / reload path, so add+remove or in-place value + # replacement (which can leave id and len unchanged) still invalidates. + _fireworks_index_cache: Optional[Tuple[int, int, List[Tuple[str, dict]]]] = None + + @classmethod + def _get_fireworks_index(cls) -> List[Tuple[str, dict]]: + model_cost = litellm.model_cost + signature = (id(model_cost), get_model_cost_mutation_generation()) + cached = cls._fireworks_index_cache + if ( + cached is not None + and cached[0] == signature[0] + and cached[1] == signature[1] + ): + return cached[2] + + index: List[Tuple[str, dict]] = [] + for key, model_info in model_cost.items(): + if not key.startswith("fireworks_ai/"): + continue + if not isinstance(model_info, dict): + continue + key_short = key[len("fireworks_ai/") :] + if key_short.startswith("accounts/fireworks/models/"): + key_short = key_short[len("accounts/fireworks/models/") :] + if not key_short: + continue + index.append((key_short, model_info)) + + cls._fireworks_index_cache = (signature[0], signature[1], index) + return index + + @staticmethod + def _matches_on_hyphen_boundary(short_name: str, key_short: str) -> bool: + """Return True if `key_short` appears in `short_name` aligned to + hyphen-separated word boundaries (or end-of-string). This avoids + spurious substring matches like `"some-model"` matching + `"awesome-model"`.""" + if short_name == key_short: + return True + if short_name.startswith(key_short + "-"): + return True + if short_name.endswith("-" + key_short): + return True + return ("-" + key_short + "-") in short_name + + def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + short_name = model + if short_name.startswith("fireworks_ai/"): + short_name = short_name[len("fireworks_ai/") :] + if short_name.startswith("accounts/fireworks/models/"): + short_name = short_name[len("accounts/fireworks/models/") :] + + candidate_keys = [ + model, + f"fireworks_ai/{short_name}", + f"fireworks_ai/accounts/fireworks/models/{short_name}", ] - # Normalize model name - remove prefix if present - normalized_model = model - if model.startswith("fireworks_ai/"): - normalized_model = model.replace("fireworks_ai/", "") - if normalized_model.startswith("accounts/fireworks/models/"): - normalized_model = normalized_model.replace( - "accounts/fireworks/models/", "" - ) + for candidate_key in candidate_keys: + model_info = litellm.model_cost.get(candidate_key) + if model_info is not None and model_info.get(capability) is not None: + return cast(Optional[bool], model_info.get(capability)) - # Check if model supports reasoning - supports_reasoning_value = any( - reasoning_model in normalized_model - for reasoning_model in reasoning_supported_models + # Fallback: preserve historical substring matching for model name + # variants (e.g. fine-tuned or regionally-suffixed versions of a + # known model). Pick the *longest* matching entry so a more specific + # known model (e.g. "qwen3-8b-instruct") wins over a less specific + # one (e.g. "qwen3-8b") when the query model is more specific still. + # Use hyphen-aligned matching to avoid false positives where a short + # known model name is an unrelated substring of a longer one. + best_match_short: Optional[str] = None + best_match_value: Optional[bool] = None + for key_short, model_info in self._get_fireworks_index(): + if model_info.get(capability) is None: + continue + if not self._matches_on_hyphen_boundary(short_name, key_short): + continue + if best_match_short is None or len(key_short) > len(best_match_short): + best_match_short = key_short + best_match_value = cast(Optional[bool], model_info.get(capability)) + + return best_match_value + + def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: + supports_function_calling_value = self._get_model_cost_capability( + model=model, capability="supports_function_calling" + ) + supports_reasoning_value = self._get_model_cost_capability( + model=model, capability="supports_reasoning" ) provider_specific_model_info: ProviderSpecificModelInfo = { @@ -288,9 +369,16 @@ class FireworksAIConfig(OpenAIGPTConfig): "supports_vision": True, # via document inlining } + if supports_function_calling_value is not None: + provider_specific_model_info["supports_function_calling"] = ( + supports_function_calling_value + ) + # Only include supports_reasoning if True if supports_reasoning_value: - provider_specific_model_info["supports_reasoning"] = True + provider_specific_model_info["supports_reasoning"] = ( + supports_reasoning_value + ) return provider_specific_model_info diff --git a/litellm/llms/reducto/__init__.py b/litellm/llms/reducto/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/reducto/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/reducto/common.py b/litellm/llms/reducto/common.py new file mode 100644 index 00000000000..4e7d96dbe87 --- /dev/null +++ b/litellm/llms/reducto/common.py @@ -0,0 +1,159 @@ +import base64 +import binascii +from collections import defaultdict +from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Tuple + +from litellm.constants import request_timeout + +REDUCTO_API_BASE = "https://platform.reducto.ai" +REDUCTO_ID_PREFIX = "reducto://" + +if TYPE_CHECKING: + from litellm.llms.base_llm.ocr.transformation import OCRPage + + +def _normalize_api_base(api_base: Optional[str]) -> str: + return (api_base or REDUCTO_API_BASE).rstrip("/") + + +def _raise_bad_request(message: str, model: str) -> NoReturn: + import litellm + + raise litellm.BadRequestError( + message=message, + model=model, + llm_provider="reducto", + ) + + +def extract_file_id_or_bytes( + source_url: str, + model: str, +) -> Tuple[Optional[str], Optional[bytes], Optional[str]]: + if source_url.startswith(REDUCTO_ID_PREFIX): + return source_url, None, None + + if source_url.startswith("http://") or source_url.startswith("https://"): + _raise_bad_request( + "Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first.", + model=model, + ) + + if not source_url.startswith("data:"): + _raise_bad_request( + "Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing.", + model=model, + ) + + try: + header, encoded = source_url.split(",", 1) + except ValueError: + _raise_bad_request("Invalid Reducto data URI provided.", model=model) + + if ";base64" not in header: + _raise_bad_request( + "Reducto only supports base64-encoded data URIs.", model=model + ) + + mime = header.removeprefix("data:").split(";")[0] or "application/octet-stream" + try: + raw_bytes = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError): + _raise_bad_request("Invalid Reducto base64 payload provided.", model=model) + + return None, raw_bytes, mime + + +def _extract_file_id_from_upload_response(response: Any) -> str: + try: + payload = response.json() + except ValueError as exc: + raise ValueError( + "Reducto /upload returned a non-JSON 200 response: {}".format(response.text) + ) from exc + file_id = (payload or {}).get("file_id") if isinstance(payload, dict) else None + if not isinstance(file_id, str) or not file_id: + raise ValueError( + "Reducto /upload returned 200 without a file_id; got payload={}".format( + payload + ) + ) + return file_id + + +def upload_bytes_sync( + raw_bytes: bytes, + mime: Optional[str], + api_key: str, + api_base: Optional[str], +) -> str: + import litellm + + response = litellm.module_level_client.post( + url="{}{}".format(_normalize_api_base(api_base), "/upload"), + headers={"Authorization": f"Bearer {api_key}"}, + files={"file": ("document", raw_bytes, mime or "application/octet-stream")}, + timeout=request_timeout, + ) + response.raise_for_status() + return _extract_file_id_from_upload_response(response) + + +async def upload_bytes_async( + raw_bytes: bytes, + mime: Optional[str], + api_key: str, + api_base: Optional[str], +) -> str: + import litellm + + response = await litellm.module_level_aclient.post( + url="{}{}".format(_normalize_api_base(api_base), "/upload"), + headers={"Authorization": f"Bearer {api_key}"}, + files={"file": ("document", raw_bytes, mime or "application/octet-stream")}, + timeout=request_timeout, + ) + response.raise_for_status() + return _extract_file_id_from_upload_response(response) + + +def build_pages_from_reducto(result: Dict[str, Any]) -> List["OCRPage"]: + from litellm.llms.base_llm.ocr.transformation import OCRPage + + chunks = result.get("chunks", []) or [] + blocks_by_page: Dict[int, List[Dict[str, Any]]] = defaultdict(list) + + for chunk in chunks: + for block in chunk.get("blocks", []) or []: + page_no = (block.get("bbox") or {}).get("page") + if page_no is None: + continue + try: + normalized_page = int(page_no) + except (TypeError, ValueError): + continue + blocks_by_page[normalized_page].append(block) + + if not blocks_by_page: + fallback_markdown = "\n\n".join( + chunk.get("content", "") for chunk in chunks if chunk.get("content") + ) + if fallback_markdown == "": + return [] + return [OCRPage(index=0, markdown=fallback_markdown)] + + pages: List["OCRPage"] = [] + for page_no, blocks in sorted(blocks_by_page.items()): + markdown = "\n\n".join( + block.get("content", "") for block in blocks if block.get("content") + ) + page_index = max(page_no - 1, 0) + page = OCRPage( + index=page_index, + markdown=markdown, + ) + # OCRPage accepts extra keys at runtime; assign blocks after construction + # so static typing does not reject provider-specific metadata. + setattr(page, "blocks", blocks) + pages.append(page) + return pages diff --git a/litellm/llms/reducto/ocr/__init__.py b/litellm/llms/reducto/ocr/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/reducto/ocr/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/reducto/ocr/transformation.py b/litellm/llms/reducto/ocr/transformation.py new file mode 100644 index 00000000000..cc338ecc484 --- /dev/null +++ b/litellm/llms/reducto/ocr/transformation.py @@ -0,0 +1,241 @@ +from typing import Any, Dict, Optional, Tuple + +import httpx + +from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig, + DocumentType, + OCRRequestData, + OCRResponse, + OCRUsageInfo, +) +from litellm.llms.reducto.common import ( + REDUCTO_API_BASE, + build_pages_from_reducto, + extract_file_id_or_bytes, + upload_bytes_async, + upload_bytes_sync, +) + + +class _BaseReductoOCRConfig(BaseOCRConfig): + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + mapped_params = dict(optional_params) + supported_params = self.get_supported_ocr_params(model=model) + for param, value in non_default_params.items(): + if param in supported_params: + mapped_params[param] = value + return mapped_params + + def validate_environment( + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + litellm_params: Optional[dict] = None, + **kwargs, + ) -> Dict: + from litellm.secret_managers.main import get_secret_str + + resolved_key = api_key or get_secret_str("REDUCTO_API_KEY") + if resolved_key is None: + raise ValueError( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + ) + + return { + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + **headers, + } + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: dict, + litellm_params: Optional[dict] = None, + **kwargs, + ) -> str: + return "{}/parse".format((api_base or REDUCTO_API_BASE).rstrip("/")) + + def _get_source_url(self, document: DocumentType, model: str) -> str: + source_url = document.get("document_url") or document.get("image_url") + if source_url is None: + raise ValueError( + "Reducto expected OCR preprocessing to produce document_url or image_url for model={}".format( + model + ) + ) + return source_url + + @staticmethod + def _resolve_credentials( + api_key: Optional[str], api_base: Optional[str] + ) -> Tuple[str, str]: + from litellm.secret_managers.main import get_secret_str + + resolved_key = api_key or get_secret_str("REDUCTO_API_KEY") + if resolved_key is None: + raise ValueError( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + ) + resolved_base = (api_base or REDUCTO_API_BASE).rstrip("/") + return resolved_key, resolved_base + + def _ensure_file_id_sync( + self, + model: str, + document: DocumentType, + api_key: Optional[str], + api_base: Optional[str], + ) -> str: + source_url = self._get_source_url(document=document, model=model) + file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model) + if file_id is not None: + return file_id + resolved_key, resolved_base = self._resolve_credentials(api_key, api_base) + return upload_bytes_sync( + raw_bytes=raw_bytes or b"", + mime=mime, + api_key=resolved_key, + api_base=resolved_base, + ) + + async def _ensure_file_id_async( + self, + model: str, + document: DocumentType, + api_key: Optional[str], + api_base: Optional[str], + ) -> str: + source_url = self._get_source_url(document=document, model=model) + file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model) + if file_id is not None: + return file_id + resolved_key, resolved_base = self._resolve_credentials(api_key, api_base) + return await upload_bytes_async( + raw_bytes=raw_bytes or b"", + mime=mime, + api_key=resolved_key, + api_base=resolved_base, + ) + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: Any, + **kwargs, + ) -> OCRResponse: + response_json = raw_response.json() + result = response_json.get("result", response_json) or {} + usage = response_json.get("usage", {}) or {} + response = OCRResponse( + pages=build_pages_from_reducto(result), + model=model, + usage_info=OCRUsageInfo( + pages_processed=usage.get("num_pages"), + credits=usage.get("credits"), + ), + object="ocr", + ) + response._hidden_params["reducto_raw"] = response_json + return response + + +class ReductoParseV3Config(_BaseReductoOCRConfig): + def get_supported_ocr_params(self, model: str) -> list: + return ["formatting", "retrieval", "settings"] + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = self._ensure_file_id_sync( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData(data={"input": file_id, **optional_params}, files=None) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = await self._ensure_file_id_async( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData(data={"input": file_id, **optional_params}, files=None) + + +class ReductoParseLegacyConfig(_BaseReductoOCRConfig): + def get_supported_ocr_params(self, model: str) -> list: + return ["enhance"] + + def _build_legacy_body(self, file_id: str, optional_params: dict) -> Dict[str, Any]: + body: Dict[str, Any] = {"document_url": file_id} + enhance = optional_params.get("enhance") + if enhance is not None: + body["options"] = {"enhance": enhance} + return body + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = self._ensure_file_id_sync( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData( + data=self._build_legacy_body( + file_id=file_id, optional_params=optional_params + ), + files=None, + ) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = await self._ensure_file_id_async( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData( + data=self._build_legacy_body( + file_id=file_id, optional_params=optional_params + ), + files=None, + ) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index ac0f07b8e0b..3f945adca0d 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -41,7 +41,7 @@ class ContextCachingEndpoints(VertexBase): """ def __init__(self) -> None: - pass + super().__init__() def _get_token_and_url_context_caching( self, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index eb67e3aa828..13aa2a5350e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -45,7 +45,7 @@ class PartnerModelPrefixes(str, Enum): class VertexAIPartnerModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() @staticmethod def is_vertex_partner_model(model: str): @@ -116,9 +116,6 @@ class VertexAIPartnerModels(VertexBase): CodestralTextCompletion, ) from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) except Exception as e: raise VertexAIError( status_code=400, @@ -133,9 +130,7 @@ class VertexAIPartnerModels(VertexBase): message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", ) try: - vertex_httpx_logic = VertexLLM() - - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 82cfe6de984..b6bf2f73b72 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -31,7 +31,7 @@ from ..vertex_llm_base import VertexBase class VertexAIGemmaModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -62,9 +62,6 @@ class VertexAIGemmaModels(VertexBase): try: import vertexai - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( VertexGemmaConfig, ) @@ -83,9 +80,8 @@ class VertexAIGemmaModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 6c6446958bc..35cd54d65f6 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -91,6 +91,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): "stream", None ) # Streaming not supported, will be faked client-side openai_request.pop("stream_options", None) # Stream options not supported + # Vertex Gemma's chatCompletions wrapper does not understand + # `context_management` (an Anthropic/Responses API concept). Strip it + # so the upstream endpoint does not 400 on the unknown field. + openai_request.pop("context_management", None) # Wrap in Vertex Gemma format return { diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 6f687dae7e8..990063bb9fb 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -4,8 +4,10 @@ Base Vertex, Google AI Studio LLM Class Handles Authentication and generating request urls for Vertex AI and Google AI Studio """ +import asyncio import json import os +import threading from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple import litellm @@ -30,6 +32,7 @@ GOOGLE_IMPORT_ERROR_MESSAGE = ( if TYPE_CHECKING: from google.auth.credentials import Credentials as GoogleCredentialsObject + from google.auth.credentials import TokenState else: GoogleCredentialsObject = Any @@ -42,10 +45,28 @@ class VertexBase: self._credentials: Optional[GoogleCredentialsObject] = None self._credentials_project_mapping: Dict[ Tuple[Optional[VERTEX_CREDENTIALS_TYPES], Optional[str]], - Tuple[GoogleCredentialsObject, str], + Tuple[GoogleCredentialsObject, Optional[str]], ] = {} self.project_id: Optional[str] = None self.async_handler: Optional[AsyncHTTPHandler] = None + # Per-credential-key asyncio.Lock for single-flight async refresh. + # Prevents thundering herd when token expires under high concurrency. + # Uses a regular dict (not WeakValueDictionary) so the lock identity is + # stable across concurrent callers — a weak reference can be GC'd + # between two coroutines arriving at the lock, breaking single-flight. + # An explicit refcount tracks the number of coroutines currently using + # each lock; the entry is pruned when the count reaches zero, so the + # dict stays bounded even in long-running high-cardinality deployments + # without depending on any private asyncio internals. + self._async_refresh_locks: Dict[tuple, asyncio.Lock] = {} + self._async_refresh_lock_refcounts: Dict[tuple, int] = {} + # Tracks in-flight background refresh tasks to avoid duplicate refreshes. + self._background_refresh_tasks: Dict[tuple, asyncio.Task] = {} + # Protects the sync get_access_token refresh path. + # Use RLock so that the reauthentication retry path (which calls + # back into get_access_token while still holding the lock) can + # re-acquire it without deadlocking the current thread. + self._sync_refresh_lock = threading.RLock() def get_vertex_region(self, vertex_region: Optional[str], model: str) -> str: import litellm @@ -77,7 +98,9 @@ class VertexBase: return vertex_region or "us-central1" def load_auth( - self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str] + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], ) -> Tuple[Any, str]: if credentials is not None: if isinstance(credentials, str): @@ -343,7 +366,241 @@ class VertexBase: except ImportError: raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) - credentials.refresh(Request()) + # Serialize all refreshes on this VertexBase across threads. + # ``credentials.refresh()`` is not safe to call concurrently on the + # same credentials object, and this method is invoked from three + # places that can run on different threads: + # - sync ``get_access_token`` (already holds ``_sync_refresh_lock``) + # - the async slow path (via ``asyncify`` in a worker thread) + # - the background proactive refresh task (via ``asyncify``) + # ``_sync_refresh_lock`` is an ``RLock`` so reentrant acquisition + # from the sync path is safe. + with self._sync_refresh_lock: + credentials.refresh(Request()) + + def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock: + """Increment the refcount and return the lock for ``credential_cache_key``. + + Every call must be paired with ``_release_async_refresh_lock`` once the + caller is done with the lock so the entry can be pruned when no other + coroutine is holding or waiting on it. + """ + lock = self._async_refresh_locks.setdefault( + credential_cache_key, asyncio.Lock() + ) + self._async_refresh_lock_refcounts[credential_cache_key] = ( + self._async_refresh_lock_refcounts.get(credential_cache_key, 0) + 1 + ) + return lock + + def _release_async_refresh_lock( + self, credential_cache_key: tuple, lock: asyncio.Lock + ) -> None: + """Decrement the refcount and drop the lock entry when it reaches zero. + + Must be called only after the caller has released ``lock`` (i.e. once + the surrounding ``async with`` has exited). asyncio is cooperative, so + the decrement-then-pop sequence below runs atomically with respect to + other coroutines. + """ + remaining = self._async_refresh_lock_refcounts.get(credential_cache_key, 0) - 1 + if remaining > 0: + self._async_refresh_lock_refcounts[credential_cache_key] = remaining + return + self._async_refresh_lock_refcounts.pop(credential_cache_key, None) + if self._async_refresh_locks.get(credential_cache_key) is lock: + self._async_refresh_locks.pop(credential_cache_key, None) + + def _try_get_cached_token( + self, + credential_cache_key: tuple, + project_id: Optional[str], + ) -> Optional[Tuple[str, str]]: + """ + Look up cached credentials and return (token, project_id) if the token + is FRESH. Returns None if not cached or not fresh. + """ + from google.auth.credentials import TokenState + + creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) + if ( + creds is not None + and self._get_token_state(creds) == TokenState.FRESH + and creds.token is not None + and isinstance(creds.token, str) + ): + resolved_project = project_id or cached_project_id + if resolved_project: + return creds.token, resolved_project + return None + + def _try_get_usable_cached_token( + self, + credential_cache_key: tuple, + project_id: Optional[str], + ) -> Optional[Tuple[str, str, "TokenState", Any, Optional[str]]]: + """ + Look up cached credentials and return usable token info for FRESH or + STALE tokens (both are still valid for outbound requests). STALE + tokens are returned along with their state and the underlying + credentials object so the caller can schedule a background refresh + without holding the per-key async lock. + """ + from google.auth.credentials import TokenState + + creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) + if creds is None: + return None + token_state = self._get_token_state(creds) + if token_state not in (TokenState.FRESH, TokenState.STALE): + return None + if creds.token is None or not isinstance(creds.token, str): + return None + resolved_project = project_id or cached_project_id + if not resolved_project: + return None + return creds.token, resolved_project, token_state, creds, cached_project_id + + def _unpack_cached_credentials( + self, credential_cache_key: tuple + ) -> Tuple[Any, Optional[str]]: + """ + Return (credentials, project_id) from the cache, or (None, None) if + not cached. Handles both tuple and legacy cache formats. + """ + if credential_cache_key not in self._credentials_project_mapping: + return None, None + cached_entry = self._credentials_project_mapping[credential_cache_key] + if isinstance(cached_entry, tuple): + return cached_entry + return cached_entry, cached_entry.quota_project_id or getattr( + cached_entry, "project_id", None + ) + + def _get_token_state(self, credentials: Any) -> "TokenState": + """ + Return the token state using google-auth's TokenState enum. + + Falls back to expired/valid checks if token_state is unavailable + (e.g. older google-auth versions or mock objects in tests). + """ + from google.auth.credentials import TokenState as _TokenState + + token_state = getattr(credentials, "token_state", None) + if isinstance(token_state, _TokenState): + return token_state + # Fallback for credentials without a real token_state (e.g. mocks) + if getattr(credentials, "expired", True): + return _TokenState.INVALID + if getattr(credentials, "valid", False): + return _TokenState.FRESH + return _TokenState.INVALID + + async def _load_and_cache_credentials( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + credential_cache_key: tuple, + ) -> Tuple[Any, Optional[str]]: + """Load credentials via load_auth (in thread) and cache the result.""" + try: + _credentials, credential_project_id = await asyncify(self.load_auth)( + credentials=credentials, + project_id=project_id, + ) + except Exception as e: + verbose_logger.exception("Failed to load vertex credentials: %s", str(e)) + raise + if _credentials is None: + raise ValueError("Could not resolve credentials") + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + return _credentials, credential_project_id + + async def _background_refresh_credentials( + self, + credentials: Any, + credential_cache_key: tuple, + credential_project_id: Optional[str], + ) -> None: + """ + Refresh credentials in the background without blocking the calling request. + + Called when the token is still valid but nearing expiry (proactive refresh). + Errors are logged but not raised — the current token is still usable. + """ + try: + verbose_logger.debug("Background proactive credential refresh") + await asyncify(self.refresh_auth)(credentials) + # Only update the cache if it still points at the credentials + # object we just refreshed. The per-key async lock is not held + # here, so a concurrent INVALID path may have already replaced + # this entry (e.g. via _handle_reauthentication_async, which + # creates a fresh credentials object). In that case our write + # would clobber the newer entry with a stale reference. + cached_creds, _ = self._unpack_cached_credentials(credential_cache_key) + if cached_creds is credentials: + self._credentials_project_mapping[credential_cache_key] = ( + credentials, + credential_project_id, + ) + except Exception: + verbose_logger.debug( + "Background credential refresh failed, will retry on next request", + exc_info=True, + ) + + async def _await_in_flight_background_refresh( + self, credential_cache_key: tuple + ) -> None: + """Wait for an in-flight background refresh to finish, if any. + + google-auth's ``Credentials.refresh()`` is not safe to invoke + concurrently on the same credentials object. Coroutines that need a + blocking refresh must first drain any background refresh that was + scheduled while a previous STALE token was being served. + """ + existing_task = self._background_refresh_tasks.get(credential_cache_key) + if existing_task is None or existing_task.done(): + return + try: + await existing_task + except Exception: + # Background refresh failures are already logged inside + # _background_refresh_credentials; the caller will fall through + # to its own blocking refresh. + pass + + def _schedule_background_refresh( + self, + credentials: Any, + credential_cache_key: tuple, + credential_project_id: Optional[str], + ) -> None: + """Kick off a single background refresh for ``credential_cache_key``. + + Skips scheduling if a refresh is already in flight. The done-callback + guards against removing a newer task that has replaced this one in the + tracking dict (done_callbacks are scheduled via ``call_soon``). + """ + existing = self._background_refresh_tasks.get(credential_cache_key) + if existing is not None and not existing.done(): + return + self._background_refresh_tasks.pop(credential_cache_key, None) + task = asyncio.create_task( + self._background_refresh_credentials( + credentials, credential_cache_key, credential_project_id + ) + ) + + def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None: + if self._background_refresh_tasks.get(credential_cache_key) is _fut: + self._background_refresh_tasks.pop(credential_cache_key, None) + + task.add_done_callback(_drop_background_refresh_task) + self._background_refresh_tasks[credential_cache_key] = task def _ensure_access_token( self, @@ -563,6 +820,65 @@ class VertexBase: # Re-raise the original error for better context raise error + async def _handle_reauthentication_async( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + credential_cache_key: Tuple, + error: Exception, + ) -> Tuple[str, str]: + """ + Async reauthentication retry that stays within the per-key async lock. + """ + verbose_logger.debug( + f"Handling async reauthentication for project_id: {project_id}. " + f"Clearing cache and retrying once." + ) + + self._credentials_project_mapping.pop(credential_cache_key, None) + + try: + _credentials, credential_project_id = ( + await self._load_and_cache_credentials( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + ) + ) + if project_id is None and isinstance(credential_project_id, str): + project_id = credential_project_id + cache_credentials = ( + json.dumps(credentials) + if isinstance(credentials, dict) + else credentials + ) + resolved_cache_key = (cache_credentials, project_id) + # Always overwrite — any pre-existing entry at the resolved key + # references the OLD credentials object we just replaced, and + # leaving it would force the next request to do a redundant + # refresh/reauth before realizing the cached creds are stale. + self._credentials_project_mapping[resolved_cache_key] = ( + _credentials, + credential_project_id, + ) + + if _credentials.token is None or not isinstance(_credentials.token, str): + raise ValueError( + "Could not resolve credentials token. Got None or non-string token (type={})".format( + type(_credentials.token).__name__ + ) + ) + if project_id is None: + raise ValueError("Could not resolve project_id") + + return _credentials.token, project_id + except Exception as retry_error: + verbose_logger.error( + f"Async reauthentication retry failed for project_id: {project_id}. " + f"Original error: {str(error)}. Retry error: {str(retry_error)}" + ) + raise error + def get_access_token( self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], @@ -646,7 +962,7 @@ class VertexBase: ) ## VALIDATE CREDENTIALS - verbose_logger.debug(f"Validating credentials for project_id: {project_id}") + verbose_logger.debug("Validating credentials") if ( project_id is None and credential_project_id is not None @@ -666,26 +982,27 @@ class VertexBase: raise ValueError("Credentials are None after loading") if _credentials.expired: - try: - verbose_logger.debug( - f"Credentials expired, refreshing for project_id: {project_id}" - ) - self.refresh_auth(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) - except Exception as e: - # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` - # in this case, we should try to reload the credentials by clearing the cache and retrying - if "Reauthentication is needed" in str(e) and not _retry_reauth: - return self._handle_reauthentication( - credentials=credentials, - project_id=project_id, - credential_cache_key=credential_cache_key, - error=e, - ) - raise e + with self._sync_refresh_lock: + # Double-check after acquiring lock + if _credentials.expired: + try: + verbose_logger.debug("Credentials expired, refreshing") + self.refresh_auth(_credentials) + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + except Exception as e: + # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` + # in this case, we should try to reload the credentials by clearing the cache and retrying + if "Reauthentication is needed" in str(e) and not _retry_reauth: + return self._handle_reauthentication( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + error=e, + ) + raise e ## VALIDATION STEP if _credentials.token is None or not isinstance(_credentials.token, str): @@ -700,6 +1017,149 @@ class VertexBase: return _credentials.token, project_id + async def get_access_token_async( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + ) -> Tuple[str, str]: + """ + Async version of get_access_token with single-flight refresh coordination. + + Prevents thundering herd: when credentials expire under high concurrency, + only one coroutine refreshes while others wait on the lock. Uses native + async refresh for service_account and authorized_user credentials. + """ + from google.auth.credentials import TokenState + + cache_credentials = ( + json.dumps(credentials) if isinstance(credentials, dict) else credentials + ) + credential_cache_key = (cache_credentials, project_id) + + # === FAST PATH (no lock) === + # If credentials are FRESH or STALE, return immediately without + # touching the per-key async lock. STALE tokens are still usable; + # we kick off a deduplicated background refresh so subsequent + # requests get a fresh token, but we must not serialize concurrent + # callers on the lock just to schedule that refresh. + usable = self._try_get_usable_cached_token(credential_cache_key, project_id) + if usable is not None: + cached_token, resolved_project, token_state, creds, cached_project_id = ( + usable + ) + if token_state == TokenState.STALE: + self._schedule_background_refresh( + creds, credential_cache_key, cached_project_id + ) + return cached_token, resolved_project + + # === SLOW PATH (per-key lock) === + lock = self._acquire_async_refresh_lock(credential_cache_key) + try: + async with lock: + # Double-check after acquiring lock — another coroutine may have refreshed. + cached = self._try_get_cached_token(credential_cache_key, project_id) + if cached is not None: + return cached + + _credentials, credential_project_id = self._unpack_cached_credentials( + credential_cache_key + ) + + # Load credentials if not cached + if _credentials is None: + _credentials, credential_project_id = ( + await self._load_and_cache_credentials( + credentials, project_id, credential_cache_key + ) + ) + + # Resolve project_id from credentials if not provided + if project_id is None and isinstance(credential_project_id, str): + project_id = credential_project_id + resolved_cache_key = (cache_credentials, project_id) + # Always overwrite — a pre-existing entry at the resolved + # key may reference stale credentials (e.g. from before a + # reauth that only repopulated the unresolved key), which + # would force the next request through an unnecessary + # refresh/reauth cycle. + self._credentials_project_mapping[resolved_cache_key] = ( + _credentials, + credential_project_id, + ) + + # Use google-auth's token_state to decide refresh strategy: + # - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of + # expiry — return it immediately and refresh in the background. + # - INVALID: token is expired or missing — must block on refresh. + token_state = self._get_token_state(_credentials) + + if token_state == TokenState.STALE: + if project_id is None: + raise ValueError("Could not resolve project_id") + current_token = _credentials.token + if current_token is None or not isinstance(current_token, str): + # Token is malformed despite STALE state — block on a full + # refresh using the same path as INVALID credentials. + token_state = TokenState.INVALID + else: + self._schedule_background_refresh( + _credentials, + credential_cache_key, + credential_project_id, + ) + return current_token, project_id + + if token_state == TokenState.INVALID: + # Drain any in-flight background refresh before invoking + # refresh_auth ourselves; google-auth's + # Credentials.refresh() is not safe to call concurrently + # on the same credentials object, and the background task + # runs outside this lock. + await self._await_in_flight_background_refresh(credential_cache_key) + cached = self._try_get_cached_token( + credential_cache_key, project_id + ) + if cached is not None: + return cached + + # Token is expired or missing — must block until refresh completes. + try: + verbose_logger.debug("Credentials expired, refreshing") + await asyncify(self.refresh_auth)(_credentials) + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + except Exception as e: + if "Reauthentication is needed" in str(e): + verbose_logger.debug( + "Reauthentication needed, clearing cache and retrying" + ) + return await self._handle_reauthentication_async( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + error=e, + ) + raise + + # Final validation + if _credentials.token is None or not isinstance( + _credentials.token, str + ): + raise ValueError( + "Could not resolve credentials token. Got None or non-string token (type={})".format( + type(_credentials.token).__name__ + ) + ) + if project_id is None: + raise ValueError("Could not resolve project_id") + + return _credentials.token, project_id + finally: + self._release_async_refresh_lock(credential_cache_key, lock) + async def _ensure_access_token_async( self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], @@ -714,13 +1174,10 @@ class VertexBase: if custom_llm_provider == "gemini": return "", "" else: - try: - return await asyncify(self.get_access_token)( - credentials=credentials, - project_id=project_id, - ) - except Exception as e: - raise e + return await self.get_access_token_async( + credentials=credentials, + project_id=project_id, + ) def set_headers( self, auth_header: Optional[str], extra_headers: Optional[dict] diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 7240d9dce57..732d5f90dc2 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -57,7 +57,7 @@ def create_vertex_url( class VertexAIModelGardenModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -89,9 +89,6 @@ class VertexAIModelGardenModels(VertexBase): import vertexai from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) except Exception as e: raise VertexAIError( status_code=400, @@ -107,9 +104,8 @@ class VertexAIModelGardenModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 6300868a641..7325c0596a6 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union import httpx @@ -26,6 +26,7 @@ from ...openai.chat.gpt_transformation import ( class XAIChatConfig(OpenAIGPTConfig): + @property def custom_llm_provider(self) -> Optional[str]: return "xai" @@ -225,21 +226,57 @@ class XAIChatConfig(OpenAIGPTConfig): verbose_logger.debug(f"Error extracting X.AI web search usage: {e}") self._fold_reasoning_tokens_into_completion(response) + self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) return response @staticmethod - def _fold_reasoning_tokens_into_completion(model_response: ModelResponse) -> None: + def _fold_reasoning_tokens_into_completion( + target: Union[ModelResponse, Usage, Dict[str, Any], None], + ) -> None: """Reconcile xAI Usage to the OpenAI invariant. xAI accounts ``reasoning_tokens`` separately from ``completion_tokens`` while still summing them into ``total_tokens``. OpenAI's contract (o1/o3) folds reasoning into ``completion_tokens``, so fold here to keep ``total = prompt + completion``. Idempotent. + + Accepts a ``ModelResponse`` (non-streaming), a ``Usage`` object, or a + raw usage ``dict`` (streaming chunk) so streaming and non-streaming + paths stay in sync. """ - usage = getattr(model_response, "usage", None) + if target is None: + return + + if isinstance(target, ModelResponse): + usage: Union[Usage, Dict[str, Any], None] = getattr(target, "usage", None) + else: + usage = target if usage is None: return + if isinstance(usage, dict): + details = usage.get("completion_tokens_details") or {} + if isinstance(details, dict): + reasoning_tokens = int(details.get("reasoning_tokens") or 0) + else: + reasoning_tokens = int(getattr(details, "reasoning_tokens", 0) or 0) + if reasoning_tokens <= 0: + return + + prompt_tokens = int(usage.get("prompt_tokens") or 0) + completion_tokens = int(usage.get("completion_tokens") or 0) + total_tokens = int(usage.get("total_tokens") or 0) + + if total_tokens == prompt_tokens + completion_tokens: + return + + # Guard against double-counting if xAI changes accounting. + if total_tokens != prompt_tokens + completion_tokens + reasoning_tokens: + return + + usage["completion_tokens"] = completion_tokens + reasoning_tokens + return + details = getattr(usage, "completion_tokens_details", None) reasoning_tokens = ( int(getattr(details, "reasoning_tokens", 0) or 0) if details else 0 @@ -284,6 +321,25 @@ class XAIChatConfig(OpenAIGPTConfig): setattr(usage, "num_sources_used", int(num_sources_used)) verbose_logger.debug(f"X.AI web search sources used: {num_sources_used}") + @staticmethod + def _normalize_openai_compatible_usage_totals( + usage: Union[Usage, Dict[str, Any], None], + ) -> None: + if usage is None: + return + if isinstance(usage, dict): + prompt_tokens = int(usage.get("prompt_tokens") or 0) + completion_tokens = int(usage.get("completion_tokens") or 0) + expected_total = prompt_tokens + completion_tokens + if int(usage.get("total_tokens") or 0) < expected_total: + usage["total_tokens"] = expected_total + return + prompt_tokens = int(usage.prompt_tokens or 0) + completion_tokens = int(usage.completion_tokens or 0) + expected_total = prompt_tokens + completion_tokens + if int(usage.total_tokens or 0) < expected_total: + usage.total_tokens = expected_total + class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): def chunk_parser(self, chunk: dict) -> ModelResponseStream: @@ -304,4 +360,8 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): # Add a dummy choice with empty delta to ensure proper processing chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] + if "usage" in chunk and chunk["usage"] is not None: + XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) + XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) + return super().chunk_parser(chunk) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 41f73ddca5e..6a4a5dd6a03 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13982,6 +13982,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "fireworks_ai", @@ -14248,6 +14263,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 6e-07, @@ -29122,6 +29152,24 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "reducto/parse-legacy": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "reducto/parse-v3": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "recraft/recraftv2": { "litellm_provider": "recraft", "mode": "image_generation", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 13381c7a6c9..09bb8057203 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1187,6 +1187,127 @@ async def get_end_user_object( return None +_END_USER_VALIDATION_NEGATIVE_TTL = 60 +_END_USER_VALIDATION_POSITIVE_TTL = 300 + + +async def resolve_and_validate_end_user_id( + raw_end_user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Optional[Span] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + route: str = "", +) -> Optional[str]: + """Optionally drop end-user ids that don't resolve to a known DB row. + + Default: pass-through. LiteLLM's documented pattern is that the `user` + field is an arbitrary caller-supplied identifier, so validation is + opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve + backwards compatibility. + + When the flag is set: accept the id when it matches any of + - LiteLLM_EndUserTable.user_id + - LiteLLM_UserTable.user_id + - LiteLLM_UserTable.user_email (case-insensitive) + + If the id doesn't match but ``litellm.max_end_user_budget_id`` is set, + we still preserve the id so the default end-user budget is applied + downstream; otherwise we return None. + + DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they + share the same cache as the rest of the auth path instead of adding new + raw Prisma queries. + """ + if raw_end_user_id is None: + return None + if not litellm.validate_end_user_id_in_db: + return raw_end_user_id + if prisma_client is None: + return raw_end_user_id + + cache_key = f"end_user_validation:{raw_end_user_id}" + cached = await user_api_key_cache.async_get_cache(key=cache_key) + if cached == "valid": + return raw_end_user_id + if cached == "invalid": + return raw_end_user_id if litellm.max_end_user_budget_id else None + + is_valid = await _end_user_id_exists_in_db( + end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value="valid" if is_valid else "invalid", + ttl=( + _END_USER_VALIDATION_POSITIVE_TTL + if is_valid + else _END_USER_VALIDATION_NEGATIVE_TTL + ), + ) + + if is_valid: + return raw_end_user_id + # Preserve id so the caller can still apply litellm.max_end_user_budget_id. + if litellm.max_end_user_budget_id: + return raw_end_user_id + return None + + +async def _end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Optional[Span] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + route: str = "", +) -> bool: + """True when the id matches an EndUser, User, or user_email row.""" + try: + end_user_obj = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_obj is not None: + return True + except litellm.BudgetExceededError: + raise + except Exception as e: + verbose_proxy_logger.debug( + f"end_user validation: get_end_user_object lookup failed: {e}" + ) + + try: + user_obj = await get_user_object( + user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=False, + user_email=end_user_id if "@" in end_user_id else None, + ) + if user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug( + f"end_user validation: get_user_object lookup failed: {e}" + ) + + return False + + @log_db_metrics async def get_tag_objects_batch( tag_names: List[str], diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 5c24aca632c..99f8ff7186f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -10,6 +10,7 @@ import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS @@ -1014,12 +1015,47 @@ def _get_customer_id_from_standard_headers( for standard_header in STANDARD_CUSTOMER_ID_HEADERS: for header_name, header_value in request_headers.items(): if header_name.lower() == standard_header.lower(): - user_id_str = str(header_value) if header_value is not None else "" - if user_id_str.strip(): + user_id_str = _coerce_user_id_to_str(header_value) + if user_id_str: return user_id_str return None +def _coerce_user_id_to_str(value: Any) -> Optional[str]: + """Return a usable end-user identifier string, or None if the value isn't one. + + Always drops non-string structured values (dict/list/tuple/set) because + stringifying them produces garbage spend-log rows like + ``"{'device_id': ...}"``. Strings that *decode* to a structured payload + are only rejected when ``litellm.validate_end_user_id_in_db`` is enabled + — operators who currently pass JSON-encoded identifiers keep their + existing behavior until they opt in. See + auth_utils.py:get_end_user_id_from_request_body for the extraction chain. + """ + if value is None: + return None + if isinstance(value, bool): + # bool is an int subclass; handle explicitly to avoid "True"/"False". + return None + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + stripped = value.strip() + if not stripped: + return None + # Reject strings that decode to a structured payload (JSON object/array) + # only when the operator has opted into end-user validation. Gating + # behind the flag preserves backwards compatibility for deployments + # that intentionally pass JSON-encoded user identifiers. + if litellm.validate_end_user_id_in_db and stripped[:1] in ("{", "["): + parsed = safe_json_loads(stripped) + if isinstance(parsed, (dict, list)): + return None + return stripped + # dict, list, tuple, set, arbitrary objects -> drop. + return None + + def get_end_user_id_from_request_body( request_body: dict, request_headers: Optional[dict] = None ) -> Optional[str]: @@ -1058,23 +1094,22 @@ def get_end_user_id_from_request_body( if isinstance(custom_header_name_to_check, list): headers_lower = {k.lower(): v for k, v in request_headers.items()} for expected_header in custom_header_name_to_check: - header_value = headers_lower.get(expected_header) - if header_value is not None: - user_id_str = str(header_value) - if user_id_str.strip(): - return user_id_str + user_id_str = _coerce_user_id_to_str(headers_lower.get(expected_header)) + if user_id_str: + return user_id_str elif isinstance(custom_header_name_to_check, str): for header_name, header_value in request_headers.items(): if header_name.lower() == custom_header_name_to_check.lower(): - user_id_str = str(header_value) if header_value is not None else "" - if user_id_str.strip(): + user_id_str = _coerce_user_id_to_str(header_value) + if user_id_str: return user_id_str # Check 3: 'user' field in request_body (commonly OpenAI) - if "user" in request_body and request_body["user"] is not None: - user_from_body_user_field = request_body["user"] - return str(user_from_body_user_field) + if "user" in request_body: + user_id_str = _coerce_user_id_to_str(request_body["user"]) + if user_id_str: + return user_id_str def _as_dict(value: Any) -> dict: # metadata / litellm_metadata can arrive as JSON strings from @@ -1083,32 +1118,30 @@ def get_end_user_id_from_request_body( if isinstance(value, dict): return value if isinstance(value, str): - from litellm.litellm_core_utils.safe_json_loads import safe_json_loads - parsed = safe_json_loads(value) return parsed if isinstance(parsed, dict) else {} return {} # Check 4: 'litellm_metadata.user' in request_body (commonly Anthropic) litellm_metadata = _as_dict(request_body.get("litellm_metadata")) - user_from_litellm_metadata = litellm_metadata.get("user") - if user_from_litellm_metadata is not None: - return str(user_from_litellm_metadata) + user_id_str = _coerce_user_id_to_str(litellm_metadata.get("user")) + if user_id_str: + return user_id_str # Check 5: 'metadata.user_id' in request_body (another common pattern) metadata_dict = _as_dict(request_body.get("metadata")) - user_id_from_metadata_field = metadata_dict.get("user_id") - if user_id_from_metadata_field is not None: - return str(user_id_from_metadata_field) + user_id_str = _coerce_user_id_to_str(metadata_dict.get("user_id")) + if user_id_str: + return user_id_str # Check 6: 'safety_identifier' in request body (OpenAI Responses API parameter) # SECURITY NOTE: safety_identifier can be set by any caller in the request body. # Only use this for end-user identification in trusted environments where you control # the calling application. For untrusted callers, prefer using headers or server-side # middleware to set the end_user_id to prevent impersonation. - if request_body.get("safety_identifier") is not None: - user_from_body_user_field = request_body["safety_identifier"] - return str(user_from_body_user_field) + user_id_str = _coerce_user_id_to_str(request_body.get("safety_identifier")) + if user_id_str: + return user_id_str return None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0cca9414b2a..6974860a22a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -44,6 +44,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, is_valid_fallback_model, + resolve_and_validate_end_user_id, ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_utils import ( @@ -1071,9 +1072,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _end_user_object = None end_user_params = {} - end_user_id = get_end_user_id_from_request_body( + raw_end_user_id = get_end_user_id_from_request_body( request_data, _safe_get_request_headers(request) ) + end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) if end_user_id: try: end_user_params["end_user_id"] = end_user_id @@ -1759,7 +1768,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached @tracer.wrap() -async def _run_centralized_common_checks( +async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_auth_obj: UserAPIKeyAuth, request: Request, request_data: dict, @@ -1837,9 +1846,23 @@ async def _run_centralized_common_checks( return parent_otel_span = user_api_key_auth_obj.parent_otel_span - end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) + # In the integrated auth flow ``_user_api_key_auth_builder`` has already + # resolved the end-user id and attached it here. Reuse that to avoid a + # second extraction pass; fall back to extracting locally when the + # function is invoked in isolation (e.g. in direct unit tests). + end_user_id = user_api_key_auth_obj.end_user_id + if end_user_id is None: + raw_end_user_id = get_end_user_id_from_request_body( + request_data, _safe_get_request_headers(request) + ) + end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) fetch_coros = [] if user_api_key_auth_obj.team_id is not None: @@ -2170,11 +2193,33 @@ async def user_api_key_auth( api_key=api_key, ) - end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) - if end_user_id is not None: - user_api_key_auth_obj.end_user_id = end_user_id + # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return + # paths (no master key, /user/auth route, JWT short-circuits) that bypass + # the end-user resolution block. If those paths produced an auth obj + # without an ``end_user_id`` set, fall back to extracting from the request + # body so spend logs are still attributed correctly. Validation honours + # ``litellm.validate_end_user_id_in_db``. + if user_api_key_auth_obj.end_user_id is None: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + raw_end_user_id = get_end_user_id_from_request_body( + request_data, _safe_get_request_headers(request) + ) + if raw_end_user_id is not None: + resolved_end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if resolved_end_user_id is not None: + user_api_key_auth_obj.end_user_id = resolved_end_user_id user_api_key_auth_obj.request_route = normalize_request_route(route) return user_api_key_auth_obj diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py new file mode 100644 index 00000000000..ab347130a30 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -0,0 +1,35 @@ +"""Rubrik guardrail integration for LiteLLM.""" + +from typing import TYPE_CHECKING + +from litellm.integrations.rubrik import RubrikLogger +from litellm.types.guardrails import SupportedGuardrailIntegrations + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> RubrikLogger: + import litellm + + rubrik_callback = RubrikLogger( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(rubrik_callback) + return rubrik_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: RubrikLogger, +} diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index 4f31c762df1..e32fee6afc5 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -178,6 +178,24 @@ async def _parse_ocr_request(request: Request) -> Dict[str, Any]: "For JSON requests, use 'document_url' or 'image_url' document types." ) + # Security: reject provider-native file IDs (e.g. reducto://) received via + # JSON. These IDs are not scoped to the LiteLLM proxy user/key, so an + # authenticated user who obtains another user's file ID could submit it + # here and receive the OCR result using the proxy's shared provider + # credentials. Force callers to upload fresh content per request via + # multipart/form-data or an inline base64 data URI, both of which produce + # a server-mediated upload bound to the current request. + if isinstance(doc, dict): + for url_field in ("document_url", "image_url"): + url_value = doc.get(url_field) + if isinstance(url_value, str) and url_value.startswith("reducto://"): + raise ValueError( + "reducto:// file IDs are not accepted through the proxy " + "OCR API; upload the file in the same request via " + "multipart/form-data with a 'file' field, or pass an " + "inline base64 data URI as the document URL." + ) + return data diff --git a/litellm/responses/sse_output_recovery.py b/litellm/responses/sse_output_recovery.py new file mode 100644 index 00000000000..5c18770a611 --- /dev/null +++ b/litellm/responses/sse_output_recovery.py @@ -0,0 +1,136 @@ +""" +Shared helpers for recovering Responses API output items from raw SSE chunks. + +The same recovery logic is needed in multiple places (e.g. the ChatGPT +Responses transformation and the LiteLLM Responses-to-Chat-Completions +bridge). Keep the implementation in a single module so a fix in one +caller automatically applies to all of them. +""" + +import json +from typing import Any, Dict, Optional + +from litellm.constants import STREAM_SSE_DONE_STRING + +_MAX_CONTENT_INDEX = 1024 + + +def parse_sse_json_chunk(chunk: str) -> Optional[Dict[str, Any]]: + """Parse a single raw SSE line into a JSON object dict. + + Returns ``None`` for empty lines, ``event:`` lines, ``[DONE]`` markers, + invalid JSON, or non-dict payloads. Centralizes the parsing step that + feeds into the recovery helpers in this module so behavior stays + consistent across all callers. + """ + # Import locally to avoid a circular import with the streaming handler. + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + stripped_chunk = ( + CustomStreamWrapper._strip_sse_data_from_chunk(chunk.strip()) or "" + ).strip() + if ( + not stripped_chunk + or stripped_chunk == STREAM_SSE_DONE_STRING + or stripped_chunk.startswith("event:") + ): + return None + try: + parsed_chunk = json.loads(stripped_chunk) + except json.JSONDecodeError: + return None + if not isinstance(parsed_chunk, dict): + return None + return parsed_chunk + + +def record_output_item_chunk( + parsed_chunk: Dict[str, Any], + output_items: Dict[int, Dict[str, Any]], +) -> None: + """Record an OUTPUT_ITEM_DONE chunk into ``output_items`` keyed by + ``output_index`` (falling back to the next free slot when missing). + """ + item = parsed_chunk.get("item") + if not isinstance(item, dict): + return + try: + output_index_raw = parsed_chunk.get("output_index") + if output_index_raw is None: + raise ValueError("missing output_index") + output_index = int(output_index_raw) + except (TypeError, ValueError): + output_index = len(output_items) + output_items[output_index] = item + + +def record_output_text_chunk( + parsed_chunk: Dict[str, Any], + output_items: Dict[int, Dict[str, Any]], + text_only_items: Dict[int, Dict[str, Any]], +) -> None: + """Record an OUTPUT_TEXT_DONE chunk as a synthetic message item in + ``text_only_items``. Real OUTPUT_ITEM_DONE events already captured in + ``output_items`` take precedence at the same ``output_index``. + """ + text = parsed_chunk.get("text") + if not isinstance(text, str): + return + + try: + output_index_raw = parsed_chunk.get("output_index") + if output_index_raw is None: + raise ValueError("missing output_index") + output_index = int(output_index_raw) + except (TypeError, ValueError): + output_index = len(text_only_items) + + if output_index in output_items: + return + + item = text_only_items.get(output_index) + if item is None: + item = { + "type": "message", + "id": parsed_chunk.get("item_id") or f"msg_{output_index}", + "role": "assistant", + "status": "completed", + "content": [], + } + text_only_items[output_index] = item + + content = item.setdefault("content", []) + if not isinstance(content, list): + return + + try: + content_index_raw = parsed_chunk.get("content_index") + if content_index_raw is None: + raise ValueError("missing content_index") + content_index = int(content_index_raw) + except (TypeError, ValueError): + content_index = len(content) + + if content_index < 0 or content_index > _MAX_CONTENT_INDEX: + return + + while len(content) <= content_index: + content.append( + { + "type": "output_text", + "text": "", + "annotations": [], + } + ) + + content_item = content[content_index] + if not isinstance(content_item, dict): + content_item = {} + content[content_index] = content_item + + content_item["type"] = "output_text" + content_item["text"] = text + if parsed_chunk.get("annotations") is not None: + content_item["annotations"] = parsed_chunk["annotations"] + else: + content_item.setdefault("annotations", []) diff --git a/litellm/router.py b/litellm/router.py index c968c819400..29025ad1437 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7778,6 +7778,38 @@ class Router: _shared_model_info = { k: v for k, v in _model_info.items() if k not in _custom_pricing_fields } + _existing_shared_mode = ( + cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {} + ).get("mode") + _deployment_mode = _shared_model_info.get("mode") + # Keep the built-in bridge mode stable for shared backend keys. + # Multiple aliases can point at the same provider/model backend, + # but their deployment-level overrides should not downgrade the + # backend from responses -> chat via last-write-wins registration. + # Only preserve in that specific direction so legitimate upgrades + # (e.g. chat -> responses) and unrelated mode changes still apply, + # and so a missing deployment mode does not silently clear the + # existing shared backend mode. + _is_responses_to_chat_downgrade = ( + _existing_shared_mode == "responses" and _deployment_mode == "chat" + ) + _would_clear_existing_mode = ( + _existing_shared_mode is not None and _deployment_mode is None + ) + if _is_responses_to_chat_downgrade or _would_clear_existing_mode: + if _deployment_mode is not None: + verbose_router_logger.warning( + "Router: preserving existing mode=%s for shared backend " + "key %s instead of the deployment-specified mode=%s " + "(prevents alias registration from downgrading the " + "shared backend mode).", + _existing_shared_mode, + _model_name, + _deployment_mode, + ) + _shared_model_info["mode"] = _existing_shared_mode + + # Always register the (possibly mode-preserved) shared backend info. _backend_alias_cost = {_model_name: _shared_model_info} if "responses/" in _model_name: _stripped_model_name = _model_name.replace("responses/", "") diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 751113400d3..0a51ce3d456 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -100,6 +100,7 @@ class SupportedGuardrailIntegrations(Enum): MCP_JWT_SIGNER = "mcp_jwt_signer" LLM_AS_A_JUDGE = "llm_as_a_judge" QOSTODIAN_NEXUS = "qostodian_nexus" + RUBRIK = "rubrik" class Role(Enum): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6084f14e2df..5082c73bf5c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -147,6 +147,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_low_reasoning_effort: Optional[bool] supports_xhigh_reasoning_effort: Optional[bool] supports_max_reasoning_effort: Optional[bool] + supports_output_config: Optional[bool] class SearchContextCostPerQuery(TypedDict, total=False): @@ -243,6 +244,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float ] # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) ocr_cost_per_page: Optional[float] # for OCR models + ocr_cost_per_credit: Optional[float] # for OCR models priced by credit annotation_cost_per_page: Optional[float] # for OCR models search_context_cost_per_query: Optional[ SearchContextCostPerQuery @@ -260,6 +262,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "chat", "audio_transcription", "responses", + "ocr", ] ] tpm: Optional[int] @@ -3219,6 +3222,7 @@ class LlmProviders(str, Enum): ANTHROPIC_TEXT = "anthropic_text" BYTEZ = "bytez" REPLICATE = "replicate" + REDUCTO = "reducto" RUNWAYML = "runwayml" AWS_POLLY = "aws_polly" HUGGINGFACE = "huggingface" diff --git a/litellm/utils.py b/litellm/utils.py index 001c89fee4c..2487d39bd0d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5387,6 +5387,16 @@ def _strip_model_name(model: str, custom_llm_provider: Optional[str]) -> str: # Global case-insensitive lookup map for model_cost (built eagerly at module import) _model_cost_lowercase_map: Optional[Dict[str, str]] = None +# Monotonic counter bumped on every model_cost mutation. Consumers that +# memoize derived state (e.g. provider-specific indices) can include this +# value in their cache key so they invalidate even when key add+remove or +# in-place value replacement leaves len/id unchanged. +_model_cost_mutation_generation: int = 0 + + +def get_model_cost_mutation_generation() -> int: + return _model_cost_mutation_generation + def _invalidate_model_cost_lowercase_map() -> None: """Invalidate the case-insensitive lookup map for model_cost. @@ -5394,8 +5404,9 @@ def _invalidate_model_cost_lowercase_map() -> None: Call this whenever litellm.model_cost is modified to ensure the map is rebuilt. Also clears related LRU caches that depend on model_cost data. """ - global _model_cost_lowercase_map + global _model_cost_lowercase_map, _model_cost_mutation_generation _model_cost_lowercase_map = None + _model_cost_mutation_generation += 1 # Clear LRU caches that depend on model_cost data get_model_info.cache_clear() @@ -5986,6 +5997,7 @@ def _get_model_info_helper( # noqa: PLR0915 tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), + ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), annotation_cost_per_page=_model_info.get( "annotation_cost_per_page", None ), @@ -9241,6 +9253,18 @@ class ProviderConfigManager: return get_vertex_ai_ocr_config(model=model) + if provider == litellm.LlmProviders.REDUCTO: + from litellm.llms.reducto.ocr.transformation import ( + ReductoParseLegacyConfig, + ReductoParseV3Config, + ) + + if model == "parse-v3": + return ReductoParseV3Config() + if model == "parse-legacy": + return ReductoParseLegacyConfig() + return None + MistralOCRConfig = getattr(sys.modules[__name__], "MistralOCRConfig") PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bda94e4768f..31a5993a240 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1011,6 +1011,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1041,6 +1042,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1071,6 +1073,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1100,6 +1103,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1129,6 +1133,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1328,6 +1333,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "global.anthropic.claude-sonnet-4-6": { @@ -1358,6 +1364,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "us.anthropic.claude-sonnet-4-6": { @@ -1388,6 +1395,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "eu.anthropic.claude-sonnet-4-6": { @@ -1417,6 +1425,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "au.anthropic.claude-sonnet-4-6": { @@ -1446,6 +1455,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "jp.anthropic.claude-sonnet-4-6": { @@ -1475,6 +1485,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "anthropic.claude-sonnet-4-20250514-v1:0": { @@ -1996,6 +2007,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 159, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -2093,6 +2105,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "azure/computer-use-preview": { @@ -9643,6 +9656,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "claude-sonnet-4-5-20250929-v1:0": { @@ -9840,6 +9854,7 @@ "us": 1.1, "fast": 6.0 }, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -9875,7 +9890,8 @@ "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -9910,7 +9926,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-opus-4-7-20260416": { "cache_creation_input_token_cost": 6.25e-06, @@ -9945,7 +9962,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", @@ -13982,6 +14000,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "fireworks_ai", @@ -14248,6 +14281,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 6e-07, @@ -28937,14 +28985,16 @@ "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-7": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", @@ -29158,6 +29208,24 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "reducto/parse-legacy": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "reducto/parse-v3": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "recraft/recraftv2": { "litellm_provider": "recraft", "mode": "image_generation", @@ -33337,6 +33405,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -33365,6 +33434,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -33478,6 +33548,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "vertex_ai/claude-sonnet-4-5@20250929": { @@ -40590,6 +40661,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "duckduckgo/search": { diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1d577213a1b..388752b032e 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1904,6 +1904,23 @@ "rerank": false } }, + "reducto": { + "display_name": "Reducto (`reducto`)", + "url": "https://docs.litellm.ai/docs/providers/reducto", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "ocr": true + } + }, "replicate": { "display_name": "Replicate (`replicate`)", "url": "https://docs.litellm.ai/docs/providers/replicate", diff --git a/pyproject.toml b/pyproject.toml index b7bae873a46..b4eb15dc38f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,8 +33,9 @@ Homepage = "https://litellm.ai" Repository = "https://github.com/BerriAI/litellm" Documentation = "https://docs.litellm.ai" -# Dependencies pinned from the published `litellm[proxy]==1.83.0` resolution. -# Docker and CI should prefer `uv.lock` rather than maintaining parallel installers. +# Optional extras retain exact pins because they are consumed by Docker images +# where exact reproducibility matters. The core SDK uses ranges so downstream +# consumers can coexist with other packages without forced downgrades. [project.optional-dependencies] proxy = [ "gunicorn==23.0.0", @@ -318,3 +319,4 @@ pytest_add_cli_args = [ [tool.coverage.run] source = ["litellm"] relative_files = true + diff --git a/tests/llm_translation/realtime/base_realtime_tests.py b/tests/llm_translation/realtime/base_realtime_tests.py index f1c42659007..1a2c6ff6a9c 100644 --- a/tests/llm_translation/realtime/base_realtime_tests.py +++ b/tests/llm_translation/realtime/base_realtime_tests.py @@ -10,7 +10,7 @@ import json import os import sys from abc import ABC, abstractmethod -from typing import Optional +from typing import Optional, Tuple, Union import pytest import websockets @@ -153,8 +153,14 @@ class BaseRealtimeTest(ABC): pass @abstractmethod - def get_initial_event_type(self) -> str: - """Return the expected initial event type (e.g., 'session.created' or 'conversation.created')""" + def get_initial_event_type(self) -> Union[str, Tuple[str, ...]]: + """Return the expected initial event type(s). + + May return a single event type (e.g. ``'session.created'``) or a tuple + of acceptable types when the upstream provider can legitimately emit + more than one initial event (e.g. xAI's Grok Voice Agent has shipped + both ``conversation.created`` and ``session.created``). + """ pass def get_skip_reason(self) -> str: @@ -229,9 +235,14 @@ class BaseRealtimeTest(ABC): # Verify initial event initial_event = websocket_client.messages_received[0] + expected_event_type = self.get_initial_event_type() + if isinstance(expected_event_type, str): + allowed_event_types: Tuple[str, ...] = (expected_event_type,) + else: + allowed_event_types = tuple(expected_event_type) assert ( - initial_event["type"] == self.get_initial_event_type() - ), f"Expected {self.get_initial_event_type()}, got {initial_event.get('type')}" + initial_event["type"] in allowed_event_types + ), f"Expected one of {allowed_event_types}, got {initial_event.get('type')}" @pytest.mark.asyncio async def test_realtime_with_query_params(self): diff --git a/tests/llm_translation/realtime/test_xai_realtime.py b/tests/llm_translation/realtime/test_xai_realtime.py index 86d0ebe3a3c..8ffcb3db30d 100644 --- a/tests/llm_translation/realtime/test_xai_realtime.py +++ b/tests/llm_translation/realtime/test_xai_realtime.py @@ -7,6 +7,7 @@ Uses the base test class to ensure consistent behavior across providers. import os import sys +from typing import Tuple import pytest @@ -20,9 +21,11 @@ class TestXAIRealtime(BaseRealtimeTest): E2E tests for xAI Realtime API. xAI's Grok Voice Agent API is OpenAI-compatible: - - Initial event: "session.created" (matches OpenAI) - - Different endpoint: wss://api.x.ai/v1/realtime + - Endpoint: wss://api.x.ai/v1/realtime - Model: grok-4-1-fast-non-reasoning + - Initial event: historically "conversation.created"; xAI has since shipped + "session.created" (matching OpenAI). Accept either to avoid spurious + failures whenever xAI flips the wire format. """ def get_model(self) -> str: @@ -31,5 +34,5 @@ class TestXAIRealtime(BaseRealtimeTest): def get_api_key_env_var(self) -> str: return "XAI_API_KEY" - def get_initial_event_type(self) -> str: - return "session.created" + def get_initial_event_type(self) -> Tuple[str, ...]: + return ("conversation.created", "session.created") diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 7f13aaf67c6..fd57b65f593 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -4225,5 +4225,9 @@ def test_gemini_google_maps_tool_simple(): assert response.choices[0].message.content is not None except (litellm.RateLimitError, litellm.InternalServerError): pass + except litellm.InternalServerError: + pytest.skip( + "Google Maps Platform returned a transient 500 (upstream flake); skipping." + ) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/proxy_unit_tests/test_reducto_ocr_route.py b/tests/proxy_unit_tests/test_reducto_ocr_route.py new file mode 100644 index 00000000000..dc658a74ee8 --- /dev/null +++ b/tests/proxy_unit_tests/test_reducto_ocr_route.py @@ -0,0 +1,137 @@ +import asyncio +import os +from unittest.mock import AsyncMock, patch + +import litellm +import pytest +from fastapi.testclient import TestClient + +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo +from litellm.proxy.proxy_server import app, initialize + + +@pytest.fixture(scope="function") +def fake_env_vars(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") + monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base") + monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base") + monkeypatch.setenv("AZURE_AI_API_KEY", "fake_azure_api_key") + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key") + monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base") + monkeypatch.setenv("AZURE_SWEDEN_API_KEY", "fake_azure_sweden_api_key") + monkeypatch.setenv("REDIS_HOST", "localhost") + + +@pytest.fixture(scope="function") +def client_no_auth(fake_env_vars): + from litellm.proxy.proxy_server import cleanup_router_config_variables + + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + cleanup_router_config_variables() + + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = os.path.join(filepath, "test_configs", "test_config_no_auth.yaml") + asyncio.run(initialize(config=config_fp, debug=True)) + + # Passthrough of api_base in the JSON body is rejected by default + # (pre_db_read_auth_checks / is_request_body_safe). This test asserts + # api_base reaches aocr(). + from litellm.proxy import proxy_server as _ps + + if _ps.general_settings is None: + _ps.general_settings = {} + _ps.general_settings["allow_client_side_credentials"] = True + + try: + yield TestClient(app) + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +def test_proxy_reducto_ocr_json_rejects_reducto_id(client_no_auth): + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "document_url", + "document_url": "reducto://proxy.pdf", + }, + "api_key": "proxy-key", + "api_base": "https://platform.reducto.ai", + }, + ) + + assert response.status_code >= 400 + assert "reducto://" in response.text + assert mock_aocr.await_count == 0 + + +def test_proxy_reducto_ocr_json_rejects_reducto_id_in_image_url(client_no_auth): + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "image_url", + "image_url": "reducto://proxy.png", + }, + }, + ) + + assert response.status_code >= 400 + assert "reducto://" in response.text + assert mock_aocr.await_count == 0 + + +def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth): + mocked_response = OCRResponse( + pages=[OCRPage(index=0, markdown="Proxy OCR")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=1, credits=1), + ) + + data_uri = "data:application/pdf;base64,JVBERi0xLjQK" + + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(return_value=mocked_response), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "document_url", + "document_url": data_uri, + }, + "api_key": "proxy-key", + "api_base": "https://platform.reducto.ai", + }, + ) + + assert response.status_code == 200 + assert mock_aocr.await_count == 1 + assert mock_aocr.await_args.kwargs["model"] == "reducto/parse-v3" + assert mock_aocr.await_args.kwargs["document"] == { + "type": "document_url", + "document_url": data_uri, + } + assert mock_aocr.await_args.kwargs["api_key"] == "proxy-key" + assert mock_aocr.await_args.kwargs["api_base"] == "https://platform.reducto.ai" + + response_body = response.json() + assert response_body["object"] == "ocr" + assert response_body["usage_info"]["credits"] == 1 + assert response_body["pages"][0]["markdown"] == "Proxy OCR" diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 697a9ebc720..d335c359aa0 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -508,6 +508,308 @@ and I learn to carry this small calm home.""" print("✓ transform_response correctly handled reasoning items and output messages") +def _make_empty_responses_api_response(model: str = "gpt-5.4"): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_from_stream", + created_at=1760144904, + error=None, + incomplete_details=None, + instructions=None, + metadata={}, + model=model, + object="response", + output=[], + parallel_tool_calls=True, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + max_output_tokens=None, + previous_response_id=None, + reasoning={"effort": "low", "summary": "detailed"}, + status="completed", + text={"format": {"type": "text"}, "verbosity": "medium"}, + truncation="disabled", + usage=ResponseAPIUsage( + input_tokens=1, + input_tokens_details=None, + output_tokens=1, + output_tokens_details=None, + total_tokens=2, + cost=None, + ), + user=None, + store=True, + background=False, + billing={"payer": "developer"}, + max_tool_calls=None, + prompt_cache_key=None, + safety_identifier=None, + service_tier="default", + top_logprobs=0, + ) + + +def _make_empty_model_response(): + from litellm.types.utils import ModelResponse, Usage + + return ModelResponse( + id="chatcmpl-test-recovered", + created=1760144904, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + +def test_transform_response_recovers_empty_output_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Recovered from SSE"}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from SSE" + + +def test_transform_response_recovers_output_item_done_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Recovered from output item","annotations":[]}]}}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from output item" + + +def test_transform_response_recovers_output_item_done_from_whitespace_padded_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + output_item_event = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "message", + "id": "msg_from_item", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Recovered from padded output item", + "annotations": [], + } + ], + }, + } + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_from_stream", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "gpt-5.4", + "output": [], + }, + } + raw_sse = "\n".join( + [ + f" data: {json.dumps(output_item_event)} ", + f"\tdata: {json.dumps(completed_event)}", + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from padded output item" + + +def test_transform_response_preserves_output_item_when_text_done_arrives_later(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Complete output item text","annotations":[]}]}}', + 'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Late text event"}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Complete output item text" + + +def test_recover_output_items_merges_text_only_items_at_distinct_indices(): + """When OUTPUT_ITEM_DONE covers some indices and OUTPUT_TEXT_DONE covers + others, both must be preserved instead of treating them as mutually + exclusive fallbacks.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_item_0","role":"assistant","status":"completed","content":[{"type":"output_text","text":"From OUTPUT_ITEM_DONE","annotations":[]}]}}', + 'data: {"type":"response.output_text.done","output_index":1,"content_index":0,"item_id":"msg_text_1","text":"From OUTPUT_TEXT_DONE only"}', + "data: [DONE]", + "", + ] + ) + + recovered = ( + LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse( + raw_sse + ) + ) + + assert len(recovered) == 2 + assert recovered[0]["id"] == "msg_item_0" + assert recovered[0]["content"][0]["text"] == "From OUTPUT_ITEM_DONE" + assert recovered[1]["id"] == "msg_text_1" + assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only" + + +def test_transform_response_prefers_completed_output_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Earlier stream text","annotations":[]}]}}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[{"type":"message","id":"msg_from_completed","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Authoritative completed text","annotations":[]}]}]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Authoritative completed text" + + def test_convert_tools_to_responses_format(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/test_litellm/integrations/rubrik_test_helpers.py new file mode 100644 index 00000000000..1bdb8cb247b --- /dev/null +++ b/tests/test_litellm/integrations/rubrik_test_helpers.py @@ -0,0 +1,23 @@ +"""Shared helpers for Rubrik plugin tests.""" + +from typing import Any, Dict + +from litellm.types.utils import GenericGuardrailAPIInputs + + +def make_tool_call_dict( + tc_id: str, name: str, arguments: str = "{}" +) -> Dict[str, Any]: + """Create a tool call dict matching the ChatCompletionMessageToolCall schema.""" + return { + "id": tc_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +def make_inputs_with_tools( + tool_calls: list, texts: list | None = None +) -> GenericGuardrailAPIInputs: + """Create GenericGuardrailAPIInputs with tool_calls.""" + return GenericGuardrailAPIInputs(texts=texts or [], tool_calls=tool_calls) diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py new file mode 100644 index 00000000000..922d2fe8a15 --- /dev/null +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -0,0 +1,1012 @@ +""" +Tests for the Rubrik LiteLLM plugin. + +Covers initialization, apply_guardrail tool blocking (all allowed, all blocked, +partial blocking, fail-open), batch logging, and Anthropic format handling. +""" + +import os +from typing import Any, Dict +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest + +from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.rubrik import RubrikLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from tests.test_litellm.integrations.rubrik_test_helpers import ( + make_inputs_with_tools, + make_tool_call_dict, +) + + +@pytest.fixture +def mock_env(): + """Set up environment variables for testing.""" + with patch.dict( + os.environ, + { + "RUBRIK_WEBHOOK_URL": "http://localhost:8080", + "RUBRIK_API_KEY": "test-api-key", + }, + ): + yield + + +@pytest.fixture +def handler(mock_env): + """Create a RubrikLogger instance for testing.""" + with patch("asyncio.create_task", Mock()): + return RubrikLogger() + + +# -- Initialization ----------------------------------------------------------- + + +class TestInitialization: + def test_init_success(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + assert ( + handler.tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + assert handler.logging_endpoint == "http://localhost:8080/v1/litellm/batch" + assert handler.key == "test-api-key" + assert isinstance(handler.tool_blocking_client, AsyncHTTPHandler) + + def test_init_with_constructor_params(self): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(api_key="ctor-key", api_base="http://ctor-host:9090") + assert handler.key == "ctor-key" + assert ( + handler.tool_blocking_endpoint + == "http://ctor-host:9090/v1/after_completion/openai/v1" + ) + + def test_init_without_url(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="Rubrik webhook URL not configured"): + RubrikLogger() + + def test_init_without_api_key(self): + with patch.dict( + os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080"}, clear=True + ): + with patch("asyncio.create_task", Mock()): + assert RubrikLogger().key is None + + def test_trailing_slash_removed(self): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): + with patch("asyncio.create_task", Mock()): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + + def test_v1_suffix_stripped_as_substring_not_charset(self): + with patch("asyncio.create_task", Mock()): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v1/after_completion/openai/v1" + ) + + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v11/v1/after_completion/openai/v1" + ) + + def test_sampling_rate_fractional(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.5 + + def test_sampling_rate_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "abc"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + + def test_sampling_rate_clamped(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "2.0"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "-0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.0 + + def test_batch_size_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "abc"}, + ): + # Should use default without crashing + assert isinstance(RubrikLogger().batch_size, int) + + def test_batch_size_valid(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "256"}, + ): + assert RubrikLogger().batch_size == 256 + + def test_init_outside_event_loop_does_not_raise(self): + """Instantiation without a running event loop must not raise RuntimeError.""" + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://localhost:8080", "RUBRIK_API_KEY": "k"}, + ): + # Do NOT patch asyncio.create_task — the real call should be + # guarded and fall back gracefully when there is no event loop. + handler = RubrikLogger() + assert handler.tool_blocking_endpoint.startswith("http://localhost:8080") + # Without a running loop at init, the periodic flush task should be + # deferred so batches still get drained once a log event arrives. + assert handler._flush_task is None + + @pytest.mark.asyncio + async def test_periodic_flush_task_started_lazily_on_first_log(self, mock_env): + """Loggers instantiated outside an event loop must still start the + periodic flush task on first use to drain low-traffic batches.""" + # Simulate sync-init by hiding the running loop from the constructor. + with patch( + "litellm.integrations.rubrik.asyncio.get_running_loop", + side_effect=RuntimeError("no running loop"), + ): + handler = RubrikLogger() + assert handler._flush_task is None + + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "id": "litellm-id", + }, + "litellm_call_id": "litellm-id", + "litellm_params": {}, + } + with patch.object(handler, "_log_batch_to_rubrik", AsyncMock()): + await handler.async_log_success_event(kwargs, None, None, None) + + assert handler._flush_task is not None + handler._flush_task.cancel() + + def test_event_hook_defaults_to_post_call_when_none_passed(self, mock_env): + """`initialize_guardrail` always passes ``event_hook=litellm_params.mode`` + (which is ``None`` when the user omits ``mode``). The logger must coerce + a None ``event_hook`` to ``post_call`` rather than leaving it as None, + which would otherwise cause the guardrail to run on every event hook.""" + from litellm.types.guardrails import GuardrailEventHooks + + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(event_hook=None) + assert handler.event_hook == GuardrailEventHooks.post_call + + def test_explicit_event_hook_preserved(self, mock_env): + from litellm.types.guardrails import GuardrailEventHooks + + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) + assert handler.event_hook == GuardrailEventHooks.pre_call + + def test_default_on_defaults_to_true_when_none_passed(self, mock_env): + """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` + (which is ``None`` when the user omits ``default_on``). The logger must + coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` + (which checks ``self.default_on is True``) silently skips the guardrail.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=None) + assert handler.default_on is True + + def test_explicit_default_on_false_preserved(self, mock_env): + """A user explicitly setting ``default_on: false`` in their guardrail + config must NOT be silently overridden to True.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=False) + assert handler.default_on is False + + def test_explicit_default_on_true_preserved(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=True) + assert handler.default_on is True + + def test_headers_with_api_key(self, handler): + assert handler._headers["Authorization"] == "Bearer test-api-key" + assert handler._headers["Content-Type"] == "application/json" + + def test_headers_without_api_key(self): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True): + with patch("asyncio.create_task", Mock()): + h = RubrikLogger() + assert "Authorization" not in h._headers + + +# -- Batch Logging ------------------------------------------------------------ + + +@pytest.mark.asyncio +class TestBatchLogging: + async def test_log_success_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_failure_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "error", + }, + } + await handler.async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_success_event_sampling_skips(self, handler): + handler.sampling_rate = 0.0 + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 0 + + async def test_flush_queue_sends_batch(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + mock_response = Mock() + mock_response.status_code = 200 + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + handler.async_httpx_client.post.assert_called_once() + assert len(handler.log_queue) == 0 + + async def test_flush_queue_preserves_events_added_during_send(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.raise_for_status = Mock() + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.flush_queue() + + assert handler.log_queue == [{"msg": "c"}] + + async def test_async_send_batch_does_not_drain_events(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.raise_for_status = Mock() + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.async_send_batch() + + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] + + async def test_log_batch_error_does_not_crash_and_preserves_events(self, handler): + """A failed batch send must not crash the caller AND must preserve the + original events in the queue so they can be retried on the next flush. + Previously the events were silently dropped on HTTP 5xx / network errors. + """ + handler.log_queue = [{"msg": "a"}] + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}] + + async def test_log_batch_network_error_preserves_events(self, handler): + """Network/timeout errors must also preserve the in-flight events.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock( + side_effect=httpx.TimeoutException("timeout") + ) + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}] + + async def test_enqueue_drops_oldest_when_queue_exceeds_max_size(self, handler): + """A sustained Rubrik webhook outage must not let the in-memory retry + queue grow without bound. Once max_queue_size is exceeded, the oldest + events are dropped to make room for new ones.""" + handler.max_queue_size = 3 + handler.batch_size = 10**6 # disable size-triggered flush + handler.flush_queue = AsyncMock() + for i in range(5): + await handler._enqueue_log_event( + kwargs={ + "standard_logging_object": { + "messages": [{"role": "user", "content": f"hi-{i}"}], + "response": "hello", + }, + }, + event_type="success", + ) + assert len(handler.log_queue) == 3 + retained = [item["messages"][0]["content"] for item in handler.log_queue] + assert retained == ["hi-2", "hi-3", "hi-4"] + + async def test_log_batch_failure_preserves_events_added_during_send(self, handler): + """Failure must preserve both the snapshot AND events appended mid-flush.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "boom" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] + + async def test_system_prompt_prepended_to_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "system": "You are a helpful assistant.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert msgs[0]["role"] == "system" + assert msgs[0]["content"] == "You are a helpful assistant." + + async def test_system_prompt_with_dict_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": {"role": "user", "content": "hi"}, + "response": "hello", + }, + "system": "Be concise.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert isinstance(msgs, list) + assert msgs[0]["role"] == "system" + assert msgs[1] == {"role": "user", "content": "hi"} + + async def test_anthropic_id_normalization(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/messages", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "litellm-call-123" + + async def test_non_anthropic_id_unchanged(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/chat/completions", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "chatcmpl-original" + + async def test_payload_deep_copied_not_mutated(self, handler): + """Verify the shared standard_logging_object is not mutated.""" + original_payload = { + "id": "original-id", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + } + kwargs = { + "standard_logging_object": original_payload, + "system": "System prompt.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + # Original payload should NOT have been mutated + assert original_payload["id"] == "original-id" + assert len(original_payload["messages"]) == 1 + + +# -- Tool Blocking (apply_guardrail) ------------------------------------------ + + +def _mock_service_response(response_json): + """Create a mock tool blocking client that returns the given JSON.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = response_json + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +def _echo_service(): + """Create a mock tool blocking client that echoes the payload back.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = kwargs.get("json", {}).get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +@pytest.mark.asyncio +class TestApplyGuardrail: + async def test_skips_requests(self, handler): + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "test_tool")]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_no_tool_calls(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_allowed(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather") + tc2 = make_tool_call_dict("call_2", "get_time") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_blocked(self, handler): + tc1 = make_tool_call_dict("call_1", "delete_table") + tc2 = make_tool_call_dict("call_2", "drop_database") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "Tool blocked by policy", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert "Tool blocked by policy" in exc_info.value.message + + async def test_partial_blocking(self, handler): + tc_blocked = make_tool_call_dict("call_A", "blocked_tool") + tc_allowed = make_tool_call_dict("call_B", "allowed_tool") + inputs = make_inputs_with_tools([tc_blocked, tc_allowed]) + + async def mock_post(*_args, **kwargs): + payload = kwargs.get("json", {}).get("response", {}) + all_tcs = payload["choices"][0]["message"]["tool_calls"] + allowed = [tc for tc in all_tcs if tc.get("id") == "call_B"] + mock_resp = Mock() + mock_resp.json.return_value = { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": allowed, + } + } + ], + } + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_service_failure_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_service_empty_choices_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + handler.tool_blocking_client = _mock_service_response({"choices": []}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_blocking_service_payload_format(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather", '{"location": "SF"}') + tc2 = make_tool_call_dict("call_2", "send_email", '{"to": "user@example.com"}') + inputs = make_inputs_with_tools([tc1, tc2]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + # Verify envelope structure + assert "request" in captured_payload + assert "response" in captured_payload + + response_data = captured_payload["response"] + message = response_data["choices"][0]["message"] + assert message["role"] == "assistant" + assert len(message["tool_calls"]) == 2 + assert message["tool_calls"][0]["id"] == "call_1" + assert message["tool_calls"][0]["function"]["name"] == "get_weather" + assert message["tool_calls"][1]["id"] == "call_2" + assert message["tool_calls"][1]["function"]["name"] == "send_email" + + async def test_request_data_included_in_envelope(self, handler): + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + logging_obj = Mock() + logging_obj.model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "litellm_params": { + "proxy_server_request": {"url": "/chat/completions"}, + }, + } + + await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + req = captured_payload["request"] + assert req["model"] == "gpt-4" + assert req["messages"] == [{"role": "user", "content": "hi"}] + + async def test_proxy_server_request_headers_stripped(self, handler): + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + logging_obj = Mock() + logging_obj.model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "litellm_params": { + "proxy_server_request": { + "url": "/chat/completions", + "method": "POST", + "headers": { + "authorization": "Bearer sk-litellm-secret", + "cookie": "session=abc", + "x-api-key": "leaked-key", + }, + "body": {"api_key": "sk-upstream-secret"}, + }, + }, + } + + await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + forwarded = captured_payload["request"]["proxy_server_request"] + assert forwarded == {"url": "/chat/completions", "method": "POST"} + + +# -- Anthropic format ---------------------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailAnthropicFormat: + """Verify blocking works correctly regardless of original provider format. + + The framework converts Anthropic tool_use blocks to OpenAI-format + tool_calls before calling apply_guardrail. + """ + + async def test_single_tool_allowed(self, handler): + tc = make_tool_call_dict( + "toolu_123", "get_weather", '{"location": "Portland, OR"}' + ) + inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_single_tool_blocked(self, handler): + tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') + inputs = make_inputs_with_tools([tc]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_text_only_response_no_blocking(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["Hello! I'm Claude."]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock() + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + assert result is inputs + mock_client.post.assert_not_called() + + async def test_service_failure_preserves_tools(self, handler): + tc = make_tool_call_dict("toolu_123", "get_weather", '{"location": "SF"}') + inputs = make_inputs_with_tools([tc]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +# -- Normalize tool calls ------------------------------------------------------ + + +class TestNormalizeToolCalls: + def test_dict_input(self): + tc = make_tool_call_dict("call_1", "test", '{"a": 1}') + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_1" + assert result[0].function.name == "test" + assert result[0].function.arguments == '{"a": 1}' + + def test_typed_object_input(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn", arguments="{}"), + ) + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_2" + assert result[0].function.name == "fn" + + def test_unsupported_type_raises(self): + with pytest.raises(TypeError, match="Cannot normalize"): + RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) + + +# -- Extract blocked tools ----------------------------------------------------- + + +class TestExtractBlockedTools: + def test_all_allowed_returns_none(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + assert result is None + + def test_some_blocked_returns_explanation(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc1 = ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="fn1", arguments="{}"), + ) + tc2 = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn2", arguments="{}"), + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "blocked fn2", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + assert result is not None + assert "blocked fn2" in result + + def test_empty_choices_raises(self): + with pytest.raises(Exception, match="empty response"): + RubrikLogger._extract_blocked_tools({"choices": []}, []) + + def test_null_tool_calls_treated_as_all_blocked(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": None, + "content": "blocked everything", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + assert result is not None + assert "blocked everything" in result + + def test_duplicate_ids_block_when_only_one_returned(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc1 = ChatCompletionMessageToolCall( + id="call_dup", + type="function", + function=Function(name="fn", arguments="{}"), + ) + tc2 = ChatCompletionMessageToolCall( + id="call_dup", + type="function", + function=Function(name="fn", arguments="{}"), + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_dup"}], + "content": "blocked duplicate", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + assert result is not None + assert "blocked duplicate" in result + + +# -- Sanitize proxy server request ------------------------------------------- + + +class TestSanitizeProxyServerRequest: + def test_drops_headers_and_body(self): + proxy_request = { + "url": "/chat/completions", + "method": "POST", + "headers": { + "authorization": "Bearer sk-litellm-secret", + "cookie": "session=abc", + "content-type": "application/json", + }, + "body": {"api_key": "sk-upstream-secret", "model": "gpt-4"}, + } + result = RubrikLogger._sanitize_proxy_server_request(proxy_request) + assert result == {"url": "/chat/completions", "method": "POST"} + + def test_none_passthrough(self): + assert RubrikLogger._sanitize_proxy_server_request(None) is None + + def test_non_dict_passthrough(self): + assert RubrikLogger._sanitize_proxy_server_request("not a dict") == "not a dict" + + def test_partial_dict(self): + result = RubrikLogger._sanitize_proxy_server_request({"url": "/v1/messages"}) + assert result == {"url": "/v1/messages"} + + +# -- Resolve model ------------------------------------------------------------- + + +class TestResolveModel: + def test_model_from_response(self): + from unittest.mock import Mock + + response = Mock() + response.model = "gpt-4" + result = RubrikLogger._resolve_model({"response": response}, {}) + assert result == "gpt-4" + + def test_model_from_call_details(self): + result = RubrikLogger._resolve_model({}, {"model": "claude-3"}) + assert result == "claude-3" + + def test_fallback_to_unknown(self): + result = RubrikLogger._resolve_model({}, {}) + assert result == "unknown" + + def test_empty_model_on_response_returns_unknown(self): + from unittest.mock import Mock + + response = Mock() + response.model = "" + result = RubrikLogger._resolve_model( + {"response": response}, {"model": "fallback"} + ) + assert result == "unknown" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c6961477a58..07ab29c5231 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2078,6 +2078,146 @@ async def test_async_success_handler_preserves_response_cost_for_pass_through_en assert slo["response_cost"] > 0 +def test_process_hidden_params_recalculates_cost_after_failure_handler_zero(): + """ + Regression: PR #21844 preserved response_cost=0 set by failure_handler on failed + router retry attempts, so a later successful response with usage logged $0 spend. + """ + from datetime import datetime + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test-retry-zero-cost", + function_id="test-retry-zero-cost", + ) + logging_obj.model_call_details["litellm_params"] = {"model": "openai/gpt-4o-mini"} + logging_obj.optional_params = {} + + err = litellm.RateLimitError( + message="rate limit", + llm_provider="openai", + model="openai/gpt-4o-mini", + ) + for _ in range(2): + logging_obj._failure_handler_helper_fn( + exception=err, + traceback_exception="", + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert logging_obj.model_call_details.get("response_cost") == 0 + + result = ModelResponse( + id="success", + choices=[{"message": {"role": "assistant", "content": "ok"}}], + usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728), + ) + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + cost = logging_obj.model_call_details.get("response_cost") + assert cost is not None and cost > 0 + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost", 0) > 0 + + +def test_process_hidden_params_preserves_zero_cost_in_hidden_params(): + """Pass-through handlers often set response_cost on result._hidden_params (including 0).""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="gemini-2.5-flash-lite", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-hidden-zero-cost", + function_id="test-hidden-zero-cost", + ) + logging_obj.model_call_details["litellm_params"] = { + "model": "gemini-2.5-flash-lite" + } + logging_obj.optional_params = {} + + result = ModelResponse( + id="batch-pending", + choices=[{"message": {"role": "assistant", "content": "pending"}}], + usage=Usage(prompt_tokens=100, completion_tokens=10, total_tokens=110), + ) + result._hidden_params = {"response_cost": 0.0} + + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + assert logging_obj.model_call_details.get("response_cost") == 0.0 + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost") == 0.0 + + +def test_process_hidden_params_uses_hidden_params_cost_after_failure_handler_zero(): + """After retry failures pin model_call_details to 0, success cost on _hidden_params wins.""" + from datetime import datetime + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test-retry-hidden-cost", + function_id="test-retry-hidden-cost", + ) + logging_obj.model_call_details["litellm_params"] = {"model": "openai/gpt-4o-mini"} + logging_obj.optional_params = {} + + err = litellm.RateLimitError( + message="rate limit", + llm_provider="openai", + model="openai/gpt-4o-mini", + ) + for _ in range(2): + logging_obj._failure_handler_helper_fn( + exception=err, + traceback_exception="", + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert logging_obj.model_call_details.get("response_cost") == 0 + + passthrough_cost = 0.00042 + result = ModelResponse( + id="success", + choices=[{"message": {"role": "assistant", "content": "ok"}}], + usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728), + ) + result._hidden_params = {"response_cost": passthrough_cost} + + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + assert logging_obj.model_call_details.get("response_cost") == passthrough_cost + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost") == passthrough_cost + + def test_function_setup_litellm_metadata_populates_metadata(): """ Test that function_setup() properly handles litellm_metadata (used by /v1/messages, diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 4495e3f4101..b2e254901f4 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from unittest.mock import patch import pytest @@ -429,6 +430,31 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): assert result["max_tokens"] == 100 +def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider(): + config = AmazonAnthropicClaudeConfig() + messages = [{"role": "user", "content": "test"}] + optional_params = {"max_tokens": 100, "output_config": {"effort": "high"}} + + with patch( + "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ) as mock_supports_factory: + result = config.transform_request( + model="us.anthropic.claude-opus-4-7", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + mock_supports_factory.assert_called_once_with( + model="us.anthropic.claude-opus-4-7", + custom_llm_provider="bedrock", + key="supports_output_config", + ) + assert result["output_config"] == {"effort": "high"} + + def test_output_format_removed_from_bedrock_invoke_request(): """ Test that output_format parameter is removed from Bedrock Invoke requests. diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 9ecdad1fcff..2e315a535f0 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -592,8 +592,15 @@ def test_remove_scope_from_cache_control(): assert request["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" -def test_bedrock_messages_forwards_output_config(): - """Bedrock Invoke /v1/messages forwards ``output_config`` for adaptive Claude models.""" +def test_bedrock_messages_strips_output_config(): + """ + Ensure output_config is stripped from the request for models that do not + support it. + + Regression test for: https://github.com/BerriAI/litellm/issues/22797 + """ + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -605,21 +612,129 @@ def test_bedrock_messages_forwards_output_config(): }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=False, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "output_config" not in result + ), "output_config should be stripped for models that don't support it" + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_preserves_output_config_for_claude_4_6(): + """ + Ensure output_config is preserved for models that support it on Bedrock Invoke. + """ + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-6-v1", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "output_config" in result + ), "output_config should be preserved for supported models" + assert result["output_config"] == {"effort": "high"} + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ) as mock_supports_factory: + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + mock_supports_factory.assert_called_with( + model="us.anthropic.claude-opus-4-7", + custom_llm_provider="bedrock", + key="supports_output_config", ) + assert result["output_config"] == {"effort": "high"} + + +def test_bedrock_messages_forwards_output_config(): + """Bedrock Invoke /v1/messages forwards ``output_config`` for supported models.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert result.get("output_config") == {"effort": "high"} - # Other params should be preserved assert result.get("max_tokens") == 4096 def test_bedrock_messages_forwards_output_config_with_output_format(): """``output_config`` is forwarded; ``output_format`` is converted to inline schema.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -636,39 +751,60 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert result.get("output_config") == {"effort": "low"} assert "output_format" not in result -def test_bedrock_messages_forwards_output_config_for_non_adaptive_model(): - """``output_config`` is forwarded for non-adaptive models so the provider's error surfaces.""" +def test_bedrock_messages_strips_output_config_with_output_format(): + """ + When both output_config and output_format are present, output_format + is converted to inline schema and output_config is stripped for + unsupported models. + """ + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] optional_params = { "max_tokens": 4096, - "output_config": {"effort": "high"}, + "output_config": {"effort": "low"}, + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-3-haiku-20240307-v1:0", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=False, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) - assert result.get("output_config") == {"effort": "high"} - assert result.get("max_tokens") == 4096 + assert "output_config" not in result + assert "output_format" not in result def test_bedrock_messages_drop_params_strips_output_config_for_pre_4_5(): @@ -701,6 +837,8 @@ def test_bedrock_messages_drop_params_strips_output_config_for_pre_4_5(): def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): """``drop_params=True`` does not strip on opus-4-7 (supports effort).""" + from unittest.mock import patch + import litellm from litellm.types.router import GenericLiteLLMParams @@ -714,13 +852,17 @@ def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): original = litellm.drop_params litellm.drop_params = True try: - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) finally: litellm.drop_params = original @@ -742,6 +884,8 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( reasoning_effort, expected_effort ): """``reasoning_effort`` maps to ``thinking`` + ``output_config.effort`` on /v1/messages.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -751,13 +895,17 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( "reasoning_effort": reasoning_effort, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert "reasoning_effort" not in result assert result.get("thinking") == {"type": "adaptive"} @@ -842,6 +990,8 @@ def test_bedrock_messages_invalid_reasoning_effort_raises_400(): def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): """Explicit ``output_config.effort`` wins over the ``reasoning_effort`` alias.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -852,13 +1002,17 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): "output_config": {"effort": "max"}, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert "reasoning_effort" not in result assert result.get("output_config") == {"effort": "max"} @@ -994,7 +1148,7 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): } result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-3-haiku-20240307-v1:0", + model="anthropic.claude-opus-4-7", messages=messages, anthropic_messages_optional_request_params=optional_params, litellm_params=GenericLiteLLMParams(), diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 2498946bb5c..90a1c24bada 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -14,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -201,3 +202,127 @@ class TestChatGPTResponsesAPITransformation: ) assert parsed.output_text == "Hello!" + + @pytest.mark.parametrize( + ("model_name", "response_model"), + [ + ("chatgpt/gpt-5.2-codex", "gpt-5.2-codex"), + ("chatgpt/gpt-5.3-codex", "gpt-5.3-codex"), + ], + ) + def test_chatgpt_non_stream_sse_response_recovers_output_items( + self, model_name: str, response_model: str + ): + config = ChatGPTResponsesAPIConfig() + response_payload = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": response_model, + "output": [], + } + streamed_output_item = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello from stream!"}], + } + sse_body = "\n".join( + [ + f"data: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': streamed_output_item})}", + f"data: {json.dumps({'type': 'response.completed', 'response': response_payload})}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + parsed = config.transform_response_api_response( + model=model_name, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert parsed.output_text == "Hello from stream!" + + def test_chatgpt_non_stream_sse_recovers_whitespace_padded_chunks(self): + """Chunks with leading whitespace before `data:` must still parse. + + `_strip_sse_data_from_chunk` only matches the prefix at position 0, + so without an outer `.strip()` such chunks would fail JSON parsing + and silently drop the contained event. + """ + config = ChatGPTResponsesAPIConfig() + response_payload = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.4", + "output": [], + } + streamed_output_item = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Recovered from padded"}], + } + sse_body = "\n".join( + [ + f" data: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': streamed_output_item})} ", + f"\tdata: {json.dumps({'type': 'response.completed', 'response': response_payload})}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + parsed = config.transform_response_api_response( + model="chatgpt/gpt-5.4", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert parsed.output_text == "Recovered from padded" + + @pytest.mark.parametrize( + "error_chunk", + [ + { + "type": "response.failed", + "response": {"error": {"message": "ChatGPT upstream failed"}}, + }, + { + "type": "error", + "error": {"message": "ChatGPT upstream failed"}, + }, + ], + ) + def test_chatgpt_non_stream_sse_response_raises_openai_error(self, error_chunk): + config = ChatGPTResponsesAPIConfig() + sse_body = "\n".join( + [ + f"data: {json.dumps(error_chunk)}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 502, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + with pytest.raises(OpenAIError) as exc_info: + config.transform_response_api_response( + model="chatgpt/gpt-5.4", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert "ChatGPT upstream failed" in str(exc_info.value) + assert exc_info.value.status_code == 502 diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 279f16a3675..a29365544df 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -6,16 +6,29 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm + sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm import supports_reasoning +from litellm import get_model_info, supports_reasoning from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message +@pytest.fixture(autouse=True) +def force_local_model_cost(monkeypatch): + """Force local model cost map usage for all tests in this file.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + # Refresh model_cost from local map + import litellm + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( @@ -62,7 +75,6 @@ def test_handle_message_content_with_tool_calls(): def test_supports_reasoning_effort(): """Test that reasoning_effort is only supported for specific Fireworks AI models.""" - # Models that support reasoning_effort supported_models = [ "fireworks_ai/accounts/fireworks/models/qwen3-8b", "fireworks_ai/accounts/fireworks/models/qwen3-32b", @@ -72,11 +84,13 @@ def test_supports_reasoning_effort(): "fireworks_ai/accounts/fireworks/models/glm-4p5", "fireworks_ai/accounts/fireworks/models/glm-4p5-air", "fireworks_ai/accounts/fireworks/models/glm-4p6", + "fireworks_ai/accounts/fireworks/models/glm-4p7", + "fireworks_ai/accounts/fireworks/models/glm-5p1", "fireworks_ai/accounts/fireworks/models/gpt-oss-120b", "fireworks_ai/accounts/fireworks/models/gpt-oss-20b", + "fireworks_ai/glm-5p1", ] - # Models that don't support reasoning_effort unsupported_models = [ "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct", "fireworks_ai/accounts/fireworks/models/mixtral-8x7b-instruct", @@ -97,19 +111,74 @@ def test_get_supported_openai_params_reasoning_effort(): """Test that reasoning_effort is only included in supported params for models that support it.""" config = FireworksAIConfig() - # Model that supports reasoning_effort supported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/qwen3-8b" + "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "reasoning_effort" in supported_params - # Model that doesn't support reasoning_effort unsupported_params = config.get_supported_openai_params( "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" ) assert "reasoning_effort" not in unsupported_params +def test_get_supported_openai_params_parallel_tool_calls(): + """Test that parallel_tool_calls is included for models that support function calling.""" + config = FireworksAIConfig() + + supported_params = config.get_supported_openai_params( + "fireworks_ai/accounts/fireworks/models/glm-4p6" + ) + assert "parallel_tool_calls" in supported_params + + unsupported_params = config.get_supported_openai_params( + "fireworks_ai/accounts/fireworks/models/glm-5p1" + ) + assert "parallel_tool_calls" not in unsupported_params + + +def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice( + monkeypatch, +): + """Test that parallel_tool_calls is gated on tools, not tool_choice.""" + config = FireworksAIConfig() + model = "fireworks_ai/test-tools-without-tool-choice" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "supports_function_calling": True, + "supports_tool_choice": False, + }, + ) + + supported_params = config.get_supported_openai_params(model) + + assert "tools" in supported_params + assert "parallel_tool_calls" in supported_params + assert "tool_choice" not in supported_params + + +def test_get_model_info_respects_explicit_fireworks_capabilities(): + """Test that get_model_info preserves explicit capability flags from the model map.""" + model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") + + assert model_info["supports_function_calling"] is False + assert model_info["supports_reasoning"] is True + assert model_info["supports_tool_choice"] is False + + +def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): + """Test that Fireworks only overrides supports_reasoning for supported models.""" + config = FireworksAIConfig() + model = "fireworks_ai/test-reasoning-false" + monkeypatch.setitem(litellm.model_cost, model, {"supports_reasoning": False}) + + info = config.get_provider_info(model) + + assert "supports_reasoning" not in info + + def test_add_transform_inline_image_block_skips_data_urls(): """ data: URLs must not have #transform=inline appended — doing so corrupts the @@ -234,6 +303,14 @@ def test_transform_messages_helper_removes_provider_specific_fields(): assert "provider_specific_fields" not in msg +def test_unmapped_model_fallback_function_calling(): + """Test that a model not in model_cost still defaults to supporting function calling for Fireworks.""" + config = FireworksAIConfig() + model = "fireworks_ai/unmapped-future-model" + info = config.get_provider_info(model) + assert info["supports_function_calling"] is True + + def test_transform_messages_helper_strips_thinking_blocks(): """thinking_blocks must not be forwarded to Fireworks chat completions.""" config = FireworksAIConfig() diff --git a/tests/test_litellm/llms/reducto/__init__.py b/tests/test_litellm/llms/reducto/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm/llms/reducto/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm/llms/reducto/test_cost.py b/tests/test_litellm/llms/reducto/test_cost.py new file mode 100644 index 00000000000..73340dc8729 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_cost.py @@ -0,0 +1,122 @@ +import litellm +import pytest + +from litellm.cost_calculator import completion_cost +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + + +def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_credit": 0.003}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="credit priced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=None, credits=10), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.03 + + +def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_credit": 0.0, + "ocr_cost_per_page": 0.5, + }, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="free credit priced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=2, credits=10), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.0 + + +def test_ocr_cost_falls_back_to_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="page priced")], + model="mistral-ocr-latest", + usage_info=OCRUsageInfo(pages_processed=2), + ) + + cost = completion_cost( + completion_response=response, + model="mistral/mistral-ocr-latest", + custom_llm_provider="mistral", + call_type="ocr", + ) + + assert cost == 1.0 + + +def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="unpriced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=None, credits=5), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.0 + + +def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="missing pages")], + model="mistral-ocr-latest", + usage_info=OCRUsageInfo(pages_processed=None), + ) + + with pytest.raises(ValueError, match="OCR response pages_processed is None"): + completion_cost( + completion_response=response, + model="mistral/mistral-ocr-latest", + custom_llm_provider="mistral", + call_type="ocr", + ) diff --git a/tests/test_litellm/llms/reducto/test_model_info.py b/tests/test_litellm/llms/reducto/test_model_info.py new file mode 100644 index 00000000000..de7a3ccba64 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_model_info.py @@ -0,0 +1,44 @@ +import uuid + +import litellm + +from litellm.utils import _invalidate_model_cost_lowercase_map + + +def test_reducto_provider_registration(): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="reducto/parse-v3" + ) + + assert model == "parse-v3" + assert custom_llm_provider == "reducto" + + +def test_get_model_info_preserves_ocr_cost_per_credit(): + test_model_name = f"reducto/test-cost-propagation-{uuid.uuid4().hex[:12]}" + previous_model_entry = litellm.model_cost.get(test_model_name) + _invalidate_model_cost_lowercase_map() + + try: + litellm.register_model( + { + test_model_name: { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.003, + } + } + ) + + model_info = litellm.get_model_info( + model=test_model_name, + custom_llm_provider="reducto", + ) + + assert model_info.get("ocr_cost_per_credit") == 0.003 + finally: + if previous_model_entry is None: + litellm.model_cost.pop(test_model_name, None) + else: + litellm.model_cost[test_model_name] = previous_model_entry + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/test_litellm/llms/reducto/test_parse_legacy.py new file mode 100644 index 00000000000..db19460baa3 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_parse_legacy.py @@ -0,0 +1,59 @@ +import json + +import litellm +import pytest + + +@pytest.fixture() +def disable_aiohttp_transport(): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_parse_legacy_wraps_enhance_under_options( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://legacy.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Legacy parse", + "blocks": [{"content": "Legacy parse", "bbox": {"page": 1}}], + } + ] + }, + } + ) + + response = await litellm.aocr( + model="reducto/parse-legacy", + document={ + "type": "file", + "file": b"%PDF-1.4 legacy", + "mime_type": "application/pdf", + }, + api_key="legacy-key", + api_base="https://platform.reducto.ai", + enhance={"agentic": [{"type": "table"}]}, + ) + + assert upload_route.called + assert parse_route.called + request_body = json.loads(parse_route.calls[0].request.read()) + assert request_body == { + "document_url": "reducto://legacy.pdf", + "options": {"enhance": {"agentic": [{"type": "table"}]}}, + } + assert response.pages[0].markdown == "Legacy parse" diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py new file mode 100644 index 00000000000..140b9737dc0 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -0,0 +1,152 @@ +import json + +import litellm +import pytest + + +def _reducto_parse_response() -> dict: + return { + "job_id": "job_123", + "usage": {"num_pages": 3, "credits": 3}, + "result": { + "chunks": [ + { + "content": "Page 1 block A", + "blocks": [ + { + "content": "Page 1 block A", + "bbox": {"page": 1}, + "kind": "text", + } + ], + }, + { + "content": "Page 2 block A", + "blocks": [ + { + "content": "Page 2 block A", + "bbox": {"page": 2}, + "kind": "table", + } + ], + }, + { + "content": "Page 1 block B", + "blocks": [ + { + "content": "Page 1 block B", + "bbox": {"page": 1}, + "kind": "text", + } + ], + }, + { + "content": "Page 3 block A", + "blocks": [ + { + "content": "Page 3 block A", + "bbox": {"page": 3}, + "kind": "figure", + } + ], + }, + ] + }, + } + + +@pytest.fixture() +def disable_aiohttp_transport(): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_parse_v3_file_upload_and_response_mapping( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://uploaded.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json=_reducto_parse_response() + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"%PDF-1.4 reducto", + "mime_type": "application/pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + formatting={"table_output_format": "html"}, + retrieval={"chunk_mode": "section"}, + settings={"ocr_system": "standard"}, + ) + + assert upload_route.called + assert parse_route.called + assert len(upload_route.calls) == 1 + assert len(parse_route.calls) == 1 + + upload_request = upload_route.calls[0].request + assert upload_request.headers["authorization"] == "Bearer test-key" + assert "application/json" not in upload_request.headers["content-type"] + upload_body = upload_request.read() + assert b'filename="document"' in upload_body + assert b"application/pdf" in upload_body + + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://uploaded.pdf" + assert parse_request_body["formatting"] == {"table_output_format": "html"} + assert parse_request_body["retrieval"] == {"chunk_mode": "section"} + assert parse_request_body["settings"] == {"ocr_system": "standard"} + + assert response.usage_info is not None + assert response.usage_info.credits == 3 + assert response.usage_info.pages_processed == 3 + assert len(response.pages) == 3 + assert response.pages[0].index == 0 + assert response.pages[0].markdown == "Page 1 block A\n\nPage 1 block B" + assert getattr(response.pages[0], "blocks")[0]["bbox"]["page"] == 1 + assert response.pages[1].markdown == "Page 2 block A" + assert response.pages[2].markdown == "Page 3 block A" + assert response._hidden_params["reducto_raw"]["usage"]["credits"] == 3 + + +@pytest.mark.asyncio +async def test_parse_v3_reducto_id_passthrough_skips_upload( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://should-not-upload.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json=_reducto_parse_response() + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "document_url", + "document_url": "reducto://already-uploaded.pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + retrieval={"chunk_mode": "section"}, + ) + + assert not upload_route.called + assert parse_route.called + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://already-uploaded.pdf" + assert parse_request_body["retrieval"]["chunk_mode"] == "section" + assert response.pages[0].markdown.startswith("Page 1 block A") diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/test_litellm/llms/reducto/test_upload.py new file mode 100644 index 00000000000..4fae90436bb --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_upload.py @@ -0,0 +1,213 @@ +import json +import os +from unittest.mock import AsyncMock, Mock + +import httpx +import litellm +import pytest + +from litellm.llms.reducto.common import ( + extract_file_id_or_bytes, + upload_bytes_async, + upload_bytes_sync, +) + + +@pytest.fixture() +def disable_aiohttp_transport(monkeypatch): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setenv("REDUCTO_API_KEY", "env-reducto-key") + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + os.environ.pop("REDUCTO_API_KEY", None) + + +@pytest.mark.asyncio +async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): + with pytest.raises(litellm.BadRequestError, match="upload the file first"): + await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "document_url", + "document_url": "https://example.com/document.pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + ) + + +@pytest.mark.asyncio +async def test_parse_v3_image_data_uri_upload_uses_image_mime( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://custom.reducto.test/upload").respond( + json={"file_id": "reducto://uploaded-image.png"} + ) + parse_route = respx_mock.post("https://custom.reducto.test/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Image OCR", + "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], + } + ] + }, + } + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"\x89PNG\r\n\x1a\npng", + "mime_type": "image/png", + }, + api_key="programmatic-key", + api_base="https://custom.reducto.test/", + ) + + assert upload_route.called + assert parse_route.called + upload_request = upload_route.calls[0].request + assert upload_request.headers["authorization"] == "Bearer programmatic-key" + assert b"image/png" in upload_request.read() + + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://uploaded-image.png" + assert response.pages[0].markdown == "Image OCR" + + +@pytest.mark.asyncio +async def test_parse_v3_uses_programmatic_api_key_over_env( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://uploaded.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Programmatic auth", + "blocks": [ + {"content": "Programmatic auth", "bbox": {"page": 1}} + ], + } + ] + }, + } + ) + + await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"%PDF-1.4 auth", + "mime_type": "application/pdf", + }, + api_key="passed-key", + api_base="https://platform.reducto.ai", + ) + + assert upload_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + assert parse_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + + +def test_upload_bytes_sync_uses_shared_client(monkeypatch): + captured = {} + + def fake_post(*, url, headers, files, timeout): + captured["url"] = url + captured["headers"] = headers + captured["files"] = files + captured["timeout"] = timeout + return httpx.Response( + 200, + json={"file_id": "reducto://sync-upload"}, + request=httpx.Request("POST", url), + ) + + sync_post = Mock(side_effect=fake_post) + monkeypatch.setattr(litellm.module_level_client, "post", sync_post) + + class ForbiddenSyncClient: + def __init__(self, *args, **kwargs): + raise AssertionError("should not construct") + + monkeypatch.setattr(httpx, "Client", ForbiddenSyncClient) + + file_id = upload_bytes_sync( + raw_bytes=b"%PDF-1.4 sync", + mime="application/pdf", + api_key="sync-key", + api_base="https://sync.reducto.test/", + ) + + assert file_id == "reducto://sync-upload" + sync_post.assert_called_once() + assert captured["url"] == "https://sync.reducto.test/upload" + assert captured["headers"] == {"Authorization": "Bearer sync-key"} + assert captured["files"]["file"] == ( + "document", + b"%PDF-1.4 sync", + "application/pdf", + ) + + +@pytest.mark.asyncio +async def test_upload_bytes_async_uses_shared_aclient(monkeypatch): + captured = {} + + async def fake_post(*, url, headers, files, timeout): + captured["url"] = url + captured["headers"] = headers + captured["files"] = files + captured["timeout"] = timeout + return httpx.Response( + 200, + json={"file_id": "reducto://async-upload"}, + request=httpx.Request("POST", url), + ) + + async_post = AsyncMock(side_effect=fake_post) + monkeypatch.setattr(litellm.module_level_aclient, "post", async_post) + + class ForbiddenAsyncClient: + def __init__(self, *args, **kwargs): + raise AssertionError("should not construct") + + monkeypatch.setattr(httpx, "AsyncClient", ForbiddenAsyncClient) + + file_id = await upload_bytes_async( + raw_bytes=b"%PDF-1.4 async", + mime="application/pdf", + api_key="async-key", + api_base="https://async.reducto.test/", + ) + + assert file_id == "reducto://async-upload" + async_post.assert_awaited_once() + assert captured["url"] == "https://async.reducto.test/upload" + assert captured["headers"] == {"Authorization": "Bearer async-key"} + assert captured["files"]["file"] == ( + "document", + b"%PDF-1.4 async", + "application/pdf", + ) + + +def test_extract_file_id_or_bytes_raises_on_malformed_data_uri(): + with pytest.raises(litellm.BadRequestError, match="Invalid Reducto data URI"): + extract_file_id_or_bytes("data:application/pdf", model="reducto/parse-v3") + + with pytest.raises(litellm.BadRequestError, match="Invalid Reducto base64 payload"): + extract_file_id_or_bytes("data:;base64,!!!not-base64", model="reducto/parse-v3") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 88aac07a0c9..2cf97081806 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -1448,3 +1449,474 @@ class TestVertexBase: aws_creds = supplier.get_aws_security_credentials(context=None, request=None) assert isinstance(aws_creds, AwsSecurityCredentials) + + @pytest.mark.asyncio + async def test_single_flight_refresh(self): + """Under high concurrency, only one coroutine should refresh expired credentials.""" + import asyncio + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "expired-token" + mock_creds.expired = True + mock_creds.expiry = None + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + refresh_call_count = 0 + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + async def slow_refresh(creds): + nonlocal refresh_call_count + refresh_call_count += 1 + await asyncio.sleep(0.05) # simulate network latency + creds.token = "refreshed-token" + creds.expired = False + + # refresh_auth is sync, but we need to count calls. + # get_access_token_async wraps it with asyncify, so the sync side_effect works. + def sync_refresh_impl(creds): + nonlocal refresh_call_count + refresh_call_count += 1 + creds.token = "refreshed-token" + creds.expired = False + + mock_refresh.side_effect = sync_refresh_impl + + # Launch 50 concurrent requests + tasks = [ + vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + for _ in range(50) + ] + results = await asyncio.gather(*tasks) + + # All should return the refreshed token + for token, project in results: + assert token == "refreshed-token" + assert project == "project-1" + + # refresh_auth should be called exactly once (single-flight) + assert ( + refresh_call_count == 1 + ), f"Expected 1 refresh call, got {refresh_call_count}" + + @pytest.mark.asyncio + async def test_async_reauthentication_uses_async_single_flight(self): + """Concurrent async reauth should reload once without using the sync path.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + stale_creds = MagicMock() + stale_creds.token = "expired-token" + stale_creds.token_state = TokenState.INVALID + stale_creds.project_id = "project-1" + stale_creds.quota_project_id = "project-1" + + refreshed_creds = MagicMock() + refreshed_creds.token = "refreshed-token" + refreshed_creds.token_state = TokenState.FRESH + refreshed_creds.project_id = "project-1" + refreshed_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + stale_creds, + "project-1", + ) + + load_call_count = 0 + + def load_auth_impl(*_args, **_kwargs): + nonlocal load_call_count + load_call_count += 1 + return refreshed_creds, "project-1" + + with ( + patch.object( + vertex_base, + "refresh_auth", + side_effect=Exception("Reauthentication is needed"), + ), + patch.object(vertex_base, "load_auth", side_effect=load_auth_impl), + patch.object(vertex_base, "get_access_token") as mock_get_access_token, + ): + results = await asyncio.gather( + *[ + vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + for _ in range(10) + ] + ) + + assert results == [("refreshed-token", "project-1")] * 10 + assert load_call_count == 1 + mock_get_access_token.assert_not_called() + + @pytest.mark.asyncio + async def test_background_refresh_when_near_expiry(self): + """When token_state is STALE (within the 3:45 REFRESH_THRESHOLD window), + return the current token immediately and refresh in the background — + zero added latency.""" + import asyncio + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + # Simulate STALE state: token is usable but near expiry. + mock_creds = MagicMock() + mock_creds.token = "near-expiry-token" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + # Should return the current (still usable) token immediately + assert token == "near-expiry-token" + + # Let the background refresh task run + await asyncio.sleep(0.05) + + assert mock_refresh.called, "Background refresh should have been triggered" + + @pytest.mark.asyncio + async def test_stale_malformed_token_blocks_on_refresh(self): + """Malformed STALE tokens should refresh instead of failing validation.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = None + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert mock_refresh.called + assert token == "refreshed-token" + assert project == "project-1" + + @pytest.mark.asyncio + async def test_fresh_token_skips_refresh(self): + """Credentials not marked expired by google-auth should not trigger refresh.""" + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "fresh-token" + mock_creds.expired = False + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + mock_creds, + "project-1", + ) + + with patch.object(vertex_base, "refresh_auth") as mock_refresh: + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert not mock_refresh.called, "Fresh token should not trigger refresh" + assert token == "fresh-token" + + @pytest.mark.asyncio + async def test_background_refresh_task_removed_after_completion(self): + """Completed background-refresh tasks must be evicted from + _background_refresh_tasks so the dict does not grow unboundedly.""" + import asyncio + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "near-expiry-token" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + # Allow the background task to complete. + await asyncio.sleep(0.1) + + # After completion the entry should have been removed by the done-callback. + assert len(vertex_base._background_refresh_tasks) == 0, ( + "Completed background refresh task was not removed from " + "_background_refresh_tasks" + ) + + @pytest.mark.asyncio + async def test_background_refresh_tasks_no_accumulation_across_many_keys(self): + """With many distinct credential keys the dict must not hold completed tasks.""" + import asyncio + import json as _json + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + num_keys = 20 + + for i in range(num_keys): + mock_creds = MagicMock() + mock_creds.token = f"token-{i}" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = f"project-{i}" + mock_creds.quota_project_id = f"project-{i}" + + credentials = {"type": "service_account", "project_id": f"project-{i}"} + + with ( + patch.object( + vertex_base, + "load_auth", + return_value=(mock_creds, f"project-{i}"), + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds, idx=i): + creds.token = f"refreshed-{idx}" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id=f"project-{i}", + custom_llm_provider="vertex_ai", + ) + + # Let all background tasks finish. + await asyncio.sleep(0.1) + + assert len(vertex_base._background_refresh_tasks) == 0, ( + f"Expected 0 tasks after all refreshes completed, " + f"found {len(vertex_base._background_refresh_tasks)}" + ) + + @pytest.mark.asyncio + async def test_async_refresh_lock_shared_while_in_use(self): + """Concurrent callers for the same key must coordinate on the same lock.""" + vertex_base = VertexBase() + key = ("creds", "project-1") + + lock_a = vertex_base._acquire_async_refresh_lock(key) + try: + async with lock_a: + lock_b = vertex_base._acquire_async_refresh_lock(key) + try: + assert lock_a is lock_b, ( + "While a coroutine still holds the lock, concurrent callers must " + "receive the same Lock instance to preserve single-flight." + ) + finally: + vertex_base._release_async_refresh_lock(key, lock_b) + finally: + vertex_base._release_async_refresh_lock(key, lock_a) + + @pytest.mark.asyncio + async def test_async_refresh_lock_pruned_after_release(self): + """get_access_token_async must drop the per-key Lock from the registry + once no coroutine is using it, so the dict stays bounded in + high-cardinality deployments. Without this, every distinct credential + leaks a Lock object for the lifetime of the process.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + for i in range(10): + mock_creds = MagicMock() + mock_creds.token = f"refreshed-{i}" + mock_creds.token_state = TokenState.FRESH + mock_creds.project_id = f"project-{i}" + mock_creds.quota_project_id = f"project-{i}" + + credentials = {"type": "service_account", "project_id": f"project-{i}"} + + with ( + patch.object( + vertex_base, + "load_auth", + return_value=(mock_creds, f"project-{i}"), + ), + patch.object(vertex_base, "refresh_auth"), + ): + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id=f"project-{i}", + custom_llm_provider="vertex_ai", + ) + + assert len(vertex_base._async_refresh_locks) == 0, ( + "expected per-key locks to be pruned once no coroutine holds or " + f"waits on them; found {len(vertex_base._async_refresh_locks)}" + ) + assert len(vertex_base._async_refresh_lock_refcounts) == 0 + + @pytest.mark.asyncio + async def test_async_refresh_lock_kept_while_waiter_pending(self): + """The prune must not run while another coroutine is still waiting on + the lock — otherwise the waiter ends up on a lock that's been replaced + in the registry and single-flight breaks.""" + vertex_base = VertexBase() + key = ("creds", "project-1") + + holder_lock = vertex_base._acquire_async_refresh_lock(key) + release_holder = asyncio.Event() + + async def hold_then_release(): + async with holder_lock: + await release_holder.wait() + vertex_base._release_async_refresh_lock(key, holder_lock) + + holder = asyncio.create_task(hold_then_release()) + await asyncio.sleep(0) # let holder grab the lock + + async def queue_for_lock(): + waiter_lock = vertex_base._acquire_async_refresh_lock(key) + try: + async with waiter_lock: + pass + finally: + vertex_base._release_async_refresh_lock(key, waiter_lock) + + waiter = asyncio.create_task(queue_for_lock()) + await asyncio.sleep(0) # let waiter queue on the lock + + assert ( + vertex_base._async_refresh_locks.get(key) is holder_lock + ), "lock with active holder/waiter must not be pruned" + + release_holder.set() + await holder + await waiter + + assert key not in vertex_base._async_refresh_locks + assert key not in vertex_base._async_refresh_lock_refcounts + + @pytest.mark.asyncio + async def test_fast_path_no_lock(self): + """Cached fresh credentials should return without acquiring the lock.""" + import datetime + + vertex_base = VertexBase() + + try: + from google.auth import _helpers as google_auth_helpers + + now = google_auth_helpers.utcnow() + except ImportError: + now = datetime.datetime.utcnow() + + mock_creds = MagicMock() + mock_creds.token = "cached-token" + mock_creds.expired = False + mock_creds.expiry = now + datetime.timedelta(minutes=30) + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + mock_creds, + "project-1", + ) + + # Spy on _acquire_async_refresh_lock to verify it's never called + with patch.object( + vertex_base, + "_acquire_async_refresh_lock", + wraps=vertex_base._acquire_async_refresh_lock, + ) as mock_get_lock: + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert token == "cached-token" + assert not mock_get_lock.called, "Fast path should not acquire lock" diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py index b16fc2bc44d..f617a8db850 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py @@ -118,7 +118,7 @@ async def test_vertex_ai_gpt_oss_simple_request(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718"), ), patch.dict( @@ -217,7 +217,7 @@ async def test_vertex_ai_gpt_oss_reasoning_effort(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718"), ), patch.dict( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py index bf6e0a5f2cd..5a86325b7fd 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py @@ -7,7 +7,6 @@ These tests verify that: 3. The completion() and responses() API work with Qwen models """ -import json import os import sys from unittest.mock import MagicMock, patch, AsyncMock @@ -179,7 +178,7 @@ async def test_vertex_ai_qwen_global_endpoint_url(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "test-project"), ), patch.dict( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py new file mode 100644 index 00000000000..b20442a032e --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py @@ -0,0 +1,220 @@ +""" +Test that VertexBase subclasses (PartnerModels, Gemma, ModelGarden) reuse +cached credentials instead of creating a new VertexLLM instance on every request. +""" + +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( + VertexAIPartnerModels, +) +from litellm.llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels +from litellm.llms.vertex_ai.vertex_model_garden.main import VertexAIModelGardenModels + + +def _mock_vertexai(): + """Return a MagicMock that satisfies the vertexai import guards.""" + m = MagicMock() + m.preview = MagicMock() + m.preview.language_models = MagicMock() + return m + + +class TestVertexBaseSubclassInit: + """All VertexBase subclasses must call super().__init__() so that + the credential cache is initialized.""" + + @pytest.mark.parametrize( + "cls", + [VertexAIPartnerModels, VertexAIGemmaModels, VertexAIModelGardenModels], + ids=["PartnerModels", "Gemma", "ModelGarden"], + ) + def test_init_calls_super(self, cls): + instance = cls() + assert hasattr(instance, "_credentials_project_mapping") + assert isinstance(instance._credentials_project_mapping, dict) + assert hasattr(instance, "access_token") + assert hasattr(instance, "project_id") + + +class TestPartnerModelsCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + partner = VertexAIPartnerModels() + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + partner, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler" + ) as mock_handler, + ): + mock_handler.completion.return_value = "response" + + partner.completion( + model="meta/llama-3.1-405b-instruct-maas", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) + + def test_credential_cache_shared_across_calls(self): + """Two successive completion() calls should hit load_auth only once.""" + partner = VertexAIPartnerModels() + + mock_creds = MagicMock() + mock_creds.token = "my-token" + mock_creds.expired = False + mock_creds.project_id = "proj" + mock_creds.quota_project_id = "proj" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + partner, "load_auth", return_value=(mock_creds, "proj") + ) as mock_load, + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler" + ) as mock_handler, + ): + mock_handler.completion.return_value = "resp" + + common_kwargs = dict( + model="meta/llama-3.1-405b-instruct-maas", + messages=[{"role": "user", "content": "hi"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="proj", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + partner.completion(**common_kwargs) + partner.completion(**common_kwargs) + + assert mock_load.call_count == 1 + + +class TestGemmaModelsCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + gemma = VertexAIGemmaModels() + + mock_gemma_config = MagicMock() + mock_gemma_config.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + gemma, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.transformation.VertexGemmaConfig", + mock_gemma_config, + ), + ): + gemma.completion( + model="gemma/gemma-3-12b-it-1234567890", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base="https://123.us-central1-1.prediction.vertexai.goog/v1/projects/proj/locations/us-central1/endpoints/456:predict", + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) + + +class TestModelGardenCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + garden = VertexAIModelGardenModels() + + mock_handler = MagicMock() + mock_handler.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + garden, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.openai_like.chat.handler.OpenAILikeChatHandler", + mock_handler, + ), + ): + garden.completion( + model="openai/5464397967697903616", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 3e3e8901706..b1c8f7234ce 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -122,17 +122,19 @@ class TestVertexGemmaCompletion: # Mock the async HTTP handler and Vertex authentication with ( patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler, + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): + mock_client = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = mock_vertex_response - mock_http_handler.return_value.post = AsyncMock(return_value=mock_response) + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client # Call litellm.acompletion() response = await litellm.acompletion( @@ -145,7 +147,7 @@ class TestVertexGemmaCompletion: ) # Verify the request sent to Vertex - call_args = mock_http_handler.return_value.post.call_args + call_args = mock_client.post.call_args assert call_args is not None, "HTTP handler was not called" request_data = call_args.kwargs["json"] @@ -210,17 +212,19 @@ class TestVertexGemmaCompletion: with ( patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler, + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "test-project"), ), ): + mock_client = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = invalid_response - mock_http_handler.return_value.post = AsyncMock(return_value=mock_response) + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client # Should raise exception (wrapped as APIConnectionError by LiteLLM) with pytest.raises(APIConnectionError) as exc_info: @@ -286,7 +290,7 @@ class TestVertexGemmaCompletion: "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): @@ -388,7 +392,7 @@ class TestVertexGemmaCompletion: "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): @@ -429,3 +433,123 @@ class TestVertexGemmaCompletion: # Verify other parameters are present assert "messages" in instance assert instance["@requestFormat"] == "chatCompletions" + + @pytest.mark.asyncio + async def test_acompletion_filters_context_management(self): + """ + Test that context_management is filtered out from the request. + + Vertex AI Gemma's chatCompletions wrapper does not understand + `context_management` (an Anthropic / OpenAI Responses API concept). + It must be stripped from the request body so the upstream endpoint + does not reject the request with an unknown-field error. + """ + mock_vertex_response = { + "deployedModelId": "1207280419999999999", + "model": "projects/993702345710/locations/us-central1/models/gemma-3-12b-it-1222199011122", + "modelDisplayName": "gemma-3-12b-it-1222199011122", + "modelVersionId": "1", + "predictions": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "ok", + "reasoning_content": None, + "role": "assistant", + "tool_calls": [], + }, + "stop_reason": None, + } + ], + "created": 1759863903, + "id": "chatcmpl-test-ctxmgmt", + "model": "google/gemma-3-12b-it", + "object": "chat.completion", + "prompt_logprobs": None, + "usage": { + "completion_tokens": 1, + "prompt_tokens": 5, + "prompt_tokens_details": None, + "total_tokens": 6, + }, + }, + } + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID"), + ), + ): + mock_client = Mock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_vertex_response + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + # Use `allowed_openai_params` so context_management actually + # reaches the transformation layer (otherwise the upstream + # validator drops it before we can prove the transformation + # strips it). This mirrors the real-world scenario where a + # caller explicitly opts in to forwarding an arbitrary param. + await litellm.acompletion( + model="vertex_ai/gemma/gemma-3-12b-it-1222199011122", + messages=[{"role": "user", "content": "Test"}], + context_management=[ + {"type": "compaction", "compact_threshold": 200000} + ], + allowed_openai_params=["context_management"], + api_base="https://test.us-central1-project.prediction.vertexai.goog/v1/projects/PROJECT_ID/locations/us-central1/endpoints/ENDPOINT_ID:predict", + vertex_project="PROJECT_ID", + vertex_location="us-central1", + ) + + call_args = mock_client.post.call_args + assert call_args is not None, "HTTP client was not called" + + request_data = call_args.kwargs["json"] + print("request body=", json.dumps(request_data, indent=4)) + instance = request_data["instances"][0] + + assert ( + "context_management" not in instance + ), "context_management should not be forwarded to Vertex Gemma" + assert instance["@requestFormat"] == "chatCompletions" + assert "messages" in instance + + def test_transform_request_strips_context_management(self): + """ + Direct unit test for VertexGemmaConfig.transform_request: verify that + `context_management` is stripped from `optional_params` regardless of + how it was supplied to the transformation layer. + """ + from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( + VertexGemmaConfig, + ) + + config = VertexGemmaConfig() + result = config.transform_request( + model="gemma-3-12b-it", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "max_tokens": 32, + "context_management": [ + {"type": "compaction", "compact_threshold": 200000} + ], + }, + litellm_params={}, + headers={}, + ) + + assert "instances" in result + instance = result["instances"][0] + assert instance["@requestFormat"] == "chatCompletions" + assert "context_management" not in instance + assert instance.get("max_tokens") == 32 diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 3ae8dfc3c0b..5c1f0f704d7 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -119,3 +119,19 @@ class TestXAIParallelToolCalls: assert result.get("parallel_tool_calls") is True assert len(result["messages"]) == 1 assert result["messages"][0]["role"] == "user" + + +class TestXAIUsageNormalization: + def test_preserves_reasoning_tokens_in_total_usage(self): + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=200) + + XAIChatConfig._normalize_openai_compatible_usage_totals(usage) + + assert usage.total_tokens == 200 + + def test_preserves_reasoning_tokens_in_streaming_usage(self): + usage = {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 200} + + XAIChatConfig._normalize_openai_compatible_usage_totals(usage) + + assert usage["total_tokens"] == 200 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..35a3bd7f657 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3016,3 +3016,340 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.max_budget == 0.0 + + +# --- resolve_and_validate_end_user_id --------------------------------------- + + +@pytest.fixture +def _validate_flag_on(monkeypatch): + """Enable opt-in DB validation for the duration of a test.""" + import litellm + + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + + +def _validation_cache(): + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + return cache + + +def _patch_validation_helpers(monkeypatch, *, end_user=None, user=None, fuzzy=None): + """Stub out the DB helpers resolve_and_validate_end_user_id delegates to.""" + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr( + auth_checks, "get_end_user_object", AsyncMock(return_value=end_user) + ) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) + monkeypatch.setattr( + auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy) + ) + + +@pytest.mark.asyncio +async def test_resolve_end_user_returns_none_for_none_input( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + assert ( + await resolve_and_validate_end_user_id( + raw_end_user_id=None, + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + is None + ) + + +@pytest.mark.asyncio +async def test_resolve_end_user_passes_through_when_flag_disabled(monkeypatch): + """Default behaviour: flag is off, arbitrary ids pass through untouched.""" + import litellm + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="codex-session-abc", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "codex-session-abc" + cache.async_set_cache.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_passes_through_when_no_prisma_client( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=None, + user_api_key_cache=cache, + ) + assert result == "alice@example.com" + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_end_user_table(_validate_flag_on, monkeypatch): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, end_user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="customer-123", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "customer-123" + cache.async_set_cache.assert_awaited_once() + kwargs = cache.async_set_cache.await_args.kwargs + assert kwargs["key"] == "end_user_validation:customer-123" + assert kwargs["value"] == "valid" + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_user_table_by_user_id( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="user-xyz", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "user-xyz" + # email fallback should not run for a non-email input + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_user_table_by_email( + _validate_flag_on, monkeypatch +): + """Email-shaped ids route through get_user_object with user_email set. + + The fuzzy lookup must happen inside get_user_object so it shares the + _should_check_db throttle and user_api_key_cache — no direct raw + Prisma calls on the auth path. + """ + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="Alice@Example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "Alice@Example.com" + auth_checks.get_user_object.assert_awaited_once() + user_kwargs = auth_checks.get_user_object.await_args.kwargs + assert user_kwargs["user_id"] == "Alice@Example.com" + assert user_kwargs["user_email"] == "Alice@Example.com" + # email branch must not bypass the cached helper with a raw fuzzy call + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_non_email_id_does_not_pass_user_email( + _validate_flag_on, monkeypatch +): + """Non-email ids skip the email fuzzy path to avoid a pointless DB hit.""" + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + await resolve_and_validate_end_user_id( + raw_end_user_id="user-xyz", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + auth_checks.get_user_object.assert_awaited_once() + user_kwargs = auth_checks.get_user_object.await_args.kwargs + assert user_kwargs["user_email"] is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_drops_codex_opaque_identifier( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) # all helpers return None + cache = _validation_cache() + + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + result = await resolve_and_validate_end_user_id( + raw_end_user_id=codex_id, + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + cache.async_set_cache.assert_awaited_once() + kwargs = cache.async_set_cache.await_args.kwargs + assert kwargs["value"] == "invalid" + + +@pytest.mark.asyncio +async def test_resolve_end_user_preserves_id_when_default_budget_configured( + _validate_flag_on, monkeypatch +): + """Don't drop unregistered ids when litellm.max_end_user_budget_id is set. + + The default end-user budget is applied downstream when the id is present + but not found in the db — dropping the id here would bypass those limits. + """ + import litellm + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-budget") + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="new-customer", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "new-customer" + + +@pytest.mark.asyncio +async def test_resolve_end_user_drops_unknown_email(_validate_flag_on, monkeypatch): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="stranger@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_uses_cached_valid_result( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + cache.async_get_cache = AsyncMock(return_value="valid") + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "alice@example.com" + auth_checks.get_end_user_object.assert_not_awaited() + auth_checks.get_user_object.assert_not_awaited() + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_uses_cached_invalid_result( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, end_user=MagicMock()) + cache = _validation_cache() + cache.async_get_cache = AsyncMock(return_value="invalid") + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="bogus", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + # Despite a matching row configured, helpers aren't called — cache wins. + auth_checks.get_end_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_swallows_db_errors_and_returns_none( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr( + auth_checks, + "get_end_user_object", + AsyncMock(side_effect=Exception("db down")), + ) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(side_effect=Exception("db down")), + ) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + # DB errors shouldn't raise through the auth path — treat as unknown. + assert result is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_reraises_budget_exceeded( + _validate_flag_on, monkeypatch +): + """BudgetExceededError from get_end_user_object must bubble up so the + auth path enforces spend limits instead of silently dropping the id.""" + import litellm + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr( + auth_checks, + "get_end_user_object", + AsyncMock( + side_effect=litellm.BudgetExceededError(current_cost=10.0, max_budget=5.0) + ), + ) + cache = _validation_cache() + + with pytest.raises(litellm.BudgetExceededError): + await resolve_and_validate_end_user_id( + raw_end_user_id="customer-over-budget", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 08035fb7173..68e1636d380 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -597,6 +597,315 @@ def test_get_end_user_id_falls_back_to_deprecated_user_header_name(): assert result == "user-legacy" +class TestCoerceUserIdToStr: + """Unit tests for the _coerce_user_id_to_str helper.""" + + def test_plain_string_is_returned_verbatim(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("alice@example.com") == "alice@example.com" + + def test_string_is_stripped(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(" bob ") == "bob" + + def test_codex_opaque_identifier_is_preserved(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + assert _coerce_user_id_to_str(codex_id) == codex_id + + def test_none_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(None) is None + + def test_empty_string_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("") is None + assert _coerce_user_id_to_str(" ") is None + + def test_dict_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + payload = { + "device_id": "abc", + "account_uuid": "", + "session_id": "c284b8cb", + } + assert _coerce_user_id_to_str(payload) is None + + def test_list_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(["a", "b"]) is None + + def test_json_encoded_dict_string_passes_through_by_default(self): + """JSON-encoded dict strings are preserved unless opt-in flag is on. + + This preserves backwards compatibility: existing deployments that + intentionally pass JSON-encoded user identifiers keep working. + """ + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + blob = ( + '{"device_id":"d5abe9199ee7759a0558974e9371e78c7b38d7621aae26d6609c1de61af6afb0",' + '"account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + assert _coerce_user_id_to_str(blob) == blob + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_dict_string_returns_none_when_validation_enabled(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + # Same broken shape we saw in spend logs, but pre-stringified to JSON. + blob = ( + '{"device_id":"d5abe9199ee7759a0558974e9371e78c7b38d7621aae26d6609c1de61af6afb0",' + '"account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + assert _coerce_user_id_to_str(blob) is None + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_list_string_passes_through_by_default(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + assert _coerce_user_id_to_str('["a","b"]') == '["a","b"]' + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_list_string_returns_none_when_validation_enabled(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + assert _coerce_user_id_to_str('["a","b"]') is None + finally: + litellm.validate_end_user_id_in_db = original + + def test_int_returns_str(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(12345) == "12345" + + def test_bool_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + # bool is an int subclass — reject explicitly, never produce "True"/"False". + assert _coerce_user_id_to_str(True) is None + assert _coerce_user_id_to_str(False) is None + + def test_brace_string_that_isnt_json_is_kept(self): + """A string starting with `{` but failing to parse stays as-is.""" + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("{not json") == "{not json" + + +class TestGetEndUserIdDropsMalformedBodyValues: + """Tests that get_end_user_id_from_request_body drops dict-shaped values + rather than stringifying them into spend logs.""" + + def test_dict_user_falls_through_to_litellm_metadata(self): + request_body = { + "user": { + "device_id": "abc", + "session_id": "c284b8cb", + }, + "litellm_metadata": {"user": "alice@example.com"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_user_with_no_other_sources_returns_none(self): + request_body = { + "user": {"device_id": "abc", "session_id": "xyz"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_json_encoded_user_string_passes_through_by_default(self): + """JSON-encoded user strings pass through unless validation is opted in. + + Gating behind ``litellm.validate_end_user_id_in_db`` keeps existing + deployments that send JSON-encoded identifiers working until they + explicitly opt into the stricter extraction. + """ + import litellm + + blob = ( + '{"device_id":"d5abe9199ee7759a","account_uuid":"",' + '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + request_body = {"user": blob} + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + finally: + litellm.validate_end_user_id_in_db = original + + assert result == blob + + def test_json_encoded_user_string_returns_none_when_validation_enabled(self): + import litellm + + request_body = { + "user": ( + '{"device_id":"d5abe9199ee7759a","account_uuid":"",' + '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ), + } + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + finally: + litellm.validate_end_user_id_in_db = original + + assert result is None + + def test_plain_string_user_is_preserved(self): + request_body = {"user": "alice@example.com"} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_codex_opaque_user_is_preserved(self): + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + request_body = {"user": codex_id} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == codex_id + + def test_int_user_is_coerced_to_string(self): + request_body = {"user": 12345} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "12345" + + def test_list_user_falls_through(self): + request_body = { + "user": ["a", "b"], + "safety_identifier": "alice@example.com", + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_safety_identifier_returns_none(self): + request_body = { + "safety_identifier": {"device_id": "abc"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_dict_metadata_user_id_returns_none(self): + request_body = { + "metadata": {"user_id": {"device_id": "abc"}}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_whitespace_user_falls_through(self): + request_body = {"user": " ", "safety_identifier": "alice@example.com"} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_user_header_falls_through_to_body(self): + """A dict-shaped value in a configured user-id header is dropped, not stringified.""" + general_settings = {"user_header_name": "x-custom-user-id"} + # A header value will normally be a str, but be defensive: the coercion + # must drop anything that isn't a usable identifier. + headers = {"x-custom-user-id": {"device_id": "abc"}} + request_body = {"user": "alice@example.com"} + + with ( + patch( + "litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", + return_value=None, + ), + patch("litellm.proxy.proxy_server.general_settings", general_settings), + ): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers=headers + ) + + assert result == "alice@example.com" + + def _make_deployment_dict( model_name: str, tpm: Optional[int] = None, rpm: Optional[int] = None ) -> dict: diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 442625c75a7..defd3bbcdcd 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3335,3 +3335,125 @@ async def test_master_key_auth_substitutes_alias_for_api_key(): finally: for k, v in _orig.items(): setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_sets_end_user_id_when_builder_skips_it(): + """Defense-in-depth: ``_user_api_key_auth_builder`` has multiple + early-return paths (master_key=None, /user/auth route, JWT + short-circuits) that bypass the end-user resolution block. The wrapper + must still attribute spend logs to the request-supplied end-user when + none of those paths set it. + + Krrish flagged the removal of this fallback as a regression risk; this + test pins the behaviour so future refactors don't silently drop it. + """ + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + # builder did NOT set end_user_id (e.g. master_key=None early return) + assert builder_token.end_user_id is None + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = json.dumps( + {"model": "gpt-4o", "user": "alice@example.com"} + ).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + # Stub the builder so the test doesn't have to traverse the full + # auth state machine; we only care about the wrapper's safety net. + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + ): + result = await user_api_key_auth(request=request, api_key="Bearer sk-test") + + # Validation flag is False by default → pass-through, raw value lands + # on the auth obj instead of being silently dropped. + assert result.end_user_id == "alice@example.com" + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder(): + """When the builder already resolved the end-user id (the primary + path), the wrapper-level safety net must not run a second resolution + pass — that would re-extract from the request body and could + overwrite a value the builder explicitly chose to set.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth( + api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id" + ) + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = json.dumps( + {"model": "gpt-4o", "user": "different-id-from-body"} + ).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + ) as mock_resolve, + ): + result = await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert result.end_user_id == "builder-resolved-id" + mock_resolve.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) diff --git a/tests/test_litellm/responses/test_sse_output_recovery.py b/tests/test_litellm/responses/test_sse_output_recovery.py new file mode 100644 index 00000000000..c8f3325a624 --- /dev/null +++ b/tests/test_litellm/responses/test_sse_output_recovery.py @@ -0,0 +1,57 @@ +"""Tests for litellm.responses.sse_output_recovery helpers.""" + +from litellm.responses.sse_output_recovery import ( + _MAX_CONTENT_INDEX, + record_output_text_chunk, +) + + +def test_text_chunk_with_oversized_content_index_is_dropped(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": _MAX_CONTENT_INDEX + 1, + "text": "ignored", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + item = text_only_items[0] + assert item["content"] == [] + + +def test_text_chunk_with_negative_content_index_is_dropped(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": -1, + "text": "ignored", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + assert text_only_items[0]["content"] == [] + + +def test_text_chunk_at_max_content_index_is_recorded(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": _MAX_CONTENT_INDEX, + "text": "kept", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + content = text_only_items[0]["content"] + assert len(content) == _MAX_CONTENT_INDEX + 1 + assert content[_MAX_CONTENT_INDEX]["text"] == "kept" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index c3f93078557..9454e03e918 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -7,8 +7,10 @@ and one has explicit zero-cost pricing in model_info, the other deployment should still use the built-in pricing. """ +import copy import os import sys +from unittest.mock import patch import pytest @@ -19,6 +21,16 @@ sys.path.insert( import litellm from litellm import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.utils import _invalidate_model_cost_lowercase_map + + +def _restore_model_cost_entries(original_entries): + for key, value in original_entries.items(): + if value is None: + litellm.model_cost.pop(key, None) + else: + litellm.model_cost[key] = value + _invalidate_model_cost_lowercase_map() def test_should_not_pollute_shared_key_with_zero_cost_pricing(): @@ -323,3 +335,70 @@ def test_responses_prefix_stripped_alias_registered_for_add_deployment(): ) is True ) + + +def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): + """ + ChatGPT aliases that share the same backend model should not be able to + downgrade the shared backend key from responses -> chat during router setup. + """ + from litellm.main import responses_api_bridge_check + + backend_model = "chatgpt/gpt-5.4" + model_keys = { + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + "chatgpt-shared-mode-base": copy.deepcopy( + litellm.model_cost.get("chatgpt-shared-mode-base") + ), + "chatgpt-shared-mode-alias": copy.deepcopy( + litellm.model_cost.get("chatgpt-shared-mode-alias") + ), + } + + try: + backend_entry = copy.deepcopy(model_keys[backend_model]) or {} + backend_entry["litellm_provider"] = "chatgpt" + backend_entry["mode"] = "responses" + litellm.model_cost[backend_model] = backend_entry + _invalidate_model_cost_lowercase_map() + + router = Router(model_list=[]) + with patch.object( + Router, "_add_deployment", lambda self, deployment: deployment + ): + router._create_deployment( + deployment_info={}, + _model_name="chatgpt/gpt-5.4", + _litellm_params={ + "model": "gpt-5.4", + "custom_llm_provider": "chatgpt", + }, + _model_info={ + "id": "chatgpt-shared-mode-base", + "mode": "responses", + }, + ) + router._create_deployment( + deployment_info={}, + _model_name="chatgpt/gpt-5.4-medium", + _litellm_params={ + "model": "gpt-5.4", + "custom_llm_provider": "chatgpt", + }, + _model_info={ + "id": "chatgpt-shared-mode-alias", + "mode": "chat", + }, + ) + + assert litellm.model_cost[backend_model]["mode"] == "responses" + assert "mode" in litellm.model_cost[backend_model] + + bridge_model_info, bridge_model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="chatgpt", + ) + assert bridge_model == "gpt-5.4" + assert bridge_model_info["mode"] == "responses" + finally: + _restore_model_cost_entries(model_keys) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index bc60375f906..de286aede93 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -754,6 +754,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_dbu_cost_per_token": {"type": "number"}, "annotation_cost_per_page": {"type": "number"}, "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_credit": {"type": "number"}, "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, "litellm_provider": {"type": "string"}, @@ -855,6 +856,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_adaptive_thinking": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, + "supports_output_config": {"type": "boolean"}, "tool_use_system_prompt_tokens": {"type": "number"}, "tpm": {"type": "number"}, "provider_specific_entry": {"type": "object"}, diff --git a/tests/test_spend_logs.py b/tests/test_spend_logs.py index fcd2bbf4a1d..c575fa07551 100644 --- a/tests/test_spend_logs.py +++ b/tests/test_spend_logs.py @@ -158,6 +158,9 @@ async def generate_team(session: aiohttp.ClientSession, org_id: str) -> dict: return await response.json() +@pytest.mark.skip( + reason="Flaky in CI: /spend/logs?request_id=... returns 500 even after a 20s wait for the spend log to be written. Same write-then-read race against the spend logs DB as test_spend_logs. Spend-log accuracy is covered by tests/test_litellm/proxy/spend_tracking/ and the proxy_spend_accuracy_tests CircleCI job." +) @pytest.mark.asyncio async def test_spend_logs_with_org_id(): """ diff --git a/tests/test_team_members.py b/tests/test_team_members.py index 415b3f07fc9..4cf85af6410 100644 --- a/tests/test_team_members.py +++ b/tests/test_team_members.py @@ -206,6 +206,9 @@ def test_error_handling(api_client): api_client.get_team_info("invalid-team-id") +@pytest.mark.skip( + reason="Flaky in CI: /team/info?team_id=... intermittently returns 404 after add_team_member calls, same race documented for test_add_multiple_members. Duplicate-prevention is covered by test_update_team_members_list_duplicate_prevention in tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py." +) def test_duplicate_user_addition(api_client, new_team): """Test that adding the same user twice is handled appropriately""" # Add user first time