diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index fcaf58099b9..6e16ed8f187 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -123,7 +123,7 @@ class _PROXY_LiteLLMManagedVectorStores( """ # Use the router to create the vector store response: Final = await llm_router.avector_store_create(model=model, **request_data) - return cast(VectorStoreCreateResponse, response) + return response # ============================================================================ # VECTOR STORE CRUD OPERATIONS diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index a5721c6f047..90a083bed2d 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -5,14 +5,16 @@ import contextvars import os from collections.abc import Coroutine, Iterable, Mapping, Sequence from functools import partial -from typing import Any, Final, Literal +from typing import Any, Final, Literal, Protocol import httpx from openai import AsyncOpenAI, OpenAI from openai.types.beta.assistant import Assistant from openai.types.beta.assistant_deleted import AssistantDeleted +from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.router import GenericLiteLLMParams from litellm.utils import ( exception_type, @@ -31,6 +33,21 @@ from .utils import get_optional_params_add_message ####### ENVIRONMENT VARIABLES ################### openai_assistants_api: Final = OpenAIAssistantsAPI() azure_assistants_api: Final = AzureAssistantsAPI() +_EXTRA_BODY: Final = TypeAdapter(object) + + +class _PopsKeys(Protocol): + def pop(self, key: str, default: None, /) -> object: ... + + +def _pop_azure_ad_token(extra_body: object) -> object: + return require_attributes(extra_body, _PopsKeys, "pop").pop("azure_ad_token", None) + + +def _popped_azure_ad_token(extra_body: object) -> str | None: + token: Final = _pop_azure_ad_token(extra_body) + return token if isinstance(token, str) else None + ### ASSISTANTS ### @@ -146,10 +163,10 @@ def get_assistants( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -331,10 +348,10 @@ def create_assistants( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -481,10 +498,10 @@ def delete_assistant( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -657,10 +674,10 @@ def create_thread( api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -804,10 +821,10 @@ def get_thread( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -982,10 +999,10 @@ def add_message( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -1134,10 +1151,10 @@ def get_messages( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token: str | None = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _popped_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") @@ -1333,10 +1350,10 @@ def run_thread( or get_secret("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) azure_ad_token = None if extra_body is not None: - azure_ad_token = extra_body.pop("azure_ad_token", None) + azure_ad_token = _pop_azure_ad_token(extra_body) else: azure_ad_token = get_secret("AZURE_AD_TOKEN") diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 62ccb07f695..1c830e45464 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils import batch_utils as core_batch_utils from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo @@ -849,14 +850,14 @@ def _get_batch_job_usage_from_response_body( return usage -def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> Mapping[str, Any]: +def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, object]) -> Mapping[str, object]: """ Get the ``result`` object from a line of an Anthropic message batch results JSONL file. Anthropic batch results lines look like: ``{"custom_id": ..., "result": {"type": "succeeded", "message": {..., "usage": {...}}}}`` """ - return batch_results_line.get("result", None) or {} + return decoded_json_object(batch_results_line.get("result", None) or {}) def _get_response_from_batch_job_output_file( @@ -866,7 +867,9 @@ def _get_response_from_batch_job_output_file( Get the response from the batch job output file """ if custom_llm_provider == "anthropic": - return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {} + return decoded_json_object( + _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {} + ) if custom_llm_provider == "bedrock": return batch_job_output_file.get("modelOutput", None) or {} _response: Final = cast( # cast-ok: batch output response comes from provider JSON diff --git a/litellm/batches/main.py b/litellm/batches/main.py index d88f549b06b..212205a9d7b 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -15,13 +15,15 @@ import contextvars import os from collections.abc import Coroutine from functools import partial -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, cast import httpx from openai.types.batch import BatchRequestCounts +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -60,10 +62,15 @@ azure_batches_instance: Final = AzureBatchesAPI() vertex_ai_batches_instance: Final = VertexAIBatchPrediction(gcs_bucket_name="") anthropic_batches_instance: Final = AnthropicBatchesHandler() xai_batches_instance: Final = XAIBatchesHandler() +_EXTRA_BODY: Final = TypeAdapter(object) base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +class _PopsKeys(Protocol): + def pop(self, key: str, default: None, /) -> object: ... + + def _resolve_timeout( optional_params: GenericLiteLLMParams, kwargs: dict[str, Any], @@ -485,9 +492,9 @@ def _handle_retrieve_batch_providers_without_provider_config( or get_secret_str("AZURE_API_KEY") ) - extra_body: Final = optional_params.get("extra_body", {}) + extra_body: Final = _EXTRA_BODY.validate_python(optional_params.get("extra_body", {})) if extra_body is not None: - extra_body.pop("azure_ad_token", None) + require_attributes(extra_body, _PopsKeys, "pop").pop("azure_ad_token", None) else: get_secret_str("AZURE_AD_TOKEN") diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index ab5d2332c21..a3cf54916f1 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -19,7 +19,7 @@ import datetime import inspect import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping -from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar +from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError from typing_extensions import TypedDict @@ -72,6 +72,8 @@ from litellm.litellm_core_utils.core_helpers import ( as_object_list, as_str_mapping, get_parent_otel_span_from_kwargs, + json_contains, + require_attributes, ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -246,6 +248,15 @@ def _stream_replay_requested(kwargs: Mapping[str, object]) -> bool: return converted_stream_requested(kwargs) and not kwargs.get("_agentic_loop_depth") +class _CacheControls(Protocol): + def get(self, key: str, default: bool, /) -> object: ... + + +def _cache_reads_disabled(kwargs: Mapping[str, object]) -> bool: + cache_controls: Final = require_attributes(kwargs.get("cache", {}), _CacheControls, "get") + return cache_controls.get("no-cache", False) is True + + def _should_defer_streaming_cache_hit_callbacks(*, cached_result: object) -> bool: """ When the cache hit is replayed as a stream, do not run success callbacks at cache-hit time. @@ -403,7 +414,7 @@ class LLMCachingHandler: if ( (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True ) and ( - kwargs.get("cache", {}).get("no-cache", False) is not True + not _cache_reads_disabled(kwargs) ): # allow users to control returning cached responses from the completion function args = args or () final_embedding_cached_response: EmbeddingResponse | None = None @@ -521,7 +532,7 @@ class LLMCachingHandler: kwargs: dict[str, Any], args: tuple[object, ...] | None = None, ) -> CachingHandlerResponse: - cached_result: Any | None = None + cached_result: object | None = None # Check if caching should be performed BEFORE doing expensive kwargs copy if _is_response_cache_excluded(model=model, kwargs=kwargs): @@ -550,7 +561,7 @@ class LLMCachingHandler: verbose_logger.debug("LiteLLM Cache: cached response has no output, treating it as a miss") return CachingHandlerResponse(cached_result=None) if cached_result is not None: - if "detail" in cached_result: + if json_contains(cached_result, "detail"): # implies an error occurred pass else: @@ -1460,6 +1471,7 @@ class LLMCachingHandler: Returns: None """ + request_kwargs: Final = dict[str, object](kwargs) litellm_params: Final = { "logger_fn": kwargs.get("logger_fn", None), "acompletion": is_async, @@ -1473,7 +1485,7 @@ class LLMCachingHandler: if litellm.cache is not None: litellm_params["preset_cache_key"] = ( - self.preset_cache_key or litellm.cache.get_preset_cache_key_from_kwargs(**kwargs) + self.preset_cache_key or litellm.cache.get_preset_cache_key_from_kwargs(**request_kwargs) ) else: litellm_params["preset_cache_key"] = None @@ -1483,11 +1495,11 @@ class LLMCachingHandler: user=kwargs.get("user", None), optional_params={}, litellm_params=litellm_params, - input=(kwargs.get("messages", "") if not is_embedding else kwargs.get("input", "")), - api_key=kwargs.get("api_key", None), + input=(request_kwargs.get("messages", "") if not is_embedding else request_kwargs.get("input", "")), + api_key=request_kwargs.get("api_key", None), original_response=str(cached_result), additional_args=None, - stream=kwargs.get("stream", False), + stream=request_kwargs.get("stream", False), custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 560a5ba06b8..501ac2037e1 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -41,6 +41,7 @@ from .base_cache import BaseCache _WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"}) _EMBEDDING_VECTOR: Final = TypeAdapter(object) +_CACHED_RESPONSE: Final = TypeAdapter(object) if TYPE_CHECKING: from litellm.router import Router @@ -175,14 +176,13 @@ class QdrantSemanticCache(BaseCache): else: raise Exception("Error while creating new collection") - def _get_cache_logic(self, cached_response: Any): + def _get_cache_logic(self, cached_response: Any) -> object: if cached_response is None: return cached_response try: - cached_response = json.loads(cached_response) # Convert string to dictionary + return _CACHED_RESPONSE.validate_python(json.loads(cached_response)) except Exception: - cached_response = ast.literal_eval(cached_response) - return cached_response + return _CACHED_RESPONSE.validate_python(ast.literal_eval(cached_response)) def _get_qdrant_cache_key_filter(self, key: str) -> dict: return { diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 7caa0c62571..db090c9b007 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator, Coroutine, Mapping from typing import Final, cast from openai.types.chat import ChatCompletionStreamOptionsParam +from typing_extensions import TypedDict import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -15,6 +16,10 @@ from .transformation import GoogleGenAIAdapter GOOGLE_GENAI_ADAPTER: Final = GoogleGenAIAdapter() +class _TranslatedCompletionKwargs(TypedDict, total=False, extra_items=object): + pass + + def _mark_failure_logged_by_inner_completion(adapter_kwargs: Mapping[str, object]) -> None: outer_logging_obj: Final = adapter_kwargs.get("litellm_logging_obj") if isinstance(outer_logging_obj, LiteLLMLoggingObj): @@ -55,7 +60,7 @@ class GenerateContentToCompletionHandler: **(extra_kwargs or {}), ) - completion_kwargs: Final = dict(completion_request) + completion_kwargs: Final = dict[str, object](completion_request) # Forward extra_kwargs that should be passed to completion call if extra_kwargs is not None: @@ -76,7 +81,7 @@ class GenerateContentToCompletionHandler: if stream_options is not None: completion_kwargs["stream_options"] = stream_options - return GenerateContentCompletionKwargs(**completion_kwargs) + return GenerateContentCompletionKwargs(**_TranslatedCompletionKwargs(**completion_kwargs)) @staticmethod async def async_generate_content_handler( diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 40c9c7bc36c..71000fec81a 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -40,6 +40,7 @@ else: # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() _GENERATE_CONTENT_RESULT: Final = TypeAdapter(object) +_GENERATE_CONTENT_PAYLOAD: Final = TypeAdapter(object) ################################################# @@ -368,8 +369,8 @@ def generate_content( response: Final = _GENERATE_CONTENT_RESULT.validate_python( base_llm_http_handler.generate_content_handler( model=setup_result.model, - contents=contents, - tools=tools, + contents=_GENERATE_CONTENT_PAYLOAD.validate_python(contents), + tools=_GENERATE_CONTENT_PAYLOAD.validate_python(tools), generate_content_provider_config=setup_result.generate_content_provider_config, generate_content_config_dict=setup_result.generate_content_config_dict, custom_llm_provider=setup_result.custom_llm_provider, @@ -465,10 +466,10 @@ async def agenerate_content_stream( return _GENERATE_CONTENT_RESULT.validate_python( await base_llm_http_handler.generate_content_handler( model=setup_result.model, - contents=contents, + contents=_GENERATE_CONTENT_PAYLOAD.validate_python(contents), generate_content_provider_config=setup_result.generate_content_provider_config, generate_content_config_dict=setup_result.generate_content_config_dict, - tools=tools, + tools=_GENERATE_CONTENT_PAYLOAD.validate_python(tools), custom_llm_provider=setup_result.custom_llm_provider, litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, @@ -556,10 +557,10 @@ def generate_content_stream( return _GENERATE_CONTENT_RESULT.validate_python( base_llm_http_handler.generate_content_handler( model=setup_result.model, - contents=contents, + contents=_GENERATE_CONTENT_PAYLOAD.validate_python(contents), generate_content_provider_config=setup_result.generate_content_provider_config, generate_content_config_dict=setup_result.generate_content_config_dict, - tools=tools, + tools=_GENERATE_CONTENT_PAYLOAD.validate_python(tools), custom_llm_provider=setup_result.custom_llm_provider, litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 249bb2cebb7..fa8105b30e8 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -37,7 +37,7 @@ from litellm.constants import ( from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import decoded_json_object, has_attributes from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider @@ -722,7 +722,7 @@ class ModelEndpoint: return await litellm.anthropic.messages.acreate(**kwargs) if route == ROUTE_CHAT: if not kwargs.get("stream"): - return await litellm.acompletion(**kwargs) + return await litellm.acompletion(**decoded_json_object(kwargs)) stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field return await litellm.acompletion( **{ @@ -731,7 +731,7 @@ class ModelEndpoint: STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options, } ) - return await litellm.aresponses(**kwargs) + return await litellm.aresponses(**decoded_json_object(kwargs)) async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response: kwargs = self._sdk_kwargs(body) diff --git a/litellm/images/main.py b/litellm/images/main.py index 79bb07fc394..b2da69bb452 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -6,6 +6,7 @@ from functools import partial from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload from pydantic import TypeAdapter, ValidationError +from typing_extensions import TypedDict if TYPE_CHECKING: from litellm.images.utils import ImageEditRequestUtils @@ -698,7 +699,7 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: init_response = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict) or isinstance(init_response, ImageResponse): ## CACHING SCENARIO if isinstance(init_response, dict): - init_response = ImageResponse(**init_response) + init_response = _image_response_from_fields(init_response) response = init_response elif asyncio.iscoroutine(init_response): response = await init_response @@ -717,6 +718,14 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: ) +class _ImageResponseFields(TypedDict, total=False, extra_items=object): + pass + + +def _image_response_from_fields(fields: Mapping[str, object]) -> ImageResponse: + return ImageResponse(**_ImageResponseFields(**fields)) + + @client def image_variation( image: FileTypes, diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index ef7442143ca..bc1a142d133 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -13,6 +13,7 @@ from types import MappingProxyType from typing import Any, Final, Protocol from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -57,6 +58,10 @@ class _SupportsItems(Protocol): def items(self) -> Iterable[tuple[object, object]]: ... +class _SpendLogKwargs(TypedDict, total=False, extra_items=object): + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + + def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]: if value is None: return @@ -243,7 +248,7 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: self._log(kwargs) - def _log(self, kwargs: Mapping[str, Any]) -> None: + def _log(self, kwargs: _SpendLogKwargs) -> None: try: payload = kwargs.get("standard_logging_object") if payload is None or _is_trace_ingest(payload): diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py index 9b685996a41..77079818908 100644 --- a/litellm/integrations/cloudzero/cz_resource_names.py +++ b/litellm/integrations/cloudzero/cz_resource_names.py @@ -17,10 +17,16 @@ """CloudZero Resource Names (CZRN) generation and validation for LiteLLM resources.""" import re +from collections.abc import Mapping from enum import Enum -from typing import Any, Final, cast +from typing import Final, Protocol, cast import litellm +from litellm.litellm_core_utils.core_helpers import require_attributes + + +class _SupportsLower(Protocol): + def lower(self) -> str: ... class CZEntityType(str, Enum): @@ -35,7 +41,7 @@ class CZRNGenerator: def __init__(self): """Initialize CZRN generator.""" - def create_from_litellm_data(self, row: dict[str, Any]) -> str: + def create_from_litellm_data(self, row: Mapping[str, object]) -> str: """Create a CZRN from LiteLLM daily spend data. CZRN format: czrn:::::: @@ -49,19 +55,21 @@ class CZRNGenerator: - cloud-local-id: model """ service_type: Final = "litellm" - provider: Final = self._normalize_provider(row.get("custom_llm_provider", "unknown")) + provider: Final = self._normalize_provider( + require_attributes(row.get("custom_llm_provider", "unknown"), _SupportsLower, "lower") + ) region: Final = "cross-region" # Use the actual entity_id (team_id or user_id) as the owner account team_id: Final = row.get("team_id", "unknown") - owner_account_id: Final = self._normalize_component(team_id) + owner_account_id: Final = self._normalize_owner(team_id) resource_type: Final = "llm-usage" # Create a unique identifier with just the model (entity info already in owner_account_id) model: Final = row.get("model", "unknown") - cloud_local_id: Final = model + cloud_local_id: Final = f"{model}" return self.create_from_components( service_type=service_type, @@ -112,7 +120,7 @@ class CZRNGenerator: return cast(tuple[str, str, str, str, str, str], match.groups()) - def _normalize_provider(self, provider: str) -> str: + def _normalize_provider(self, provider: _SupportsLower) -> str: """Normalize provider names to standard CZRN format.""" # Map common provider names to CZRN standards provider_map: Final = { @@ -141,16 +149,14 @@ class CZRNGenerator: return "unknown" # Convert to lowercase unless uppercase is allowed - if not allow_uppercase: - component = component.lower() + return self._clean_component(component if allow_uppercase else component.lower()) - # Replace invalid characters with hyphens - component = re.sub(r"[^a-zA-Z0-9-]", "-", component) + def _normalize_owner(self, owner: object) -> str: + if not owner: + return "unknown" + return self._clean_component(require_attributes(owner, _SupportsLower, "lower").lower()) - # Remove consecutive hyphens - component = re.sub(r"-+", "-", component) - - # Remove leading/trailing hyphens - component = component.strip("-") - - return component or "unknown" + def _clean_component(self, component: str) -> str: + hyphenated: Final = re.sub(r"[^a-zA-Z0-9-]", "-", component) + collapsed: Final = re.sub(r"-+", "-", hyphenated) + return collapsed.strip("-") or "unknown" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 385522c2435..b582d6b8c4f 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -759,7 +759,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac response: object, plan: AgenticLoopPlan, kwargs: dict, - ) -> Any: + ) -> object: """ Post-process the response returned by the agentic-loop follow-up call. diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 0e52ca0b62f..dfa3cb97f5f 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -5,14 +5,14 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d import re from collections.abc import Mapping from pathlib import Path -from typing import Any, Final +from typing import Any, Final, Protocol import yaml from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes class _PromptFileJson(TypedDict): @@ -27,6 +27,14 @@ class PromptJsonEntry(TypedDict): metadata: ReadOnly[dict[str, object]] +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _metadata_section(metadata: Mapping[str, object], name: str) -> _SupportsGet: + return require_attributes(metadata.get(name, {}), _SupportsGet, "get") + + def strip_version_suffix(prompt_id: str) -> str | None: base, separator, version = prompt_id.rpartition(".v") if separator and base and version.isdigit(): @@ -51,8 +59,8 @@ class PromptTemplate: restricted_keys: Final = ["model", "input", "output"] self.model = self.metadata.get("model") self.input_schema = self.metadata.get("input", {}).get("schema", {}) - self.output_format = self.metadata.get("output", {}).get("format") - self.output_schema = self.metadata.get("output", {}).get("schema", {}) + self.output_format = _metadata_section(self.metadata, "output").get("format") + self.output_schema = _metadata_section(self.metadata, "output").get("schema", {}) self.optional_params = {} for key in self.metadata: if key not in restricted_keys: diff --git a/litellm/integrations/focus/destinations/gcs_destination.py b/litellm/integrations/focus/destinations/gcs_destination.py index 2f7052597ff..82f1a65c217 100644 --- a/litellm/integrations/focus/destinations/gcs_destination.py +++ b/litellm/integrations/focus/destinations/gcs_destination.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Mapping from datetime import timezone -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase @@ -21,14 +22,14 @@ class FocusGCSDestination(GCSBucketBase, FocusDestination): self, *, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, str] | None = None, ) -> None: - config = config or {} - bucket_name: Final = config.get("bucket_name") + resolved_config: Final[Mapping[str, str]] = config or {} + bucket_name: Final = resolved_config.get("bucket_name") if not bucket_name: raise ValueError("bucket_name must be provided for GCS destination") super().__init__(bucket_name=bucket_name) - service_account_json: Final = config.get("service_account_json") + service_account_json: Final = resolved_config.get("service_account_json") if service_account_json is not None: self.path_service_account_json = service_account_json self.prefix = prefix.rstrip("/") diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 0f7dcdbcc81..f04604ddad8 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -6,7 +6,7 @@ import re import uuid from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timezone, tzinfo -from typing import Any, Final, Protocol, cast +from typing import Final, Protocol, cast import httpx from pydantic import ConfigDict, Field, TypeAdapter @@ -109,9 +109,58 @@ class LLMResponse(LiteLLMBaseModel): user_metadata: dict[str, object] | None = None +class _GalileoRecord(TypedDict, total=False): + latency_ms: ReadOnly[int] + status_code: ReadOnly[int] + input_text: ReadOnly[str] + output_text: ReadOnly[str] + node_type: ReadOnly[str] + model: ReadOnly[str] + num_input_tokens: ReadOnly[int] + num_output_tokens: ReadOnly[int] + num_total_tokens: ReadOnly[int] + cost: ReadOnly[float | None] + output_logprobs: ReadOnly[Mapping[str, object] | None] + created_at: ReadOnly[str] + tags: ReadOnly[Sequence[str] | None] + user_metadata: ReadOnly[Mapping[str, object] | None] + messages: ReadOnly[Sequence[object]] + + +class _GalileoLoggingKwargs(TypedDict, total=False, extra_items=object): + standard_logging_object: ReadOnly[GalileoStandardLoggingFields | None] + + +def _record_messages(messages: object) -> object: + return messages.get("messages") if isinstance(messages, dict) else messages + + +def _galileo_record(response: LLMResponse, messages: object) -> _GalileoRecord: + record: Final[_GalileoRecord] = { + "latency_ms": response.latency_ms, + "status_code": response.status_code, + "input_text": response.input_text, + "output_text": response.output_text, + "node_type": response.node_type, + "model": response.model, + "num_input_tokens": response.num_input_tokens, + "num_output_tokens": response.num_output_tokens, + "num_total_tokens": response.num_total_tokens, + "cost": response.cost, + "output_logprobs": response.output_logprobs, + "created_at": response.created_at, + "tags": response.tags, + "user_metadata": response.user_metadata, + } + record_messages: Final = _record_messages(messages) + if isinstance(record_messages, list) and record_messages: + return {**record, "messages": record_messages} + return record + + class GalileoObserve(CustomLogger): def __init__(self) -> None: - self.in_memory_records: list[Mapping[str, object]] = [] + self.in_memory_records: list[_GalileoRecord] = [] self.batch_size = 1 self.api_key = os.getenv("GALILEO_API_KEY") self.project_id = os.getenv("GALILEO_PROJECT_ID") @@ -278,7 +327,7 @@ class GalileoObserve(CustomLogger): return created_at @staticmethod - def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, object]: + def _token_metrics_from_record(record: _GalileoRecord) -> dict[str, object]: num_input_tokens: Final = int(record.get("num_input_tokens") or 0) num_output_tokens: Final = int(record.get("num_output_tokens") or 0) num_total_tokens = int(record.get("num_total_tokens") or 0) @@ -296,7 +345,7 @@ class GalileoObserve(CustomLogger): @staticmethod def _record_to_v2_span( - record: Mapping[str, Any], + record: _GalileoRecord, *, trace_id: str, span_id: str, @@ -322,12 +371,13 @@ class GalileoObserve(CustomLogger): **GalileoObserve._token_metrics_from_record(record), }, } - if record.get("tags"): - span["tags"] = record["tags"] + tags: Final = record.get("tags") + if tags: + span["tags"] = tags return span @staticmethod - def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, object]: + def _record_to_v2_trace(record: _GalileoRecord) -> dict[str, object]: trace_id: Final = str(uuid.uuid4()) span_id: Final = str(uuid.uuid4()) created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", "")) @@ -347,7 +397,7 @@ class GalileoObserve(CustomLogger): "spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)], } - def _build_traces_payload(self, records: Sequence[Mapping[str, object]]) -> dict[str, object]: + def _build_traces_payload(self, records: Sequence[_GalileoRecord]) -> dict[str, object]: payload: Final[dict[str, object]] = { "traces": [self._record_to_v2_trace(record) for record in records], "logging_method": "api_direct", @@ -668,7 +718,7 @@ class GalileoObserve(CustomLogger): async def async_log_success_event( self, - kwargs: Mapping[str, object], + kwargs: _GalileoLoggingKwargs, response_obj: object, start_time: object, end_time: object, @@ -686,7 +736,7 @@ class GalileoObserve(CustomLogger): async def _async_log_success_event_impl( self, - kwargs: Mapping[str, Any], + kwargs: _GalileoLoggingKwargs, response_obj: object, start_time: object, end_time: object, @@ -752,12 +802,7 @@ class GalileoObserve(CustomLogger): created_at=GalileoObserve._format_created_at(start_ts), ) - request_dict: Final = request_record.model_dump() - if isinstance(messages, dict): - messages = messages.get("messages") - if isinstance(messages, list) and messages: - request_dict["messages"] = messages - self.in_memory_records.append(request_dict) + self.in_memory_records.append(_galileo_record(request_record, messages)) verbose_logger.debug("Galileo Logger: queued record, in_memory=%d", len(self.in_memory_records)) # Bound the buffer so persistent flush failures cannot grow it diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index adaacd615a8..c51c513b633 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -3,13 +3,14 @@ GitLab prompt manager with configurable prompts folder. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, TypeVar +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.litellm_core_utils.core_helpers import require_attributes if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -40,6 +41,18 @@ class GitLabCachedPrompt(TypedDict): optional_params: ReadOnly[Mapping[str, object]] +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _input_schema(metadata: Mapping[str, object]) -> object: + return require_attributes(metadata.get("input", {}), _SupportsGet, "get").get("schema", {}) + + +def _optional_params(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} + + def encode_prompt_id(raw_id: str) -> str: """Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'""" if raw_id.startswith(GITLAB_PREFIX): @@ -68,8 +81,8 @@ class GitLabPromptTemplate: self.model = model or metadata.get("model") self.temperature = metadata.get("temperature") self.max_tokens = metadata.get("max_tokens") - self.input_schema = metadata.get("input", {}).get("schema", {}) - self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} + self.input_schema = _input_schema(metadata) + self.optional_params = _optional_params(metadata) def __repr__(self): return f"GitLabPromptTemplate(id='{self.template_id}', model='{self.model}')" @@ -84,7 +97,7 @@ class GitLabTemplateManager: def __init__( self, - gitlab_config: dict[str, Any], + gitlab_config: Mapping[str, Any], prompt_id: str | None = None, ref: str | None = None, gitlab_client: GitLabClient | None = None, @@ -297,7 +310,7 @@ class GitLabPromptManager(CustomPromptManagement): def __init__( self, - gitlab_config: dict[str, object], + gitlab_config: Mapping[str, object], prompt_id: str | None = None, ref: str | None = None, # tag/branch/SHA override gitlab_client: GitLabClient | None = None, diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index bd68ffe8cd6..842eaaf847d 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -22,6 +22,7 @@ from litellm.integrations.langfuse.langfuse_mock_client import ( from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, filter_exceptions_from_params, + json_contains, reconstruct_model_name, safe_deep_copy, ) @@ -1311,11 +1312,11 @@ def _end_grounding_span(*, tracing: "LangfuseTracing", parent: "LangfuseObservat ).end() -def log_requester_metadata(clean_metadata: Mapping[str, Any]): +def log_requester_metadata(clean_metadata: Mapping[str, object]): returned_metadata: Final = {} - requester_metadata: Final = clean_metadata.get("requester_metadata") or {} + requester_metadata: Final[object] = clean_metadata.get("requester_metadata") or {} for k, v in clean_metadata.items(): - if k not in requester_metadata: + if not json_contains(requester_metadata, k): returned_metadata[k] = v returned_metadata.update({"requester_metadata": requester_metadata}) diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index 1ceb9aa763f..d4fefe54333 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -2,15 +2,28 @@ import json from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol from pydantic import TypeAdapter from litellm import _logging +from litellm.litellm_core_utils.core_helpers import require_attributes _DECODED_JSON: Final = TypeAdapter(object) +class _SpanIdentifiers(Protocol): + @property + def trace_id(self) -> str | None: ... + + @property + def id(self) -> str | None: ... + + +def _trace_and_span_ids(span: _SpanIdentifiers) -> tuple[str | None, str | None]: + return span.trace_id, span.id + + def normalize_provider_name(provider: str | None) -> str | None: """ Normalize LiteLLM provider names to standardized string names. @@ -77,7 +90,7 @@ def extract_opik_metadata( def extract_span_identifiers( - current_span_data: Any, + current_span_data: object, ) -> tuple[str | None, str | None]: """ Extract trace_id and parent_span_id from current_span_data. @@ -95,7 +108,7 @@ def extract_span_identifiers( return (current_span_data.get("trace_id"), current_span_data.get("id")) try: - return current_span_data.trace_id, current_span_data.id + return _trace_and_span_ids(require_attributes(current_span_data, _SpanIdentifiers, "trace_id", "id")) except AttributeError: _logging.verbose_logger.warning("Unexpected current_span_data format: %s", type(current_span_data)) return None, None diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index f95b5611831..628f35edd50 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -2,8 +2,8 @@ Handler for transforming interactions API requests to litellm.responses requests. """ -from collections.abc import AsyncIterator, Callable, Coroutine, Iterator -from typing import Any, Final +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping +from typing import Final import litellm from litellm.interactions.litellm_responses_transformation.streaming_iterator import ( @@ -72,6 +72,7 @@ class LiteLLMResponsesInteractionsHandler: model=model, input=input, optional_params=optional_params, + custom_llm_provider=custom_llm_provider, **kwargs, ) @@ -108,10 +109,11 @@ class LiteLLMResponsesInteractionsHandler: async def async_interactions_api_handler( self, - responses_request: dict[str, Any], + responses_request: Mapping[str, object], model: str, input: InteractionInput | None, optional_params: InteractionsAPIOptionalRequestParams, + custom_llm_provider: str | None = None, **kwargs, ) -> InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]: """Async handler for interactions API requests.""" @@ -128,7 +130,7 @@ class LiteLLMResponsesInteractionsHandler: litellm_custom_stream_wrapper=responses_response, request_input=input, optional_params=optional_params, - custom_llm_provider=responses_request.get("custom_llm_provider"), + custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) @@ -139,5 +141,5 @@ class LiteLLMResponsesInteractionsHandler: return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( responses_response=responses_api_response, model=model, - custom_llm_provider=responses_request.get("custom_llm_provider"), + custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 53451b04515..93799c72209 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -99,6 +99,11 @@ def _post_hook_overridden(callback: CustomLogger) -> bool: _STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +_HOOK_RESPONSE: Final = TypeAdapter(object) + + +def _plan_response_override(plan: AgenticLoopPlan) -> object | None: + return plan.response_override def _stashed_stream_options(kwargs: Mapping[str, object]) -> Mapping[str, object] | None: @@ -396,8 +401,10 @@ async def _execute_chat_completion_agentic_plan( ) if _post_hook_overridden(callback): try: - response_followup = await callback.async_post_agentic_loop_response_hook( - response=response_followup, plan=plan, kwargs=kwargs + response_followup = _HOOK_RESPONSE.validate_python( + await callback.async_post_agentic_loop_response_hook( + response=response_followup, plan=plan, kwargs=kwargs + ) ) except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") @@ -646,8 +653,8 @@ async def run_chat_completion_agentic_hooks( kwargs=hook_kwargs, ) - if plan.response_override is not None: - return plan.response_override + if (response_override := _plan_response_override(plan)) is not None: + return response_override if plan.terminate: verbose_logger.debug( "Agentic chat loop terminated by callback=%s reason=%s", type(callback).__name__, plan.stop_reason diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 8dec2941e02..71b20f92cf7 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -102,7 +102,7 @@ _LAX_BOOL: Final = TypeAdapter(bool) def as_str_mapping(value: object) -> Mapping[str, object] | None: - return value if isinstance(value, Mapping) else None + return value if isinstance(value, (dict, Mapping)) else None def as_object_list(value: object) -> Sequence[object] | None: @@ -180,7 +180,24 @@ def require_attributes(value: object, shape: type[_Shape], *names: str) -> _Shap raise AttributeError(f"'{type(value).__name__}' object has no attribute '{missing}'", name=missing, obj=value) +class KeywordArguments(Protocol): + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def keyword_arguments(value: object, target: type[object]) -> KeywordArguments: + if has_attributes(value, KeywordArguments, "keys"): + return value + raise TypeError( + f"{target.__module__}.{target.__qualname__}() argument after ** must be a mapping, not {type(value).__name__}" + ) + + def iter_items(value: object) -> Iterable[object]: + listed: Final = as_object_list(value) + if listed is not None: + return listed if isinstance(value, _SupportsIter): return value if isinstance(value, _SupportsGetItem): @@ -746,17 +763,22 @@ def process_response_headers( return {**additional_headers, **({"request-id": request_id} if request_id is not None else {})} -def preserve_upstream_non_openai_attributes( - model_response: "ModelResponseStream", original_chunk: "ModelResponseStream" -): +def preserve_upstream_non_openai_attributes(model_response: "ModelResponseStream", original_chunk: object): """ Preserve non-OpenAI attributes from the original chunk. """ # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings expected_keys: Final = frozenset(type(model_response).model_fields) | frozenset({"usage"}) + chunk_model: Final = original_chunk if isinstance(original_chunk, BaseModel) else None + if chunk_model is None: + raise AttributeError( + f"type object '{type(original_chunk).__name__}' has no attribute 'model_fields'", + name="model_fields", + obj=type(original_chunk), + ) candidate_keys: Final = tuple( key - for key in (*type(original_chunk).model_fields, *(original_chunk.__pydantic_extra__ or {})) + for key in (*type(chunk_model).model_fields, *(chunk_model.__pydantic_extra__ or {})) if key not in expected_keys ) if not candidate_keys: @@ -766,7 +788,7 @@ def preserve_upstream_non_openai_attributes( key, cast( # cast-ok: Pydantic extras are scalar-checked before assignment object, - getattr(original_chunk, key), + getattr(chunk_model, key), ), ) for key in candidate_keys @@ -774,7 +796,7 @@ def preserve_upstream_non_openai_attributes( if all(value is None or isinstance(value, (str, int, float, bool)) for _, value in candidate_values): _assign_attributes(model_response, candidate_values) return - _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) + _assign_attributes(model_response, chunk_model.model_dump(include=set(candidate_keys)).items()) def _assign_attributes(target: BaseModel, values: Iterable[tuple[str, object]]) -> None: diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 6d39f272f3e..16cd076db39 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, should_report_bug, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, json_contains, subscript from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string from litellm.llms.bedrock.common_utils import bedrock_model_reaches_converse_by_default from litellm.llms.bigmodel.common_utils import raise_bigmodel_business_error @@ -1859,26 +1859,32 @@ def _map_together_ai_exception( extra_information: str, ) -> None: try: - error_response = json.loads(error_str) + error_response = _DECODED_ERROR_BODY.validate_python(json.loads(error_str)) except Exception: error_response = {"error": error_str} - if "error" in error_response and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"]: + if json_contains(error_response, "error") and json_contains( + subscript(error_response, "error"), "`inputs` tokens + `max_new_tokens` must be <=" + ): raise ContextWindowExceededError( - message=f"TogetherAIException - {error_response['error']}", + message=f"TogetherAIException - {subscript(error_response, 'error')}", model=model, llm_provider="together_ai", response=getattr(original_exception, "response", None), ) - elif "error" in error_response and "invalid private key" in error_response["error"]: + elif json_contains(error_response, "error") and json_contains( + subscript(error_response, "error"), "invalid private key" + ): raise AuthenticationError( - message=f"TogetherAIException - {error_response['error']}", + message=f"TogetherAIException - {subscript(error_response, 'error')}", llm_provider="together_ai", model=model, response=getattr(original_exception, "response", None), ) - elif "error" in error_response and "INVALID_ARGUMENT" in error_response["error"]: + elif json_contains(error_response, "error") and json_contains( + subscript(error_response, "error"), "INVALID_ARGUMENT" + ): raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", + message=f"TogetherAIException - {subscript(error_response, 'error')}", model=model, llm_provider="together_ai", response=getattr(original_exception, "response", None), @@ -1890,13 +1896,13 @@ def _map_together_ai_exception( llm_provider="together_ai", ) elif ( - "error" in error_response - and "API key doesn't match expected format." in error_response["error"] - or "error_type" in error_response - and error_response["error_type"] == "validation" + json_contains(error_response, "error") + and json_contains(subscript(error_response, "error"), "API key doesn't match expected format.") + or json_contains(error_response, "error_type") + and subscript(error_response, "error_type") == "validation" ): raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", + message=f"TogetherAIException - {subscript(error_response, 'error')}", model=model, llm_provider="together_ai", response=getattr(original_exception, "response", None), @@ -1910,7 +1916,7 @@ def _map_together_ai_exception( ) elif original_exception.status_code == 422: raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", + message=f"TogetherAIException - {subscript(error_response, 'error')}", model=model, llm_provider="together_ai", response=getattr(original_exception, "response", None), diff --git a/litellm/litellm_core_utils/json_validation_rule.py b/litellm/litellm_core_utils/json_validation_rule.py index b6b16aff7a7..030bee48530 100644 --- a/litellm/litellm_core_utils/json_validation_rule.py +++ b/litellm/litellm_core_utils/json_validation_rule.py @@ -1,9 +1,11 @@ import json -from typing import Any, Final +from collections.abc import Mapping +from typing import Final from pydantic import TypeAdapter from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict _DECODED_RESPONSE: Final = TypeAdapter(object) @@ -76,7 +78,7 @@ def normalize_json_schema_types( return schema -def normalize_tool_schema(tool: dict[str, Any]) -> dict[str, Any]: +def normalize_tool_schema(tool: Mapping[str, object]) -> Mapping[str, object]: """ Normalize a tool's parameter schema to use standard JSON Schema lowercase types. @@ -92,10 +94,14 @@ def normalize_tool_schema(tool: dict[str, Any]) -> dict[str, Any]: normalized_tool: Final = tool.copy() # Normalize function parameters if present - if "function" in tool and isinstance(tool["function"], dict): - normalized_tool["function"] = tool["function"].copy() - if "parameters" in tool["function"]: - normalized_tool["function"]["parameters"] = normalize_json_schema_types(tool["function"]["parameters"]) + if "function" not in tool: + return normalized_tool + function: Final = tool["function"] + if is_str_keyed_dict(function): + normalized_function: Final = function.copy() + normalized_tool["function"] = normalized_function + if "parameters" in function: + normalized_function["parameters"] = normalize_json_schema_types(function["parameters"]) return normalized_tool diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8f3101915be..5d21b2762cd 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -25,6 +25,7 @@ from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream from openai import Stream as OpenAIStream from pydantic import BaseModel, JsonValue, TypeAdapter +from typing_extensions import TypedDict import litellm from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase @@ -91,8 +92,11 @@ from litellm.litellm_core_utils.classifier_logging import ( ) from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, + attribute_or_none, + decoded_json_object, finite_cost, get_provider_response_headers_from_hidden_params, + has_attributes, is_expected_client_error, native_messages_reported_cost, provider_reported_usage_cost, @@ -103,6 +107,7 @@ from litellm.litellm_core_utils.core_helpers import ( require_attributes, set_response_cost_in_hidden_params, single_log_line, + subscript, ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params @@ -669,6 +674,10 @@ class UsageObservation(Enum): ESTIMATED = "estimated" +class _UsageFields(TypedDict, total=False, extra_items=object): + pass + + class _A2AUsageSlot(Protocol): usage: Mapping[str, object] | None @@ -685,6 +694,31 @@ class _ParamsLookup(Protocol): def get(self, key: str, default: object, /) -> object: ... +class _HiddenParamsSlot(Protocol): + get: Callable[[str], object] + __setitem__: Callable[[str, object], None] + + +class _HiddenParamsCopier(Protocol): + copy: Callable[[], object] + + +class _MetadataDefaults(Protocol): + setdefault: Callable[[str, Mapping[str, object]], object] + + +class _MetadataSlot(Protocol): + __setitem__: Callable[[str, object], None] + + +class _MetadataUpdate(Protocol): + update: Callable[[Mapping[str, object]], None] + + +class _ExceptionHeaders(Protocol): + headers: object + + def _callback_control_params(litellm_params: object) -> dict[str, object]: mapping: Final = as_str_mapping(litellm_params) if isinstance(mapping, dict): @@ -1682,8 +1716,10 @@ class Logging(LiteLLMLoggingBaseClass): ) # log raw request to provider (like LangFuse) -- if opted in. if self.log_raw_request_response is True or litellm.log_raw_request_response is True: - _litellm_params: Final = self.model_call_details.get("litellm_params", {}) - _metadata: Final = _litellm_params.get("metadata", {}) or {} + _litellm_params: Final = require_attributes( + self._model_call_details_view.get("litellm_params", dict[str, object]()), _ParamsLookup, "get" + ) + _metadata: Final = _litellm_params.get("metadata", dict[str, object]()) or dict[str, object]() try: self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")), @@ -1697,6 +1733,8 @@ class Logging(LiteLLMLoggingBaseClass): error=None, ) if should_redact_message_logging(self.model_call_details): + if not has_attributes(_metadata, _MetadataSlot, "__setitem__"): + raise TypeError(f"'{type(_metadata).__name__}' object does not support item assignment") _metadata["raw_request"] = REDACTED_BY_LITELLM else: curl_command: Final = self._get_request_curl_command( @@ -1705,15 +1743,21 @@ class Logging(LiteLLMLoggingBaseClass): additional_args=additional_args, data=additional_args.get("complete_input_dict", {}), ) - _metadata["raw_request"] = redact_string(str(curl_command)) + redacted_curl_command: Final = redact_string(str(curl_command)) + if not has_attributes(_metadata, _MetadataSlot, "__setitem__"): + raise TypeError(f"'{type(_metadata).__name__}' object does not support item assignment") + _metadata["raw_request"] = redacted_curl_command except Exception as e: self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( error=str(e), ) - _metadata["raw_request"] = redact_string( + unlogged_raw_request: Final = redact_string( f"Unable to Log \ raw request: {e}" ) + if not has_attributes(_metadata, _MetadataSlot, "__setitem__"): + raise TypeError(f"'{type(_metadata).__name__}' object does not support item assignment") + _metadata["raw_request"] = unlogged_raw_request if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: self.logger_fn( @@ -2344,7 +2388,7 @@ class Logging(LiteLLMLoggingBaseClass): def _record_zero_cost_diagnostic( self, result: object, - response_cost: float | None, + response_cost: object, *, calculation_failed: bool = False, litellm_model_name: str | None = None, @@ -2375,7 +2419,7 @@ class Logging(LiteLLMLoggingBaseClass): def _zero_cost_finding( self, result: object, - response_cost: float | None, + response_cost: object, *, calculation_failed: bool, litellm_model_name: str | None, @@ -3079,19 +3123,21 @@ class Logging(LiteLLMLoggingBaseClass): elif self.relays_llm_passthrough and isinstance(result, Response): from litellm.utils import ProviderConfigManager - provider_config: Final = ProviderConfigManager.get_provider_passthrough_config( - provider=self.model_call_details.get("custom_llm_provider", ""), - model=self.model, - ) - if provider_config is not None: - logging_result = provider_config.logging_non_streaming_response( + passthrough_provider: Final = self._model_call_details_view.get("custom_llm_provider", "") + if isinstance(passthrough_provider, str): + provider_config: Final = ProviderConfigManager.get_provider_passthrough_config( + provider=passthrough_provider, model=self.model, - custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), - httpx_response=result, - request_data=self.model_call_details.get("request_data", {}), - logging_obj=self, - endpoint=self.model_call_details.get("endpoint", ""), ) + if provider_config is not None: + logging_result = provider_config.logging_non_streaming_response( + model=self.model, + custom_llm_provider=passthrough_provider, + httpx_response=result, + request_data=self.model_call_details.get("request_data", {}), + logging_obj=self, + endpoint=self.model_call_details.get("endpoint", ""), + ) return logging_result def _surface_response_headers_from_result(self, logging_result: object) -> None: @@ -3138,19 +3184,25 @@ class Logging(LiteLLMLoggingBaseClass): """ if logging_result is None: return - hidden_params: Final = getattr(logging_result, HIDDEN_PARAMS_ATTR, None) + hidden_params: Final = attribute_or_none(logging_result, HIDDEN_PARAMS_ATTR) if not hidden_params: return if self.model_call_details.get("litellm_params") is None: return - metadata_hidden_params: Final = hidden_params.copy() + metadata_hidden_params: Final = require_attributes( + require_attributes(hidden_params, _HiddenParamsCopier, "copy").copy(), _HiddenParamsSlot, "get" + ) response_cost: Final[object] = self.model_call_details.get("response_cost") if metadata_hidden_params.get("response_cost") is None and response_cost is not None: metadata_hidden_params["response_cost"] = response_cost - litellm_params: Final = self.model_call_details["litellm_params"] - metadata: Final = litellm_params.get("metadata") or {} + litellm_params: Final = require_attributes( + self._model_call_details_view["litellm_params"], _HiddenParamsSlot, "get" + ) + metadata: Final = litellm_params.get("metadata") or dict[str, object]() litellm_params["metadata"] = metadata + if not has_attributes(metadata, _MetadataSlot, "__setitem__"): + raise TypeError(f"'{type(metadata).__name__}' object does not support item assignment") metadata["hidden_params"] = metadata_hidden_params def _process_hidden_params_and_response_cost( @@ -3166,7 +3218,9 @@ class Logging(LiteLLMLoggingBaseClass): hidden_params: Final = getattr(logging_result, HIDDEN_PARAMS_ATTR, {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: - self.model_call_details["litellm_params"].setdefault("metadata", {}) + require_attributes( + self._model_call_details_view["litellm_params"], _MetadataDefaults, "setdefault" + ).setdefault("metadata", {}) if self.model_call_details["litellm_params"]["metadata"] is None: self.model_call_details["litellm_params"]["metadata"] = {} self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = getattr( @@ -3179,8 +3233,9 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = 0.0 self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result)) elif "response_cost" in hidden_params: - self.model_call_details["response_cost"] = hidden_params["response_cost"] - self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"]) + hidden_response_cost: Final = subscript(hidden_params, "response_cost") + self.model_call_details["response_cost"] = hidden_response_cost + self._record_zero_cost_diagnostic(logging_result, hidden_response_cost) self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result)) elif (existing_cost := self.model_call_details.get("response_cost")) is not None and ( existing_cost != 0 or self.call_type == CallTypes.pass_through.value # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped @@ -4659,10 +4714,23 @@ class Logging(LiteLLMLoggingBaseClass): elif self.model_call_details.get("combined_usage_object") is None: self.model_call_details["response_cost"] = 0 - if hasattr(exception, "headers") and isinstance(exception.headers, dict): + raised_exception: Final[object] = exception + exception_headers: Final = ( + as_str_mapping(raised_exception.headers) + if has_attributes(raised_exception, _ExceptionHeaders, "headers") + else None + ) + if isinstance(exception_headers, dict): self.model_call_details.setdefault("litellm_params", {}) - metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} - metadata.update(exception.headers) + metadata: Final = require_attributes( + require_attributes(self._model_call_details_view["litellm_params"], _ParamsLookup, "get").get( + "metadata", dict[str, object]() + ) + or dict[str, object](), + _MetadataUpdate, + "update", + ) + metadata.update(exception_headers) ## STANDARDIZED LOGGING PAYLOAD @@ -6412,7 +6480,7 @@ def build_custom_logger_compatible_class( return callback # Get global BitBucket config - gitlab_config: Final = getattr(litellm, "global_gitlab_config", None) + gitlab_config: Final = litellm.global_gitlab_config if hasattr(litellm, "global_gitlab_config") else None if gitlab_config is None: raise ValueError("Gitlab configuration not found. Please set litellm.global_gitlab_config first.") @@ -7198,11 +7266,15 @@ class StandardLoggingPayloadSetup: if isinstance(reported, Usage): return reported normalized: Final = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj) - return Usage(**normalized) if normalized else Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + return ( + Usage(**_UsageFields(**decoded_json_object(normalized))) + if normalized + else Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + ) @staticmethod def get_usage_as_dict( - response_obj: dict[str, Any] | None, + response_obj: Mapping[str, object] | None, combined_usage_object: Usage | None = None, ) -> dict[str, Any]: """ @@ -7226,19 +7298,23 @@ class StandardLoggingPayloadSetup: reported: Final = response_obj.get("usage_info") if isinstance(reported, BaseModel): return reported.model_dump(exclude_none=True) + reported_fields: Final = as_str_mapping(reported) return ( - {key: value for key, value in reported.items() if value is not None} - if isinstance(reported, dict) + {key: value for key, value in reported_fields.items() if value is not None} + if isinstance(reported_fields, dict) else dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) ) if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() - if isinstance(_raw, dict): - if ResponseAPILoggingUtils.is_response_api_usage(_raw): - return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() - if InteractionsUsageObjectTransformation.is_interactions_usage_object(_raw): - return InteractionsUsageObjectTransformation.transform_interactions_usage_object(_raw).model_dump() - return _raw + raw_fields: Final = as_str_mapping(_raw) + if isinstance(raw_fields, dict): + if ResponseAPILoggingUtils.is_response_api_usage(raw_fields): + return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(raw_fields).model_dump() + if InteractionsUsageObjectTransformation.is_interactions_usage_object(raw_fields): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object( + raw_fields + ).model_dump() + return raw_fields if isinstance(_raw, Usage): return _raw.model_dump() return _empty diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 1f6a66f699a..4ff82f705d9 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1429,7 +1429,7 @@ def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | residency: Final = data_residency.lower() if residency not in _VALID_DATA_RESIDENCIES: return 1.0 - multiplier: Final = model_info.get(f"regional_processing_uplift_multiplier_{residency}") + multiplier: Final[object] = model_info.get(f"regional_processing_uplift_multiplier_{residency}") if multiplier is None: return 1.0 if not isinstance(multiplier, (str, Buffer, SupportsFloat, SupportsIndex)): diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 0fed2cb2fce..20d1205c005 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -18,6 +18,7 @@ from litellm.constants import ( STREAM_FINISH_CHOICE_FIELDS, VLLM_CHOICE_FIELDS, ) +from litellm.litellm_core_utils.core_helpers import iter_items, subscript from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( extract_reasoning_content, @@ -68,6 +69,14 @@ _MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields. _ANNOTATION_ENTRIES: Final = TypeAdapter(tuple[object, ...]) +class _ImageResponseFields(TypedDict, total=False, extra_items=object): + pass + + +def _image_response_from_fields(fields: Mapping[str, object]) -> ImageResponse: + return ImageResponse(**_ImageResponseFields(**fields)) + + def _is_provider_specific_annotation(annotation: object) -> bool: match annotation: case {"type": "url_citation"}: @@ -639,7 +648,7 @@ class LiteLLMResponseObjectHandler: model_response_dict: Final = model_response_object.model_dump() model_response_dict.update(response_object) - model_response_object = ImageResponse(**model_response_dict) + model_response_object = _image_response_from_fields(model_response_dict) return model_response_object @staticmethod @@ -683,12 +692,12 @@ class LiteLLMResponseObjectHandler: choices_list: Final[list[TextChoices]] = [] # Convert each choice to TextChoices - for choice in response["choices"]: + for choice in iter_items(subscript(response, "choices")): text_choices = TextChoices() - text_choices["text"] = choice["message"]["content"] - text_choices["index"] = choice["index"] + text_choices["text"] = subscript(subscript(choice, "message"), "content") + text_choices["index"] = subscript(choice, "index") text_choices["logprobs"] = transformed_logprobs - text_choices["finish_reason"] = choice["finish_reason"] + text_choices["finish_reason"] = subscript(choice, "finish_reason") choices_list.append(text_choices) text_completion_response["choices"] = choices_list diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ae36484b294..704af2b9ebd 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -13,7 +13,7 @@ from types import MappingProxyType from typing import Annotated, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment -from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError, model_validator from pydantic_core import ErrorDetails, PydanticCustomError import litellm @@ -22,7 +22,12 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, require_attributes +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, + attribute_or_none, + require_attributes, +) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -1455,22 +1460,25 @@ def convert_to_gemini_tool_call_invoke( raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e}") -def _json_children(node: JsonValue) -> Iterable[JsonValue]: - if isinstance(node, dict): - return node.values() - if isinstance(node, list): - return (child for child in node if not isinstance(child, list)) +def _json_children(node: object) -> Iterable[object]: + mapping: Final = as_str_mapping(node) + if mapping is not None: + return mapping.values() + items: Final = as_object_list(node) + if items is not None: + return (child for child in items if not isinstance(child, list)) return () -def _next_json_level(level: tuple[JsonValue, ...], _: None) -> tuple[JsonValue, ...]: +def _next_json_level(level: tuple[object, ...], _: None) -> tuple[object, ...]: return tuple(itertools.chain.from_iterable(_json_children(node) for node in level)) -def _is_part_reference(node: JsonValue) -> bool: - if not isinstance(node, dict): +def _is_part_reference(node: object) -> bool: + mapping: Final = as_str_mapping(node) + if mapping is None: return False - ref: Final[JsonValue] = node.get("$ref") + ref: Final = mapping.get("$ref") return isinstance(ref, str) and ref != "" @@ -1482,7 +1490,7 @@ def _may_spell_ref_key(raw_json: str) -> bool: return '"$ref"' in raw_json or _REF_CHAR_ESCAPE.search(raw_json) is not None -def _contains_part_reference(payload: JsonValue) -> bool: +def _contains_part_reference(payload: object) -> bool: """Gemini reads a {"$ref": ""} object in a function_response as a pointer to one of its parts, unless it sits under a list nested directly in a list @@ -1739,7 +1747,7 @@ def convert_to_gemini_tool_call_result( try: if content_str.strip().startswith("{") or content_str.strip().startswith("["): # Try to parse as JSON (for Computer Use structured responses) - parsed: Final[JsonValue] = json.loads(content_str) + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(content_str))) if isinstance(parsed, dict) and not (_may_spell_ref_key(content_str) and _contains_part_reference(parsed)): response_data = parsed # Use the parsed JSON directly else: diff --git a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py index d2531e70d6b..39baaed25e6 100644 --- a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py +++ b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py @@ -1,6 +1,6 @@ import json from datetime import datetime -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -28,6 +28,10 @@ class _ChatTemplateFileResult(TypedDict): chat_template: NotRequired[ReadOnly[str]] +class _AddedToken(TypedDict, total=False, extra_items=object): + content: ReadOnly[str] + + def strftime_now(fmt: str) -> str: """ Custom function for templates that need current date/time formatting (e.g., gpt-oss) @@ -159,7 +163,7 @@ async def aget_chat_template_file(hf_model_name: str) -> _ChatTemplateFileResult _aget_chat_template_file = aget_chat_template_file -def extract_token_value(token_value: None | str | dict[str, Any]) -> str: +def extract_token_value(token_value: None | str | _AddedToken) -> str: """ Extract token string from various formats (string, dict, etc.) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 01aa7296b0f..6d40d963d76 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -26,7 +26,7 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, json_contains from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.base_llm.realtime.transformation import ( @@ -220,9 +220,13 @@ def _unscreened_audio_error_event(client_event_id: object) -> Mapping[str, objec } +def _decode_json(payload: str) -> object: + return _DECODED_JSON.validate_python(json.loads(payload)) + + def _decode_json_object(payload: str) -> Mapping[str, object]: """Decode a realtime frame into its top-level field mapping.""" - return json.loads(payload) + return decoded_json_object(_decode_json(payload)) _GUARDRAIL_TURN_DETECTION_GATE: Final[Mapping[str, object]] = MappingProxyType( @@ -511,7 +515,7 @@ class RealTimeStreaming: self._held_client_content: tuple[str, ...] = () self._held_client_content_bytes: int = 0 self._held_client_content_deadline: float | None = None - self._replay_scope: Final = client_websocket_scope(websocket) + self._replay_scope: Final = client_websocket_scope(self.websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker @@ -533,7 +537,7 @@ class RealTimeStreaming: self._translation_usage_finalized = False # Detect whether the client is explicitly opting into the beta protocol. - self._client_wants_beta = self.detect_beta_header(websocket) + self._client_wants_beta = self.detect_beta_header(self.websocket) self._backend_uses_beta_protocol = not translation_session and ( self._client_wants_beta if backend_uses_beta_protocol is None else backend_uses_beta_protocol ) @@ -674,7 +678,7 @@ class RealTimeStreaming: if not self._should_store_message(message_obj): return try: - event_type: Final = message_obj.get("type", "") + event_type: Final = decoded_json_object(message_obj).get("type", "") if event_type in self._SESSION_EVENT_TYPES: typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) else: @@ -1020,7 +1024,7 @@ class RealTimeStreaming: sent = True continue try: - msg_obj = _decode_json_object(msg) + msg_obj = as_str_mapping(_decode_json(msg)) except (json.JSONDecodeError, TypeError): msg_obj = None if isinstance(msg_obj, dict) and self.provider_config.is_setup_message(msg_obj): @@ -1162,13 +1166,15 @@ class RealTimeStreaming: return _session_update_frame(self._held_session_update or _NO_SESSION_FIELDS) async def _hold_session_update(self, message: str) -> None: - client_event: Final = _decode_json_object(message) + client_event: Final = _decode_json(message) if self._held_session_update_bytes + len(message.encode("utf-8")) > RealTimeStreaming._MAX_BUFFERED_BYTES: await self.websocket.send_text( - json.dumps(_HELD_SESSION_UPDATE_TOO_LARGE.error_event(client_event.get("event_id"))) + json.dumps( + _HELD_SESSION_UPDATE_TOO_LARGE.error_event(decoded_json_object(client_event).get("event_id")) + ) ) return - session: Final = client_event.get("session") + session: Final = decoded_json_object(client_event).get("session") update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS held: Final = _merge_session_fields(self._held_session_update or _NO_SESSION_FIELDS, update) self._held_session_update = held @@ -1182,7 +1188,7 @@ class RealTimeStreaming: _session_update_frame(session), self.model, None ) return next( - (frame for frame in preview if isinstance(frame, str) and "setup" in _decode_json_object(frame)), None + (frame for frame in preview if isinstance(frame, str) and json_contains(_decode_json(frame), "setup")), None ) def _leaves_applied_setup_unchanged(self, update: Mapping[str, object]) -> bool: @@ -1194,7 +1200,7 @@ class RealTimeStreaming: return ( applied_setup is not None and merged_setup is not None - and _decode_json_object(applied_setup) == _decode_json_object(merged_setup) + and _decode_json(applied_setup) == _decode_json(merged_setup) ) async def _acknowledge_session_setup(self, session: Mapping[str, object]) -> None: @@ -1522,8 +1528,8 @@ class RealTimeStreaming: ``return_new_content_delta_events`` modality lookup, ...). """ try: - message_obj: Final = _decode_json_object(transformed_message) - if "setup" in message_obj: + message_obj: Final = _decode_json(transformed_message) + if json_contains(message_obj, "setup"): self.session_configuration_request = transformed_message except (json.JSONDecodeError, TypeError): return @@ -1852,7 +1858,7 @@ class RealTimeStreaming: def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: - event: Final = _decode_json_object(raw_response) + event: Final = as_str_mapping(_decode_json(raw_response)) except (json.JSONDecodeError, TypeError): return None return event if isinstance(event, dict) else None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index c13c9bc31fb..3760d7f3059 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -351,7 +351,7 @@ def _redacted_web_search_action(action: object) -> object: def _redacted_web_search_call(item: object) -> object: if isinstance(item, BaseModel): - action: Final = getattr(item, "action", None) + action: Final = attribute_or_none(item, "action") return item if action is None else item.model_copy(update={"action": _redacted_web_search_action(action)}) mapping: Final = _as_mapping(item) if mapping is None or mapping.get("action") is None: diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index e2b651f8885..ee70152636f 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypedDic from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required +from typing_extensions import TypedDict as ExtraItemsTypedDict import litellm from litellm._logging import verbose_logger @@ -19,6 +20,7 @@ from litellm.litellm_core_utils.core_helpers import ( chunk_usage, finite_cost, provider_reported_usage_cost, + require_attributes, stamp_litellm_bill_on_usage, ) from litellm.litellm_core_utils.hidden_params import ( @@ -82,6 +84,18 @@ if TYPE_CHECKING: ) +class _UsageFields(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + +def _usage_fields(fields: Mapping[str, object]) -> _UsageFields: + return _UsageFields(**fields) + + +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + class _ThinkingBlockFragment(TypedDict, total=False): type: str | None data: str | None @@ -946,12 +960,12 @@ class ChunkProcessor: usage_chunk = chunk.hidden_params.get("usage", None) if isinstance(usage_chunk, dict): - return Usage(**usage_chunk) + return Usage(**_usage_fields(usage_chunk)) if isinstance(usage_chunk, UnreportedUsage): return None if usage_chunk is None or isinstance(usage_chunk, Usage): return usage_chunk - return Usage(**usage_chunk.model_dump()) + return Usage(**_usage_fields(require_attributes(usage_chunk, _SupportsModelDump, "model_dump").model_dump())) def _calculate_usage_per_chunk( self, @@ -1275,7 +1289,7 @@ class ChunkProcessor: if value is not None } - returned_usage = Usage(**returned_usage.model_dump(), **provider_pricing_fields) + returned_usage = Usage(**_usage_fields(returned_usage.model_dump()), **provider_pricing_fields) return returned_usage diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 079b29d5509..62ba7bb0e04 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -27,11 +27,17 @@ from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_F from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, + as_str_mapping, attribute_or_none, chunk_usage, decoded_json_object, + has_attributes, + json_contains, provider_reported_usage_cost, + require_attributes, stamp_litellm_bill_on_usage, + subscript, + without, ) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, @@ -195,6 +201,10 @@ def _response_choice_index(response_obj: Mapping[str, object]) -> int: return _upstream_choice_index(response_obj.get("original_chunk")) +def _delta_choice_index(completion_fields: Mapping[str, object], response_fields: Mapping[str, object]) -> object: + return completion_fields.get("index", response_fields.get("choice_index")) + + def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: first_choice: Final = _upstream_first_choice(upstream_chunk) upstream_extras: Final = first_choice.model_extra if first_choice is not None else None @@ -231,13 +241,21 @@ def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) +class _KeywordFields(TypedDict, total=False, extra_items=object): + pass + + +def _keyword_fields(fields: Mapping[str, object]) -> _KeywordFields: + return _KeywordFields(**fields) + + def _delta_from_upstream(upstream_delta: "BaseModel | _UpstreamDeltaLike | None") -> Delta | None: if upstream_delta is None: return Delta() if not isinstance(upstream_delta, BaseModel): return None try: - return Delta(**dict(iter(upstream_delta))) + return Delta(**_keyword_fields(dict(iter(upstream_delta)))) except Exception: return None @@ -275,7 +293,7 @@ def print_verbose(print_statement: object): @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, object] + response_obj: Mapping[str, object] @dataclass(frozen=True, slots=True) @@ -286,19 +304,8 @@ class _ProviderChunkEarlyReturn: _ProviderChunkResult = _ProviderChunkParsed | _ProviderChunkEarlyReturn -class _PredibaseStreamData(TypedDict): - token: NotRequired[Mapping[str, str]] - details: Mapping[str, str] - generated_text: str | None - error: str | None - - -class _NlpCloudStreamData(TypedDict): - generated_text: str - - -class _AlephAlphaStreamData(TypedDict): - completions: Sequence[Mapping[str, str]] +class _Replaceable(Protocol): + def replace(self, old: str, new: str, /) -> object: ... class _DeltaDumpDict(TypedDict): @@ -365,6 +372,32 @@ class _VertexChunkLike(Protocol): candidates: Sequence[_VertexCandidateLike] +class _TokenTotalsLike(Protocol): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class _ParamsLookupLike(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +class _IdentifiedChunkLike(Protocol): + id: str + + +class _Decodable(Protocol): + def decode(self, encoding: str, /) -> str: ... + + +class _VllmRequestOutputLike(Protocol): + outputs: object + + +class _VllmCompletionOutputLike(Protocol): + text: object + + class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None upstream_provider: str | None = None @@ -783,28 +816,33 @@ class CustomStreamWrapper: self.holding_chunk = "" return hold, curr_chunk - def handle_predibase_chunk(self, chunk): + def handle_predibase_chunk(self, chunk: object): try: if not isinstance(chunk, str): - chunk = chunk.decode("utf-8") # DO NOT REMOVE this: This is required for HF inference API + Streaming + chunk = require_attributes(chunk, _Decodable, "decode").decode( + "utf-8" + ) # DO NOT REMOVE this: This is required for HF inference API + Streaming text = "" is_finished = False finish_reason = "" print_verbose(f"chunk: {chunk}") if chunk.startswith("data:"): - data_json: Final[_PredibaseStreamData] = json.loads(chunk[5:]) + data_json: Final = _DECODED_JSON.validate_python(json.loads(chunk[5:])) print_verbose(f"data json: {data_json}") - if "token" in data_json and "text" in data_json["token"]: - text = data_json["token"]["text"] - if data_json.get("details", False) and data_json["details"].get("finish_reason", False): + if json_contains(data_json, "token") and json_contains(subscript(data_json, "token"), "text"): + text = subscript(subscript(data_json, "token"), "text") + data_lookup: Final = require_attributes(data_json, _ParamsLookupLike, "get") + if data_lookup.get("details", False) and require_attributes( + subscript(data_json, "details"), _ParamsLookupLike, "get" + ).get("finish_reason", False): is_finished = True - finish_reason = data_json["details"]["finish_reason"] - elif data_json.get("generated_text", False): # if full generated text exists, then stream is complete + finish_reason = subscript(subscript(data_json, "details"), "finish_reason") + elif data_lookup.get("generated_text", False): # if full generated text exists, then stream is complete text = "" # don't return the final bos token is_finished = True finish_reason = "stop" - elif data_json.get("error", False): - raise Exception(data_json.get("error")) + elif data_lookup.get("error", False): + raise Exception(data_lookup.get("error", None)) return { "text": text, "is_finished": is_finished, @@ -828,11 +866,11 @@ class CustomStreamWrapper: if self.model and "dolphin" in self.model: chunk = self.process_chunk(chunk=chunk) else: - data_json: Final[_NlpCloudStreamData] = json.loads(chunk) - chunk = data_json["generated_text"] + data_json: Final = _DECODED_JSON.validate_python(json.loads(chunk)) + chunk = subscript(data_json, "generated_text") text = chunk - if "[DONE]" in text: - text = text.replace("[DONE]", "") + if json_contains(text, "[DONE]"): + text = require_attributes(text, _Replaceable, "replace").replace("[DONE]", "") is_finished = True finish_reason = "stop" return { @@ -845,9 +883,9 @@ class CustomStreamWrapper: def handle_aleph_alpha_chunk(self, chunk): chunk = chunk.decode("utf-8") - data_json: Final[_AlephAlphaStreamData] = json.loads(chunk) + data_json: Final = _DECODED_JSON.validate_python(json.loads(chunk)) try: - text: Final = data_json["completions"][0]["completion"] + text: Final = subscript(subscript(subscript(data_json, "completions"), 0), "completion") is_finished: Final = True finish_reason: Final = "stop" return { @@ -1014,20 +1052,17 @@ class CustomStreamWrapper: def model_response_creator( self, - chunk: dict | None = None, + chunk: dict[str, object] | None = None, hidden_params: Mapping[str, object] | None = None, delta: Delta | None = None, ) -> ModelResponseStream: _model: Final = self._served_openrouter_model() or self._cached_model_name _logging_obj_llm_provider: Final = self._cached_logging_llm_provider - if chunk is None: - args: dict[str, Any] = {"model": _model} - else: + if chunk is not None: chunk.pop("model", None) - args = {"model": _model} - if chunk: - args.update({k: v for k, v in chunk.items() if k != "stream"}) + chunk_fields: Final[Mapping[str, object]] = {k: v for k, v in chunk.items() if k != "stream"} if chunk else {} + args: Final = _keyword_fields({"model": _model, **chunk_fields}) model_response: Final = ( ModelResponseStream(**args) @@ -1108,7 +1143,8 @@ class CustomStreamWrapper: provider_specific_fields: Final = getattr(original_chunk, "provider_specific_fields", None) if provider_specific_fields is not None: model_response.provider_specific_fields = provider_specific_fields - _apply_provider_fields(model_response, provider_specific_fields) + if model_response.provider_specific_fields is not None: + _apply_provider_fields(model_response, model_response.provider_specific_fields) return model_response def is_chunk_non_empty( @@ -1169,10 +1205,8 @@ class CustomStreamWrapper: model_response.choices[0].delta["role"] = "assistant" self.sent_first_chunk = True elif self.sent_first_chunk is True and hasattr(model_response.choices[0].delta, "role"): - _initial_delta: Final = model_response.choices[0].delta.model_dump() - - _initial_delta.pop("role", None) - set_extra_attribute(model_response.choices[0], "delta", Delta(**_initial_delta)) + _initial_delta: Final = without(model_response.choices[0].delta.model_dump(), "role") + set_extra_attribute(model_response.choices[0], "delta", Delta(**_keyword_fields(_initial_delta))) return model_response def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: @@ -1275,7 +1309,7 @@ class CustomStreamWrapper: self, completion_obj: dict[str, Any], model_response: ModelResponseStream, - response_obj: dict[str, Any], + response_obj: Mapping[str, Any], plain_delta_created: bool = False, ): from litellm.litellm_core_utils.core_helpers import ( @@ -1300,19 +1334,15 @@ class CustomStreamWrapper: if hold is False: ## check if openai/azure chunk - original_chunk: Final = response_obj.get("original_chunk", None) + original_chunk: Final[object] = response_obj.get("original_chunk", None) if original_chunk: - upstream_chunk: Final[_UpstreamChunkLike] = original_chunk + upstream_chunk: Final = require_attributes(original_chunk, _UpstreamChunkLike, "choices") if len(upstream_chunk.choices) > 0: - original_chunk_for_reuse: Final[object] = cast( # cast-ok: runtime type-check before access - object, - original_chunk, - ) upstream_choice: Final = ( - original_chunk_for_reuse.choices[0] + original_chunk.choices[0] if plain_delta_created - and isinstance(original_chunk_for_reuse, (ChatCompletionChunk, ModelResponseStream)) - and len(original_chunk_for_reuse.choices) == 1 + and isinstance(original_chunk, (ChatCompletionChunk, ModelResponseStream)) + and len(original_chunk.choices) == 1 else None ) upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent @@ -1343,13 +1373,13 @@ class CustomStreamWrapper: for choice in upstream_chunk.choices: try: if isinstance(choice, BaseModel): - choice_json = choice.model_dump() - choice_json.pop( - "finish_reason", None - ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choice_json.pop("logprobs", None) choices.append( - StreamingChoices(**choice_json, logprobs=getattr(choice, "logprobs", None)) + StreamingChoices( + **_keyword_fields( + without(choice.model_dump(), "finish_reason", "logprobs") + ), # for mistral etc. which return a value in their last chunk (not-openai compatible). + logprobs=getattr(choice, "logprobs", None), + ) ) except Exception: choices.append(StreamingChoices()) @@ -1377,10 +1407,10 @@ class CustomStreamWrapper: self.sent_first_chunk = True if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] - model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) + model_response.choices[0].delta = Delta(**_keyword_fields(completion_obj)) + _index: Final = _delta_choice_index(completion_obj, response_obj) if _index is not None: - model_response.choices[0].index = _index + set_extra_attribute(model_response.choices[0], "index", _index) _apply_choice_fields(model_response.choices[0], chunk_choice_fields) if sibling_finish_reason is not None: @@ -1436,7 +1466,7 @@ class CustomStreamWrapper: # Preserve custom attributes from original chunk (applies to both # empty and non-empty delta final chunks). - _original_chunk: Final = response_obj.get("original_chunk", None) + _original_chunk: Final[object] = response_obj.get("original_chunk", None) if _original_chunk is not None: preserve_upstream_non_openai_attributes( model_response=model_response, @@ -1504,7 +1534,7 @@ class CustomStreamWrapper: def _dispatch_provider_chunk( self, - chunk: Any, + chunk: object, model_response: ModelResponseStream, completion_obj: dict[str, object], ) -> _ProviderChunkResult: @@ -1532,19 +1562,23 @@ class CustomStreamWrapper: chunk.choices[0].finish_reason = None return _ProviderChunkEarlyReturn(chunk) + chunk_fields: Final = as_str_mapping(chunk) + generic_chunk: Final = chunk_fields if isinstance(chunk_fields, dict) else None if ( - isinstance(chunk, dict) - and generic_chunk_has_all_required_fields(chunk=chunk) # check if chunk is a generic streaming chunk + # check if chunk is a generic streaming chunk + generic_chunk is not None and generic_chunk_has_all_required_fields(chunk=generic_chunk) ) or (self.custom_llm_provider and self.custom_llm_provider in litellm._custom_providers): if self.received_finish_reason is not None: - _chunk_has_content: Final = isinstance(chunk, dict) and ( - bool(chunk.get("text", "")) - or chunk.get("tool_use") is not None + _chunk_has_content: Final = generic_chunk is not None and ( + bool(generic_chunk.get("text", "")) + or generic_chunk.get("tool_use") is not None # Usage-only final chunks are valid and needed to surface # finish_reason/usage to downstream translators. - or chunk.get("usage") is not None + or generic_chunk.get("usage") is not None ) - if not _chunk_has_content and (not isinstance(chunk, dict) or "provider_specific_fields" not in chunk): + if not _chunk_has_content and ( + generic_chunk is None or "provider_specific_fields" not in generic_chunk + ): raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] @@ -1582,7 +1616,12 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider and self.custom_llm_provider == "vllm": - completion_obj["content"] = chunk[0].outputs[0].text + vllm_output: Final = require_attributes( + subscript(require_attributes(subscript(chunk, 0), _VllmRequestOutputLike, "outputs").outputs, 0), + _VllmCompletionOutputLike, + "text", + ) + completion_obj["content"] = vllm_output.text elif ( self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha" ): # aleph alpha doesn't provide streaming @@ -1689,7 +1728,13 @@ class CustomStreamWrapper: if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None: self.received_finish_reason = response_obj["finish_reason"] if response_obj["usage"] is not None: - _text_completion_usage: Final[Usage] = response_obj["usage"] + _text_completion_usage: Final = require_attributes( + _DECODED_JSON.validate_python(response_obj["usage"]), + _TokenTotalsLike, + "prompt_tokens", + "completion_tokens", + "total_tokens", + ) setattr( model_response, "usage", @@ -1711,7 +1756,13 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] if "usage" in response_obj is not None: - _codestral_usage: Final[Usage] = response_obj["usage"] + _codestral_usage: Final = require_attributes( + _DECODED_JSON.validate_python(response_obj["usage"]), + _TokenTotalsLike, + "prompt_tokens", + "completion_tokens", + "total_tokens", + ) setattr( model_response, "usage", @@ -1813,14 +1864,14 @@ class CustomStreamWrapper: ) return _ProviderChunkParsed(response_obj) - def chunk_creator(self, chunk: Any): + def chunk_creator(self, chunk: object): if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields) if self.custom_llm_provider != "cached_response" and isinstance( chunk, (ChatCompletionChunk, ModelResponseStream) ): self._ollama_response_ids.observe(chunk) - if hasattr(chunk, "id"): + if has_attributes(chunk, _IdentifiedChunkLike, "id"): self.response_id = chunk.id if self._ollama_response_ids.response_id is not None: self.response_id = self._ollama_response_ids.response_id @@ -1840,10 +1891,10 @@ class CustomStreamWrapper: ), delta=prebuilt_delta, ) - response_obj: dict[str, Any] = {} + response_obj: Mapping[str, object] = {} try: # return this for all models - completion_obj: Final[dict[str, Any]] = {"content": ""} + completion_obj: Final[dict[str, object]] = {"content": ""} dispatch_result: Final = self._dispatch_provider_chunk( chunk=chunk, model_response=model_response, @@ -1855,24 +1906,23 @@ class CustomStreamWrapper: model_response.model = self._served_openrouter_model() or self.model ## FUNCTION CALL PARSING - original_chunk: Final[_UpstreamChunkLike | None] = ( - response_obj.get("original_chunk") if response_obj is not None else None - ) + original_chunk: Final = response_obj.get("original_chunk") if response_obj is not None else None if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints # enter this branch when no content has been passed in response - if hasattr(original_chunk, "id"): + if has_attributes(original_chunk, _IdentifiedChunkLike, "id"): model_response = self.set_model_id(original_chunk.id, model_response) if stored_attribute(original_chunk, "provider_specific_fields") is not None: model_response = self.copy_model_response_level_provider_specific_fields( original_chunk, model_response ) - if original_chunk.choices and len(original_chunk.choices) > 0: - delta: Final = original_chunk.choices[0].delta + upstream_chunk: Final = require_attributes(original_chunk, _UpstreamChunkLike, "choices") + if upstream_chunk.choices and len(upstream_chunk.choices) > 0: + delta: Final = upstream_chunk.choices[0].delta if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): try: - model_response.system_fingerprint = original_chunk.system_fingerprint + model_response.system_fingerprint = upstream_chunk.system_fingerprint ## AZURE - check if arguments is not None if delta.function_call is not None: if delta.function_call.arguments is None: @@ -1922,7 +1972,7 @@ class CustomStreamWrapper: if "tool_calls" in completion_obj and _entry_len(completion_obj, "tool_calls") > 0: if self.is_function_call is True: # user passed in 'functions' param - completion_obj["function_call"] = completion_obj["tool_calls"][0]["function"] + completion_obj["function_call"] = subscript(subscript(completion_obj["tool_calls"], 0), "function") completion_obj["tool_calls"] = None self.tool_call = True @@ -1931,7 +1981,7 @@ class CustomStreamWrapper: if isinstance(chunk_usage, Usage): model_response.usage = chunk_usage elif isinstance(chunk_usage, BaseModel): - model_response.usage = Usage(**chunk_usage.model_dump()) + model_response.usage = Usage(**_keyword_fields(chunk_usage.model_dump())) ## RETURN ARG if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: @@ -2126,7 +2176,11 @@ class CustomStreamWrapper: else: asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler - litellm_params: Final[Mapping[str, object]] = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: Final = require_attributes( + _DECODED_JSON.validate_python(self.logging_obj.model_call_details.get("litellm_params", {})), + _ParamsLookupLike, + "get", + ) if self.logging_obj.is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) @@ -3015,9 +3069,9 @@ class CustomStreamWrapper: _strip_sse_data_from_chunk = strip_sse_data_from_chunk -def _usage_from_payload(payload: Mapping[str, Any]) -> Usage: +def _usage_from_payload(payload: Mapping[str, object]) -> Usage: """Build ``Usage`` from a wire-shaped usage dict, forwarding every key.""" - return litellm.Usage(**payload) + return litellm.Usage(**_keyword_fields(payload)) def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index f6ff6d3bc47..d9dd4b2b395 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -2,11 +2,12 @@ Common utilities for A2A (Agent-to-Agent) Protocol """ -from collections.abc import Awaitable, Callable, Mapping -from typing import Any, Final +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final from pydantic import BaseModel +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -67,7 +68,15 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: return "\n".join(conversation_parts) -def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_depth: int = 10) -> str: +def _join_with_spaces(texts: Sequence[object]) -> str: + strings: Final = [text for text in texts if isinstance(text, str)] + if len(strings) == len(texts): + return " ".join(strings) + first_bad: Final = next(position for position, text in enumerate(texts) if not isinstance(text, str)) + raise TypeError(f"sequence item {first_bad}: expected str instance, {type(texts[first_bad]).__name__} found") + + +def extract_text_from_a2a_message(message: object, depth: int = 0, max_depth: int = 10) -> str: """ Extract text content from A2A message parts. @@ -82,19 +91,20 @@ def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_d if message is None or depth >= max_depth: return "" - parts: Final = message.get("parts", []) - text_parts: Final[list[str]] = [] + parts: Final = iter_items(decoded_json_object(message).get("parts", [])) + text_parts: Final[list[object]] = [] for part in parts: - if part.get("kind") == "text": - text_parts.append(part.get("text", "")) + part_fields = decoded_json_object(part) + if part_fields.get("kind") == "text": + text_parts.append(part_fields.get("text", "")) # Handle nested parts if they exist - elif "parts" in part: - nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth) + elif "parts" in part_fields: + nested_text = extract_text_from_a2a_message(part_fields, depth + 1, max_depth) if nested_text: text_parts.append(nested_text) - return " ".join(text_parts) + return _join_with_spaces(text_parts) def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_depth: int = 10) -> str: diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index 9253ca389ee..9ad5a585e66 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -11,7 +11,8 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final from aiohttp import ClientResponse -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -31,6 +32,16 @@ _JSON_OBJECTS: Final = TypeAdapter( ) +class _ChatCompletionBody(TypedDict, extra_items=object): + id: ReadOnly[str] + created: ReadOnly[int] + model: ReadOnly[str] + object: ReadOnly[str] + choices: ReadOnly[JsonValue] + system_fingerprint: ReadOnly[NotRequired[str | None]] + service_tier: ReadOnly[NotRequired[str | None]] + + class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): def get_complete_url( self, @@ -73,7 +84,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - _json_response: Final = await raw_response.json() + _json_response: Final[_ChatCompletionBody] = await raw_response.json() model_response.id = _json_response.get("id") model_response.choices = [ Choices.model_validate(choice) for choice in _JSON_OBJECTS.validate_python(_json_response.get("choices")) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index c777f1f6ac9..1d24cfab131 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Final, NoReturn, cast +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Final, NoReturn, Protocol, cast, runtime_checkable import httpx from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch @@ -31,6 +31,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.core_helpers import ( + as_object_list, as_str_mapping, coerce_token_limit, drops_unsupported_params, @@ -450,6 +451,46 @@ def _enum_conflicts_with_declared_type(schema: Mapping[str, object]) -> bool: return check is not None and not all(check(value) for value in enum_values) +def _as_json_object(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + +def _filtered_output_subschema(value: object) -> object: + subschema: Final = _as_json_object(value) + return value if subschema is None else AnthropicConfig.filter_anthropic_output_schema(subschema) + + +def _filtered_output_subschemas(value: object) -> object: + subschemas: Final = _as_json_object(value) + if subschemas is None: + return value + return {name: _filtered_output_subschema(subschema) for name, subschema in subschemas.items()} + + +def _filtered_output_branches(branches: Sequence[object]) -> Sequence[object]: + return [_filtered_output_subschema(branch) for branch in branches] + + +@runtime_checkable +class _SupportsAdd(Protocol): + def __add__(self, other: object, /) -> object: ... + + +def _concatenated(left: object, right: str | Sequence[object]) -> object: + if isinstance(left, _SupportsAdd): + return left + right + raise TypeError(f"unsupported operand type(s) for +: '{type(left).__name__}' and '{type(right).__name__}'") + + +def _extended_any_of(result: Mapping[str, object], branches: Sequence[object]) -> object: + filtered: Final = _filtered_output_branches(branches) + return _concatenated(result["anyOf"], filtered) if "anyOf" in result else filtered + + +def _described_with_note(description: object, note: str) -> object: + return _concatenated(_concatenated(description, " "), note) if description else note + + # Single, internal-only key on ``litellm_params`` used to thread the per- # request reverse map (sanitized -> original) from request build to response # parsing. ``litellm_params`` is never serialized to a provider; ``optional_ @@ -927,7 +968,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return params @staticmethod - def filter_anthropic_output_schema(schema: dict[str, Any]) -> dict[str, Any]: + def filter_anthropic_output_schema(schema: Mapping[str, object]) -> Mapping[str, object]: """ Filter out unsupported fields from JSON schema for Anthropic's output_format API. @@ -1019,16 +1060,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): note_value = json.dumps(value) if isinstance(value, (dict, list)) else value constraint_descriptions.append(label.format(note_value)) - result: Final[dict[str, Any]] = {} + result: Final[dict[str, object]] = {} # Update description with removed constraint info if constraint_descriptions: - existing_desc: Final = schema.get("description", "") constraint_note: Final = "Note: " + ", ".join(constraint_descriptions) + "." - if existing_desc: - result["description"] = existing_desc + " " + constraint_note - else: - result["description"] = constraint_note + result["description"] = _described_with_note(schema.get("description"), constraint_note) drops_conflicting_type: Final = _enum_conflicts_with_declared_type(schema) @@ -1041,26 +1078,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Already handled above continue - if key in ("properties", "$defs", "definitions") and isinstance(value, dict): - nested_schemas: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict - result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in nested_schemas.items()} - elif key == "items" and isinstance(value, dict): - nested_schema: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict - result[key] = AnthropicConfig.filter_anthropic_output_schema(nested_schema) - elif key == "anyOf" and isinstance(value, list): - result["anyOf"] = result.get("anyOf", []) + [ - AnthropicConfig.filter_anthropic_output_schema(item) for item in value - ] - elif key == "allOf" and isinstance(value, list): - result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] + branches = as_object_list(value) + if key in ("properties", "$defs", "definitions"): + result[key] = _filtered_output_subschemas(value) + elif key == "items": + result[key] = _filtered_output_subschema(value) + elif key == "anyOf" and branches is not None: + result["anyOf"] = _extended_any_of(result, branches) + elif key == "allOf" and branches is not None: + result[key] = _filtered_output_branches(branches) elif key == "$ref" and isinstance(value, str): result[key] = schema_fragment_ref(value) - elif key == "oneOf" and isinstance(value, list): + elif key == "oneOf" and branches is not None: # Anthropic rejects oneOf ("Schema type 'oneOf' is not supported"); # the Anthropic SDK rewrites it to anyOf, so do the same. - result["anyOf"] = result.get("anyOf", []) + [ - AnthropicConfig.filter_anthropic_output_schema(item) for item in value - ] + result["anyOf"] = _extended_any_of(result, branches) else: result[key] = value diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 937afeae402..e32754adf03 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -2570,7 +2570,7 @@ def changed_server_tool_call_arguments( def write_anthropic_hosted_web_call_arguments(holder: Message | Delta, arguments_by_tool_id: Mapping[str, str]) -> None: - fields: Final = holder.provider_specific_fields + fields: Final[Mapping[str, object]] = holder.provider_specific_fields or {} if not arguments_by_tool_id or not fields or fields.get("web_search_calls") is None: return holder.provider_specific_fields = { # rebind-ok: the hosted copy of a rewritten call must not keep its old input @@ -2598,7 +2598,7 @@ def _jsonable_hosted_web_call(item: object, arguments_by_tool_id: Mapping[str, s call: Final = _with_hosted_web_call_arguments(item, arguments_by_tool_id) if not isinstance(call, ResponseFunctionWebSearch): return call - return _WEB_SEARCH_CALL_ADAPTER.dump_python(call, mode="json", exclude_none=True) + return _DECODED_JSON.validate_python(_WEB_SEARCH_CALL_ADAPTER.dump_python(call, mode="json", exclude_none=True)) def _with_hosted_web_call_arguments(item: object, arguments_by_tool_id: Mapping[str, str]) -> object: @@ -2613,7 +2613,7 @@ def _with_hosted_web_call_arguments(item: object, arguments_by_tool_id: Mapping[ def without_anthropic_server_tool_calls(holder: Message, tool_ids: frozenset[str]) -> None: - fields: Final = holder.provider_specific_fields or {} + fields: Final[Mapping[str, object]] = holder.provider_specific_fields or {} holder.provider_specific_fields = { # rebind-ok: a denied call must leave every copy of the message it was on **fields, "server_tool_calls": [ diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 0c564a194dc..e68360e69b9 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -4,12 +4,17 @@ import json from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter import litellm -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + attribute_items, + normalize_drop_params, + require_attributes, +) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.llms.anthropic.messages.convertible_blocks import tool_result_text_part from litellm.llms.anthropic.messages.utils import ( @@ -254,6 +259,11 @@ if TYPE_CHECKING: ToolResultContent: TypeAlias = str | list[ToolMessageContentPart] +class _CarriesUsage(Protocol): + @property + def usage(self) -> Usage: ... + + def target_supports_mid_conversation_system(model: str | None, custom_llm_provider: str | None) -> bool: if not model: return False @@ -681,7 +691,7 @@ class LiteLLMAnthropicMessagesAdapter: ## ASSISTANT MESSAGE ## assistant_message_str: str | None = None - assistant_content_list: list[dict[str, Any]] = [] # For content blocks with cache_control + assistant_content_list: list[ChatCompletionTextObject] = [] # For content blocks with cache_control has_cache_control_in_text = False tool_calls: list[ChatCompletionAssistantToolCall] = [] thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] @@ -694,7 +704,7 @@ class LiteLLMAnthropicMessagesAdapter: assistant_message_str = str(content) elif isinstance(content, dict): if content.get("type") == "text": - text_block: dict[str, object] = { + text_block: ChatCompletionTextObject = { "type": "text", "text": content.get("text", ""), } @@ -753,7 +763,7 @@ class LiteLLMAnthropicMessagesAdapter: ): # Use list format if any text block has cache_control, otherwise use string if has_cache_control_in_text and len(assistant_content_list) > 0: - assistant_content: Any = assistant_content_list + assistant_content: str | Sequence[ChatCompletionTextObject] | None = assistant_content_list elif len(assistant_content_list) > 0 and not has_cache_control_in_text: # Concatenate text blocks into string when no cache_control assistant_content = "".join(block.get("text", "") for block in assistant_content_list) @@ -1795,7 +1805,7 @@ class LiteLLMAnthropicMessagesAdapter: else translated_finish_reason ) # extract usage - usage: Final[Usage] = getattr(response, "usage") + usage: Final = require_attributes(response, _CarriesUsage, "usage").usage message_usage: Final = self.translate_openai_usage_to_anthropic_usage(usage) polyfill_iterations: Final = polyfill_result.iterations_usage if polyfill_result is not None else None anthropic_usage: Final[AnthropicUsage] = ( diff --git a/litellm/llms/anthropic/messages/native/mcp_handler.py b/litellm/llms/anthropic/messages/native/mcp_handler.py index 8fe18d79f24..cefd6c6a6d9 100644 --- a/litellm/llms/anthropic/messages/native/mcp_handler.py +++ b/litellm/llms/anthropic/messages/native/mcp_handler.py @@ -46,6 +46,10 @@ def _get_stop_reason(response: AnthropicMessagesResponse) -> str | None: return stop_reason if isinstance(stop_reason, str) else None +def _requests_stream(kwargs: Mapping[str, object]) -> bool: + return bool(kwargs.get("stream", False)) + + def _build_tool_result_message(tool_results: Sequence[Mapping[str, object]]) -> AnthropicMessagesUserMessageParam: """Turn executed tool results into the user message Anthropic expects.""" return AnthropicMessagesUserMessageParam( @@ -120,7 +124,8 @@ async def anthropic_messages_with_mcp( should_auto_execute: Final = LiteLLM_Proxy_MCP_Handler.should_auto_execute_tools( mcp_tools_with_litellm_proxy=mcp_references ) - stream: Final = bool(kwargs.pop("stream", False)) + stream: Final = _requests_stream(kwargs) + kwargs.pop("stream", None) base_call_args: Final[Mapping[str, object]] = { "max_tokens": max_tokens, diff --git a/litellm/llms/anthropic/messages/native/streaming_iterator.py b/litellm/llms/anthropic/messages/native/streaming_iterator.py index 52b3f1042f8..1195652dfd2 100644 --- a/litellm/llms/anthropic/messages/native/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/streaming_iterator.py @@ -28,6 +28,7 @@ from litellm.llms.anthropic.messages.utils import anthropic_sse_frame from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) +from litellm.responses.mcp.parked_stream_dispatch import DeferredStreamRelease from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -647,8 +648,12 @@ class BaseAnthropicMessagesStreamingIterator: disconnect (GeneratorExit / CancelledError) re-raises past it. Without this dispatch that window loses the spend row entirely. """ - deferred_cb: Final = getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None) - deferred_args: Final = getattr(self.litellm_logging_obj, "_deferred_stream_complete_args", None) + deferred_cb: Final[DeferredStreamRelease | None] = getattr( + self.litellm_logging_obj, "_on_deferred_stream_complete", None + ) + deferred_args: Final[tuple[object, ...] | None] = getattr( + self.litellm_logging_obj, "_deferred_stream_complete_args", None + ) if deferred_cb is None or deferred_args is None: return self.litellm_logging_obj._on_deferred_stream_complete = None diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 5dae6cc9fc3..5ed1030aba9 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -29,6 +29,7 @@ import httpx from typing_extensions import ReadOnly from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, subscript from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.azure_ai.agents.transformation import ( @@ -81,19 +82,6 @@ class _AzureMessagesData(TypedDict, total=False): data: ReadOnly[list[_AzureMessage]] -class _CreatedObject(TypedDict): - id: ReadOnly[str] - - -class _RunError(TypedDict, total=False): - message: ReadOnly[str] - - -class _RunStatus(TypedDict, total=False): - status: ReadOnly[str] - last_error: ReadOnly[_RunError] - - class _SSEDelta(TypedDict, total=False): content: ReadOnly[list[_AzureContentItem]] @@ -114,6 +102,14 @@ class _AsyncAgentRequest(Protocol): ) -> Awaitable[httpx.Response]: ... +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class AzureAIAgentsHandler: """ Handler for Azure AI Agent Service. @@ -133,21 +129,21 @@ class AzureAIAgentsHandler: def _build_thread_url(self, api_base: str, api_version: str) -> str: return f"{join_api_base(api_base, 'threads')}?api-version={api_version}" - def _build_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: + def _build_messages_url(self, api_base: str, thread_id: object, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" - def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: + def _build_runs_url(self, api_base: str, thread_id: object, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/runs')}?api-version={api_version}" - def _build_run_status_url(self, api_base: str, thread_id: str, run_id: str, api_version: str) -> str: + def _build_run_status_url(self, api_base: str, thread_id: object, run_id: object, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") run_status_route: Final = f"threads/{encoded_thread_id}/runs/{encoded_run_id}" return f"{join_api_base(api_base, run_status_route)}?api-version={api_version}" - def _build_list_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: + def _build_list_messages_url(self, api_base: str, thread_id: object, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" @@ -217,7 +213,7 @@ class AzureAIAgentsHandler: model: str, content: str, model_response: ModelResponse, - thread_id: str, + thread_id: object, messages: list[dict[str, object]], annotations: list[_TransformedAnnotation] | None = None, ) -> ModelResponse: @@ -354,10 +350,10 @@ class AzureAIAgentsHandler: api_base: str, api_version: str, agent_id: str, - thread_id: str | None, + thread_id: object, messages: list[dict[str, object]], optional_params: dict, - ) -> tuple[str, str, list[_TransformedAnnotation] | None]: + ) -> tuple[object, str, list[_TransformedAnnotation] | None]: """Execute the agent flow synchronously. Returns (thread_id, content, annotations).""" # Step 1: Create thread if not provided @@ -365,8 +361,7 @@ class AzureAIAgentsHandler: verbose_logger.debug("Creating thread at: %s", self._build_thread_url(api_base, api_version)) response = make_request("POST", self._build_thread_url(api_base, api_version), {}) self._check_response(response, [200, 201], "Failed to create thread") - thread_data: Final[_CreatedObject] = response.json() - thread_id = thread_data["id"] + thread_id = subscript(_json_body(response), "id") verbose_logger.debug("Created thread: %s", thread_id) # At this point thread_id is guaranteed to be a string @@ -386,8 +381,7 @@ class AzureAIAgentsHandler: response = make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) self._check_response(response, [200, 201], "Failed to create run") - run_data: Final[_CreatedObject] = response.json() - run_id: Final = run_data["id"] + run_id: Final = subscript(_json_body(response), "id") verbose_logger.debug("Created run: %s", run_id) # Step 4: Poll for completion @@ -396,15 +390,15 @@ class AzureAIAgentsHandler: response = make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") - status_data: _RunStatus = response.json() + status_data = decoded_json_object(_json_body(response)) status = status_data.get("status") verbose_logger.debug("Run status: %s", status) if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_data: _RunStatus = response.json() - error_msg = error_data.get("last_error", {}).get("message", "Unknown error") + error_data = decoded_json_object(_json_body(response)) + error_msg = decoded_json_object(error_data.get("last_error", {})).get("message", "Unknown error") raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") time.sleep(self.config.POLL_INTERVAL_SECONDS) @@ -482,10 +476,10 @@ class AzureAIAgentsHandler: api_base: str, api_version: str, agent_id: str, - thread_id: str | None, + thread_id: object, messages: list[dict[str, object]], optional_params: dict, - ) -> tuple[str, str, list[_TransformedAnnotation] | None]: + ) -> tuple[object, str, list[_TransformedAnnotation] | None]: """Execute the agent flow asynchronously. Returns (thread_id, content, annotations).""" # Step 1: Create thread if not provided @@ -493,8 +487,7 @@ class AzureAIAgentsHandler: verbose_logger.debug("Creating thread at: %s", self._build_thread_url(api_base, api_version)) response = await make_request("POST", self._build_thread_url(api_base, api_version), {}) self._check_response(response, [200, 201], "Failed to create thread") - thread_data: Final[_CreatedObject] = response.json() - thread_id = thread_data["id"] + thread_id = subscript(_json_body(response), "id") verbose_logger.debug("Created thread: %s", thread_id) # At this point thread_id is guaranteed to be a string @@ -514,8 +507,7 @@ class AzureAIAgentsHandler: response = await make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) self._check_response(response, [200, 201], "Failed to create run") - run_data: Final[_CreatedObject] = response.json() - run_id: Final = run_data["id"] + run_id: Final = subscript(_json_body(response), "id") verbose_logger.debug("Created run: %s", run_id) # Step 4: Poll for completion @@ -524,15 +516,15 @@ class AzureAIAgentsHandler: response = await make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") - status_data: _RunStatus = response.json() + status_data = decoded_json_object(_json_body(response)) status = status_data.get("status") verbose_logger.debug("Run status: %s", status) if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_data: _RunStatus = response.json() - error_msg = error_data.get("last_error", {}).get("message", "Unknown error") + error_data = decoded_json_object(_json_body(response)) + error_msg = decoded_json_object(error_data.get("last_error", {})).get("message", "Unknown error") raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") await asyncio.sleep(self.config.POLL_INTERVAL_SECONDS) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index a4777be146f..59254ac6c0a 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -4,7 +4,9 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.vector_store.transformation import ( @@ -30,6 +32,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM): """ @@ -230,13 +234,14 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM # Build attributes with all available metadata # Exclude system fields and already-processed fields - attributes = {} - for key, value in result.items(): - if key not in ["id", "content", "contentVector", "@search.score"]: - attributes[key] = value + metadata = { + key: value + for key, value in decoded_json_object(_DECODED_JSON.validate_python(result)).items() + if key not in ["id", "content", "contentVector", "@search.score"] + } # Always include document_id in attributes - attributes["document_id"] = document_id + attributes = {**metadata, "document_id": document_id} result_obj = VectorStoreSearchResult( score=score, diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 3afc685b3dd..a705cfe20ad 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -1,7 +1,10 @@ import json from abc import abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, cast +from typing_extensions import TypedDict + import litellm from litellm.constants import STREAM_SSE_DONE_STRING @@ -27,10 +30,18 @@ def _indexed_thinking_block( return {**block, "index": index} +class _DeltaFields(TypedDict, total=False, extra_items=object): + pass + + +def _message_fields_without_thinking_blocks(message: Message) -> Mapping[str, object]: + return message.model_dump(exclude={"thinking_blocks"}) + + def _streaming_delta(message: Message) -> Delta: thinking_blocks: Final = message.thinking_blocks if hasattr(message, "thinking_blocks") else None return Delta( - **message.model_dump(exclude={"thinking_blocks"}), + **_DeltaFields(**_message_fields_without_thinking_blocks(message)), thinking_blocks=( None if thinking_blocks is None diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index b87a783e244..08f986de85a 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -73,8 +73,8 @@ _NO_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) def with_attribution_headers( attribution_headers: Mapping[str, str], - headers: dict[str, str] | None, # mutable-ok: returned as-is when there is nothing to add -) -> dict[str, str] | None: # mutable-ok: becomes the request's outbound headers + headers: dict[str, object] | None, # mutable-ok: returned as-is when there is nothing to add +) -> dict[str, object] | None: # mutable-ok: becomes the request's outbound headers """ `headers` plus any attribution header the caller didn't already set (names compared case-insensitively). Builds a new dict; `headers` is never mutated. diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index d7b643eb7f4..3e48c6f2ff2 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -235,7 +235,7 @@ class BaseVideoConfig(ABC): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, - ) -> dict[str, str]: + ) -> object: pass @abstractmethod diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index dc14ef005b5..f9e834e9d37 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -7,10 +7,11 @@ https://docs.twelvelabs.io/docs/models/pegasus import json import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx -from pydantic import ConfigDict, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -43,6 +44,7 @@ class _PegasusResponse(TypedDict): _PEGASUS_RESPONSE: Final = TypeAdapter(_PegasusResponse) +_OPTIONAL_PARAM: Final = TypeAdapter(object) class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): @@ -65,7 +67,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): def map_openai_params( self, - non_default_params: dict, + non_default_params: Mapping[str, object], optional_params: dict, model: str, drop_params: bool, @@ -79,7 +81,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): optional_params["responseFormat"] = self._normalize_response_format(value) return optional_params - def _normalize_response_format(self, value: Any) -> object: + def _normalize_response_format(self, value: object) -> object: """Normalize response_format to TwelveLabs format. TwelveLabs expects: @@ -110,7 +112,11 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): return value # Otherwise return the dict as is return value - return type_to_response_format_param(response_format=value) or value + if value is None: + return None + if isinstance(value, type) and issubclass(value, BaseModel): + return type_to_response_format_param(response_format=value) or value + raise TypeError(f"Unsupported response_format type - {value}") def transform_request( self, @@ -134,7 +140,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): # Handle responseFormat - transform to TwelveLabs format if "responseFormat" in optional_params: - response_format: Final = optional_params["responseFormat"] + response_format: Final = _OPTIONAL_PARAM.validate_python(optional_params["responseFormat"]) transformed_format: Final = self._normalize_response_format(response_format) if transformed_format: request_data["responseFormat"] = transformed_format diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 6f972986e56..5685a7b3ef7 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -1,8 +1,9 @@ import types from collections.abc import Mapping -from typing import Any, Final +from typing import Final from openai.types.image import Image +from typing_extensions import TypedDict from litellm.llms.bedrock.common_utils import get_cached_model_info from litellm.types.llms.bedrock import ( @@ -19,6 +20,10 @@ from litellm.types.llms.bedrock import ( from litellm.types.utils import ImageResponse +class _NovaCanvasTaskFields(TypedDict, total=False, extra_items=object): + pass + + class AmazonNovaCanvasConfig: """ Reference: https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/model-catalog/serverless/amazon.nova-canvas-v1:0 @@ -82,10 +87,12 @@ class AmazonNovaCanvasConfig: image_generation_config = {**image_generation_config, **optional_params} if task_type == "TEXT_IMAGE": - text_to_image_params: dict[str, Any] = image_generation_config.pop("textToImageParams", {}) + text_to_image_params: Mapping[str, object] = image_generation_config.pop("textToImageParams", {}) text_to_image_params = {"text": text, **text_to_image_params} try: - text_to_image_params_typed: Final = AmazonNovaCanvasTextToImageParams(**text_to_image_params) + text_to_image_params_typed: Final = AmazonNovaCanvasTextToImageParams( + **_NovaCanvasTaskFields(**text_to_image_params) + ) except Exception as e: raise ValueError( f"Error transforming text to image params: {e}. Got params: {text_to_image_params}, Expected params: {AmazonNovaCanvasTextToImageParams.__annotations__}" @@ -104,7 +111,7 @@ class AmazonNovaCanvasConfig: imageGenerationConfig=image_generation_config_typed, ) if task_type == "COLOR_GUIDED_GENERATION": - color_guided_generation_params: dict[str, Any] = image_generation_config.pop( + color_guided_generation_params: Mapping[str, object] = image_generation_config.pop( "colorGuidedGenerationParams", {} ) color_guided_generation_params = { @@ -113,7 +120,7 @@ class AmazonNovaCanvasConfig: } try: color_guided_generation_params_typed: Final = AmazonNovaCanvasColorGuidedGenerationParams( - **color_guided_generation_params + **_NovaCanvasTaskFields(**color_guided_generation_params) ) except Exception as e: raise ValueError( @@ -133,10 +140,12 @@ class AmazonNovaCanvasConfig: imageGenerationConfig=image_generation_config_typed, ) if task_type == "INPAINTING": - inpainting_params: dict[str, Any] = image_generation_config.pop("inpaintingParams", {}) + inpainting_params: Mapping[str, object] = image_generation_config.pop("inpaintingParams", {}) inpainting_params = {"text": text, **inpainting_params} try: - inpainting_params_typed: Final = AmazonNovaCanvasInpaintingParams(**inpainting_params) + inpainting_params_typed: Final = AmazonNovaCanvasInpaintingParams( + **_NovaCanvasTaskFields(**inpainting_params) + ) except Exception as e: raise ValueError( f"Error transforming inpainting params: {e}. Got params: {inpainting_params}, Expected params: {AmazonNovaCanvasInpaintingParams.__annotations__}" diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index a20eae0b59d..be23e450b19 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -8,13 +8,13 @@ from dataclasses import dataclass from functools import partial, reduce from itertools import accumulate, chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias, TypeVar +from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypeVar, runtime_checkable from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, subscript from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -171,10 +171,33 @@ def _reads_flat_texts_only(guardrail: "CustomGuardrail") -> bool: return is_compression_guardrail(guardrail) or not guardrail.supports_structured_messages() -_StringHolder = tuple[Any, str | int] +_StringHolder: TypeAlias = tuple[object, str | int] +_StringLeaf: TypeAlias = tuple[_StringHolder, str] _LeafIdentity: TypeAlias = tuple[int, str | int] +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str | int, value: object, /) -> None: ... + + +@runtime_checkable +class _SupportsDelItem(Protocol): + def __delitem__(self, key: str | int, /) -> None: ... + + +def _assignable(container: object) -> _SupportsSetItem: + if isinstance(container, _SupportsSetItem): + return container + raise TypeError(f"'{type(container).__name__}' object does not support item assignment") + + +def _deletable(container: object) -> _SupportsDelItem: + if isinstance(container, _SupportsDelItem): + return container + raise TypeError(f"'{type(container).__name__}' object doesn't support item deletion") + + @dataclass(frozen=True, slots=True) class _Assign: holder: _StringHolder @@ -194,7 +217,7 @@ _OBJECT_BY_REFERENCE: Final = TypeAdapter(InstanceOf[Mapping[str, object]]) _ARRAY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[list[object]]) -def _collect_strings(node: object, holders: list[_StringHolder]) -> None: +def _collect_strings(node: object, leaves: list[_StringLeaf]) -> None: """ Record a (container, key) holder for every non-empty string value nested under an arbitrary JSON node, so prompt content a caller hides in fields @@ -209,28 +232,28 @@ def _collect_strings(node: object, holders: list[_StringHolder]) -> None: for key, value in current.items(): if isinstance(value, str): if value: - holders.append((current, key)) + leaves.append(((current, key), value)) else: stack.append(value) elif isinstance(current, list): for index, value in enumerate(current): if isinstance(value, str): if value: - holders.append((current, index)) + leaves.append(((current, index), value)) else: stack.append(value) -def _collect_block_text(block: dict, holders: list[_StringHolder]) -> None: +def _collect_block_text(block: dict, leaves: list[_StringLeaf]) -> None: text: Final = block.get("text") if isinstance(text, str) and text: - holders.append((block, "text")) + leaves.append(((block, "text"), text)) -def _request_field_holders(body: Mapping[str, object]) -> tuple[_StringHolder, ...]: - holders: Final = list[_StringHolder]() - _collect_strings(body.get("additionalModelRequestFields"), holders) - return tuple(holders) +def _request_field_leaves(body: Mapping[str, object]) -> tuple[_StringLeaf, ...]: + leaves: Final = list[_StringLeaf]() + _collect_strings(body.get("additionalModelRequestFields"), leaves) + return tuple(leaves) def _extract_converse_texts( @@ -267,15 +290,15 @@ def _extract_converse_texts( that model. Scanning it stays the fail-closed default. Its strings come last, in the same order as they follow the latest user message in the rows. """ - system_holders: Final = list[_StringHolder]() + system_leaves: Final = list[_StringLeaf]() if not skip_system and not scan_only_tool_results: for block in body.get("system") or []: if isinstance(block, dict): - _collect_block_text(block, system_holders) + _collect_block_text(block, system_leaves) - request_field_holders: Final = () if scan_only_tool_results else _request_field_holders(body) + request_field_leaves: Final = () if scan_only_tool_results else _request_field_leaves(body) - message_holders: Final = list[_StringHolder]() + message_leaves: Final = list[_StringLeaf]() for message in body.get("messages") or []: if not isinstance(message, dict): continue @@ -286,20 +309,19 @@ def _extract_converse_texts( continue if scan_only_tool_results and "toolResult" not in block: continue - _collect_block_text(block, message_holders) + _collect_block_text(block, message_leaves) tool_use = block.get("toolUse") if isinstance(tool_use, dict): - _collect_strings(tool_use.get("input"), message_holders) + _collect_strings(tool_use.get("input"), message_leaves) tool_result = block.get("toolResult") if isinstance(tool_result, dict): for inner in tool_result.get("content") or []: if isinstance(inner, dict): - _collect_block_text(inner, message_holders) - _collect_strings(inner.get("json"), message_holders) + _collect_block_text(inner, message_leaves) + _collect_strings(inner.get("json"), message_leaves) - holders: Final = [*system_holders, *message_holders, *request_field_holders] - texts: Final = [container[key] for container, key in holders] - return texts, holders + leaves: Final = (*system_leaves, *message_leaves, *request_field_leaves) + return [text for _, text in leaves], [holder for holder, _ in leaves] @dataclass(frozen=True, slots=True) @@ -464,9 +486,9 @@ def _apply(writes: Sequence[_Assign | _Remove]) -> None: for write in writes: match write: case _Assign(holder=(container, key), value=value): - container[key] = value + _assignable(container)[key] = value case _Remove(holder=(container, key)): - del container[key] + del _deletable(container)[key] def _block_text_holder(block: object) -> _ContentHolder | None: @@ -545,7 +567,7 @@ def _edits_drop_a_text_change( if writes is None: return True _apply(writes) - edited: Final = [container[key] for container, key in holders] + edited: Final = [subscript(container, key) for container, key in holders] return any( returned != before and after == before for before, returned, after in zip(texts, guardrailed_texts, edited) ) @@ -562,26 +584,25 @@ def _extract_converse_output_texts( ``citationsContent.content[].text`` -- while leaving structural values such as reasoning signatures and citation sources untouched. """ - holders: Final[list[_StringHolder]] = [] + leaves: Final[list[_StringLeaf]] = [] for block in content_blocks: if not isinstance(block, dict): continue - _collect_block_text(block, holders) + _collect_block_text(block, leaves) tool_use = block.get("toolUse") if isinstance(tool_use, dict): - _collect_strings(tool_use.get("input"), holders) + _collect_strings(tool_use.get("input"), leaves) reasoning = block.get("reasoningContent") if isinstance(reasoning, dict): reasoning_text = reasoning.get("reasoningText") if isinstance(reasoning_text, dict): - _collect_block_text(reasoning_text, holders) + _collect_block_text(reasoning_text, leaves) citations = block.get("citationsContent") if isinstance(citations, dict): for cited in citations.get("content") or []: if isinstance(cited, dict): - _collect_block_text(cited, holders) - texts: Final = [container[key] for container, key in holders] - return texts, holders + _collect_block_text(cited, leaves) + return [text for _, text in leaves], [holder for holder, _ in leaves] def _write_back_texts( @@ -598,7 +619,7 @@ def _write_back_texts( for idx, (container, key) in enumerate(holders): if idx >= len(guardrailed_texts): break - container[key] = guardrailed_texts[idx] + _assignable(container)[key] = guardrailed_texts[idx] _GroupKey: TypeAlias = str | tuple[str, int] diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 25498fef46f..c59baf2e07c 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -18,6 +18,7 @@ from pydantic import ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import decoded_json_object, subscript from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection from litellm.llms.bedrock.common_utils import BedrockError @@ -918,7 +919,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_content_start_event( self, - event: dict, + event: object, current_response_id: str | None, current_output_item_id: str | None, current_conversation_id: str | None, @@ -942,7 +943,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Returns: Tuple of (events, response_id, output_item_id, conversation_id, delta_type) """ - content_start: Final = event["contentStart"] + content_start: Final = decoded_json_object(subscript(event, "contentStart")) role: Final = content_start.get("role") if role != "ASSISTANT": @@ -1228,7 +1229,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_content_end_event( self, - event: dict, + event: object, current_output_item_id: str | None, current_response_id: str | None, current_delta_type: str | None, @@ -1247,7 +1248,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Returns: Tuple of (events, reset_delta_chunks) """ - content_end: Final = event["contentEnd"] + content_end: Final = subscript(event, "contentEnd") verbose_logger.debug("Handling contentEnd: %s", content_end) if self._assistant_final_text_active: @@ -1338,7 +1339,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_prompt_end_event( self, - event: dict, + event: object, current_response_id: str | None, current_conversation_id: str | None, ) -> tuple[ @@ -1400,7 +1401,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_tool_use_event( self, - event: dict, + event: object, current_response_id: str | None, current_conversation_id: str | None, ) -> tuple[list[OpenAIRealtimeEvents], str, str]: @@ -1416,7 +1417,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Tuple of (events, response_id, conversation_id) """ verbose_logger.debug("Handling toolUse") - tool_use: Final = BedrockToolUse.model_validate(event["toolUse"]) + tool_use: Final = BedrockToolUse.model_validate(subscript(event, "toolUse")) verbose_logger.debug("Tool use event: %s (ID: %s)", tool_use.toolName, tool_use.toolUseId) response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" @@ -1668,7 +1669,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_conversation_id, current_delta_type, ) = self.transform_content_start_event( - event, + _JSON_VALUE.validate_python(event), current_response_id, current_output_item_id, current_conversation_id, @@ -1710,7 +1711,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): elif "contentEnd" in event: events, current_delta_chunks = self.transform_content_end_event( - event, + _JSON_VALUE.validate_python(event), current_output_item_id, current_response_id, current_delta_type, @@ -1735,7 +1736,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): elif "toolUse" in event: events, current_response_id, current_conversation_id = self.transform_tool_use_event( - event, current_response_id, current_conversation_id + _JSON_VALUE.validate_python(event), current_response_id, current_conversation_id ) returned_messages.extend(events) @@ -1745,7 +1746,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_output_item_id, current_response_id, current_delta_type, - ) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id) + ) = self.transform_prompt_end_event( + _JSON_VALUE.validate_python(event), current_response_id, current_conversation_id + ) returned_messages.extend(events) return { diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index b2b3a9b0faf..afcbe61cd75 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -4,9 +4,11 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable from urllib.parse import urlparse import httpx +from pydantic import TypeAdapter +from typing_extensions import TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.core_helpers import decoded_json_object, keyword_arguments from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -61,6 +63,13 @@ def _vector_search_configuration(retrieval_config: object) -> Mapping[str, objec return decoded_json_object(decoded_json_object(retrieval_config).get("vectorSearchConfiguration", {})) +_DECODED_JSON: Final = TypeAdapter(object) + + +class _KBResponseFields(TypedDict, total=False, extra_items=object): + pass + + class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS Bedrock Knowledge Bases.""" @@ -361,7 +370,11 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: try: - response_data: Final = BedrockKBResponse(**response.json()) + response_data: Final = BedrockKBResponse( + **_KBResponseFields( + **keyword_arguments(_DECODED_JSON.validate_python(response.json()), BedrockKBResponse) + ) + ) results: Final[list[VectorStoreSearchResult]] = [] for item in response_data.get("retrievalResults", []) or []: content: BedrockKBContent | None = item.get("content") diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index 636760b093d..c22bdf9d887 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -1,10 +1,20 @@ import json from typing import Final +from pydantic import TypeAdapter +from typing_extensions import TypedDict + import litellm +from litellm.litellm_core_utils.core_helpers import keyword_arguments from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig from litellm.types.llms.databricks import GenericStreamingChunk +_DECODED_JSON: Final = TypeAdapter(object) + + +class _StreamChunkFields(TypedDict, total=False, extra_items=object): + pass + class CodestralTextCompletionConfig(OpenAITextCompletionConfig): """ @@ -100,7 +110,11 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): "finish_reason": finish_reason, } - original_chunk: Final = litellm.ModelResponseStream(**chunk_data_dict) + original_chunk: Final = litellm.ModelResponseStream( + **_StreamChunkFields( + **keyword_arguments(_DECODED_JSON.validate_python(chunk_data_dict), litellm.ModelResponseStream) + ) + ) _choices: Final = chunk_data_dict.get("choices", []) or [] if len(_choices) == 0: return { diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 4ba5f9bd57b..f5d084ea800 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -397,8 +397,12 @@ class CohereV2ModelResponseIterator: if index != -1: str_line = str_line[index:] - data_json: Final = json.loads(str_line) - return self.chunk_parser(chunk=data_json) + data_json: Final = _DECODED_JSON.validate_python(json.loads(str_line)) + if isinstance(data_json, dict): + return self.chunk_parser(chunk=data_json) + raise ValueError( + f"Failed to parse v2 chunk: {type(data_json).__name__!r} object has no attribute 'get', chunk: {data_json}" + ) # Async iterator def __aiter__(self): diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index 056b5bce742..ca096e369e1 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -2,12 +2,14 @@ Legacy /v1/embedding transformation logic for Bedrock Cohere. """ -from collections.abc import Sized +from collections.abc import Mapping, Sized from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE +from litellm.litellm_core_utils.core_helpers import as_str_mapping, subscript from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import header_value from litellm.types.llms.bedrock import ( @@ -24,6 +26,17 @@ class _SupportsEncode(Protocol): def encode(self, text: str, /) -> Sized: ... +_DECODED_JSON: Final = TypeAdapter(object) + + +def _embed_response_body(response_json: object) -> Mapping[str, object]: + body: Final = as_str_mapping(response_json) + if body is None: + subscript(response_json, "embeddings") + raise TypeError(f"'{type(response_json).__name__}' object is not subscriptable") + return body + + def _bedrock_billed_input_tokens(response: httpx.Response) -> int | None: count: Final = header_value(response.headers, "x-amzn-bedrock-input-token-count") return int(count) if count is not None and count.isdigit() else None @@ -119,7 +132,7 @@ class CohereEmbeddingConfig: encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: - response_json: Final = response.json() + response_json: Final = _DECODED_JSON.validate_python(response.json()) ## LOGGING logging_obj.post_call( input=input, @@ -127,13 +140,14 @@ class CohereEmbeddingConfig: additional_args={"complete_input_dict": data}, original_response=response_json, ) + billed_input_tokens: Final = _bedrock_billed_input_tokens(response) if "images" in data else None return self._populate_embedding_response( - response_json=response_json, + response_json=dict(_embed_response_body(response_json)), model_response=model_response, model=model, encoding=encoding, input=input, - billed_input_tokens=_bedrock_billed_input_tokens(response) if "images" in data else None, + billed_input_tokens=billed_input_tokens, ) def transform_response( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index efb8b302664..f65c4205944 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -632,6 +632,7 @@ _OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) _DECODED_JSON: Final = TypeAdapter(object) +_AGENTIC_HOOK_RESULT: Final = TypeAdapter(object) def _decoded_body_headers(response: httpx.Response) -> httpx.Headers: @@ -906,7 +907,7 @@ class BaseLLMHTTPHandler: stream: bool | None = False, fake_stream: bool = False, api_key: str | None = None, - headers: dict[str, Any] | None = None, + headers: dict[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, provider_config: BaseConfig | None = None, shared_session: Optional["ClientSession"] = None, @@ -1407,7 +1408,7 @@ class BaseLLMHTTPHandler: api_key: str | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, aembedding: bool | None = False, - headers: dict[str, Any] | None = None, + headers: dict[str, object] | None = None, ) -> EmbeddingResponse: provider_config: Final = ProviderConfigManager.get_provider_embedding_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -1588,7 +1589,7 @@ class BaseLLMHTTPHandler: api_key: str | None = None, api_base: str | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, - litellm_params: dict[str, Any] | None = None, + litellm_params: dict[str, object] | None = None, ) -> RerankResponse: # get config from model, custom llm provider headers = provider_config.validate_environment( @@ -6146,7 +6147,7 @@ class BaseLLMHTTPHandler: fingerprint: str, stream: bool = False, callback: Optional["CustomLogger"] = None, - ) -> AnthropicMessagesResponse | AsyncIterator[object]: + ) -> object: from litellm.anthropic_interface import messages as anthropic_messages patch: Final = plan.request_patch or AgenticLoopRequestPatch() @@ -6187,7 +6188,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup["max_agentic_loops"] = max_loops kwargs_for_followup["_agentic_loop_fingerprints"] = [*fingerprints, fingerprint] - response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( + response: Final[AnthropicMessagesResponse | AsyncIterator[object]] = await anthropic_messages.acreate( **{ "max_tokens": max_tokens, "messages": patch.messages, @@ -6200,9 +6201,7 @@ class BaseLLMHTTPHandler: if callback is not None: try: - response = await callback.async_post_agentic_loop_response_hook( - response=response, plan=plan, kwargs=kwargs - ) + return await callback.async_post_agentic_loop_response_hook(response=response, plan=plan, kwargs=kwargs) except Exception as e: _call_id: Final = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( @@ -6227,7 +6226,7 @@ class BaseLLMHTTPHandler: fingerprints: Sequence[str], fingerprint: str, callback: Optional["CustomLogger"] = None, - ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: + ) -> object: patch: Final = plan.request_patch or AgenticLoopRequestPatch() if patch.messages is None: raise ValueError("Agentic loop plan missing patched responses input") @@ -6263,7 +6262,7 @@ class BaseLLMHTTPHandler: try: with internal_sub_call(logging_obj.litellm_call_id): - response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses( + response: Final[ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] = await litellm.aresponses( model=patch.model or model, input=patch.messages, # pyright: ignore[reportArgumentType] # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime **optional_params, @@ -6272,7 +6271,7 @@ class BaseLLMHTTPHandler: if callback is not None: try: - response = await callback.async_post_agentic_loop_response_hook( + return await callback.async_post_agentic_loop_response_hook( response=response, plan=plan, kwargs=kwargs ) except Exception as e: @@ -6501,21 +6500,23 @@ class BaseLLMHTTPHandler: callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan ) if not build_plan_overridden: - agentic_result: object = await bill_first_round_on_failure( - callback.async_run_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + agentic_result: object = _AGENTIC_HOOK_RESULT.validate_python( + await bill_first_round_on_failure( + callback.async_run_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ), logging_obj=logging_obj, + response=response, stream=stream, - kwargs=kwargs_with_provider, - ), - logging_obj=logging_obj, - response=response, - stream=stream, + ) ) return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface) @@ -6536,8 +6537,9 @@ class BaseLLMHTTPHandler: stream=stream, ) - if plan.response_override is not None: - return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface) + response_override: object = plan.response_override + if response_override is not None: + return self._maybe_wrap_in_fake_stream(response_override, logging_obj, api_surface) if plan.terminate: verbose_logger.debug( "Agentic loop terminated by callback=%s reason=%s", @@ -6545,10 +6547,12 @@ class BaseLLMHTTPHandler: plan.stop_reason, ) return self._maybe_wrap_in_fake_stream( - await callback.async_post_agentic_loop_response_hook( - response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), - plan=plan, - kwargs=kwargs_with_provider, + _AGENTIC_HOOK_RESULT.validate_python( + await callback.async_post_agentic_loop_response_hook( + response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), + plan=plan, + kwargs=kwargs_with_provider, + ) ), logging_obj, api_surface, @@ -12396,10 +12400,10 @@ class BaseLLMHTTPHandler: def generate_content_handler( self, model: str, - contents: Any, + contents: object, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, generate_content_config_dict: dict, - tools: Any, + tools: object, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, @@ -12517,10 +12521,10 @@ class BaseLLMHTTPHandler: async def async_generate_content_handler( self, model: str, - contents: Any, + contents: object, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, generate_content_config_dict: dict, - tools: Any, + tools: object, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index cb56717993e..fae2c49a978 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -1,14 +1,30 @@ import json import time import types -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sized from typing import Final import httpx +from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import iter_items, json_contains, subscript from litellm.utils import Choices, Message, ModelResponse, Usage +_DECODED_JSON: Final = TypeAdapter(object) +_STR: Final = TypeAdapter(str) + + +def _has_length(value: object) -> bool: + if isinstance(value, Sized): + return len(value) > 0 + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + +def _completion_message(item: object) -> Message: + completion: Final = subscript(item, "completion") + return Message(content=_STR.validate_python(completion) if _has_length(completion) else None) + class AlephAlphaError(Exception): def __init__(self, status_code, message): @@ -249,22 +265,19 @@ def completion( ) print_verbose(f"raw model_response: {response.text}") ## RESPONSE OBJECT - completion_response: Final = response.json() - if "error" in completion_response: + completion_response: Final = _DECODED_JSON.validate_python(response.json()) + if json_contains(completion_response, "error"): raise AlephAlphaError( - message=completion_response["error"], + message=subscript(completion_response, "error"), status_code=response.status_code, ) else: try: choices_list: Final = [] - for idx, item in enumerate(completion_response["completions"]): - if len(item["completion"]) > 0: - message_obj = Message(content=item["completion"]) - else: - message_obj = Message(content=None) + for idx, item in enumerate(iter_items(subscript(completion_response, "completions"))): + message_obj = _completion_message(item) choice_obj = Choices( - finish_reason=item["finish_reason"], + finish_reason=subscript(item, "finish_reason"), index=idx + 1, message=message_obj, ) diff --git a/litellm/llms/edenai/videos/transformation.py b/litellm/llms/edenai/videos/transformation.py index 0cb9c5afaec..97235a0ffb8 100644 --- a/litellm/llms/edenai/videos/transformation.py +++ b/litellm/llms/edenai/videos/transformation.py @@ -142,7 +142,7 @@ class EdenAIVideoConfig(OpenAIVideoConfig): raw_response: httpx.Response, logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str | None = None, - ) -> dict[str, str]: # mutable-ok: inherited contract + ) -> object: raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising return super().transform_video_list_response( raw_response=raw_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 55a75423042..ec8de5af5e2 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast import httpx from pydantic import ConfigDict, TypeAdapter +from typing_extensions import TypedDict import litellm from litellm._logging import verbose_logger @@ -65,6 +66,10 @@ _STR: Final = TypeAdapter(str) _DECODED_JSON: Final = TypeAdapter(object) +class _FunctionFields(TypedDict, total=False, extra_items=object): + pass + + def _map_reasoning_effort(value: object) -> object: effort: Final[object] = cast(Mapping[str, object], value).get("effort") if isinstance(value, Mapping) else value if effort is True: @@ -718,7 +723,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ if tool_calls is not None and message.content is not None and message.tool_calls is None: try: - function: Final = Function(**json.loads(message.content)) + function: Final = Function(**_FunctionFields(**_JSON_DICT.validate_python(json.loads(message.content)))) if function.name != RESPONSE_FORMAT_TOOL_NAME and function.name in [ tool["function"]["name"] for tool in tool_calls ]: diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 7bc783b960f..59058b03d67 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -452,36 +453,37 @@ class GigaChatConfig(BaseConfig): choices: Final[list[Choices]] = [] # mutable-ok: accumulator for building response choices for choice in response_json.get("choices", []): message_data = choice.get("message", {}) - finish_reason = choice.get("finish_reason", "stop") + finish_reason = _JSON_VALUE.validate_python(choice.get("finish_reason", "stop")) - if finish_reason == "function_call" and message_data.get("function_call"): - func_call = message_data["function_call"] + function_call = ( + decoded_json_object(_JSON_VALUE.validate_python(message_data)).get("function_call") + if finish_reason == "function_call" + else None + ) + if function_call: + func_call = decoded_json_object(function_call) args = func_call.get("arguments", {}) + args_mapping = as_str_mapping(args) if is_structured_output: - if isinstance(args, dict): - content = json.dumps(args, ensure_ascii=False) - else: - content = str(args) - message_data["content"] = content - message_data.pop("function_call", None) - message_data.pop("functions_state_id", None) + message_data["content"] = ( + json.dumps(args_mapping, ensure_ascii=False) if args_mapping is not None else str(args) + ) finish_reason = "stop" else: - if isinstance(args, dict): - args = json.dumps(args, ensure_ascii=False) + arguments = json.dumps(args_mapping, ensure_ascii=False) if args_mapping is not None else args message_data["tool_calls"] = [ { "id": f"call_{uuid.uuid4().hex[:24]}", "type": "function", "function": { "name": func_call.get("name", ""), - "arguments": args, + "arguments": arguments, }, } ] - message_data.pop("function_call", None) finish_reason = "tool_calls" + message_data.pop("function_call", None) message_data.pop("functions_state_id", None) diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 30ca7147865..5a988edc144 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -9,12 +9,13 @@ from __future__ import annotations import types from collections.abc import Mapping -from typing import Final +from typing import Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders +from litellm.litellm_core_utils.core_helpers import iter_items, json_contains, subscript from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -26,6 +27,28 @@ from litellm.types.utils import EmbeddingResponse from ..authenticator import get_access_token _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsDelItem(Protocol): + def __delitem__(self, key: str, /) -> None: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _add_prompt_tokens(total_tokens: float, embedding: object) -> float: + if not (json_contains(embedding, "usage") and json_contains(subscript(embedding, "usage"), "prompt_tokens")): + return total_tokens + prompt_tokens: Final = subscript(subscript(embedding, "usage"), "prompt_tokens") + if isinstance(prompt_tokens, int | float): + return total_tokens + prompt_tokens + raise TypeError( + f"unsupported operand type(s) for +=: '{type(total_tokens).__name__}' and '{type(prompt_tokens).__name__}'" + ) class GigaChatEmbeddingError(BaseLLMException): @@ -144,7 +167,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): "model": "Embeddings" } """ - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) # Log response logging_obj.post_call( @@ -155,16 +178,19 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): ) # Calculate total tokens from individual embeddings - total_tokens = 0 - if "data" in response_json: - for emb in response_json["data"]: - if "usage" in emb and "prompt_tokens" in emb["usage"]: - total_tokens += emb["usage"]["prompt_tokens"] + total_tokens: float = 0 + if json_contains(response_json, "data"): + for emb in iter_items(subscript(response_json, "data")): + total_tokens = _add_prompt_tokens(total_tokens, emb) # Remove usage from individual embeddings (not part of OpenAI format) - if "usage" in emb: + if json_contains(emb, "usage"): + if not isinstance(emb, _SupportsDelItem): + raise TypeError(f"'{type(emb).__name__}' object does not support item deletion") del emb["usage"] # Set overall usage + if not isinstance(response_json, _SupportsSetItem): + raise TypeError(f"'{type(response_json).__name__}' object does not support item assignment") response_json["usage"] = { "prompt_tokens": total_tokens, "total_tokens": total_tokens, diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index bb836eb1b35..e3210482b2a 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -1,13 +1,15 @@ import json import os import time -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, Sequence, Sized from copy import deepcopy from typing import TYPE_CHECKING, Any, Final, Protocol import httpx +from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import iter_items, subscript from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -44,6 +46,27 @@ def _header_value(headers: Mapping[str, str], key: str, default: str) -> str: return headers.get(key, default) +_DECODED_JSON: Final = TypeAdapter(object) + + +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + +def _generated_text(value: object) -> str: + if isinstance(value, str): + return value + raise AttributeError(f"'{type(value).__name__}' object has no attribute 'strip'", name="strip", obj=value) + + +def _add_logprob(total: float, logprob: object) -> float: + if isinstance(logprob, (int, float)): + return total + logprob + raise TypeError(f"unsupported operand type(s) for +=: '{type(total).__name__}' and '{type(logprob).__name__}'") + + tgi_models_cache = None conv_models_cache = None @@ -407,33 +430,41 @@ class HuggingFaceEmbeddingConfig(BaseConfig): headers=None, ) - if len(completion_response[0]["generated_text"]) > 0: - model_response.choices[0].message.content = output_parser(completion_response[0]["generated_text"]) + if _json_len(_DECODED_JSON.validate_python(completion_response[0]["generated_text"])) > 0: + model_response.choices[0].message.content = output_parser( + _generated_text(_DECODED_JSON.validate_python(completion_response[0]["generated_text"])) + ) ## GETTING LOGPROBS + FINISH REASON if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]: model_response.choices[0].finish_reason = completion_response[0]["details"]["finish_reason"] sum_logprob = 0 - for token in completion_response[0]["details"]["tokens"]: - if token["logprob"] is not None: - sum_logprob += token["logprob"] + for token in iter_items(_DECODED_JSON.validate_python(completion_response[0]["details"]["tokens"])): + logprob = subscript(token, "logprob") + if logprob is not None: + sum_logprob = _add_logprob(sum_logprob, logprob) setattr(model_response.choices[0].message, "_logprob", sum_logprob) if "best_of" in optional_params and optional_params["best_of"] > 1: if "details" in completion_response[0] and "best_of_sequences" in completion_response[0]["details"]: choices_list: Final = [] - for idx, item in enumerate(completion_response[0]["details"]["best_of_sequences"]): + for idx, item in enumerate( + iter_items( + _DECODED_JSON.validate_python(completion_response[0]["details"]["best_of_sequences"]) + ) + ): sum_logprob = 0 - for token in item["tokens"]: - if token["logprob"] is not None: - sum_logprob += token["logprob"] - if len(item["generated_text"]) > 0: + for token in iter_items(subscript(item, "tokens")): + logprob = subscript(token, "logprob") + if logprob is not None: + sum_logprob = _add_logprob(sum_logprob, logprob) + if _json_len(subscript(item, "generated_text")) > 0: message_obj = Message( - content=output_parser(item["generated_text"]), + content=output_parser(_generated_text(subscript(item, "generated_text"))), logprobs=sum_logprob, ) else: message_obj = Message(content=None) choice_obj = Choices( - finish_reason=item["finish_reason"], + finish_reason=subscript(item, "finish_reason"), index=idx + 1, message=message_obj, ) @@ -442,8 +473,13 @@ class HuggingFaceEmbeddingConfig(BaseConfig): elif task == "text-classification": model_response.choices[0].message.content = json.dumps(completion_response) else: - if isinstance(completion_response, list) and len(completion_response[0]["generated_text"]) > 0: - model_response.choices[0].message.content = output_parser(completion_response[0]["generated_text"]) + if ( + isinstance(completion_response, list) + and _json_len(_DECODED_JSON.validate_python(completion_response[0]["generated_text"])) > 0 + ): + model_response.choices[0].message.content = output_parser( + _generated_text(_DECODED_JSON.validate_python(completion_response[0]["generated_text"])) + ) ## CALCULATING USAGE prompt_tokens = 0 try: diff --git a/litellm/llms/minimax/videos/transformation.py b/litellm/llms/minimax/videos/transformation.py index aff99542d9b..fd014a3d521 100644 --- a/litellm/llms/minimax/videos/transformation.py +++ b/litellm/llms/minimax/videos/transformation.py @@ -803,7 +803,7 @@ class MinimaxVideoConfig(BaseVideoConfig): params.update(extra_query) return f"{api_base}/v2/query/video_generation", params - def transform_video_list_response( # pyright: ignore[reportIncompatibleMethodOverride] # base declares dict[str, str] but the payload is a heterogeneous list body + def transform_video_list_response( self, raw_response: httpx.Response, logging_obj: "Logging", diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py index 0cca299de93..8301bee9b90 100644 --- a/litellm/llms/nlp_cloud/chat/transformation.py +++ b/litellm/llms/nlp_cloud/chat/transformation.py @@ -1,9 +1,10 @@ import json import time -from collections.abc import Mapping +from collections.abc import Mapping, Sized from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, @@ -23,6 +24,15 @@ else: LoggingClass = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + class NLPCloudConfig(BaseConfig): """ Reference: https://docs.nlpcloud.com/#generation @@ -199,7 +209,7 @@ class NLPCloudConfig(BaseConfig): ) else: try: - if len(completion_response["generated_text"]) > 0: + if _json_len(_DECODED_JSON.validate_python(completion_response["generated_text"])) > 0: model_response.choices[0].message.content = completion_response["generated_text"] except Exception: raise NLPCloudError( diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py index 8367853673f..8f8111be3ce 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py @@ -6,11 +6,12 @@ request carries them exactly once over both plaintext and TLS. """ from collections.abc import Iterator, Mapping, Sequence -from typing import Any, Final, TypeAlias +from typing import Final, Protocol, SupportsIndex, SupportsInt, TypeAlias, runtime_checkable import grpc from google.protobuf.message import Message from pydantic import TypeAdapter +from typing_extensions import Buffer from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema as schema from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException @@ -38,16 +39,29 @@ def open_channel(api_base: str, use_ssl: bool) -> grpc.Channel: return grpc.insecure_channel(api_base) -def recognition_config(config: Mapping[str, Any]) -> Message: +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +def _native_int(value: object) -> int: + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)): + return int(value) + raise TypeError( + f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'" + ) + + +def recognition_config(config: Mapping[str, object]) -> Message: endpointing: Final = _endpointing_config(config.get("endpointing_config")) return schema.RecognitionConfig( encoding=schema.LINEAR_PCM, - sample_rate_hertz=int(config["sample_rate_hertz"]), + sample_rate_hertz=_native_int(config["sample_rate_hertz"]), language_code=config["language_code"], - audio_channel_count=int(config["audio_channel_count"]), + audio_channel_count=_native_int(config["audio_channel_count"]), enable_automatic_punctuation=bool(config.get("enable_automatic_punctuation", True)), enable_word_time_offsets=bool(config.get("enable_word_time_offsets", False)), - max_alternatives=int(config.get("max_alternatives", 1)), + max_alternatives=_native_int(config.get("max_alternatives", 1)), model=config.get("model", "") or "", verbatim_transcripts=bool(config.get("verbatim_transcripts", False)), profanity_filter=bool(config.get("profanity_filter", False)), diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 1e3e6186dd1..22e09584b1d 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -8,7 +8,7 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm._uuid import uuid -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import decoded_json_object, map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_images_from_message, @@ -47,6 +47,15 @@ else: LiteLLMLoggingObj = Any _REQUEST_PARAMS: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) + + +def _json_document(value: object) -> str | bytes | bytearray: + if isinstance(value, (str, bytes, bytearray)): + return value + raise TypeError(f"the JSON object must be str, bytes or bytearray, not {value.__class__.__name__}") + + _TOOL_CALL_ARGUMENTS: Final = TypeAdapter(dict[str, object]) @@ -372,7 +381,7 @@ class OllamaChatConfig(BaseConfig): # remap 'thinking' to 'reasoning_content' response_json_message["reasoning_content"] = response_json_message["thinking"] del response_json_message["thinking"] - elif response_json_message.get("content") is not None: + elif decoded_json_object(_DECODED_JSON.validate_python(response_json_message)).get("content") is not None: # parse reasoning content from content from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( parse_content_for_reasoning, @@ -387,14 +396,20 @@ class OllamaChatConfig(BaseConfig): and litellm_params.get("function_name") is not None and response_json_message is not None ): - function_call: Final = json.loads(response_json_message["content"]) + function_call: Final = decoded_json_object( + _DECODED_JSON.validate_python( + json.loads(_json_document(_DECODED_JSON.validate_python(response_json_message["content"]))) + ) + ) message: Final = litellm.Message( content=None, tool_calls=[ { "id": f"call_{uuid.uuid4()}", "function": { - "name": function_call.get("name", litellm_params.get("function_name")), + "name": function_call.get( + "name", _DECODED_JSON.validate_python(litellm_params.get("function_name")) + ), "arguments": json.dumps(function_call.get("arguments", function_call)), }, "type": "function", diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index b4dc57ff5ed..a473ef18630 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -1,7 +1,9 @@ import json -from collections.abc import Callable +from collections.abc import Callable, Sized from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter + import litellm from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client @@ -15,6 +17,13 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj oobabooga_config: Final = OobaboogaConfig() +_DECODED_JSON: Final = TypeAdapter(object) + + +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") def completion( @@ -149,7 +158,7 @@ def embedding( } ] - num_tokens: Final = len(completion_response["data"][0]["embedding"]) + num_tokens: Final = _json_len(_DECODED_JSON.validate_python(completion_response["data"][0]["embedding"])) # Adding metadata to response setattr( model_response, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index f7bd46536a1..c4c72ab6be6 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -8,15 +8,21 @@ import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextvars import ContextVar from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast, overload from urllib.parse import urlparse import httpx from pydantic import TypeAdapter, ValidationError +from typing_extensions import TypedDict import litellm from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason +from litellm.litellm_core_utils.core_helpers import ( + drops_unsupported_params, + map_finish_reason, + require_attributes, + subscript, +) from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( extract_reasoning_content, handle_invalid_parallel_tool_calls, @@ -147,6 +153,14 @@ _CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _DECODED_JSON: Final = TypeAdapter(object) +class _ModelResponseStreamFields(TypedDict, total=False, extra_items=object): + pass + + +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: str, /) -> object: ... + + def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None: value: Final = litellm_params.get(key) if litellm_params is not None else None return value if isinstance(value, str) else None @@ -808,7 +822,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): convert_tool_call_to_json_mode=json_mode, ): # to support response_format on claude models - json_mode_content_str: str | None = str(new_tool_calls[0]["function"].get("arguments", "")) or None + tool_call_function = require_attributes( + subscript(new_tool_calls[0], "function"), _SupportsGetWithDefault, "get" + ) + json_mode_content_str: str | None = str(tool_call_function.get("arguments", "")) or None if json_mode_content_str is not None: translated_message = Message(content=json_mode_content_str) finish_reason = "stop" @@ -1086,7 +1103,7 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): choices: Final = self._map_reasoning_to_reasoning_content(raw_choices) fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint") - kwargs: Final[dict[str, Any]] = { + kwargs: Final[dict[str, object]] = { "id": chunk.get("id"), "system_fingerprint": fingerprint, "object": "chat.completion.chunk", @@ -1099,6 +1116,6 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): service_tier: Final = chunk.get("service_tier") if isinstance(service_tier, str) and service_tier: kwargs["service_tier"] = service_tier - return ModelResponseStream(**kwargs) + return ModelResponseStream(**_ModelResponseStreamFields(**kwargs)) except Exception as e: raise e diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index bc8aaab3185..a6dbc79254e 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -4,7 +4,9 @@ Helper util for handling openai-specific cost calculation """ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal, SupportsFloat, SupportsIndex + +from typing_extensions import Buffer from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -100,8 +102,14 @@ def _video_resolution_to_cost_field_suffix(resolution: str) -> str | None: return safe +def _native_float(value: object) -> float: + if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(value) + raise TypeError(f"float() argument must be a string or a real number, not '{type(value).__name__}'") + + def video_output_cost_per_second( - model_info: Mapping[str, Any], + model_info: Mapping[str, object], video_resolution: str | None, ) -> float | None: """ @@ -118,10 +126,10 @@ def video_output_cost_per_second( tier_key: Final = f"output_cost_per_second_{suffix}" tier_rate: Final = model_info.get(tier_key) if tier_rate is not None: - return float(tier_rate) + return _native_float(tier_rate) out: Final = model_info.get("output_cost_per_second") if out is not None: - return float(out) + return _native_float(out) return None diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 0bba58096ef..92fdffebe34 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -3,9 +3,12 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter +from typing_extensions import TypedDict import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import keyword_arguments from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -28,6 +31,13 @@ else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +class _ImageResponseFields(TypedDict, total=False, extra_items=object): + pass + + class OpenAIImageEditConfig(BaseImageEditConfig): """ Base configuration for OpenAI image edit API. @@ -149,10 +159,10 @@ class OpenAIImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: """No transform applied since outputs are in OpenAI spec already""" try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) - return ImageResponse(**raw_response_json) + return ImageResponse(**_ImageResponseFields(**keyword_arguments(raw_response_json, ImageResponse))) def validate_environment( self, diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 504dfb50015..0200d9acbb0 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -7,6 +7,8 @@ from typing import Final import httpx from openai import AsyncOpenAI, OpenAI +from pydantic import BaseModel +from typing_extensions import TypedDict import litellm from litellm.types.utils import FileTypes, ImageResponse, LlmProviders @@ -18,6 +20,14 @@ from ..common_utils import OpenAIError, upstream_error_response from ..workload_identity import build_async_openai_client, build_openai_client +class _ImageResponseFields(TypedDict, total=False, extra_items=object): + pass + + +def _dumped_fields(model: BaseModel) -> Mapping[str, object]: + return model.model_dump() + + class OpenAIImageVariationsHandler: def get_sync_client( self, @@ -94,7 +104,7 @@ class OpenAIImageVariationsHandler: raw_response: Final = await client.images.with_raw_response.create_variation(**data) response: Final = raw_response.parse() - response_json: Final = response.model_dump() + response_json: Final = _dumped_fields(response) ## LOGGING logging_obj.post_call( @@ -109,7 +119,7 @@ class OpenAIImageVariationsHandler: ## RESPONSE OBJECT return provider_config.transform_response_image_variation( model=model, - model_response=ImageResponse(**response_json), + model_response=ImageResponse(**_ImageResponseFields(**response_json)), raw_response=httpx.Response( status_code=200, request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object @@ -215,7 +225,7 @@ class OpenAIImageVariationsHandler: raw_response: Final = client.images.with_raw_response.create_variation(**json_data) response: Final = raw_response.parse() - response_json: Final = response.model_dump() + response_json: Final = _dumped_fields(response) ## LOGGING logging_obj.post_call( @@ -230,7 +240,7 @@ class OpenAIImageVariationsHandler: ## RESPONSE OBJECT return provider_config.transform_response_image_variation( model=model, - model_response=ImageResponse(**response_json), + model_response=ImageResponse(**_ImageResponseFields(**response_json)), raw_response=httpx.Response( status_code=200, request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 8e9f2a0a0bd..77c0223b375 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -36,7 +36,7 @@ from collections.abc import Iterable, Mapping, MutableSequence, Sequence from dataclasses import dataclass from itertools import accumulate, chain, groupby, repeat from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar, Union, cast, runtime_checkable from openai.types.responses.response_output_refusal import ResponseOutputRefusal from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -52,7 +52,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i OpenAiResponsesToChatCompletionStreamIterator, tool_call_dict_from_output_item, ) -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -694,6 +694,24 @@ def _is_textless_tool_row(message: Mapping[str, object], texts: tuple[str, ...] return _is_string_tool_row(message) and (not message.get("content") or texts == ()) +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _assignable(item: object) -> _SupportsSetItem: + if isinstance(item, _SupportsSetItem): + return item + raise TypeError(f"'{type(item).__name__}' object does not support item assignment") + + +def _summarized_encrypted_content(item: Mapping[str, object]) -> str | None: + if compaction_summary(item) is None: + return None + encrypted_content: Final = item["encrypted_content"] + return encrypted_content if isinstance(encrypted_content, str) else None + + def _rewritten_whole_text(item: object, field: str, guardrail_response: str) -> object | None: if isinstance(stream_item_field(item, field), str): return guardrail_response @@ -1134,7 +1152,7 @@ class OpenAIResponsesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam + messages: Sequence[Mapping[str, object]], responses: Sequence[str], task_mappings: Sequence[tuple[int, int | None]], ) -> None: @@ -1149,26 +1167,27 @@ class OpenAIResponsesHandler(BaseTranslation): for guardrail_response, mapping in zip(responses, task_mappings): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) + item = messages[msg_idx] - if compaction_summary(messages[msg_idx]) is not None: - messages[msg_idx]["encrypted_content"] = with_compaction_summary( # rebind-ok: in-place input rewrite - messages[msg_idx]["encrypted_content"], guardrail_response - ) + encrypted_content = _summarized_encrypted_content(item) + if encrypted_content is not None: + _assignable(item)["encrypted_content"] = with_compaction_summary(encrypted_content, guardrail_response) continue - field = _scanned_text_field(messages[msg_idx]) - content = messages[msg_idx].get(field, None) + field = _scanned_text_field(item) + content = item.get(field, None) if content is None: continue + content_parts = as_object_list(content) if content_idx_optional is None: - rewritten = _rewritten_whole_text(messages[msg_idx], field, guardrail_response) + rewritten = _rewritten_whole_text(item, field, guardrail_response) if rewritten is not None: - messages[msg_idx][field] = rewritten + _assignable(item)[field] = rewritten - elif isinstance(content, list): + elif content_parts is not None: # Replace specific text item in list content - self._write_content_part_text(content[content_idx_optional], guardrail_response) + self._write_content_part_text(content_parts[content_idx_optional], guardrail_response) async def process_output_response( self, @@ -1299,7 +1318,7 @@ class OpenAIResponsesHandler(BaseTranslation): request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, - ) -> list[Any]: + ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. diff --git a/litellm/llms/openai/transcriptions/whisper_transformation.py b/litellm/llms/openai/transcriptions/whisper_transformation.py index 947d6175122..21fdf202a52 100644 --- a/litellm/llms/openai/transcriptions/whisper_transformation.py +++ b/litellm/llms/openai/transcriptions/whisper_transformation.py @@ -3,7 +3,10 @@ from collections.abc import Mapping from typing import Final from httpx import Headers, Response +from pydantic import TypeAdapter +from typing_extensions import TypedDict +from litellm.litellm_core_utils.core_helpers import json_contains, keyword_arguments from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, @@ -19,6 +22,12 @@ from litellm.types.utils import FileTypes, TranscriptionResponse from ..common_utils import OpenAIError +_DECODED_JSON: Final = TypeAdapter(object) + + +class _TranscriptionResponseFields(TypedDict, total=False, extra_items=object): + pass + def _header_value(headers: Mapping[str, str], key: str, default: str) -> str: return headers.get(key, default) @@ -131,14 +140,16 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): raw_response: Response, ) -> TranscriptionResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except json.JSONDecodeError: content_type: Final = _header_value(raw_response.headers, "content-type", "").lower() if "application/json" in content_type: raise return TranscriptionResponse(text=raw_response.text) - if any(key in raw_response_json for key in TranscriptionResponse.model_fields): - return TranscriptionResponse(**raw_response_json) + if any(json_contains(raw_response_json, key) for key in TranscriptionResponse.model_fields): + return TranscriptionResponse( + **_TranscriptionResponseFields(**keyword_arguments(raw_response_json, TranscriptionResponse)) + ) else: raise ValueError("Invalid response format. Received response does not match the expected format.") diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index d108802f7f9..7fee55ade1c 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -1,8 +1,11 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import TypeAdapter +from typing_extensions import TypedDict import litellm +from litellm.litellm_core_utils.core_helpers import keyword_arguments from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.openai.workload_identity import resolve_openai_bearer_token @@ -27,6 +30,12 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +class _VectorStoreResponseFields(TypedDict, total=False, extra_items=object): + pass + class OpenAIVectorStoreConfig(BaseVectorStoreConfig): ASSISTANTS_HEADER_KEY = "OpenAI-Beta" @@ -118,8 +127,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: try: - response_json: Final = response.json() - return VectorStoreSearchResponse(**response_json) + response_json: Final = _DECODED_JSON.validate_python(response.json()) + return VectorStoreSearchResponse( + **_VectorStoreResponseFields(**keyword_arguments(response_json, VectorStoreSearchResponse)) + ) except Exception as e: raise self.get_error_class( error_message=str(e), @@ -149,8 +160,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: try: - response_json: Final = response.json() - return VectorStoreCreateResponse(**response_json) + response_json: Final = _DECODED_JSON.validate_python(response.json()) + return VectorStoreCreateResponse( + **_VectorStoreResponseFields(**keyword_arguments(response_json, VectorStoreCreateResponse)) + ) except Exception as e: raise self.get_error_class( error_message=str(e), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index abd157a9a08..4d3efc29fa7 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -2,12 +2,21 @@ import mimetypes from collections.abc import Mapping from io import BufferedReader, BytesIO from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from urllib.parse import quote import httpx from httpx._types import FileContent, FileTypes, RequestFiles +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + decoded_json_object, + iter_items, + json_contains, + require_attributes, + subscript, +) from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields from litellm.llms.base_llm.videos.transformation import BaseVideoConfig @@ -39,6 +48,59 @@ else: BaseLLMException = Any +_REQUEST_VALUE: Final = TypeAdapter(object) + + +def _dict_fields(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + +def _model_label(model: object) -> str | None: + return format(model) if model else None + + +def _encoded_video_id(video_id: object, provider: str, model: object) -> object: + if not video_id: + return video_id + return encode_video_id_with_provider( + require_attributes(video_id, str, "startswith"), + provider, + _model_label(model), + ) + + +def _with_encoded_id(video: object, provider: str) -> object: + fields: Final = _dict_fields(video) + if fields is None or "id" not in fields: + return video + return {**fields, "id": _encoded_video_id(fields["id"], provider, fields.get("model"))} + + +def _edge_model(data: object, index: int) -> object: + if not data: + return None + edge: Final = _dict_fields(subscript(data, index)) + return None if edge is None else edge.get("model") + + +def _decoded_character(character: object) -> object: + fields: Final = _dict_fields(character) + if fields is None: + return character + character_id: Final = fields.get("id") + if not isinstance(character_id, str): + return character + return {**fields, "id": extract_original_character_id(character_id)} + + class OpenAIVideoConfig(BaseVideoConfig): """ Configuration class for OpenAI video generation. @@ -157,25 +219,11 @@ class OpenAIVideoConfig(BaseVideoConfig): `character_`, convert it back to the original provider id before forwarding upstream. """ - raw_characters: Final = request_dict.get("characters") - if not isinstance(raw_characters, list): + raw_characters: Final = as_object_list(_REQUEST_VALUE.validate_python(request_dict.get("characters"))) + if raw_characters is None: return request_dict - decoded_characters: Final[list[Any]] = [] - for character in raw_characters: - if not isinstance(character, dict): - decoded_characters.append(character) - continue - - character_id = character.get("id") - if isinstance(character_id, str): - decoded_character = dict(character) - decoded_character["id"] = extract_original_character_id(character_id) - decoded_characters.append(decoded_character) - else: - decoded_characters.append(character) - - request_dict["characters"] = decoded_characters + request_dict["characters"] = [_decoded_character(character) for character in raw_characters] return request_dict def transform_video_create_response( @@ -343,41 +391,27 @@ class OpenAIVideoConfig(BaseVideoConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, - ) -> dict[str, str]: - response_data: Final = raw_response.json() + ) -> object: + response_data: Final = _json_body(raw_response) + if not custom_llm_provider or not json_contains(response_data, "data"): + return response_data - if custom_llm_provider and "data" in response_data: - for video_obj in response_data.get("data", []): - if isinstance(video_obj, dict) and "id" in video_obj: - video_obj["id"] = encode_video_id_with_provider( - video_obj["id"], - custom_llm_provider, - video_obj.get("model"), - ) + fields: Final = decoded_json_object(response_data) + data: Final = fields.get("data", []) + encoded_items: Final = [_with_encoded_id(item, custom_llm_provider) for item in iter_items(data)] + encoded_data: Final = encoded_items if as_object_list(data) is not None else data - # Encode pagination cursor IDs so they remain consistent - # with the wrapped data[].id format - data_list: Final = response_data.get("data", []) - if response_data.get("first_id"): - first_model = None - if data_list and isinstance(data_list[0], dict): - first_model = data_list[0].get("model") - response_data["first_id"] = encode_video_id_with_provider( - response_data["first_id"], - custom_llm_provider, - first_model, - ) - if response_data.get("last_id"): - last_model = None - if data_list and isinstance(data_list[-1], dict): - last_model = data_list[-1].get("model") - response_data["last_id"] = encode_video_id_with_provider( - response_data["last_id"], - custom_llm_provider, - last_model, - ) - - return response_data + # Encode pagination cursor IDs so they remain consistent + # with the wrapped data[].id format + first_id: Final = fields.get("first_id") + encoded_first: Final[Mapping[str, object]] = ( + {"first_id": _encoded_video_id(first_id, custom_llm_provider, _edge_model(data, 0))} if first_id else {} + ) + last_id: Final = fields.get("last_id") + encoded_last: Final[Mapping[str, object]] = ( + {"last_id": _encoded_video_id(last_id, custom_llm_provider, _edge_model(data, -1))} if last_id else {} + ) + return {**fields, "data": encoded_data, **encoded_first, **encoded_last} def transform_video_delete_request( self, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index a450947afdf..5a5d038979b 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -5,14 +5,16 @@ For handling OpenAI-like chat completions, like IBM WatsonX, etc. """ import json -from collections.abc import Callable, Mapping, Sequence -from typing import Final, TypedDict +from collections.abc import Callable +from typing import Final import httpx -from typing_extensions import ReadOnly +from pydantic import TypeAdapter +from typing_extensions import TypedDict import litellm from litellm import LlmProviders +from litellm.litellm_core_utils.core_helpers import keyword_arguments from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -26,22 +28,16 @@ from litellm.utils import CustomStreamWrapper, ProviderConfigManager from ..common_utils import OpenAILikeBase, OpenAILikeError from .transformation import OpenAILikeChatConfig - -class _OpenAILikeChatCompletion(TypedDict, total=False): - """The chat-completion JSON body an OpenAI-like provider returns for a non-streamed call.""" - - id: ReadOnly[str] - choices: ReadOnly[Sequence[Mapping[str, object]]] - created: ReadOnly[int] - model: ReadOnly[str] - system_fingerprint: ReadOnly[str] - usage: ReadOnly[Mapping[str, object]] - object: ReadOnly[str] +_DECODED_JSON: Final = TypeAdapter(object) -def _fake_streamed_model_response(payload: _OpenAILikeChatCompletion) -> ModelResponse: +class _ModelResponseFields(TypedDict, total=False, extra_items=object): + pass + + +def _fake_streamed_model_response(payload: object) -> ModelResponse: """Build the single response a fake-streamed provider call replays as one chunk.""" - return ModelResponse(**payload) + return ModelResponse(**_ModelResponseFields(**keyword_arguments(payload, ModelResponse))) async def make_call( @@ -63,7 +59,7 @@ async def make_call( if streaming_decoder is not None: completion_stream = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = _fake_streamed_model_response(response.json()) + model_response: Final = _fake_streamed_model_response(_DECODED_JSON.validate_python(response.json())) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator( @@ -103,7 +99,7 @@ def make_sync_call( if streaming_decoder is not None: completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = _fake_streamed_model_response(response.json()) + model_response: Final = _fake_streamed_model_response(_DECODED_JSON.validate_python(response.json())) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator( diff --git a/litellm/llms/openai_like/embedding/handler.py b/litellm/llms/openai_like/embedding/handler.py index 540feeca6e7..b064505b445 100644 --- a/litellm/llms/openai_like/embedding/handler.py +++ b/litellm/llms/openai_like/embedding/handler.py @@ -10,9 +10,11 @@ from typing import Final import httpx from pydantic import TypeAdapter +from typing_extensions import TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import keyword_arguments from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -40,6 +42,11 @@ _TRANSPORT_HEADER_NAMES: Final = frozenset( ) _HEADER_NAME: Final = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+") _HEADER_VALUE: Final = re.compile(r"[\t\x20-\x7e]*") +_DECODED_JSON: Final = TypeAdapter(object) + + +class _EmbeddingResponseFields(TypedDict, total=False, extra_items=object): + pass def _header_text(name: str, value: object) -> str | None: @@ -108,7 +115,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): response.raise_for_status() - response_json: Final = response.json() + response_json: Final = _DECODED_JSON.validate_python(response.json()) except httpx.HTTPStatusError as e: raise OpenAILikeError( status_code=e.response.status_code, @@ -126,7 +133,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): additional_args={"complete_input_dict": data}, original_response=response_json, ) - return EmbeddingResponse(**response_json) + return EmbeddingResponse(**_EmbeddingResponseFields(**keyword_arguments(response_json, EmbeddingResponse))) except Exception as e: ## LOGGING logging_obj.post_call( diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index b30713fa119..94e64dfb95b 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -5,10 +5,11 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal from httpx import Headers, Response +from pydantic import TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_TOKENS -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import json_contains, map_finish_reason, subscript from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, @@ -28,6 +29,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + class PredibaseConfig(BaseConfig): """ @@ -163,13 +166,13 @@ class PredibaseConfig(BaseConfig): additional_args={"complete_input_dict": request_data}, ) try: - completion_response: Final = raw_response.json() + completion_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise PredibaseError(message=raw_response.text, status_code=422) - if "error" in completion_response: + if json_contains(completion_response, "error"): raise PredibaseError( - message=str(completion_response["error"]), + message=str(subscript(completion_response, "error")), status_code=raw_response.status_code, ) elif not isinstance(completion_response, dict): diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ef44b6e4c0f..421bd769778 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -41,6 +41,11 @@ class _RAGFlowCreateResponseRoot(RootModel[_RAGFlowCreateResponse]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) +class _RAGFlowStatus(TypedDict): + code: ReadOnly[NotRequired[int]] + message: ReadOnly[NotRequired[str]] + + _RAGFLOW_EMPTY_DATA: Final[_RAGFlowDataset] = {} @@ -197,7 +202,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): } """ try: - response_json: Final = response.json() + response_json: Final[_RAGFlowStatus] = response.json() # Check for RAGFlow error response if response_json.get("code") != 0: diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index 62dfac43bd7..4553a02e451 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH @@ -33,6 +34,10 @@ _TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_i _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +class _ReplicatePrediction(TypedDict): + urls: ReadOnly[NotRequired[Mapping[str, str]]] + + class ReplicateConfig(BaseConfig): """ Reference: https://replicate.com/meta/llama-2-70b-chat/api @@ -292,7 +297,7 @@ class ReplicateConfig(BaseConfig): "urls":{"cancel":"https://api.replicate.com/v1/predictions/gqsmqmp1pdrj00cknr08dgmvb4/cancel","get":"https://api.replicate.com/v1/predictions/gqsmqmp1pdrj00cknr08dgmvb4","stream":"https://stream-b.svc.rno2.c.replicate.net/v1/streams/eot4gbydowuin4snhncydwxt57dfwgsc3w3snycx5nid7oef7jga"} } """ - response_json: Final = response.json() + response_json: Final[_ReplicatePrediction] = response.json() prediction_url: Final = response_json.get("urls", {}).get("get") if prediction_url is None: raise ReplicateError( diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 5f6bef4246e..88f953091b7 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -11,9 +11,11 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from httpx._models import Headers, Response +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import as_object_list, json_contains, subscript from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, @@ -33,6 +35,14 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +def _concatenated(prefix: str, value: object) -> str: + if isinstance(value, str): + return prefix + value + raise TypeError(f'can only concatenate str (not "{type(value).__name__}") to str') + class SagemakerConfig(BaseConfig): """ @@ -199,7 +209,7 @@ class SagemakerConfig(BaseConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - completion_response: Final = raw_response.json() + completion_response: Final = _DECODED_JSON.validate_python(raw_response.json()) ## LOGGING logging_obj.post_call( input=messages, @@ -212,15 +222,19 @@ class SagemakerConfig(BaseConfig): ## RESPONSE OBJECT try: - if isinstance(completion_response, list): - completion_response_choices = completion_response[0] - else: - completion_response_choices = completion_response + completion_response_list: Final = as_object_list(completion_response) + completion_response_choices: Final = ( + completion_response if completion_response_list is None else completion_response_list[0] + ) completion_output = "" - if "generation" in completion_response_choices: - completion_output += completion_response_choices["generation"] - elif "generated_text" in completion_response_choices: - completion_output += completion_response_choices["generated_text"] + if json_contains(completion_response_choices, "generation"): + completion_output = _concatenated( + completion_output, subscript(completion_response_choices, "generation") + ) + elif json_contains(completion_response_choices, "generated_text"): + completion_output = _concatenated( + completion_output, subscript(completion_response_choices, "generated_text") + ) # check if the prompt template is part of output, if so - filter it out if completion_output.startswith(prompt) and "" in prompt: diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 928dcc3bc52..161cf908a6f 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict import httpx from typing_extensions import ReadOnly -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_process_openai_file_message, convert_to_anthropic_tool_result, @@ -354,13 +354,15 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): if isinstance(block, Mapping) and block.get("type") == "text" ) elif role == "assistant": - tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + tool_calls: object = ( + msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + ) thinking_blocks = _signed_thinking_blocks(msg) if tool_calls: content_blocks: list[dict[str, object]] = list(thinking_blocks) # mutable-ok: JSON wire blocks if content: content_blocks.append({"type": "text", "text": content}) - for tc in tool_calls: + for tc in iter_items(tool_calls): func = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", {}) tc_id = tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", "") func_name = func.get("name", "") if isinstance(func, dict) else getattr(func, "name", "") diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index 30d872f248f..cefb195f4a5 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -13,11 +13,13 @@ import re from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, NoReturn +from typing import Any, Final, NoReturn, SupportsFloat, SupportsIndex from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError +from typing_extensions import Buffer +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -172,6 +174,12 @@ def _nested_form_param(param: str, fields: tuple[_BracketedField, ...]) -> Mappi } +def _as_float(value: object) -> float: + if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(value) + raise TypeError(f"float() argument must be a string or a real number, not {type(value).__name__!r}") + + def _is_soniox_context(value: object) -> bool: try: _STRUCTURED_CONTEXT.validate_python(value) @@ -345,7 +353,7 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): response_format: str | None = None, ) -> TranscriptionResponse: """Shared response-building logic (also used by the handler).""" - transcription_meta: dict[str, Any] = {} + transcription_meta: object = {} transcript: dict[str, Any] if isinstance(payload, dict) and "transcript" in payload: @@ -385,9 +393,10 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): response["task"] = "transcribe" # Best-effort metadata fields matching OpenAI's verbose_json shape. - if transcription_meta.get("audio_duration_ms") is not None: + audio_duration_ms: Final = decoded_json_object(transcription_meta).get("audio_duration_ms") + if audio_duration_ms is not None: try: - response["duration"] = float(transcription_meta["audio_duration_ms"]) / 1000.0 + response["duration"] = _as_float(audio_duration_ms) / 1000.0 except (TypeError, ValueError): pass @@ -403,10 +412,12 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): words: Final[list[dict[str, object]]] = [] for token in tokens: word_entry: dict[str, object] = {"word": token.get("text", "")} - if token.get("start_ms") is not None: - word_entry["start"] = float(token["start_ms"]) / 1000.0 - if token.get("end_ms") is not None: - word_entry["end"] = float(token["end_ms"]) / 1000.0 + start_ms: object = token.get("start_ms") + if start_ms is not None: + word_entry["start"] = _as_float(start_ms) / 1000.0 + end_ms: object = token.get("end_ms") + if end_ms is not None: + word_entry["end"] = _as_float(end_ms) / 1000.0 words.append(word_entry) if words: response["words"] = words diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 974996c5e48..65277973d39 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -6,13 +6,13 @@ Handles transformation between OpenAI-compatible format and Stability AI API for API Reference: https://platform.stability.ai/docs/api-reference """ -from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from typing_extensions import ReadOnly, TypedDict +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object, json_contains, subscript from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -31,12 +31,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any - -class _StabilityImageEditResponse(TypedDict, total=False): - image: ReadOnly[str] - finish_reason: ReadOnly[str] - seed: ReadOnly[int] - errors: ReadOnly[Sequence[str]] +_DECODED_JSON: Final = TypeAdapter(object) class StabilityImageEditConfig(BaseImageEditConfig): @@ -268,7 +263,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} """ try: - response_data: Final[_StabilityImageEditResponse] = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Stability AI response: {e}", @@ -277,14 +272,15 @@ class StabilityImageEditConfig(BaseImageEditConfig): ) # Check for errors in response - if "errors" in response_data: + if json_contains(response_json, "errors"): raise self.get_error_class( - error_message=f"Stability AI error: {response_data['errors']}", + error_message=f"Stability AI error: {subscript(response_json, 'errors')}", status_code=raw_response.status_code, headers=raw_response.headers, ) # Check finish_reason + response_data: Final = decoded_json_object(response_json) finish_reason: Final = response_data.get("finish_reason", "") if finish_reason == "CONTENT_FILTERED": raise self.get_error_class( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index d0de71e0305..aa275f254b4 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -7,7 +7,7 @@ from enum import Enum from functools import lru_cache from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType -from typing import Any, Final, Literal, TypeVar, cast, get_type_hints +from typing import Final, Literal, TypeVar, cast, get_type_hints import httpx from pydantic import JsonValue, TypeAdapter, ValidationError @@ -1384,6 +1384,13 @@ class VertexAIModelInfo(BaseLLMModelInfo): raise NotImplementedError("Vertex AI models are not supported yet") +_COUNT_TOKENS_REQUEST_FIELDS: Final = frozenset({"model", "contents", "system_instruction", "tools"}) + + +class _ForwardedCountTokensParams(TypedDict, total=False, extra_items=object): + pass + + class VertexAITokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" @@ -1446,7 +1453,7 @@ class VertexAITokenCounter(BaseTokenCounter): model_to_use: str, messages: list[dict[str, object]] | None, contents: list[dict[str, object]] | None, - deployment: dict[str, Any] | None = None, + deployment: dict[str, object] | None = None, request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, @@ -1458,6 +1465,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.common_utils import context_management_with_edits from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( @@ -1467,19 +1475,19 @@ class VertexAITokenCounter(BaseTokenCounter): VertexAIPartnerModels, ) - deployment = deployment or {} - count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_deployment: Final = deployment or {} + raw_litellm_params: Final = copy.deepcopy(resolved_deployment.get("litellm_params", {})) if not VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and not uses_vertex_gemini_api( - model_to_use, count_tokens_params_request + model_to_use, decoded_json_object(raw_litellm_params) if raw_litellm_params else None ): return None without_count_route: Final = ( VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None ) - model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {} - if without_count_route and messages and model_info.get("custom_tokenizer") is not None: + model_info: Final[object] = resolved_deployment.get("model_info") or {} + if without_count_route and messages and decoded_json_object(model_info).get("custom_tokenizer") is not None: return None if without_count_route: return TokenCountResponse( @@ -1497,7 +1505,7 @@ class VertexAITokenCounter(BaseTokenCounter): translated_count: Final = await self._count_translated_claude_request( model_to_use=model_to_use, messages=messages, - deployment=deployment, + deployment=resolved_deployment, request_model=request_model, tools=tools, system=system, @@ -1513,7 +1521,7 @@ class VertexAITokenCounter(BaseTokenCounter): partner_models_handler: Final = VertexAIPartnerModels() # Extract vertex-specific params from litellm_params - partner_litellm_params: Final[dict[str, object]] = count_tokens_params_request + partner_litellm_params: Final = dict(decoded_json_object(raw_litellm_params)) vertex_project = partner_litellm_params.get("vertex_project") or partner_litellm_params.get( "vertex_ai_project" ) @@ -1585,11 +1593,12 @@ class VertexAITokenCounter(BaseTokenCounter): ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter + gemini_litellm_params: Final = decoded_json_object(raw_litellm_params) request_view: Final = await inference_request_view( messages or [], system, messages_format, - count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), + count_context_edits(context_management, tools, gemini_litellm_params, model_to_use), ) if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) @@ -1600,7 +1609,7 @@ class VertexAITokenCounter(BaseTokenCounter): return count_context_error_response(count_context, request_model, model_to_use) blocks_error: Final = unconvertible_messages_error( request_view.messages, messages_format - ) or malformed_tool_result_text_count_error(messages, messages_format, deployment) + ) or malformed_tool_result_text_count_error(messages, messages_format, resolved_deployment) if blocks_error is not None: return count_context_error_response(blocks_error, request_model, model_to_use) @@ -1608,22 +1617,21 @@ class VertexAITokenCounter(BaseTokenCounter): contents if contents is not None else await gemini_count_contents( - request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request + request_view.messages, messages_format, model_to_use, "vertex_ai", gemini_litellm_params ) ) if isinstance(resolved_contents, VertexGeminiCountContextError): return count_context_error_response(resolved_contents, request_model, model_to_use) - count_tokens_params: Final = { - "model": model_to_use, - "contents": resolved_contents, - "system_instruction": count_context.system_instruction, - "tools": count_context.tools, - } - count_tokens_params_request.update(count_tokens_params) try: result = await VertexAITokenCounter().acount_tokens( - **count_tokens_params_request, + **_ForwardedCountTokensParams( + **{k: v for k, v in gemini_litellm_params.items() if k not in _COUNT_TOKENS_REQUEST_FIELDS} + ), + model=model_to_use, + contents=resolved_contents, + system_instruction=count_context.system_instruction, + tools=count_context.tools, ) except litellm.APIError as error: return provider_error_response(error, request_model, model_to_use) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 1f03de4625c..7a7afcb83f1 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -10,7 +10,7 @@ from contextlib import aclosing from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import IO, Any, Final, TypedDict +from typing import IO, Final, TypedDict from urllib.parse import quote, unquote import httpx @@ -827,7 +827,7 @@ def _openai_batch_jsonl_entry_to_vertex_anthropic_row( def _openai_batch_jsonl_entry_to_vertex_rows( - openai_entry: dict[str, Any], + openai_entry: Mapping[str, object], map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], batch_model: str | None = None, deployment_params: Mapping[str, object] = MappingProxyType({}), @@ -842,7 +842,7 @@ def _openai_batch_jsonl_entry_to_vertex_rows( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ - entry_body: Final = openai_entry.get("body") + entry_body: Final = as_str_mapping(openai_entry.get("body")) model: Final = batch_model or (str(entry_body.get("model") or "") if isinstance(entry_body, dict) else "") is_anthropic_model: Final = is_vertex_anthropic_batch_model(model) if _is_embeddings_batch_entry(openai_entry): @@ -855,11 +855,10 @@ def _openai_batch_jsonl_entry_to_vertex_rows( if not isinstance(entry_body, dict): raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object") - openai_request_body: Final = openai_entry["body"] chat_request_body: Final = ( - responses_batch_body_to_chat_body(openai_request_body, custom_llm_provider="vertex_ai") + responses_batch_body_to_chat_body(entry_body, custom_llm_provider="vertex_ai") if _is_responses_batch_entry(openai_entry) - else openai_request_body + else entry_body ) chat_messages: Final = map_developer_role_to_system_role(chat_request_body.get("messages", [])) if is_anthropic_model: diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index ec915e10fcd..56843d2e1a1 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -7,6 +7,7 @@ from openai.types.image import Image from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import json_contains, subscript from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -20,23 +21,24 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj _PREDICTIONS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class VertexImageGeneration(VertexLLM): def process_image_generation_response( self, - json_response: Mapping[str, object], + json_response: object, model_response: ImageResponse, model: str | None = None, ) -> ImageResponse: - if "predictions" not in json_response: + if not json_contains(json_response, "predictions"): raise litellm.InternalServerError( message=f"image generation response does not contain 'predictions', got {json_response}", llm_provider="vertex_ai", model=model, ) - predictions: Final = _PREDICTIONS.validate_python(json_response["predictions"]) + predictions: Final = _PREDICTIONS.validate_python(subscript(json_response, "predictions")) response_data: Final[list[Image]] = [] for prediction in predictions: @@ -168,7 +170,7 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final[Mapping[str, object]] = response.json() + json_response: Final = _DECODED_JSON.validate_python(response.json()) return self.process_image_generation_response(json_response, model_response, model) async def aimage_generation( @@ -272,11 +274,11 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final[Mapping[str, object]] = response.json() + json_response: Final = _DECODED_JSON.validate_python(response.json()) return self.process_image_generation_response(json_response, model_response, model) - def is_image_generation_response(self, json_response: dict[str, Any]) -> bool: + def is_image_generation_response(self, json_response: Mapping[str, object]) -> bool: if "predictions" in json_response: - if "bytesBase64Encoded" in json_response["predictions"][0]: + if json_contains(subscript(json_response["predictions"], 0), "bytesBase64Encoded"): return True return False diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 7635cd32a06..5781481dde8 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -20,7 +20,10 @@ import json from collections.abc import Awaitable, Callable from typing import Final +from pydantic import TypeAdapter + from litellm import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( @@ -30,6 +33,8 @@ from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +_DECODED_JSON: Final = TypeAdapter(object) + class VertexAIRealtimeConfig(GeminiRealtimeConfig): """ @@ -217,7 +222,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): ``setup`` closes the connection with a 1007 policy error. """ json_message: Final = json.loads(message) - msg_type: Final = json_message.get("type") + msg_type: Final = decoded_json_object(_DECODED_JSON.validate_python(json_message)).get("type") if msg_type == "session.update": if session_configuration_request is None: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py index d11be440172..e26b26af278 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py @@ -131,7 +131,8 @@ class VertexAILlama3Config(OpenAIGPTConfig): ## RESPONSE OBJECT try: - completion_response: Final = OpenAIChatCompletionResponse(**raw_response.json()) + payload: Final[OpenAIChatCompletionResponse] = raw_response.json() + completion_response: Final = OpenAIChatCompletionResponse(**payload) except Exception as e: response_headers: Final = getattr(raw_response, "headers", None) raise VertexAIError( diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index cc1e7d59caa..283b013917b 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -8,7 +8,10 @@ from collections.abc import Mapping from typing import Final import httpx +from pydantic import TypeAdapter +from typing_extensions import TypedDict +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -18,6 +21,12 @@ from litellm.types.utils import EmbeddingResponse from ..common_utils import get_volcengine_base_url, get_volcengine_headers +_DECODED_JSON: Final = TypeAdapter(object) + + +class _VolcengineEmbeddingFields(TypedDict, total=False, extra_items=object): + pass + class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): """ @@ -159,9 +168,10 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): ) -> EmbeddingResponse: """Transform Volcengine response to EmbeddingResponse""" try: - response_json: Final = raw_response.json() + decoded_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise ValueError(f"Failed to parse Volcengine response as JSON: {e}") + response_json: Final = decoded_json_object(decoded_response) # Volcengine response format matches OpenAI format closely # Just need to ensure all required fields are present @@ -177,7 +187,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): transformed_response["id"] = response_json["id"] # Create EmbeddingResponse from transformed data - return EmbeddingResponse(**transformed_response) + return EmbeddingResponse(**_VolcengineEmbeddingFields(**transformed_response)) def validate_environment( self, diff --git a/litellm/llms/volcengine/videos/transformation.py b/litellm/llms/volcengine/videos/transformation.py index 5af1080b9a0..ca4d3546636 100644 --- a/litellm/llms/volcengine/videos/transformation.py +++ b/litellm/llms/volcengine/videos/transformation.py @@ -874,7 +874,7 @@ class VolcEngineVideoConfig(BaseVideoConfig): **(extra_query or {}), } - def transform_video_list_response( # pyright: ignore[reportIncompatibleMethodOverride] # base declares dict[str, str] but the payload is a heterogeneous list body + def transform_video_list_response( self, raw_response: httpx.Response, logging_obj: "Logging", diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py index 6d3686e9f85..5d2aa418173 100644 --- a/litellm/llms/watsonx/common_utils.py +++ b/litellm/llms/watsonx/common_utils.py @@ -1,10 +1,12 @@ from typing import Final, cast import httpx +from pydantic import TypeAdapter import litellm from litellm import verbose_logger from litellm.caching import InMemoryCache +from litellm.litellm_core_utils.core_helpers import subscript from litellm.litellm_core_utils.prompt_templates import factory as ptf from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -23,6 +25,13 @@ class WatsonXAIError(BaseLLMException): iam_token_cache: Final = InMemoryCache() +_IAM_TOKEN_RESPONSE: Final = TypeAdapter(object) + + +def _ttl_with_buffer(expires_in: object) -> float: + if isinstance(expires_in, (int, float)): + return expires_in - 10 + raise TypeError(f"unsupported operand type(s) for -: '{type(expires_in).__name__}' and 'int'") def get_watsonx_iam_url(): @@ -30,7 +39,7 @@ def get_watsonx_iam_url(): def generate_iam_token(api_key=None, **params) -> str: - result: str | None = iam_token_cache.get_cache(api_key) + result: object = iam_token_cache.get_cache(api_key) if result is None: headers: Final = {} @@ -58,13 +67,13 @@ def generate_iam_token(api_key=None, **params) -> str: ) response: Final = litellm.module_level_client.post(url=iam_token_url, data=data, headers=headers) response.raise_for_status() - json_data: Final = response.json() + json_data: Final = _IAM_TOKEN_RESPONSE.validate_python(response.json()) - result = json_data["access_token"] + result = subscript(json_data, "access_token") iam_token_cache.set_cache( key=api_key, value=result, - ttl=json_data["expires_in"] - 10, # leave some buffer + ttl=_ttl_with_buffer(subscript(json_data, "expires_in")), ) return cast(str, result) diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 9ea062ea2f8..fd1ecaf29dd 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -4,7 +4,9 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock @@ -30,6 +32,9 @@ else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): """ Reference: https://cloud.ibm.com/apidocs/watsonx-ai#text-generation @@ -305,12 +310,15 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): completion_tokens: Final = json_resp["results"][0]["generated_token_count"] model_response.choices[0].message.content = generated_text model_response.choices[0].finish_reason = map_finish_reason(json_resp["results"][0]["stop_reason"]) - if json_resp.get("created_at"): + created_at: Final = decoded_json_object(_DECODED_JSON.validate_python(json_resp)).get("created_at") + if created_at: + if not isinstance(created_at, str): + raise TypeError("fromisoformat: argument must be str") try: - created_datetime = datetime.fromisoformat(json_resp["created_at"]) + created_datetime = datetime.fromisoformat(created_at) except ValueError: # datetime.fromisoformat cannot handle 'Z' in Python 3.10 - created_datetime = datetime.fromisoformat(f"{json_resp['created_at'].rstrip('Z')}+00:00") + created_datetime = datetime.fromisoformat(f"{created_at.rstrip('Z')}+00:00") model_response.created = int(created_datetime.timestamp()) else: model_response.created = int(time.time()) diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index f53b71a5a9a..da8d673ce9f 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,15 +1,17 @@ import copy from collections.abc import AsyncIterator, Iterator, Mapping from types import MappingProxyType -from typing import Any, Final +from typing import Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter +from typing_extensions import Buffer import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, @@ -36,6 +38,28 @@ _RESPONSE_BODY: Final = TypeAdapter(dict[object, object], config=ConfigDict(hide _FILTERED_TOOL: Final = TypeAdapter(object) +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +def _token_count(value: object) -> int: + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)): + return int(value) + raise TypeError( + f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'" + ) + + +def _reasoning_tokens(details: object) -> int: + if not details: + return 0 + details_fields: Final = as_str_mapping(details) + if details_fields is not None and isinstance(details, dict): + return _token_count(details_fields.get("reasoning_tokens") or 0) + return _token_count(attribute_or_none(details, "reasoning_tokens") or 0) + + def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) if usage is None or reported_cost is None: @@ -324,7 +348,7 @@ class XAIChatConfig(OpenAIGPTConfig): @staticmethod def fold_reasoning_tokens_into_completion( - target: ModelResponse | Usage | dict[str, Any] | None, + target: ModelResponse | Usage | dict[str, object] | None, ) -> None: """Reconcile xAI Usage to the OpenAI invariant. @@ -341,24 +365,20 @@ class XAIChatConfig(OpenAIGPTConfig): return if isinstance(target, ModelResponse): - usage: Usage | dict[str, Any] | None = getattr(target, "usage", None) + usage: Usage | dict[str, object] | 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) + reasoning_tokens = _reasoning_tokens(usage.get("completion_tokens_details")) 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) + prompt_tokens = _token_count(usage.get("prompt_tokens") or 0) + completion_tokens = _token_count(usage.get("completion_tokens") or 0) + total_tokens = _token_count(usage.get("total_tokens") or 0) if total_tokens == prompt_tokens + completion_tokens: return @@ -370,8 +390,8 @@ class XAIChatConfig(OpenAIGPTConfig): 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 + details = attribute_or_none(usage, "completion_tokens_details") + reasoning_tokens = _token_count(attribute_or_none(details, "reasoning_tokens") or 0) if details else 0 if reasoning_tokens <= 0: return @@ -407,15 +427,15 @@ class XAIChatConfig(OpenAIGPTConfig): @staticmethod def _normalize_openai_compatible_usage_totals( - usage: Usage | dict[str, Any] | None, + usage: Usage | dict[str, object] | 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) + prompt_tokens = _token_count(usage.get("prompt_tokens") or 0) + completion_tokens = _token_count(usage.get("completion_tokens") or 0) expected_total = prompt_tokens + completion_tokens - if int(usage.get("total_tokens") or 0) < expected_total: + if _token_count(usage.get("total_tokens") or 0) < expected_total: usage["total_tokens"] = expected_total return prompt_tokens = int(usage.prompt_tokens or 0) diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 9228deda26b..b56b547a949 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -8,7 +8,9 @@ OpenAPI spec: https://you.com/specs/openapi_search_v1.yaml from typing import Final, TypedDict import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, subscript from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.search.transformation import ( @@ -18,6 +20,8 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_DECODED_JSON: Final = TypeAdapter(object) + class _YouComSearchRequestRequired(TypedDict): """Required fields for You.com Search API request.""" @@ -169,23 +173,26 @@ class YouComSearchConfig(BaseSearchConfig): - snippets[0] → SearchResult.snippet (falls back to `description`) - page_age → SearchResult.date """ - response_json: Final = raw_response.json() - raw_results: Final = response_json.get("results") or {} + response_json: Final = decoded_json_object(_DECODED_JSON.validate_python(raw_response.json())) + raw_results: Final = decoded_json_object(response_json.get("results") or {}) web_results: Final = raw_results.get("web") or [] news_results: Final = raw_results.get("news") or [] results: Final[list[SearchResult]] = [] - for item in list(web_results) + list(news_results): + for entry in list(iter_items(web_results)) + list(iter_items(news_results)): + item = decoded_json_object(entry) snippets = item.get("snippets") or [] - snippet = snippets[0] if snippets else item.get("description", "") + snippet = subscript(snippets, 0) if snippets else item.get("description", "") results.append( - SearchResult( - title=item.get("title", ""), - url=item.get("url", ""), - snippet=snippet, - date=item.get("page_age"), - last_updated=None, + SearchResult.model_validate( + { + "title": item.get("title", ""), + "url": item.get("url", ""), + "snippet": snippet, + "date": item.get("page_age"), + "last_updated": None, + } ) ) diff --git a/litellm/main.py b/litellm/main.py index df8e8f29483..c86a888aa96 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3657,8 +3657,10 @@ def _complete_openrouter(ctx: CompletionDispatchContext) -> _CompletionDispatchR or get_secret_str("OR_API_KEY") ) - openrouter_site_url: Final = get_secret("OR_SITE_URL") or "https://litellm.ai" - openrouter_app_name: Final = get_secret("OR_APP_NAME") or "liteLLM" + openrouter_site_url: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + get_secret("OR_SITE_URL") or "https://litellm.ai" + ) + openrouter_app_name: Final = _OPAQUE_VALUE_ADAPTER.validate_python(get_secret("OR_APP_NAME") or "liteLLM") openrouter_headers: Final = { "HTTP-Referer": openrouter_site_url, @@ -3762,8 +3764,10 @@ def _complete_vercel_ai_gateway( api_key = api_key or litellm.api_key or get_secret("VERCEL_AI_GATEWAY_API_KEY") - vercel_site_url: Final = get_secret("VERCEL_SITE_URL") or "https://litellm.ai" - vercel_app_name: Final = get_secret("VERCEL_APP_NAME") or "liteLLM" + vercel_site_url: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + get_secret("VERCEL_SITE_URL") or "https://litellm.ai" + ) + vercel_app_name: Final = _OPAQUE_VALUE_ADAPTER.validate_python(get_secret("VERCEL_APP_NAME") or "liteLLM") vercel_headers: Final = { "http-referer": vercel_site_url, @@ -7349,8 +7353,10 @@ def embedding( or get_secret_str("OR_API_KEY") ) - openrouter_site_url: Final = get_secret("OR_SITE_URL") or "https://litellm.ai" - openrouter_app_name: Final = get_secret("OR_APP_NAME") or "liteLLM" + openrouter_site_url: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + get_secret("OR_SITE_URL") or "https://litellm.ai" + ) + openrouter_app_name: Final = _OPAQUE_VALUE_ADAPTER.validate_python(get_secret("OR_APP_NAME") or "liteLLM") openrouter_headers: Final = { "HTTP-Referer": openrouter_site_url, diff --git a/litellm/models/base.py b/litellm/models/base.py index 3105fc9b91d..728fea1cff7 100644 --- a/litellm/models/base.py +++ b/litellm/models/base.py @@ -2,14 +2,34 @@ Base model class for domain models. """ +from collections.abc import Iterable, Mapping from datetime import datetime -from typing import Any +from typing import Final, Protocol from pydantic import ConfigDict +from typing_extensions import TypedDict +from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes from litellm.types.llms.base import LiteLLMBaseModel +class _RecordFields(TypedDict, total=False, extra_items=object): + pass + + +class _ModelDumpRecord(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _DictRecord(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +def _as_dict(record: object) -> Mapping[str, object] | None: + fields: Final = as_str_mapping(record) + return fields if isinstance(record, dict) else None + + class DomainModel(LiteLLMBaseModel): """Base class for all domain models.""" @@ -23,17 +43,18 @@ class DomainModel(LiteLLMBaseModel): updated_at: datetime | None = None @classmethod - def from_db_record(cls, record: Any) -> "DomainModel": + def from_db_record(cls, record: Mapping[str, object] | Iterable[tuple[str, object]] | None) -> "DomainModel": """Create a domain model from a database record.""" if record is None: raise ValueError("Cannot create domain model from None record") - if isinstance(record, dict): - return cls(**record) - if hasattr(record, "model_dump") and callable(record.model_dump): - return cls(**record.model_dump()) - if hasattr(record, "dict") and callable(record.dict): - return cls(**record.dict()) - return cls(**dict(record)) + fields: Final = _as_dict(record) + if fields is not None: + return cls(**_RecordFields(**fields)) + if has_attributes(record, _ModelDumpRecord, "model_dump") and callable(record.model_dump): + return cls(**_RecordFields(**record.model_dump())) + if has_attributes(record, _DictRecord, "dict") and callable(record.dict): + return cls(**_RecordFields(**record.dict())) + return cls(**_RecordFields(**dict(record))) def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]: """Convert domain model to a dictionary for database operations.""" diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index c5a63b0f990..e428017badd 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -13,7 +13,7 @@ the JSON, then run `npm run gen:api` in ui/litellm-dashboard and commit schema.d import json import re import sys -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -43,7 +43,7 @@ def _stabilize_multi_method_route_ids(routes) -> None: """FastAPI derives route IDs from a set of methods; make snapshots stable.""" for route in routes: - methods = sorted(getattr(route, "methods", None) or []) + methods: Sequence[str] = sorted(getattr(route, "methods", None) or []) if len(methods) <= 1 or not getattr(route, "path_format", None): continue diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index f0a4d47f1e4..bff270b8983 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,7 +23,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, keyword_arguments from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.llms.azure.common_utils import caused_by_entra_rejection @@ -145,19 +145,20 @@ def _to_jsonrpc_object(chunk: object) -> object: return chunk -def _build_message_send_params(params: dict[str, Any]) -> "MessageSendParams": +def _build_message_send_params(params: object) -> "MessageSendParams": """Build MessageSendParams from wire (0.3) or A2A 1.0 JSON-RPC params.""" from a2a.compat.v0_3.types import MessageSendParams + fields: Final = dict(**keyword_arguments(params, MessageSendParams)) try: - return MessageSendParams(**params) + return MessageSendParams.model_validate(fields) except ValidationError: from a2a.compat.v0_3.conversions import pb2_v10, to_compat_send_message_request from google.protobuf.json_format import ParseDict, ParseError pb: Final = pb2_v10.SendMessageRequest() try: - ParseDict(params, pb, ignore_unknown_fields=True) + ParseDict(fields, pb, ignore_unknown_fields=True) except ParseError as e: raise ValueError(f"Invalid message/send params: {e}") from e return to_compat_send_message_request(pb, "").params @@ -1027,7 +1028,7 @@ async def invoke_agent_a2a( from a2a.compat.v0_3.types import SendMessageRequest try: - message_send_params: Final = _build_message_send_params(params) + message_send_params: Final = _build_message_send_params(_DECODED_JSON.validate_python(params)) except (ValidationError, ValueError) as e: return _jsonrpc_error(request_id, -32602, f"Invalid params: {e}") diff --git a/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py b/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py index a69dd6ee025..0e9fe830f3f 100644 --- a/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py +++ b/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py @@ -68,7 +68,13 @@ async def update_deprecation_alert_state( ) -> DeprecationAlertState: _require_admin(user, write=True) state: Final = DeprecationAlertState( - **data.model_dump(), updated_at=datetime.now(timezone.utc), updated_by=user.user_id + model=data.model, + deprecation_date=data.deprecation_date, + status=data.status, + last_seen_date=data.last_seen_date, + last_seen_requests=data.last_seen_requests, + updated_at=datetime.now(timezone.utc), + updated_by=user.user_id, ) await repository.set_param(_state_key(data.model, data.deprecation_date), state.model_dump(mode="json")) return state diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4c4e506ce38..6b39511d89b 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -30,6 +30,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import json_contains, subscript from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy._types import ( @@ -164,6 +165,10 @@ class _UserInfoResponse(Protocol): def json(self) -> dict[str, object]: ... +class _JWKSResponse(Protocol): + def json(self) -> object: ... + + @dataclass(frozen=True, slots=True) class JWTIdentity: user_id: str | None @@ -218,6 +223,10 @@ def _userinfo_document(response: _UserInfoResponse) -> dict[str, object]: return response.json() +def _jwks_document(response: _JWKSResponse) -> object: + return response.json() + + def jwks_unavailable_exception(error: JWKSUnreachableError) -> ProxyException: return ProxyException( message=( @@ -960,12 +969,12 @@ class JWTHandler: ) try: - response_json: Final = response.json() + response_json: Final = _jwks_document(response) except Exception as e: verbose_proxy_logger.error("Error parsing response: %s. Original Response: %s", e, response.text) raise Exception(f"Error parsing response: {e}. Check server logs for original response.") - keys: Final = response_json["keys"] if "keys" in response_json else response_json + keys: Final = subscript(response_json, "keys") if json_contains(response_json, "keys") else response_json return cast(JWKKeyValue, keys) # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> JWTKeyItem: diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index e1d6b2e090c..e828afe999c 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -3,14 +3,17 @@ import base64 import json import os +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import NON_LLM_CONNECTION_TIMEOUT +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import HTTPHandler if TYPE_CHECKING: @@ -23,6 +26,23 @@ AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' fe _LICENSE_VERDICT: Final = TypeAdapter(object) +class _LicenseFields(TypedDict, total=False, extra_items=object): + pass + + +def _license_fields(license_data: object) -> Mapping[str, object]: + fields: Final = as_str_mapping(license_data) + if fields is None: + raise TypeError(f"'{type(license_data).__name__}' object is not subscriptable") + return fields + + +def _expiration_text(expiration_date: object) -> str: + if isinstance(expiration_date, str): + return expiration_date + raise TypeError(f"strptime() argument 1 must be str, not {type(expiration_date).__name__}") + + class LicenseCheck: """ - Check if license in env @@ -201,15 +221,15 @@ class LicenseCheck: ) # Decode and parse the data - license_data: Final = json.loads(message.decode()) + license_data: Final = _license_fields(_LICENSE_VERDICT.validate_python(json.loads(message.decode()))) # Check expiration date - expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") + expiration_date: Final = datetime.strptime(_expiration_text(license_data["expiration_date"]), "%Y-%m-%d") if expiration_date < datetime.now(): self.airgapped_license_data = None return False, "License has expired" - self.airgapped_license_data = EnterpriseLicenseData(**license_data) + self.airgapped_license_data = EnterpriseLicenseData(**_LicenseFields(**license_data)) return True diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index 558c9139f5b..1b6db2b0d5c 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -2,7 +2,7 @@ import builtins import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Final, Literal, Protocol +from typing import Final, Literal, Protocol import click import requests @@ -33,6 +33,22 @@ class _JsonBodyView(TypedDict): body: ReadOnly[object] +class _KeyImportFields(TypedDict, total=False, extra_items=object): + pass + + +_KEY_IMPORT_FIELD_NAMES: Final = ( + "models", + "aliases", + "spend", + "key_alias", + "team_id", + "user_id", + "budget_id", + "config", +) + + def _cli_context(view: _CliContextView) -> _CliContext: return view.obj @@ -302,25 +318,9 @@ def _display_dry_run_table(source_keys: Sequence[Mapping[str, object]]) -> None: rich.print(table) -def _prepare_key_import_data(key: Mapping[str, object]) -> dict[str, Any]: +def _prepare_key_import_data(key: Mapping[str, object]) -> _KeyImportFields: """Prepare key data for import by extracting relevant fields.""" - import_data: Final = {} - - # Copy relevant fields if they exist - for field in [ - "models", - "aliases", - "spend", - "key_alias", - "team_id", - "user_id", - "budget_id", - "config", - ]: - if key.get(field): - import_data[field] = key[field] - - return import_data + return _KeyImportFields(**{field: key[field] for field in _KEY_IMPORT_FIELD_NAMES if key.get(field)}) def _import_keys_to_destination( diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index 538b42a8c83..ae3751cff69 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -10,12 +10,17 @@ from typing import TYPE_CHECKING, Any, Final, Literal import click import rich import yaml +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import json_contains, subscript + # local imports from ... import Client from ._cli_context import cli_context_values +_DECODED_YAML: Final = TypeAdapter(object) + if TYPE_CHECKING: from rich.console import JustifyMethod @@ -385,10 +390,10 @@ def _print_summary_table(provider_counts): def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, object]]: """Load and validate the model list from a YAML file.""" with open(yaml_file, "r") as f: - data: Final = yaml.safe_load(f) - if not data or "model_list" not in data: + data: Final = _DECODED_YAML.validate_python(yaml.safe_load(f)) + if not data or not json_contains(data, "model_list"): raise click.ClickException("YAML file must contain a 'model_list' key with a list of models.") - model_list: Final = data["model_list"] + model_list: Final = subscript(data, "model_list") if not isinstance(model_list, list): raise click.ClickException("'model_list' must be a list of model definitions.") return model_list diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1b004423181..e585cd7b1ff 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -73,6 +73,7 @@ from litellm.litellm_core_utils.core_helpers import ( is_expected_client_error, proxy_request_headers, redact_nested_match_and_regex_keys, + require_attributes, stream_cost_injection_requested, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer @@ -175,7 +176,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError from litellm.types.router_weights import RouterWeights, validate_router_weights -_LateResponseT = TypeVar("_LateResponseT", bound=Response) +_LateResponseT = TypeVar("_LateResponseT") _LlmCallT = TypeVar("_LlmCallT") _ResponseT: Final = TypeVar("_ResponseT") _ChunkT: Final = TypeVar("_ChunkT") @@ -1365,7 +1366,17 @@ class LateSseFraming: """How a response the keepalive already opened ends once its upstream fails late.""" failure: Callable[[Exception], tuple[str, ...]] - body: Callable[[Response], tuple[bytes, ...]] + body: Callable[[object], tuple[bytes, ...]] + + +class _RenderedBody(Protocol): + @property + def body(self) -> bytes | memoryview: ... + + +class _CarriesStatusCode(Protocol): + @property + def status_code(self) -> int: ... def _openai_late_failure_frames(exc: Exception) -> tuple[str, ...]: @@ -1373,8 +1384,11 @@ def _openai_late_failure_frames(exc: Exception) -> tuple[str, ...]: return _sse_error_frames(error_obj) -def _openai_late_body_frames(produced: Response) -> tuple[bytes, ...]: - return b"data: " + (bytes(produced.body) or b"{}") + b"\n\n", b"data: [DONE]\n\n" +def _openai_late_body_frames(produced: object) -> tuple[bytes, ...]: + return ( + b"data: " + (bytes(require_attributes(produced, _RenderedBody, "body").body) or b"{}") + b"\n\n", + b"data: [DONE]\n\n", + ) def _anthropic_late_failure_frames(exc: Exception) -> tuple[str, ...]: @@ -1393,10 +1407,14 @@ def _responses_late_failure_frames(exc: Exception) -> tuple[str, ...]: return _responses_error_frames(sse_error_payload(exc)[1]) -def _responses_late_body_frames(produced: Response) -> tuple[bytes, ...]: - body: Final = bytes(produced.body) +def _responses_late_body_frames(produced: object) -> tuple[bytes, ...]: + body: Final = bytes(require_attributes(produced, _RenderedBody, "body").body) error_obj: Final = sse_error_object(body or b"{}") or {"message": body.decode()} - with_status: Final = {**error_obj, "code": error_obj.get("code") or str(produced.status_code)} + with_status: Final = { + **error_obj, + "code": error_obj.get("code") + or str(require_attributes(produced, _CarriesStatusCode, "status_code").status_code), + } return tuple(frame.encode() for frame in _responses_error_frames(with_status)) @@ -1620,7 +1638,7 @@ def ttft_keepalive_interval(request_data: Mapping[str, object], llm_router: Rout return resolve_ttft_keepalive_interval(deployments, litellm.sse_keepalive_ping_interval_seconds) -async def _aclose_late_response(produced: Response) -> None: +async def _aclose_late_response(produced: object) -> None: """Run the cleanup Starlette would have run, for a response it never called. Closing an already-closed async generator is a no-op, so this is safe to call @@ -1643,7 +1661,7 @@ async def _aclose_late_response(produced: Response) -> None: verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) -async def _relay_late_response(produced: Response, framing: LateSseFraming) -> AsyncGenerator[bytes, None]: +async def _relay_late_response(produced: object, framing: LateSseFraming) -> AsyncGenerator[bytes, None]: """Replay a Response that was built after a keepalive had already opened the wire.""" if not isinstance(produced, StreamingResponse): # The status line is already on the wire, so a non-streaming body, an error @@ -3303,7 +3321,7 @@ class ProxyBaseLLMRequestProcessing: contents: list[object] | None = None, # Add contents parameter skip_pre_call_logic: bool = False, first_chunk_framing: FirstChunkSseFraming = OPENAI_FIRST_CHUNK_FRAMING, - ) -> Any: + ) -> object: """ Common request processing logic for both chat completions and responses API endpoints """ @@ -3571,20 +3589,21 @@ class ProxyBaseLLMRequestProcessing: ) # Non-streaming response - fall through to normal response handling elif select_data_generator: - selected_data_generator = select_data_generator( + produced_generator: Final[AsyncGenerator[str, None]] = select_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=self.data, request=request, ) - if route_type in ("aresponses", "aget_responses"): - selected_data_generator = ( - ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( - original_stream_response=response, - wrapped_generator=selected_data_generator, - user_api_key_dict=user_api_key_dict, - ) + selected_data_generator = ( + ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=response, + wrapped_generator=produced_generator, + user_api_key_dict=user_api_key_dict, ) + if route_type in ("aresponses", "aget_responses") + else produced_generator + ) if selected_data_generator is not None: return await create_response( generator=selected_data_generator, @@ -3826,9 +3845,9 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def _wrap_responses_stream_for_container_ownership( original_stream_response: object, - wrapped_generator: Any, + wrapped_generator: AsyncIterator[_ChunkT], user_api_key_dict: UserAPIKeyAuth, - ): + ) -> AsyncGenerator[_ChunkT, None]: """Forward SSE chunks and record container ownership before the terminal chunk goes out. Streaming ``/v1/responses`` short-circuits out of diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 0aa8d2291ef..82120023b21 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -93,6 +93,10 @@ class _CallbackNotDispatchable: _CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable +class _PresidioOptions(TypedDict, total=False, extra_items=object): + pass + + def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError: """ Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched. @@ -193,15 +197,15 @@ def initialize_callbacks_on_proxy( if presidio_logging_only is not None: presidio_logging_only = bool(presidio_logging_only) # validate boolean given - _presidio_params = {} + _presidio_params: Mapping[str, object] = {} if "presidio" in callback_specific_params and isinstance(callback_specific_params["presidio"], dict): _presidio_params = callback_specific_params["presidio"] - params: dict[str, Any] = { + params: Mapping[str, object] = { "logging_only": presidio_logging_only, **_presidio_params, } - pii_masking_object = OPTIONAL_PresidioPIIMasking(**params) + pii_masking_object = OPTIONAL_PresidioPIIMasking(**_PresidioOptions(**params)) imported_list.append(pii_masking_object) elif isinstance(callback, str) and callback == "llamaguard_moderations": try: diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 644a35ef4ef..5b7ce5c8f3e 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -3,7 +3,7 @@ import ctypes import os import re from collections.abc import Mapping -from typing import Final, Literal, cast +from typing import Final, Literal, TypeVar, cast, overload from pydantic import TypeAdapter, ValidationError @@ -201,7 +201,18 @@ def decrypt_bearer_token(token: str, prefix: str) -> str | None: return None -def encrypt_value_helper(value: str, new_encryption_key: str | None = None): +_Passthrough: Final = TypeVar("_Passthrough") + + +@overload +def encrypt_value_helper(value: str, new_encryption_key: str | None = None) -> str: ... + + +@overload +def encrypt_value_helper(value: _Passthrough, new_encryption_key: str | None = None) -> _Passthrough: ... + + +def encrypt_value_helper(value: object, new_encryption_key: str | None = None) -> object: signing_key: Final = new_encryption_key or get_salt_key() try: diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 9e3071f020e..ae91cbb1f2a 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1200,7 +1200,7 @@ class ResetBudgetJob: if updated_keys: await self._write_key_reset_updates(updated_keys=updated_keys) for k in updated_keys: - token = getattr(k.row, "token", None) + token = k.row.token if token: await self._invalidate_spend_counter(key_counter(str(token))) @@ -1317,7 +1317,7 @@ class ResetBudgetJob: if updated_users: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: - user_id = getattr(u.row, "user_id", None) + user_id = u.row.user_id if user_id: await self._invalidate_spend_counter(user_counter(str(user_id))) if user_id == LITELLM_PROXY_BUDGET_NAME: @@ -1438,7 +1438,7 @@ class ResetBudgetJob: if updated_teams: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: - team_id = getattr(t.row, "team_id", None) + team_id = t.row.team_id if team_id: await self._invalidate_spend_counter(team_counter(str(team_id))) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index d2b69f24397..ccabd9dcde6 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -18,6 +18,7 @@ from collections.abc import Mapping, Sequence from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter @@ -149,7 +150,13 @@ class _UpsertBatch(Protocol): class _UpsertBatchManager(Protocol): async def __aenter__(self) -> _UpsertBatch: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... class _UpsertTransaction(Protocol): @@ -161,7 +168,13 @@ class _UpsertTransaction(Protocol): class _UpsertTransactionManager(Protocol): async def __aenter__(self) -> _UpsertTransaction: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... def _seed_totals_query(prisma_client: "PrismaClient") -> _SeedTotalsQuery: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 12267120965..5a9a52447b7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,7 +18,7 @@ from contextlib import asynccontextmanager from contextvars import ContextVar from datetime import datetime, timedelta, timezone from functools import reduce -from types import MappingProxyType +from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload from urllib.parse import quote, unquote @@ -249,21 +249,43 @@ class SpendUpsertTable(Protocol): class _SpendBatch(Protocol): - litellm_usertable: BatchTable - litellm_endusertable: SpendUpsertTable - litellm_dailymodelusage: SpendUpsertTable + @property + def litellm_usertable(self) -> BatchTable: ... + + @property + def litellm_endusertable(self) -> SpendUpsertTable: ... + + @property + def litellm_dailymodelusage(self) -> SpendUpsertTable: ... def execute_raw(self, query: LiteralString, *args: object) -> None: ... - litellm_verificationtoken: BatchTable - litellm_teamtable: BatchTable - litellm_teammembership: BatchTable - litellm_organizationtable: BatchTable - litellm_organizationmembership: BatchTable - litellm_projecttable: BatchTable - litellm_tagtable: BatchTable - litellm_agentstable: BatchTable - litellm_modelaccessgroupbudgettable: BatchTable + @property + def litellm_verificationtoken(self) -> BatchTable: ... + + @property + def litellm_teamtable(self) -> BatchTable: ... + + @property + def litellm_teammembership(self) -> BatchTable: ... + + @property + def litellm_organizationtable(self) -> BatchTable: ... + + @property + def litellm_organizationmembership(self) -> BatchTable: ... + + @property + def litellm_projecttable(self) -> BatchTable: ... + + @property + def litellm_tagtable(self) -> BatchTable: ... + + @property + def litellm_agentstable(self) -> BatchTable: ... + + @property + def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... _EntitySpendTable: TypeAlias = Literal[ @@ -288,7 +310,13 @@ def _entity_spend_table(batcher: _SpendBatch, table_accessor: _EntitySpendTable) class SpendBatchManager(Protocol): async def __aenter__(self) -> _SpendBatch: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... class _SpendTransaction(Protocol): @@ -300,7 +328,13 @@ class _SpendTransaction(Protocol): class _SpendTransactionManager(Protocol): async def __aenter__(self) -> _SpendTransaction: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... _DailySpendTransactionT = TypeVar("_DailySpendTransactionT", bound=BaseDailySpendTransaction) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index c2d58512fea..33d973cccfb 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,11 +1,12 @@ import re -from collections.abc import Awaitable, Callable, Iterator +from collections.abc import Awaitable, Callable, Iterator, Mapping from http import HTTPStatus from typing import Final, Protocol, TypeVar from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, @@ -93,6 +94,28 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException)) +class _CodedError(Protocol): + @property + def code(self) -> object: ... + + +def _error_code(error: _CodedError) -> object: + return error.code + + +class _FrameGlobals(Protocol): + @property + def f_globals(self) -> Mapping[str, object]: ... + + +class _PrefixMatchable(Protocol): + def startswith(self, prefix: str, /) -> bool: ... + + +def _frame_module_name(frame: _FrameGlobals) -> _PrefixMatchable: + return require_attributes(frame.f_globals.get("__name__", ""), _PrefixMatchable, "startswith") + + def _prisma_error_code(e: Exception) -> str | None: import prisma.engine.errors @@ -317,7 +340,7 @@ class PrismaDBExceptionHandler: if not isinstance(e, _exception_types(prisma.errors.PrismaError)): return False - if isinstance(e, prisma.errors.DataError) and e.code == _PRISMA_POOL_TIMEOUT_CODE: + if isinstance(e, prisma.errors.DataError) and _error_code(e) == _PRISMA_POOL_TIMEOUT_CODE: return True if PrismaDBExceptionHandler.postgres_sqlstate(e) == "53300": return True @@ -433,7 +456,7 @@ class PrismaDBExceptionHandler: return False tb = e.__traceback__ if hasattr(e, "__traceback__") else None while tb is not None: - if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): + if _frame_module_name(tb.tb_frame).startswith("prisma.engine"): return True tb = tb.tb_next return False diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index f8cdabe9296..1a5f15fe294 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from itertools import groupby +from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter, ValidationError @@ -42,7 +43,13 @@ class ModelUsageBatch(Protocol): class ModelUsageBatchManager(Protocol): async def __aenter__(self) -> ModelUsageBatch: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index a7f18769461..e0be2e1e5ef 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -35,12 +35,170 @@ from litellm.proxy.db.token_auth import ( from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: + from collections.abc import Sequence + from types import TracebackType + from prisma import Prisma from prisma import models as prisma_models from prisma.types import HttpConfig + from typing_extensions import Self from litellm.repositories.prisma_protocols import TableActions + _DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) + _DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) + + class PrismaBatchActions(Protocol): + def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> None: ... + + def create_many(self, data: Sequence[Mapping[str, object]], *, skip_duplicates: bool | None = None) -> None: ... + + def delete(self, where: Mapping[str, object], include: Mapping[str, object] | None = None) -> None: ... + + def update( + self, + data: Mapping[str, object], + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> None: ... + + def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> None: ... + + def update_many(self, data: Mapping[str, object], where: Mapping[str, object]) -> None: ... + + def delete_many(self, where: Mapping[str, object] | None = None) -> None: ... + + class PrismaBatchSession(Protocol): + @property + def litellm_agentstable(self) -> PrismaBatchActions: ... + + @property + def litellm_budgettable(self) -> PrismaBatchActions: ... + + @property + def litellm_dailymodelusage(self) -> PrismaBatchActions: ... + + @property + def litellm_dailyteamspend(self) -> PrismaBatchActions: ... + + @property + def litellm_dailytoolspend(self) -> PrismaBatchActions: ... + + @property + def litellm_endusertable(self) -> PrismaBatchActions: ... + + @property + def litellm_modelaccessgroupbudgettable(self) -> PrismaBatchActions: ... + + @property + def litellm_organizationmembership(self) -> PrismaBatchActions: ... + + @property + def litellm_organizationtable(self) -> PrismaBatchActions: ... + + @property + def litellm_projecttable(self) -> PrismaBatchActions: ... + + @property + def litellm_proxymodeltable(self) -> PrismaBatchActions: ... + + @property + def litellm_tagtable(self) -> PrismaBatchActions: ... + + @property + def litellm_teammembership(self) -> PrismaBatchActions: ... + + @property + def litellm_teamtable(self) -> PrismaBatchActions: ... + + @property + def litellm_usertable(self) -> PrismaBatchActions: ... + + @property + def litellm_verificationtoken(self) -> PrismaBatchActions: ... + + def execute_raw(self, query: str, *args: object) -> None: ... + + async def commit(self) -> None: ... + + async def __aenter__(self) -> Self: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + exc_tb: TracebackType | None, + /, + ) -> None: ... + + class PrismaTransaction(Protocol): + @property + def litellm_accessgrouptable(self) -> TableActions[prisma_models.LiteLLM_AccessGroupTable]: ... + + @property + def litellm_agentstable(self) -> TableActions[prisma_models.LiteLLM_AgentsTable]: ... + + @property + def litellm_budgettable(self) -> TableActions[prisma_models.LiteLLM_BudgetTable]: ... + + @property + def litellm_invitationlink(self) -> TableActions[prisma_models.LiteLLM_InvitationLink]: ... + + @property + def litellm_mcpservertable(self) -> TableActions[prisma_models.LiteLLM_MCPServerTable]: ... + + @property + def litellm_mcpuserenvvars(self) -> TableActions[prisma_models.LiteLLM_MCPUserEnvVars]: ... + + @property + def litellm_objectpermissiontable(self) -> TableActions[prisma_models.LiteLLM_ObjectPermissionTable]: ... + + @property + def litellm_organizationtable(self) -> TableActions[prisma_models.LiteLLM_OrganizationTable]: ... + + @property + def litellm_proxymodeltable(self) -> TableActions[prisma_models.LiteLLM_ProxyModelTable]: ... + + @property + def litellm_teamtable(self) -> TableActions[prisma_models.LiteLLM_TeamTable]: ... + + @property + def litellm_usertable(self) -> TableActions[prisma_models.LiteLLM_UserTable]: ... + + @property + def litellm_verificationtoken(self) -> TableActions[prisma_models.LiteLLM_VerificationToken]: ... + + @property + def litellm_workflowevent(self) -> TableActions[prisma_models.LiteLLM_WorkflowEvent]: ... + + @property + def litellm_workflowrun(self) -> TableActions[prisma_models.LiteLLM_WorkflowRun]: ... + + def batch_(self) -> PrismaBatchSession: ... + + async def execute_raw(self, query: str, *args: object) -> int: ... + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + + class PrismaTransactionManager(Protocol): + async def __aenter__(self) -> PrismaTransaction: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + exc_tb: TracebackType | None, + /, + ) -> None: ... + + _ENVIRONMENT_VARIABLES_REVISION_MIGRATION: Final = "20261015120000_add_environment_variables_revision_trigger" __all__ = ( @@ -925,6 +1083,12 @@ class PrismaWrapper: async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + def tx( + self, *, max_wait: int | timedelta = _DEFAULT_TX_MAX_WAIT, timeout: int | timedelta = _DEFAULT_TX_TIMEOUT + ) -> PrismaTransactionManager: ... + + def batch_(self) -> PrismaBatchSession: ... + @property def litellm_deprecatedverificationtoken( self, diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index c71e2206f45..e6a436a9b1e 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -19,8 +19,12 @@ if TYPE_CHECKING: from prisma import models as prisma_models from prisma.types import HttpConfig + from litellm.proxy.db.prisma_client import PrismaBatchSession, PrismaTransactionManager from litellm.repositories.prisma_protocols import TableActions + _DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) + _DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -311,6 +315,12 @@ class RoutingPrismaWrapper: async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + def tx( + self, *, max_wait: int | timedelta = _DEFAULT_TX_MAX_WAIT, timeout: int | timedelta = _DEFAULT_TX_TIMEOUT + ) -> PrismaTransactionManager: ... + + def batch_(self) -> PrismaBatchSession: ... + @property def litellm_deprecatedverificationtoken( self, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 92181d546eb..e2cebc691a5 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -19,6 +19,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from itertools import chain, groupby +from types import TracebackType from typing import TYPE_CHECKING, Final, Literal, Protocol from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -175,7 +176,8 @@ class _ToolRollupTable(Protocol): class _ToolRollupBatch(Protocol): - litellm_dailytoolspend: _ToolRollupTable + @property + def litellm_dailytoolspend(self) -> _ToolRollupTable: ... def execute_raw(self, query: LiteralString, *args: object) -> None: ... @@ -183,7 +185,13 @@ class _ToolRollupBatch(Protocol): class _ToolRollupBatchManager(Protocol): async def __aenter__(self) -> _ToolRollupBatch: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... def _tool_rollup_batch(prisma_client: PrismaClient) -> _ToolRollupBatchManager: diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py index e45ca7e10f2..1665cb4caaf 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_archive.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -17,7 +17,7 @@ import zlib from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import ClassVar, Final +from typing import ClassVar, Final, Protocol import yaml from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -64,6 +64,15 @@ _YAML_PAIRS: Final = TypeAdapter( ) +class _YamlValueNode(Protocol): + @property + def value(self) -> object: ... + + +def _yaml_node_value(node: _YamlValueNode) -> object: + return node.value + + class _AgentSkillFrontmatterLoader(yaml.BaseLoader): yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state None: list(_AGENT_SKILL_YAML_SCALARS.items()) @@ -259,7 +268,7 @@ def agent_skill_frontmatter(manifest: bytes) -> Mapping[str, str]: return MappingProxyType({}) if not isinstance(node, yaml.MappingNode) or not _valid_frontmatter_node(node): return MappingProxyType({}) - pairs: Final = _YAML_PAIRS.validate_python(node.value) + pairs: Final = _YAML_PAIRS.validate_python(_yaml_node_value(node)) fields: Final = { key_text: value_text for key, value in pairs @@ -273,7 +282,7 @@ def agent_skill_frontmatter(manifest: bytes) -> Mapping[str, str]: def _frontmatter_string(node: yaml.Node) -> str | None: if not isinstance(node, yaml.ScalarNode) or node.tag != "tag:yaml.org,2002:str": return None - value: Final = _DECODED_YAML.validate_python(node.value) + value: Final = _yaml_node_value(node) return value if isinstance(value, str) else None @@ -282,18 +291,18 @@ def _valid_frontmatter_node(node: yaml.Node, depth: int = 0) -> bool: return False if isinstance(node, yaml.ScalarNode): pattern: Final = _AGENT_SKILL_YAML_SCALARS.get(node.tag) - scalar: Final = _DECODED_YAML.validate_python(node.value) + scalar: Final = _yaml_node_value(node) return isinstance(scalar, str) and ( node.tag == "tag:yaml.org,2002:str" or (pattern is not None and pattern.fullmatch(scalar) is not None) ) if isinstance(node, yaml.SequenceNode): - values: Final = _YAML_NODES.validate_python(node.value) + values: Final = _YAML_NODES.validate_python(_yaml_node_value(node)) return node.tag == "tag:yaml.org,2002:seq" and all( _valid_frontmatter_node(value, depth + 1) for value in values ) if not isinstance(node, yaml.MappingNode) or node.tag != "tag:yaml.org,2002:map": return False - pairs: Final = _YAML_PAIRS.validate_python(node.value) + pairs: Final = _YAML_PAIRS.validate_python(_yaml_node_value(node)) keys: Final = tuple(_frontmatter_string(key) for key, _ in pairs) return ( None not in keys diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 030f8d7a46a..850a8508428 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -6,12 +6,14 @@ ########################################################################## import asyncio -from typing import Final, cast +from typing import Final from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -25,6 +27,7 @@ from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMFineTuningJob router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) from litellm.types.llms.openai import LiteLLMFineTuningJobCreate @@ -50,7 +53,7 @@ def set_fine_tuning_config(config): # Function to search for specific custom_llm_provider and return its configuration def get_fine_tuning_provider_config( - custom_llm_provider: str, + custom_llm_provider: object, ): global fine_tuning_config if fine_tuning_config is None: @@ -61,6 +64,13 @@ def get_fine_tuning_provider_config( return None +async def _decoded_request_body(request: Request) -> object: + try: + return _DECODED_JSON.validate_python(await request.json()) + except Exception: + return {} + + @router.post( "/v1/fine_tuning/jobs", dependencies=[Depends(user_api_key_auth)], @@ -163,7 +173,7 @@ async def create_fine_tuning_job( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = cast(LiteLLMFineTuningJob, await llm_router.acreate_fine_tuning_job(**data)) + response = await llm_router.acreate_fine_tuning_job(**data) response.training_file = training_file response.validation_file = fine_tuning_request.validation_file set_hidden_param(response, "unified_file_id", unified_file_id) @@ -285,12 +295,8 @@ async def retrieve_fine_tuning_job( route_type=CallTypes.aretrieve_fine_tuning_job.value, ) - try: - request_body = await request.json() - except Exception: - request_body = {} - - custom_llm_provider = request_body.get("custom_llm_provider", None) or custom_llm_provider + request_body: Final = decoded_json_object(await _decoded_request_body(request)) + requested_provider: Final = request_body.get("custom_llm_provider", None) or custom_llm_provider ## CHECK IF MANAGED FILE ID unified_finetuning_job_id: str | Literal[False] = False @@ -305,16 +311,13 @@ async def retrieve_fine_tuning_job( status_code=500, detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = cast( - LiteLLMFineTuningJob, - await llm_router.aretrieve_fine_tuning_job( - **data, - ), + response = await llm_router.aretrieve_fine_tuning_job( + **data, ) set_hidden_param(response, "unified_finetuning_job_id", unified_finetuning_job_id) - elif custom_llm_provider: + elif requested_provider: # get configs for custom_llm_provider - llm_provider_config: Final = get_fine_tuning_provider_config(custom_llm_provider=custom_llm_provider) + llm_provider_config: Final = get_fine_tuning_provider_config(custom_llm_provider=requested_provider) if llm_provider_config is not None: data.update(llm_provider_config) @@ -560,11 +563,7 @@ async def cancel_fine_tuning_job( route_type=CallTypes.acancel_fine_tuning_job.value, ) - try: - request_body = await request.json() - except Exception: - request_body = {} - + request_body: Final = decoded_json_object(await _decoded_request_body(request)) custom_llm_provider: Final = request_body.get("custom_llm_provider", None) ## CHECK IF MANAGED FILE ID @@ -580,11 +579,8 @@ async def cancel_fine_tuning_job( status_code=500, detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = cast( - LiteLLMFineTuningJob, - await llm_router.acancel_fine_tuning_job( - **data, - ), + response = await llm_router.acancel_fine_tuning_job( + **data, ) set_hidden_param(response, "unified_finetuning_job_id", unified_finetuning_job_id) else: diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index c773e88c769..03f83616fea 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -11,7 +11,7 @@ import os import ssl from collections.abc import AsyncGenerator, AsyncIterable, Iterator, Mapping, Sequence from ssl import SSLContext -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias from fastapi import HTTPException from pydantic import BaseModel @@ -23,7 +23,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import attribute_items, iter_items +from litellm.litellm_core_utils.core_helpers import attribute_items, iter_items, require_attributes, subscript from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_calls, @@ -60,6 +60,19 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_ToolCallKind: TypeAlias = Literal["server_tool_call", "tool_call"] +_OutputTarget: TypeAlias = tuple[Literal["content"], None] | tuple[_ToolCallKind, int] + + +class _ToolCallFunction(Protocol): + arguments: str + + +class _FunctionToolCall(Protocol): + @property + def function(self) -> _ToolCallFunction: ... + + class CatoNetworksGuardrailMissingSecrets(Exception): pass @@ -522,14 +535,16 @@ class CatoNetworksGuardrail(CustomGuardrail): ) @staticmethod - def _output_fragments(message: Message) -> Sequence[tuple[tuple[str, int | None], str]]: + def _output_fragments(message: Message) -> Sequence[tuple[_OutputTarget, str]]: """Assistant text the proxy returns to the caller: ``content`` plus every tool call's ``function.arguments`` string, Claude's server tool calls first, each tagged with where a redaction must be written back. ``content`` is only included when present so a tool-call-only choice keeps its ``None`` content (the text-vs-tool-call signal downstream consumers rely on) while its arguments are still inspected.""" - content: Final = () if message.content is None else ((("content", None), message.content),) + content: Final[tuple[tuple[_OutputTarget, str], ...]] = ( + () if message.content is None else ((("content", None), message.content),) + ) return ( *content, *CatoNetworksGuardrail._tool_call_fragments("server_tool_call", anthropic_server_tool_calls(message)), @@ -537,7 +552,9 @@ class CatoNetworksGuardrail(CustomGuardrail): ) @staticmethod - def _tool_call_fragments(kind: str, tool_calls: Sequence[object]) -> tuple[tuple[tuple[str, int | None], str], ...]: + def _tool_call_fragments( + kind: _ToolCallKind, tool_calls: Sequence[object] + ) -> tuple[tuple[tuple[_ToolCallKind, int], str], ...]: return tuple( ((kind, idx), arguments) for idx, arguments in enumerate( @@ -547,14 +564,16 @@ class CatoNetworksGuardrail(CustomGuardrail): ) @staticmethod - def _apply_output_fragment(message: Any, target: tuple[str, int | None], redacted: str) -> None: + def _apply_output_fragment(message: Message, target: _OutputTarget, redacted: str) -> None: match target: case ("content", _): message.content = redacted case ("server_tool_call", int() as idx): write_anthropic_server_tool_call_arguments(message, idx, redacted) case (_, idx): - message.tool_calls[idx].function.arguments = redacted + require_attributes( + subscript(message.tool_calls, idx), _FunctionToolCall, "function" + ).function.arguments = redacted @staticmethod def _responses_output_field(item: object, key: str) -> str | Sequence[object] | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 55f7abd40a8..51b849751c8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -13,7 +13,7 @@ from typing import Final, Literal from urllib.parse import urlparse import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -22,6 +22,8 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, vali from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider +_DECODED_JSON: Final = TypeAdapter(object) + # ============================================================================= # Result Types - Used by Starlark code to return guardrail decisions # ============================================================================= @@ -200,7 +202,7 @@ class JsonSchemaNode(TypedDict, total=False): properties: ReadOnly[Mapping[str, "JsonSchemaNode"]] -def json_parse(text: str) -> JsonValue: +def json_parse(text: str) -> object: """ Parse a JSON string into a Python object. @@ -211,7 +213,7 @@ def json_parse(text: str) -> JsonValue: Parsed Python object, or None if parsing fails """ try: - return json.loads(text) + return _DECODED_JSON.validate_python(json.loads(text)) except (json.JSONDecodeError, TypeError) as e: verbose_proxy_logger.debug("Starlark json_parse error: %s", e) return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 1b9bfc725db..8c698934796 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -8,7 +8,7 @@ import os from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional import httpx @@ -186,7 +186,6 @@ class EnkryptAIGuardrails(CustomGuardrail): details: Final = response.get("details", {}) detected_attacks: Final[list[str]] = [] - attack_details: Final[dict[str, Any]] = {} for key, value in summary.items(): # Check if attack is detected @@ -195,13 +194,14 @@ class EnkryptAIGuardrails(CustomGuardrail): if key == "toxicity": if isinstance(value, list) and len(value) > 0: detected_attacks.append(key) - attack_details[key] = details.get(key, {}) else: if value == 1: detected_attacks.append(key) - attack_details[key] = details.get(key, {}) - return {"attacks_detected": detected_attacks, "attack_details": attack_details} + return { + "attacks_detected": detected_attacks, + "attack_details": {key: details.get(key, {}) for key in detected_attacks}, + } def _determine_guardrail_status(self, response_json: EnkryptAIResponse) -> GuardrailStatus: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..85136c46a67 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -21,7 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, @@ -46,6 +46,7 @@ _CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues) _SENT_ROW: Final = TypeAdapter(Mapping[str, object]) _SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]]) _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DECODED_JSON: Final = TypeAdapter(object) # Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*). _HEADER_VALUE_ALLOWLIST: Final = frozenset( @@ -242,6 +243,16 @@ def _structured_rows_to_write_back( ) +def _guardrail_response_from_json(response_json: object) -> GenericGuardrailAPIResponse: + if isinstance(response_json, dict): + return GenericGuardrailAPIResponse.from_dict(response_json) + raise AttributeError( + f"{type(response_json).__name__!r} object has no attribute 'get'", + name="get", + obj=response_json, + ) + + class GenericGuardrailAPI(CustomGuardrail): """ Generic Guardrail API integration for LiteLLM. @@ -396,7 +407,7 @@ class GenericGuardrailAPI(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"], error: Exception, - http_status_code: int | None = None, + http_status_code: object = None, ) -> GenericGuardrailAPIInputs: status_suffix: Final = f" http_status_code={http_status_code}" if http_status_code else "" error_type: Final = "HTTPStatusError" if isinstance(error, httpx.HTTPStatusError) else type(error).__name__ @@ -482,7 +493,7 @@ class GenericGuardrailAPI(CustomGuardrail): ) -> GenericGuardrailAPIInputs: unreachable_fail_open: Final = is_unreachable and self.unreachable_fallback == "fail_open" if unreachable_fail_open or not self.fail_on_error: - http_status_code: Final = getattr(getattr(error, "response", None), "status_code", None) + http_status_code: Final = attribute_or_none(attribute_or_none(error, "response"), "status_code") return self._fail_open_passthrough( inputs=inputs, input_type=input_type, @@ -608,11 +619,11 @@ class GenericGuardrailAPI(CustomGuardrail): ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final = _DECODED_JSON.validate_python(response.json()) verbose_proxy_logger.debug("Generic Guardrail API response: %s", response_json) - guardrail_response: Final = GenericGuardrailAPIResponse.from_dict(response_json) + guardrail_response: Final = _guardrail_response_from_json(response_json) # Handle the response if guardrail_response.action == "BLOCKED": diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 9463023854e..85bf39c481b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -135,7 +135,7 @@ _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None _MCP_JWT_CALL_TYPES: Final = frozenset({"call_mcp_tool", "list_mcp_tools"}) # Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at). -_jwks_cache: Final[dict[str, tuple[Sequence[Mapping[str, object]], float]]] = {} +_jwks_cache: Final[dict[str, tuple[object, float]]] = {} _JWKS_CACHE_TTL: Final = 3600 # 1 hour @@ -181,7 +181,7 @@ def _compute_kid(public_key: RSAPublicKey) -> str: return hashlib.sha256(der_bytes).hexdigest()[:16] -async def _fetch_jwks(jwks_uri: str, timeout: float | httpx.Timeout | None = None) -> Sequence[Mapping[str, object]]: +async def _fetch_jwks(jwks_uri: str, timeout: float | httpx.Timeout | None = None) -> object: """ Fetch and cache a JWKS from the given URI. @@ -202,7 +202,7 @@ async def _fetch_jwks(jwks_uri: str, timeout: float | httpx.Timeout | None = Non client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) resp: Final = await client.get(jwks_uri, headers={"Accept": "application/json"}, timeout=timeout) resp.raise_for_status() - jwks_body: Final[Mapping[str, Sequence[Mapping[str, object]]]] = resp.json() + jwks_body: Final = decoded_json_object(_DECODED_JSON.validate_python(resp.json())) fetched_keys: Final = jwks_body.get("keys", []) _jwks_cache[jwks_uri] = (fetched_keys, now) return fetched_keys diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 1790f335548..16a820e1c0b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -9,7 +9,7 @@ import httpx from typing_extensions import NotRequired, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, json_contains +from litellm.litellm_core_utils.core_helpers import attribute_or_none, decoded_json_object, iter_items, json_contains from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -445,10 +445,10 @@ class PurviewGuardrailBase: tool_calls: Final = tool_calls_with_server_tool_calls(message) if tool_calls: for tc in tool_calls: - fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) + fn = tc.get("function") if isinstance(tc, dict) else attribute_or_none(tc, "function") if fn is None: continue - arguments = fn.get("arguments") if isinstance(fn, dict) else getattr(fn, "arguments", None) + arguments = fn.get("arguments") if isinstance(fn, dict) else attribute_or_none(fn, "arguments") if isinstance(arguments, str) and arguments.strip(): args.append(arguments) diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 152515c50ed..1f7f93bb9cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -1,6 +1,6 @@ # litellm/proxy/guardrails/guardrail_hooks/pangea.py import os -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sized from typing import TYPE_CHECKING, Final from fastapi import HTTPException @@ -12,7 +12,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, subscript from litellm.llms.anthropic.common_utils import ( message_dict_with_server_tool_calls, message_dict_with_server_tool_calls_split_out, @@ -40,6 +40,12 @@ if TYPE_CHECKING: _DECODED_JSON: Final = TypeAdapter(object) +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + def _choice_with_message(choice: object, rewrite: Callable[[object], object]) -> object: mapping: Final = validated_object_mapping(choice) if mapping is None or "message" not in mapping: @@ -59,9 +65,9 @@ class _TextCompletionRequest: return [{"role": "user", "content": self.body["prompt"]}] # This mutates the original dict, but we'll still return it anyways - def update_original_body(self, prompt_messages: list[dict]) -> dict[str, object]: - assert len(prompt_messages) == 1 - self.body["prompt"] = prompt_messages[0]["content"] + def update_original_body(self, prompt_messages: object) -> dict[str, object]: + assert _json_len(prompt_messages) == 1 + self.body["prompt"] = subscript(subscript(prompt_messages, 0), "content") return self.body @@ -117,7 +123,7 @@ class PangeaHandler(CustomGuardrail): self.api_base, ) - async def _call_pangea_ai_guard(self, api: str, payload: dict, hook_name: str) -> dict: + async def _call_pangea_ai_guard(self, api: str, payload: dict, hook_name: str) -> Mapping[str, object]: """ Makes the API call to the Pangea AI Guard endpoint. The function itself will raise an error in the case that a response @@ -153,8 +159,7 @@ class PangeaHandler(CustomGuardrail): ) response.raise_for_status() - result: Final = response.json() - decoded_result: Final = decoded_json_object(_DECODED_JSON.validate_python(result)) + decoded_result: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) verdict: Final = decoded_json_object(decoded_result.get("result", {})) if verdict.get("blocked"): @@ -172,7 +177,7 @@ class PangeaHandler(CustomGuardrail): "Pangea Guardrail (%s): Request passed. Response: %s", hook_name, verdict.get("detectors") ) - return result + return decoded_result async def _async_pre_call_hook( self, @@ -200,14 +205,15 @@ class PangeaHandler(CustomGuardrail): ai_guard_response = await self._call_pangea_ai_guard("v1beta/guard", ai_guard_payload, "async_pre_call_hook") add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) - if not ai_guard_response.get("result", {}).get("transformed"): + pangea_result: Final = decoded_json_object(ai_guard_response.get("result", {})) + if not pangea_result.get("transformed"): return - output: Final = ai_guard_response.get("result", {}).get("output", {}) + output: Final = pangea_result.get("output", {}) if call_type == "text_completion" or call_type == "atext_completion": - data = transformer.update_original_body(output["messages"]) + data = transformer.update_original_body(subscript(output, "messages")) else: - data["messages"] = output["messages"] + data["messages"] = subscript(output, "messages") return data @log_guardrail_information @@ -290,12 +296,14 @@ class PangeaHandler(CustomGuardrail): ai_guard_response = await self._call_pangea_ai_guard("v1beta/guard", ai_guard_payload, "async_pre_call_hook") add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) - if not ai_guard_response.get("result", {}).get("transformed"): + pangea_result: Final = decoded_json_object(ai_guard_response.get("result", {})) + if not pangea_result.get("transformed"): return - output: Final = ai_guard_response.get("result", {}).get("output", {}) + output: Final = pangea_result.get("output", {}) response.choices = [ # pyright: ignore[reportAttributeAccessIssue] # Pangea returns untyped choice dicts - _choice_with_message(choice, message_dict_with_server_tool_calls_split_out) for choice in output["choices"] + _choice_with_message(choice, message_dict_with_server_tool_calls_split_out) + for choice in iter_items(subscript(output, "choices")) ] return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 73ccd82d37f..f4b153d86ba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -28,7 +28,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.core_helpers import decoded_json_object, json_contains from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, tool_calls_with_server_tool_calls, @@ -145,6 +145,7 @@ class _ResponsesInputItem(LiteLLMBaseModel): _ResponsesInput: TypeAlias = str | tuple[_ResponsesInputItem, ...] | None _RESPONSES_INPUT: Final[TypeAdapter[_ResponsesInput]] = TypeAdapter(_ResponsesInput) _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_DECODED_JSON: Final = TypeAdapter(object) _NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) @@ -394,7 +395,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): metadata: Mapping[str, object] | None = None, call_id: object = None, tool_event: _AirsToolEvent | None = None, - ) -> dict[str, object]: + ) -> Mapping[str, object]: """Call PANW Prisma AIRS API to scan content or a tool_event.""" if tool_event is None and not content.strip(): @@ -506,30 +507,32 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) response.raise_for_status() - result: Final[dict[str, object]] = response.json() + result: Final = _DECODED_JSON.validate_python(response.json()) # Validate response format - if "action" not in result: + if not json_contains(result, "action"): verbose_proxy_logger.error("PANW Prisma AIRS: Invalid API response format: %s", result) return {"action": "block", "category": "api_error"} + body: Final = decoded_json_object(result) + # Check for profile-related errors from PANW API - if result.get("action") == "block" and "error" in result: - error_msg: Final = str(result.get("error", "")).lower() + if body.get("action") == "block" and "error" in body: + error_msg: Final = str(body.get("error", "")).lower() if "profile" in error_msg and ( "not found" in error_msg or "required" in error_msg or "invalid" in error_msg ): verbose_proxy_logger.error( "PANW Prisma AIRS: Profile configuration error. Ensure your API key has a linked profile in Strata Cloud Manager, or provide 'profile_name' or 'profile_id' in config/metadata. PANW API response: %s", - result, + body, ) verbose_proxy_logger.debug( "PANW Prisma AIRS: Scan result - Action: %s, Category: %s", - result.get("action"), - result.get("category", "unknown"), + body.get("action"), + body.get("category", "unknown"), ) - return result + return body except httpx.HTTPStatusError as e: status: Final = e.response.status_code @@ -839,7 +842,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): def _handle_api_error_with_logging( self, - scan_result: dict[str, object], + scan_result: Mapping[str, object], data: dict[str, object], start_time: datetime, event_type: GuardrailEventHooks, @@ -1363,7 +1366,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): assembled_model_response: ModelResponse, request_data: dict, start_time: datetime, - ) -> tuple[bool, ModelResponse, dict[str, object]]: + ) -> tuple[bool, ModelResponse, Mapping[str, object]]: """ Scan assembled streaming response and apply masking if needed. Returns (content_was_modified, response, scan_result). diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index f2697c0aa11..09f92f9c8ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -5,11 +5,11 @@ from dataclasses import dataclass from functools import partial, reduce from itertools import chain, groupby from types import MappingProxyType -from typing import Any, Final, Literal, TypeAlias, TypedDict, TypeVar +from typing import Final, Literal, TypeAlias, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter -from typing_extensions import ReadOnly, Required +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache @@ -22,7 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, keyword_arguments from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, @@ -472,19 +472,13 @@ def _rewritten_tool_config(tool_config: object, removed_names: frozenset[str], f } -class _ToolPermissionRuleFields(TypedDict, total=False): - """The config-file shape a :class:`ToolPermissionRule` is built from.""" - - id: ReadOnly[Required[str]] - tool_name: ReadOnly[str | None] - tool_type: ReadOnly[str | None] - decision: ReadOnly[Required[Literal["allow", "deny"]]] - allowed_param_patterns: ReadOnly[dict[str, str] | None] +class _ToolPermissionRuleFields(TypedDict, total=False, extra_items=object): + pass -def _rule_from_fields(fields: _ToolPermissionRuleFields) -> ToolPermissionRule: +def _rule_from_fields(fields: object) -> ToolPermissionRule: """Validate one config-file rule entry into a :class:`ToolPermissionRule`.""" - return ToolPermissionRule(**fields) + return ToolPermissionRule.model_validate(_ToolPermissionRuleFields(**keyword_arguments(fields, ToolPermissionRule))) def _is_tool_use_block(block: object) -> bool: @@ -529,7 +523,7 @@ class ToolPermissionGuardrail(CustomGuardrail): self.default_action, ) - def _load_rules(self, rules: list[Any] | None) -> None: + def _load_rules(self, rules: Sequence[object] | None) -> None: """Parse ``rules`` and (re)build the compiled target/pattern lookups. ``self.rules`` plus ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` @@ -948,10 +942,10 @@ class ToolPermissionGuardrail(CustomGuardrail): return is_allowed, None, message @staticmethod - def _get_mapping_value(item: object, key: str) -> Any: + def _get_mapping_value(item: object, key: str) -> object: if isinstance(item, dict): return item.get(key) - return getattr(item, key, None) + return attribute_or_none(item, key) @staticmethod def _legacy_function_call_id(choice_index: int) -> str: diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 2f2294db522..cafd79626ef 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -10,6 +10,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast from pydantic import BaseModel, TypeAdapter, ValidationError +from typing_extensions import TypedDict import litellm from litellm import Router @@ -17,6 +18,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, GUARDRAIL_ROTATION_ATTEMPTS from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import as_str_mapping, keyword_arguments, subscript from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, @@ -102,6 +104,10 @@ def _as_json_array(value: object) -> list[object] | None: return _JSON_ARRAY.validate_python(value) if isinstance(value, list) else None +class _CategoryThresholdFields(TypedDict, total=False, extra_items=object): + pass + + def _extra_param_items(litellm_params: LitellmParams) -> Iterable[tuple[str, object]]: return ( litellm_params.model_dump(exclude_none=True).items() @@ -775,7 +781,11 @@ class InMemoryGuardrailHandler: litellm_params = litellm_params_data if "category_thresholds" in litellm_params_data and litellm_params_data["category_thresholds"]: - lakera_category_thresholds: Final = LakeraCategoryThresholds(**litellm_params_data["category_thresholds"]) + lakera_category_thresholds: Final = LakeraCategoryThresholds( + **_CategoryThresholdFields( + **keyword_arguments(subscript(litellm_params_data, "category_thresholds"), LakeraCategoryThresholds) + ) + ) litellm_params.category_thresholds = lakera_category_thresholds if litellm_params.api_key and litellm_params.api_key.startswith("os.environ/"): @@ -871,7 +881,8 @@ class InMemoryGuardrailHandler: params: Final = guardrail.get("litellm_params") if not isinstance(params, (dict, LitellmParams)): return - litellm_params: Final = LitellmParams(**params) if isinstance(params, dict) else params + params_mapping: Final = as_str_mapping(params) + litellm_params: Final = LitellmParams(**params_mapping) if params_mapping is not None else params guardrail_name: Final = guardrail.get("guardrail_name", "Unknown") for custom_guardrail_callback in self._tracked_callbacks(guardrail_id): for error in ( diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index b4d343b3932..625e1e2df7e 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -26,6 +26,7 @@ from urllib.parse import parse_qs, urlsplit from fastapi import HTTPException from pydantic import Field, TypeAdapter, ValidationError +from typing_extensions import TypedDict import litellm from litellm._internal_context import with_service_target @@ -105,6 +106,10 @@ _DECODED_JSON: Final = TypeAdapter(object) _WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) +class _FileContentFetchKwargs(TypedDict, total=False, extra_items=object): + pass + + def _batch_create_routing_model(data: Mapping[str, object]) -> str | None: body_model: Final = data.get("model") if isinstance(body_model, str) and body_model: @@ -537,7 +542,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): custom_llm_provider: str, data: dict, user_api_key_dict: UserAPIKeyAuth | None = None, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: """ Map proxy-facing file IDs to provider file IDs and credentials. @@ -954,7 +959,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content = await litellm.afile_content( file_id=provider_file_id, user_api_key_dict=user_api_key_dict, - **fetch_kwargs, + **_FileContentFetchKwargs(**fetch_kwargs), ) file_content_bytes: Final = getattr(file_content, "content", None) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b1eb8f7d8ae..8abc91a72ed 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -54,7 +54,10 @@ from litellm.integrations.custom_guardrail import merged_guardrail_entries, with from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, attribute_or_none, + decoded_json_object, is_codex_user_agent, + iter_items, + keyword_arguments, normalize_drop_params, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor @@ -1251,6 +1254,10 @@ def convert_key_logging_metadata_to_callback( return team_callback_settings_obj +class _CallbackModelFields(TypedDict, total=False, extra_items=object): + pass + + def get_validated_callback_metadata(item: object, *, source: str) -> AddTeamCallback | None: if not isinstance(item, Mapping): raise TypeError(f"{source} metadata.logging must be a list of objects, got an item of {type(item).__name__}") @@ -1280,9 +1287,9 @@ class KeyAndTeamLoggingSettings: return None @staticmethod - def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): + def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth) -> object: if user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata: - return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging") + return decoded_json_object(decrypt_callback_vars(user_api_key_dict.team_metadata)).get("logging") return None @@ -1291,14 +1298,14 @@ def get_dynamic_logging_metadata( ) -> TeamCallbackMetadata | None: callback_settings_obj: TeamCallbackMetadata | None = None key_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - team_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( + team_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( user_api_key_dict ) ######################################################################################### # Key-based callbacks ######################################################################################### if key_dynamic_logging_settings is not None: - for item in key_dynamic_logging_settings: + for item in iter_items(key_dynamic_logging_settings): callback = get_validated_callback_metadata(item=item, source="key-level") if callback is None: continue @@ -1310,7 +1317,7 @@ def get_dynamic_logging_metadata( # Team-based callbacks ######################################################################################### elif team_dynamic_logging_settings is not None: - for item in team_dynamic_logging_settings: + for item in iter_items(team_dynamic_logging_settings): callback = get_validated_callback_metadata(item=item, source="team-level") if callback is None: continue @@ -1331,9 +1338,11 @@ def get_dynamic_logging_metadata( } } """ - team_metadata: Final = decrypt_callback_vars(user_api_key_dict.team_metadata) + team_metadata: Final = decoded_json_object(decrypt_callback_vars(user_api_key_dict.team_metadata)) callback_settings: Final = team_metadata.get("callback_settings", None) or {} - callback_settings_obj = TeamCallbackMetadata(**callback_settings) + callback_settings_obj = TeamCallbackMetadata( + **_CallbackModelFields(**keyword_arguments(callback_settings, TeamCallbackMetadata)) + ) verbose_proxy_logger.debug("Team callback settings activated: %s", callback_settings_obj) ######################################################################################### # Enter here when configured on the config.yaml file. @@ -1465,7 +1474,7 @@ def resolve_tenant_otel_destinations( disabled: Final = _dynamically_disabled_backends(user_api_key_dict, request_headers) callbacks: Final = tuple( callback - for item in entries + for item in iter_items(entries) if (callback := get_validated_callback_metadata(item=item, source="otel-destination")) is not None if callback.callback_name.lower() not in disabled ) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 4cbc2155af0..10bdd1d2353 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,7 +1,7 @@ import asyncio from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass -from types import MappingProxyType +from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status @@ -98,11 +98,11 @@ class _KeyRecord(Protocol): class _AccessGroupTable(Protocol): async def find_unique(self, where: Mapping[str, object]) -> _AccessGroupRecord | None: ... - async def find_many(self, order: Mapping[str, object]) -> Sequence[_AccessGroupRecord]: ... + async def find_many(self, *, order: Mapping[str, object]) -> Sequence[_AccessGroupRecord]: ... async def create(self, data: Mapping[str, object]) -> _AccessGroupRecord: ... - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _AccessGroupRecord: ... + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _AccessGroupRecord | None: ... async def delete(self, where: Mapping[str, object]) -> object: ... @@ -118,9 +118,9 @@ class _TeamTable(Protocol): class _KeyTable(Protocol): async def find_unique(self, where: Mapping[str, object]) -> _KeyRecord | None: ... - async def find_many(self, where: Mapping[str, object]) -> Sequence[_KeyRecord]: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_KeyRecord]: ... - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... class _AgentRecord(Protocol): @@ -132,9 +132,9 @@ class _AgentRecord(Protocol): class _AgentTable(Protocol): - async def find_many(self, where: Mapping[str, object]) -> Sequence[_AgentRecord]: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_AgentRecord]: ... - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... class _HasSomeFilter(TypedDict): @@ -170,7 +170,13 @@ class _AccessGroupTx(Protocol): class _AccessGroupTransaction(Protocol): async def __aenter__(self) -> _AccessGroupTx: ... - async def __aexit__(self, *exc_info: object) -> None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> None: ... def _access_group_transaction(prisma_client: PrismaClient) -> _AccessGroupTransaction: @@ -731,6 +737,8 @@ async def update_access_group( from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + if record is None: + raise AttributeError("'NoneType' object has no attribute 'dict'", name="dict", obj=None) await _cache_access_group_record(record) await _patch_team_caches_add_access_group(teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) await _patch_team_caches_remove_access_group( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 21ac82754b4..3307f3050d9 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1122,13 +1122,13 @@ async def _target_in_caller_admin_org( caller_user_id: str, target_user_id: str, ) -> bool: - admin_memberships: Final = await prisma_client.db.litellm_organizationmembership.find_many( + admin_memberships: Final = await _organization_membership_table(prisma_client).find_many( where={"user_id": caller_user_id, "user_role": LitellmUserRoles.ORG_ADMIN.value} ) if not admin_memberships: return False admin_org_ids: Final = {membership.organization_id for membership in admin_memberships} - target_memberships: Final = await prisma_client.db.litellm_organizationmembership.find_many( + target_memberships: Final = await _organization_membership_table(prisma_client).find_many( where={"user_id": target_user_id} ) return any(membership.organization_id in admin_org_ids for membership in target_memberships) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b59606a8dcf..90779f732c7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -20,7 +20,7 @@ import secrets import traceback from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone -from types import MappingProxyType +from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, cast import fastapi @@ -245,13 +245,20 @@ class _UserRowLike(Protocol): class _TxTables(Protocol): - litellm_proxymodeltable: TableActions[object] + @property + def litellm_proxymodeltable(self) -> TableActions[object]: ... class _ModelReencryptTransaction(Protocol): async def __aenter__(self) -> _TxTables: ... - async def __aexit__(self, *exc_info: object) -> None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> None: ... class _ModelParamsUpdate(TypedDict): diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index da84c2d0013..378d3b34d12 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,12 +1,11 @@ """`/management/v1/spend_logs` facets.""" -from collections.abc import Sequence from datetime import datetime, timezone from functools import partial -from typing import Annotated, Final, Literal, Protocol +from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, Query, Request -from typing_extensions import ReadOnly, TypedDict +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth @@ -21,7 +20,7 @@ from litellm.proxy.list_api.common import ( reject_unknown_query_params, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import query_raw_rows from litellm.types.proxy.management_endpoints.management_v1 import ( FacetListResponse, PageMeta, @@ -35,18 +34,7 @@ router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) # query, so both reads of the same table stop at the same depth. SPEND_LOGS_FACET_SCAN_CAP: Final = 10000 - -class _FacetRow(TypedDict, total=False): - end_user: ReadOnly[str | None] - user: ReadOnly[str | None] - - -class _FacetQuery(Protocol): - async def __call__(self, query: str, *args: object) -> Sequence[_FacetRow]: ... - - -def _facet_query(prisma_client: PrismaClient) -> _FacetQuery: - return prisma_client.db.query_raw +_FACET_VALUES: Final = TypeAdapter(list[str]) def _as_utc(value: datetime) -> datetime: @@ -141,8 +129,8 @@ async def _list_spend_log_facet( f" ORDER BY {column_sql} ASC" f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" ) - rows: Final = await _facet_query(prisma_client)(facet_sql, *params) - values: Final[list[str]] = [value for row in rows if (value := row.get(column))] + rows: Final = await query_raw_rows(prisma_client.db, facet_sql, *params) + values: Final = _FACET_VALUES.validate_python([value for row in rows if (value := row.get(column))]) has_more: Final = len(values) > page_size return FacetListResponse( diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 3a2de12d108..abc2fe4651a 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -20,13 +20,14 @@ import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol +from typing import TYPE_CHECKING, Annotated, Final, Protocol from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import json_contains, require_attributes from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import publish_model_rows_written @@ -104,6 +105,14 @@ def _http_error(status_code: int, message: str) -> HTTPException: return HTTPException(status_code=status_code, detail=detail) +class _AppendableAccessGroups(Protocol): + def append(self, value: str, /) -> None: ... + + +class _RemovableAccessGroups(Protocol): + def remove(self, value: str, /) -> None: ... + + class _DeploymentRow(Protocol): @property def model_id(self) -> str: ... @@ -376,7 +385,7 @@ async def _missing_models_after_read_through( return tuple(still_missing) -def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]: +def add_access_group_to_deployment(model_info: dict[str, object], access_group: str) -> tuple[dict[str, object], bool]: """ Add an access group to a deployment's model_info. @@ -385,16 +394,16 @@ def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str access_group: The access group name to add Returns: - Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified) + tuple[dict[str, object], bool]: (updated_model_info, was_modified) """ access_groups: Final = model_info.get("access_groups", []) # Check if access group already exists - if access_group in access_groups: + if json_contains(access_groups, access_group): return model_info, False # Add the access group - access_groups.append(access_group) + require_attributes(access_groups, _AppendableAccessGroups, "append").append(access_group) model_info["access_groups"] = access_groups return model_info, True @@ -729,7 +738,9 @@ async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> return deployment.model_info -def remove_access_group_from_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]: +def remove_access_group_from_deployment( + model_info: dict[str, object], access_group: str +) -> tuple[dict[str, object], bool]: """ Remove an access group from a deployment's model_info. @@ -738,16 +749,16 @@ def remove_access_group_from_deployment(model_info: dict[str, Any], access_group access_group: The access group name to remove Returns: - Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified) + tuple[dict[str, object], bool]: (updated_model_info, was_modified) """ access_groups: Final = model_info.get("access_groups", []) # Check if access group exists - if access_group not in access_groups: + if not json_contains(access_groups, access_group): return model_info, False # Remove the access group - access_groups.remove(access_group) + require_attributes(access_groups, _RemovableAccessGroups, "remove").remove(access_group) model_info["access_groups"] = access_groups return model_info, True diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0b4f266f838..6d4dc01a8c1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -553,6 +553,21 @@ def _effective_complexity_router_config( ).effective +def _encrypted_litellm_params( + dumped: Mapping[str, object], + incoming_params: GenericLiteLLMParams | None, + existing_params: GenericLiteLLMParams | None, +) -> Mapping[str, object]: + return { + name: ( + _effective_complexity_router_config(incoming_params, existing_params) + if name == "complexity_router_config" + else encrypt_value_helper(value) + ) + for name, value in dumped.items() + } + + def _dump_without_none(model: BaseModel) -> Mapping[str, object]: return model.model_dump(exclude_none=True) @@ -1511,14 +1526,11 @@ def update_db_model( # update litellm params if updated_patch.litellm_params: # Encrypt any sensitive values - encrypted_params: Final = { - k: ( - _effective_complexity_router_config(updated_patch.litellm_params, db_model.litellm_params) - if k == "complexity_router_config" - else encrypt_value_helper(v) - ) - for k, v in updated_patch.litellm_params.model_dump(exclude_none=True, exclude_unset=True).items() - } + encrypted_params: Final = _encrypted_litellm_params( + updated_patch.litellm_params.model_dump(exclude_none=True, exclude_unset=True), + updated_patch.litellm_params, + db_model.litellm_params, + ) merged_litellm_params.update(encrypted_params) @@ -3471,14 +3483,9 @@ async def update_model( ### ENCRYPT PARAMS ### encrypted_params: Final = MappingProxyType( - { - k: ( - _effective_complexity_router_config(model_params.litellm_params, deployment.litellm_params) - if k == "complexity_router_config" - else encrypt_value_helper(value=v) - ) - for k, v in _new_litellm_params_dict.items() - } + _encrypted_litellm_params( + _new_litellm_params_dict, model_params.litellm_params, deployment.litellm_params + ) ) ### MERGE WITH EXISTING DATA ### diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index a185c0dda35..853f29ecb62 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -206,7 +206,19 @@ class _ObjectPermissionTxClient(Protocol): class _BudgetTxClient(Protocol): - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaBudgetTable | None": ... + async def update( + self, *, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaBudgetTable | None": ... + + +class _OrganizationTxClient(Protocol): + async def update( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> "PrismaOrganizationTable | None": ... class _TransactionTables(Protocol): @@ -217,7 +229,7 @@ class _TransactionTables(Protocol): def litellm_budgettable(self) -> "_BudgetTxClient": ... @property - def litellm_organizationtable(self) -> "_OrganizationTableClient": ... + def litellm_organizationtable(self) -> "_OrganizationTxClient": ... class _TransactionManager(Protocol): @@ -228,6 +240,7 @@ class _TransactionManager(Protocol): exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: "TracebackType | None", + /, ) -> bool | None: ... diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 1d35d34cbcf..24675406567 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -164,7 +164,9 @@ def _resolve_team_callbacks(team_metadata: object) -> TeamCallbackMetadata: for entry in logging_entries if isinstance(logging_entries, list) else (): if not isinstance(entry, dict): continue - callback = get_validated_callback_metadata(item=entry, source="team-level read") + callback = get_validated_callback_metadata( + item=_DECODED_JSON.validate_python(entry), source="team-level read" + ) if callback is None: continue resolved = convert_key_logging_metadata_to_callback(data=callback, team_callback_settings_obj=resolved) diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index 27f8b90a586..f29acc3b380 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -5,6 +5,7 @@ import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from types import TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast from fastapi import HTTPException @@ -68,8 +69,11 @@ if TYPE_CHECKING: class _UserEnvVarsTransactionClient(Protocol): - litellm_mcpuserenvvars: "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]" - litellm_mcpservertable: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]" + @property + def litellm_mcpuserenvvars(self) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]": ... + + @property + def litellm_mcpservertable(self) -> "TableActions[prisma_db_models.LiteLLM_MCPServerTable]": ... async def execute_raw(self, query: str, *args: object) -> int: ... @@ -77,7 +81,13 @@ class _UserEnvVarsTransactionClient(Protocol): class _UserEnvVarsTransaction(Protocol): async def __aenter__(self) -> _UserEnvVarsTransactionClient: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... @dataclass(frozen=True, slots=True) @@ -139,6 +149,10 @@ def _blank_to_none(value: object) -> str | None: return value.strip() or None +def _submitted_or(submitted: Mapping[str, object], field: str, fallback: str | None) -> object: + return submitted.get(field, fallback) + + # Token-exchange settings with dedicated columns that also exist on # ``MCPCredentials`` as a legacy shape (rows and REST callers that predate the # columns). Every write lifts blob values into the columns and strips them from @@ -1379,8 +1393,10 @@ async def update_mcp_server( "credentials": safe_dumps( _bind_submitted_oauth_client( safe_submitted, - _blank_to_none(cleared_auth_fields.get("issuer", prepared_data.get("issuer", old_issuer))), - _blank_to_none(prepared_data.get("url", existing.url if existing else None)), + _blank_to_none( + cleared_auth_fields.get("issuer", _submitted_or(prepared_data, "issuer", old_issuer)) + ), + _blank_to_none(_submitted_or(prepared_data, "url", existing.url if existing else None)), ) ) } diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index ecc40b6e98f..f36609aa8b3 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -6,7 +6,17 @@ import time from collections.abc import AsyncIterator, Callable, Mapping from contextlib import asynccontextmanager from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import ( + TYPE_CHECKING, + Final, + Literal, + Optional, + Protocol, + SupportsIndex, + SupportsInt, + overload, + runtime_checkable, +) from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx @@ -15,6 +25,7 @@ from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Resp from mcp.client.auth.utils import get_client_metadata_scopes, union_scopes from mcp.shared.auth import OAuthMetadata from pydantic import ConfigDict, Field, SecretStr, TypeAdapter, ValidationError +from typing_extensions import Buffer from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -673,10 +684,24 @@ def _validate_token_response( ) +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +class _OAuthTokenResponseFields(Protocol): + @overload + def get(self, key: Literal["access_token", "refresh_token"], /) -> str | None: ... + @overload + def get(self, key: Literal["expires_in"], /) -> object: ... + @overload + def get(self, key: Literal["scope"], default: str | None, /) -> object: ... + + async def _store_per_user_token_server_side( server: MCPServer, user_id: str, - token_response: dict[str, Any], + token_response: _OAuthTokenResponseFields, identity_binding_proof: str | None = None, cimd_client_id: str | None = None, requested_scope: str | None = None, @@ -699,7 +724,11 @@ async def _store_per_user_token_server_side( raw_expires: Final = token_response.get("expires_in") try: - expires_in: int | None = int(raw_expires) if raw_expires is not None else None + expires_in: int | None = ( + int(raw_expires) + if isinstance(raw_expires, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)) + else None + ) except (TypeError, ValueError): expires_in = None diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index cc02fc43ef3..3d55f63c73d 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -29,7 +29,7 @@ from dataclasses import dataclass, replace from functools import lru_cache from itertools import chain, groupby from types import EllipsisType, MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse import anyio @@ -73,6 +73,7 @@ from litellm.constants import ( MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.mcp_client.client import ( @@ -469,14 +470,6 @@ class MCPServerConfig(TypedDict, total=False): rpm: ReadOnly[int | None] -class _ProtectedResourceMetadataPayload(TypedDict, total=False): - """The RFC 9728 protected-resource metadata document fields this gateway reads.""" - - authorization_servers: Sequence[object] - scopes_supported: Sequence[str] - scopes: Sequence[str] - - class _AuthorizationServerMetadataPayload(TypedDict, total=False): """The RFC 8414 / OpenID Discovery authorization-server metadata fields this gateway reads.""" @@ -1910,6 +1903,10 @@ def _listing_task_was_cancelled(error: asyncio.CancelledError) -> bool: return error.args != (SESSION_CLEANUP_TIMEOUT_MESSAGE,) +class _InternalNetworkSettings(Protocol): + def get(self, key: Literal["mcp_internal_ip_ranges"], /) -> list[str] | None: ... + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -5369,7 +5366,7 @@ class MCPServerManager: try: response: Final = await self._fetch_oauth_discovery_url(resource_metadata_url, server_url) response.raise_for_status() - data: Final[_ProtectedResourceMetadataPayload] = response.json() + decoded: Final = _UNTYPED_VALUE.validate_python(response.json()) except SSRFError as exc: verbose_logger.warning( "MCP OAuth discovery: refusing to fetch resource metadata from %s " @@ -5387,11 +5384,13 @@ class MCPServerManager: ) return [], None - raw_servers: Final = data.get("authorization_servers") - if isinstance(raw_servers, list): - authorization_servers = [entry for entry in raw_servers if isinstance(entry, str) and entry.strip() != ""] - else: - authorization_servers = [] + data: Final = decoded_json_object(decoded) + raw_servers: Final = as_object_list(data.get("authorization_servers")) + authorization_servers: Final = ( + [entry for entry in raw_servers if isinstance(entry, str) and entry.strip() != ""] + if raw_servers is not None + else [] + ) scopes: Final = self._extract_scopes(data.get("scopes_supported") or data.get("scopes")) @@ -5631,13 +5630,14 @@ class MCPServerManager: "aws_session_name": credentials_dict.get("aws_session_name"), } - def _extract_scopes(self, scopes_value: str | Sequence[object] | None) -> list[str] | None: + def _extract_scopes(self, scopes_value: object) -> list[str] | None: if isinstance(scopes_value, str): scopes = [s.strip() for s in scopes_value.split() if s.strip()] return scopes or None - if isinstance(scopes_value, list): - scopes = [s for s in scopes_value if isinstance(s, str) and s.strip()] - return scopes or None + scope_entries: Final = as_object_list(scopes_value) + if scope_entries is not None: + listed_scopes: Final = [s for s in scope_entries if isinstance(s, str) and s.strip()] + return listed_scopes or None return None async def _fetch_tools_with_timeout( @@ -7354,7 +7354,7 @@ class MCPServerManager: servers.append(server) return servers - def _get_general_settings(self) -> dict[str, Any]: + def _get_general_settings(self) -> _InternalNetworkSettings: """Get general_settings, importing lazily to avoid circular imports.""" try: from litellm.proxy.proxy_server import ( diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 8422d516d8c..afa3a373f3c 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -2080,7 +2080,7 @@ async def _list_tools_before_first_call( async def execute_mcp_tool( name: str, - arguments: dict[str, object], + arguments: Mapping[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 815c1d6270c..165fc23d754 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -15,7 +15,7 @@ import httpx2 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pydantic import TypeAdapter, ValidationError from starlette.datastructures import Headers -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT @@ -118,6 +118,8 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( _CLIENT_FORWARDED_TOKEN_AUTH_TYPES: Final = frozenset((MCPAuth.true_passthrough, MCPAuth.oauth_delegate)) +_NO_TOOL_ARGUMENTS: Final[Mapping[str, Never]] = MappingProxyType({}) + def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: reference: Final = uuid4().hex @@ -356,7 +358,7 @@ if MCP_AVAILABLE: status_code=403, detail={"error": "forbidden", "message": f"{tool_name} requires mcp_tool_search_enabled on the key"}, ) - tool_arguments: Final = data.get("arguments") or {} + tool_arguments: Final = data.get("arguments") or _NO_TOOL_ARGUMENTS if tool_name == AGENT_SEARCH_TOOL_NAME: return await handle_agent_search( query=str(tool_arguments.get("query", "")), diff --git a/litellm/proxy/mcp_server/tool_search.py b/litellm/proxy/mcp_server/tool_search.py index 96bcf8fbacf..e65fcf61f49 100644 --- a/litellm/proxy/mcp_server/tool_search.py +++ b/litellm/proxy/mcp_server/tool_search.py @@ -10,13 +10,13 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain, islice from types import EllipsisType, MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict +from typing import TYPE_CHECKING, Final, Protocol, SupportsIndex, SupportsInt, TypedDict, runtime_checkable import anyio from anyio import to_process from anyio.lowlevel import RunVar from pydantic import JsonValue, SecretStr, ValidationError -from typing_extensions import ReadOnly, Required, assert_never +from typing_extensions import Buffer, ReadOnly, Required, assert_never import litellm from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K @@ -131,7 +131,14 @@ async def _tool_argument_validation_error( return "Tool argument validation worker failed" -def coerce_top_k(value: Any, default: int = 5) -> int: +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +def coerce_top_k(value: object, default: int = 5) -> int: + if not isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)): + return default try: return int(value) except (TypeError, ValueError): @@ -699,7 +706,7 @@ async def handle_mcp_proxy_tool( async def handle_mcp_tool_call( tool_name: str, - arguments: dict[str, Any], + arguments: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, client_ip: str | None = None, mcp_servers: list[str] | None = None, diff --git a/litellm/proxy/middleware/admission_control_middleware.py b/litellm/proxy/middleware/admission_control_middleware.py index 83503d0c3d4..1b1af7fff95 100644 --- a/litellm/proxy/middleware/admission_control_middleware.py +++ b/litellm/proxy/middleware/admission_control_middleware.py @@ -3,7 +3,7 @@ import os from collections.abc import Callable, Mapping from dataclasses import dataclass from functools import lru_cache -from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable +from typing import Annotated, Final, Literal, Protocol, TypeAlias, runtime_checkable from pydantic import Field, TypeAdapter, ValidationError from starlette.responses import JSONResponse @@ -133,6 +133,12 @@ class AdmissionControlState: return self._metrics +class _RoutedScope(Protocol): + def __getitem__(self, key: Literal["path"], /) -> str: ... + + def get(self, key: Literal["root_path"], default: str, /) -> str: ... + + class _AdmissionLease: def __init__(self, state: AdmissionControlState) -> None: self.state: Final = state @@ -207,7 +213,7 @@ class AdmissionControlMiddleware: state.record_release() -def _get_route_path(scope: Scope) -> str: +def _get_route_path(scope: _RoutedScope) -> str: """Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching.""" path: Final[str] = scope["path"] root_path: Final[str] = scope.get("root_path", "") diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index 95e5c9d1293..14b99b0e25d 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -19,7 +19,7 @@ import re import threading from collections.abc import Callable, Sequence from enum import Enum -from typing import Final, Protocol, runtime_checkable +from typing import Final, Literal, Protocol, runtime_checkable from starlette.types import ASGIApp, Message, Receive, Scope, Send @@ -196,6 +196,26 @@ def _extract_model_id(headers: Sequence[tuple[bytes, bytes]]) -> str | None: ) +class _RequestLineFields(Protocol): + def get(self, key: Literal["path", "method"], default: str, /) -> str: ... + + +class _ResponseStartFields(Protocol): + def __getitem__(self, key: Literal["status"], /) -> int: ... + + def get( + self, key: Literal["headers"], default: Sequence[tuple[bytes, bytes]], / + ) -> Sequence[tuple[bytes, bytes]]: ... + + +def _request_line_fields(scope: _RequestLineFields) -> _RequestLineFields: + return scope + + +def _response_start_fields(message: _ResponseStartFields) -> _ResponseStartFields: + return message + + class BillableRequestMetricsMiddleware: """ Pure ASGI middleware that classifies each request once and fans the result @@ -263,7 +283,10 @@ class BillableRequestMetricsMiddleware: await self.app(scope, receive, send) return - classification: Final = classify_billable_request(scope.get("path", ""), scope.get("method", "POST")) + request_line: Final = _request_line_fields(scope) + classification: Final = classify_billable_request( + request_line.get("path", ""), request_line.get("method", "POST") + ) if classification is None: await self.app(scope, receive, send) return @@ -275,8 +298,9 @@ class BillableRequestMetricsMiddleware: async def send_wrapper(message: Message) -> None: nonlocal status_code, model_id if message["type"] == "http.response.start": - status_code = message["status"] - model_id = _extract_model_id(message.get("headers", [])) + response_start: Final = _response_start_fields(message) + status_code = response_start["status"] + model_id = _extract_model_id(response_start.get("headers", [])) await send(message) await self.app(scope, receive, send_wrapper) diff --git a/litellm/proxy/middleware/per_request_root_path_middleware.py b/litellm/proxy/middleware/per_request_root_path_middleware.py index 3afa268b182..fdc374a58d7 100644 --- a/litellm/proxy/middleware/per_request_root_path_middleware.py +++ b/litellm/proxy/middleware/per_request_root_path_middleware.py @@ -14,7 +14,7 @@ Opt-in: with ``SERVER_ROOT_PATHS`` unset the middleware is not added at all. import os from collections.abc import Sequence from contextvars import ContextVar -from typing import Final +from typing import Final, Literal, Protocol from starlette.types import ASGIApp, Receive, Scope, Send @@ -89,6 +89,14 @@ def get_server_root_paths() -> tuple[str, ...]: return normalize_root_paths(configured.split(",")) +class _RootPathFields(Protocol): + def get(self, key: Literal["path", "root_path"], default: str, /) -> str: ... + + +def _root_path_fields(scope: _RootPathFields) -> _RootPathFields: + return scope + + class PerRequestRootPathMiddleware: """Sets ``scope["root_path"]`` to the configured prefix matching the request path on a whole-segment boundary. ``scope["path"]`` is left @@ -104,7 +112,7 @@ class PerRequestRootPathMiddleware: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] in ("http", "websocket"): - path: Final[str] = scope.get("path", "") + path: Final[str] = _root_path_fields(scope).get("path", "") for prefix in self.root_paths: if path == prefix or path.startswith(prefix + "/"): scope["root_path"] = prefix # rebind-ok: ASGI middleware contract; Router and base_url read it @@ -113,7 +121,7 @@ class PerRequestRootPathMiddleware: # existing value when nothing matched — i.e. FastAPI's scalar # SERVER_ROOT_PATH) so code that emits URLs off the request path # picks the same prefix the router will resolve the request under. - token: Final = _request_root_path_var.set(str(scope.get("root_path", ""))) + token: Final = _request_root_path_var.set(str(_root_path_fields(scope).get("root_path", ""))) try: await self.app(scope, receive, send) finally: diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py index 329a3c38d91..06649d72153 100644 --- a/litellm/proxy/middleware/request_size_limit_middleware.py +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -1,7 +1,7 @@ import json -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from types import MappingProxyType -from typing import Final +from typing import Final, Literal, Protocol from fastapi import HTTPException from starlette._utils import get_route_path @@ -14,6 +14,18 @@ MaxRequestSizeGetter = Callable[[], int | float | None] RequestSizeLimitEnabledGetter = Callable[[], bool] +class _RequestHeadersFields(Protocol): + def get(self, key: Literal["headers"], /) -> Iterable[tuple[bytes, bytes]] | None: ... + + +class _RequestBodyFields(Protocol): + def get(self, key: Literal["body"], default: bytes, /) -> bytes: ... + + +def _request_body_fields(message: _RequestBodyFields) -> _RequestBodyFields: + return message + + class RequestEntityTooLarge(HTTPException): def __init__(self, max_request_size_mb: float) -> None: super().__init__(status_code=413, detail=f"Request size is too large. Max size is {max_request_size_mb} MB") @@ -100,7 +112,7 @@ class RequestSizeLimitMiddleware: if message["type"] != "http.request": return message - received_body_bytes += len(message.get("body", b"")) + received_body_bytes += len(_request_body_fields(message).get("body", b"")) if received_body_bytes > max_request_size_bytes: raise RequestEntityTooLarge(max_request_size_mb) return message @@ -128,7 +140,7 @@ def _mb_to_bytes(max_request_size_mb: float | None) -> int | None: return int(max_request_size_mb * 1024 * 1024) -def _get_content_length(scope: Scope) -> int | None: +def _get_content_length(scope: _RequestHeadersFields) -> int | None: headers: Final = dict(scope.get("headers") or []) raw_content_length: Final = headers.get(b"content-length") if raw_content_length is None: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index e5e68a930c9..a81fced9021 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -25,7 +25,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.llms.base_llm.managed_resources.utils import is_base64_encoded_unified_id from litellm.proxy._types import ProxyException from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -1567,7 +1567,7 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: if not prisma_client: return for attr in ("output_file_id", "error_file_id"): - raw_id = getattr(response, attr, None) + raw_id = attribute_or_none(response, attr) if not raw_id or is_base64_encoded_unified_file_id(raw_id): continue try: diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py index 3b809688d6c..fb67cc93e37 100644 --- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py +++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast from fastapi.responses import StreamingResponse @@ -12,6 +12,14 @@ if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging +class _CallIdFields(Protocol): + def get(self, key: Literal["litellm_call_id"], default: str, /) -> str: ... + + +def _call_id_fields(data: _CallIdFields) -> _CallIdFields: + return data + + class FileContentStreamingHandler: @staticmethod def resolve_streaming_request_params( @@ -77,7 +85,7 @@ class FileContentStreamingHandler: async for chunk in stream_iterator: yield chunk await proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" + litellm_call_id=_call_id_fields(data).get("litellm_call_id", ""), status="success" ) except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 8b8d9ec81be..df51740b553 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -11,7 +11,7 @@ import itertools import traceback from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass -from typing import Any, BinaryIO, Final, NoReturn, TypedDict, cast, get_args +from typing import Any, BinaryIO, Final, NoReturn, Protocol, TypedDict, cast, get_args import httpx from fastapi import ( @@ -35,7 +35,7 @@ from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, ) -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, has_attributes from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import build_list_page @@ -185,6 +185,11 @@ _LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject]) _REQUEST_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +class _CarriesPrismaClient(Protocol): + @property + def prisma_client(self) -> PrismaClient: ... + + class UploadedFileInfo(TypedDict): filename: ReadOnly[str | None] content_type: ReadOnly[str | None] @@ -1067,8 +1072,11 @@ async def get_file_content( ) # Check if file is stored in a storage backend (check DB) - if hasattr(managed_files_obj, "prisma_client") and getattr(managed_files_obj, "prisma_client", None): - prisma_client: Final[PrismaClient] = getattr(managed_files_obj, "prisma_client") + if ( + has_attributes(managed_files_obj, _CarriesPrismaClient, "prisma_client") + and managed_files_obj.prisma_client + ): + prisma_client: Final = managed_files_obj.prisma_client db_file: Final = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": file_id} ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a2e2e36d458..f1a9e282bc9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2510,21 +2510,23 @@ async def _relay_router_model( user_api_key_dict: UserAPIKeyAuth, ) -> Response: try: - result: Final = await llm_router.allm_passthrough_route( - model=model, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict, request), + result: Final = _DECODED_JSON.validate_python( + await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict, request), + ) ) except httpx.HTTPStatusError as upstream_error: return await _relay_upstream_response(upstream_error.response) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index c1fc6479421..dafbf6f74b6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE from litellm.litellm_core_utils.core_helpers import ( + json_contains, native_messages_reported_cost, override_stamped_response_cost, response_cost_in_hidden_params, @@ -58,6 +59,7 @@ else: EndpointType = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) def _event_payload(event: str) -> Mapping[str, object] | None: @@ -609,10 +611,10 @@ class AnthropicPassthroughLoggingHandler: from litellm.types.utils import Choices, SpecialEnums try: - _json_response: Final = httpx_response.json() + _json_response: Final = _DECODED_JSON.validate_python(httpx_response.json()) # Only handle successful batch job creation (POST requests with 201 status) - if httpx_response.status_code == 200 and "id" in _json_response: + if httpx_response.status_code == 200 and json_contains(_json_response, "id"): # Transform Anthropic response to LiteLLM batch format anthropic_batches_config: Final = AnthropicBatchesConfig() litellm_batch_response: Final = anthropic_batches_config.transform_retrieve_batch_response( @@ -626,7 +628,7 @@ class AnthropicPassthroughLoggingHandler: litellm_batch_response.status = "validating" # Extract batch ID from the response - batch_id: Final = _json_response.get("id", "") + batch_id: Final = litellm_batch_response.id # Get model from request body (batch response doesn't include model) request_body = request_body or {} diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 42a7c458ef2..00e2aeb9e6e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,7 +6,16 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Callable, + Coroutine, + Iterable, + Mapping, + MutableMapping, + Sequence, +) from dataclasses import dataclass from datetime import datetime from functools import reduce @@ -278,6 +287,11 @@ async def set_env_variables_in_header(custom_headers: dict | None) -> dict | Non return headers +class _AdapterCompletionRouter(Protocol): + @property + def aadapter_completion(self) -> Callable[..., Coroutine[object, object, object]]: ... + + async def chat_completion_pass_through_endpoint( fastapi_response: Response, request: Request, @@ -368,7 +382,8 @@ async def chat_completion_pass_through_endpoint( and data["model"] is not None and llm_router.registry().resolve(data["model"], kind="recognized") ): - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + adapter_router = require_attributes(llm_router, _AdapterCompletionRouter, "aadapter_completion") + llm_response = asyncio.create_task(adapter_router.aadapter_completion(**data)) elif ( llm_router is not None and data["model"] not in router_model_names @@ -377,11 +392,13 @@ async def chat_completion_pass_through_endpoint( or len(llm_router.registry().patterns().patterns) > 0 ) ): # check for wildcard routes or default deployment before checking deployment_names - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + adapter_router = require_attributes(llm_router, _AdapterCompletionRouter, "aadapter_completion") + llm_response = asyncio.create_task(adapter_router.aadapter_completion(**data)) elif llm_router is not None and data["model"] in llm_router.registry().names( kind="deployment_names" ): # model in router deployments, calling a specific deployment on the router (lowest priority) - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data, specific_deployment=True)) + adapter_router = require_attributes(llm_router, _AdapterCompletionRouter, "aadapter_completion") + llm_response = asyncio.create_task(adapter_router.aadapter_completion(**data, specific_deployment=True)) elif user_model is not None: # `litellm --model ` llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) else: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f5db97e7a3..b01dcfd2f03 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -394,6 +394,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_parent_otel_span_from_kwargs, iter_items, + json_contains, normalize_drop_params, require_attributes, ) @@ -2532,6 +2533,7 @@ class _InvitationAcceptTxManager(Protocol): exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, + /, ) -> bool | None: ... @@ -13806,7 +13808,7 @@ async def moderations( llm_router=llm_router, user_model=user_model, ) - response: Final = await llm_call + response: Final = _UNTYPED_VALUE.validate_python(await llm_call) ### ALERTING ### asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) @@ -13856,14 +13858,20 @@ async def moderations( ) +class _SpeechBytes(Protocol): + def aiter_bytes(self, chunk_size: int | None = None) -> Awaitable[AsyncIterator[bytes]]: ... + + async def _audio_speech_chunk_generator( - _response: HttpxBinaryResponseContent, + _response: object, ) -> AsyncGenerator[bytes, None]: # chunk_size has a big impact on latency, it can't be too small or too large # too small: latency is high # too large: latency is low, but memory usage is high # 8192 is a good compromise - _generator: Final = await _response.aiter_bytes(chunk_size=AUDIO_SPEECH_CHUNK_SIZE) + _generator: Final = await require_attributes(_response, _SpeechBytes, "aiter_bytes").aiter_bytes( + chunk_size=AUDIO_SPEECH_CHUNK_SIZE + ) async for chunk in _generator: yield chunk @@ -13924,7 +13932,7 @@ async def audio_speech( llm_router=llm_router, user_model=user_model, ) - response: Final = await llm_call + response: Final = _UNTYPED_VALUE.validate_python(await llm_call) ### ALERTING ### asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) @@ -14088,7 +14096,7 @@ async def audio_transcriptions( llm_router=llm_router, user_model=user_model, ) - response: Final = await llm_call + response: Final = _UNTYPED_VALUE.validate_python(await llm_call) except Exception as e: raise e finally: @@ -14665,7 +14673,7 @@ async def get_assistants( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.aget_assistants(**data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.aget_assistants(**data)) ### ALERTING ### asyncio.create_task( @@ -14741,7 +14749,7 @@ async def create_assistant( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.acreate_assistants(**data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.acreate_assistants(**data)) ### ALERTING ### asyncio.create_task( @@ -14816,7 +14824,9 @@ async def delete_assistant( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.adelete_assistant(assistant_id=assistant_id, **data) + response: Final = _UNTYPED_VALUE.validate_python( + await llm_router.adelete_assistant(assistant_id=assistant_id, **data) + ) ### ALERTING ### asyncio.create_task( @@ -14891,7 +14901,7 @@ async def create_threads( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.acreate_thread(**data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.acreate_thread(**data)) ### ALERTING ### asyncio.create_task( @@ -14964,7 +14974,7 @@ async def get_thread( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.aget_thread(thread_id=thread_id, **data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.aget_thread(thread_id=thread_id, **data)) ### ALERTING ### asyncio.create_task( @@ -15041,7 +15051,7 @@ async def add_messages( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.a_add_message(thread_id=thread_id, **data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.a_add_message(thread_id=thread_id, **data)) ### ALERTING ### asyncio.create_task( @@ -15114,7 +15124,7 @@ async def get_messages( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - response: Final = await llm_router.aget_messages(thread_id=thread_id, **data) + response: Final = _UNTYPED_VALUE.validate_python(await llm_router.aget_messages(thread_id=thread_id, **data)) ### ALERTING ### asyncio.create_task( @@ -15194,7 +15204,9 @@ async def run_thread( if "stream" in data and data["stream"] is True: # use generate_responses to stream responses async def produce_run_stream() -> StreamingResponse | JSONResponse: - run_stream: Final = await router.arun_thread(thread_id=thread_id, **data) + run_stream: Final = _UNTYPED_VALUE.validate_python( + await router.arun_thread(thread_id=thread_id, **data) + ) return await create_response( generator=async_assistants_data_generator( user_api_key_dict=user_api_key_dict, @@ -15222,7 +15234,7 @@ async def run_thread( on_late_failure=audit_late_failure, ) - response: Final = await router.arun_thread(thread_id=thread_id, **data) + response: Final = _UNTYPED_VALUE.validate_python(await router.arun_thread(thread_id=thread_id, **data)) ### ALERTING ### asyncio.create_task( @@ -16175,11 +16187,11 @@ async def _count_request_tokens( tools: Final = request.tools system: Final = request.system - deployment: dict[str, Any] | None = None if routed is None else dict(routed) + deployment: dict[str, object] | None = None if routed is None else dict(routed) litellm_model_name = None model_info: Final = None if routed is None else routed.get("model_info", {}) if deployment is not None: - litellm_model_name = deployment.get("litellm_params", {}).get("model") + litellm_model_name = _with_get(deployment.get("litellm_params", _EMPTY_MAPPING)).get("model") deployment = { # rebind-ok: merge the named credential into a copy **deployment, "litellm_params": _with_chat_timeout( @@ -16187,8 +16199,8 @@ async def _count_request_tokens( ), } # remove the custom_llm_provider_prefix in the litellm_model_name - if "/" in litellm_model_name: - litellm_model_name = litellm_model_name.split("/", 1)[1] + if json_contains(litellm_model_name, "/"): + litellm_model_name = _split_once(litellm_model_name, "/")[1] model_to_use: str = ( litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model @@ -17071,6 +17083,14 @@ class _SupportsLower(Protocol): def lower(self) -> object: ... +class _SupportsSplit(Protocol): + def split(self, sep: str, maxsplit: int, /) -> Sequence[object]: ... + + +def _split_once(value: object, sep: str) -> Sequence[object]: + return require_attributes(value, _SupportsSplit, "split").split(sep, 1) + + def _lowered(value: object) -> object: return require_attributes(value, _SupportsLower, "lower").lower() @@ -17239,18 +17259,19 @@ async def _apply_search_filter_to_models( def _is_byok_outside_caller_teams(model_info_dict: object) -> bool: return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) - def _model_matches_search(m: dict[str, Any]) -> bool: + def _model_matches_search(m: Mapping[str, object]) -> bool: # Team BYOK models persist an internal `model_name` # (e.g. `model_name_{team_id}_{uuid}`) and expose the user-facing # name via `model_info.team_public_model_name`. Match both so the # name shown in the UI is searchable. - if search_lower in (m.get("model_name") or "").lower(): + if json_contains(_lowered(m.get("model_name") or ""), search_lower): return True - team_public_model_name: Final = (m.get("model_info") or {}).get("team_public_model_name") or "" - if search_lower in team_public_model_name.lower(): + team_public_model_name: Final = ( + _with_get(m.get("model_info") or _EMPTY_MAPPING).get("team_public_model_name") or "" + ) + if json_contains(_lowered(team_public_model_name), search_lower): return True - litellm_params: Final[Mapping[str, object]] = m.get("litellm_params") or MappingProxyType({}) - litellm_model: Final = litellm_params.get("model") + litellm_model: Final = _with_get(m.get("litellm_params") or _EMPTY_MAPPING).get("model") return isinstance(litellm_model, str) and search_lower in litellm_model.lower() # Filter models in router by search term, dropping BYOK rows that @@ -19481,7 +19502,7 @@ async def async_queue_request( return StreamingResponse( async_data_generator( user_api_key_dict=user_api_key_dict, - response=await router.schedule_acompletion(**data), + response=_UNTYPED_VALUE.validate_python(await router.schedule_acompletion(**data)), request_data=data, request=request, ), @@ -19500,7 +19521,7 @@ async def async_queue_request( on_late_failure=audit_late_failure, ) - response: Final = await router.schedule_acompletion(**data) + response: Final = _UNTYPED_VALUE.validate_python(await router.schedule_acompletion(**data)) fastapi_response.headers.update({"x-litellm-priority": str(queued_priority)}) return response except Exception as e: @@ -20147,7 +20168,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): ) if user_obj and hasattr(user_obj, "__dict__"): - user_obj.__dict__.pop("password", None) + vars(user_obj).pop("password", None) # The password just changed via an invitation/reset link; any UI session # minted under the old password may be in hostile hands. Revoke them all — diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 47050f01644..38041ed1862 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse, StreamingResponse from pydantic import TypeAdapter from starlette.datastructures import UploadFile +from typing_extensions import TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -91,6 +92,10 @@ _EMPTY_TEAM_PROVIDER_CREDENTIALS: Final[Mapping[str, object]] = MappingProxyType _TEAM_CREDENTIAL_PROVIDER: Final = "openai" +class _ForwardedRequestFields(TypedDict, total=False, extra_items=object): + pass + + def _as_string_keyed_mapping(value: object) -> Mapping[str, object] | None: if isinstance(value, Mapping): return value @@ -904,7 +909,7 @@ async def rag_ingest( vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None if vertex_target_error is not None: raise HTTPException(status_code=400, detail={"error": vertex_target_error}) - request_data: Final[dict[str, Any]] = await add_litellm_data_to_request( + request_data: Final[dict[str, object]] = await add_litellm_data_to_request( data={}, request=request, general_settings=general_settings, @@ -955,7 +960,7 @@ async def rag_ingest( file_url=file_url, file_id=ingest_file_id, router=llm_router, - **request_data, + **_ForwardedRequestFields(**request_data), ) if ingest_response.get("status") == "failed": raise HTTPException( diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 4905751a85d..b80220b1aac 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -6,6 +6,7 @@ from typing import Final import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -23,6 +24,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( from litellm.types.utils import CallTypes router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) @router.post( @@ -85,7 +87,7 @@ async def rerank( llm_router=llm_router, user_model=user_model, ) - response: Final = await llm_call + response: Final = _DECODED_JSON.validate_python(await llm_call) ### ALERTING ### asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 11a0b947abe..369d167f623 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2686,7 +2686,7 @@ async def _responses_websocket_session( llm_router=llm_router, user_model=user_model, ) - outcome: Final = await llm_call + outcome: Final = _DECODED_JSON.validate_python(await llm_call) except asyncio.CancelledError: if refused_turns.first is not None: _ = bill_socket_failure(refused_turns.first) diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 0a13e3a9e99..11be7a736ed 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -21,6 +21,13 @@ if TYPE_CHECKING: _RESPONSE_POLLING_TARGET: Final = "response_polling" _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_LIST: Final = TypeAdapter(list[object]) +_CACHED_VALUE: Final = TypeAdapter(object) + + +def _json_document(value: object) -> str | bytes | bytearray: + if isinstance(value, (str, bytes, bytearray)): + return value + raise TypeError(f"the JSON object must be str, bytes or bytearray, not {value.__class__.__name__}") class ResponsePollingHandler: @@ -156,13 +163,13 @@ class ResponsePollingHandler: cache_key: Final = self.get_cache_key(polling_id) # Get current state - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) + cached_state: Final = _CACHED_VALUE.validate_python(await self.redis_cache.async_get_cache(cache_key)) if not cached_state: verbose_proxy_logger.warning("No cached state found for polling_id: %s", polling_id) return # Parse existing ResponsesAPIResponse from cache - state: Final = _JSON_OBJECT.validate_python(json.loads(cached_state)) + state: Final = _JSON_OBJECT.validate_python(json.loads(_json_document(cached_state))) # Update status (using OpenAI native status values) if status: @@ -239,10 +246,10 @@ class ResponsePollingHandler: return None cache_key: Final = self.get_cache_key(polling_id) - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) + cached_state: Final = _CACHED_VALUE.validate_python(await self.redis_cache.async_get_cache(cache_key)) if cached_state: - return _JSON_OBJECT.validate_python(json.loads(cached_state)) + return _JSON_OBJECT.validate_python(json.loads(_json_document(cached_state))) return None diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index a2c4066a14f..dee8c8d5c60 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -7,6 +7,7 @@ from datetime import datetime from typing import Any, Final, TypeAlias from fastapi import APIRouter, Depends, HTTPException +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import UI_SESSION_TOKEN_TEAM_ID @@ -68,7 +69,18 @@ async def _refresh_router_search_tools() -> None: verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e) -def _with_loaded_tools_where_undecryptable(db_search_tools: Sequence[dict[str, object]]) -> list[dict[str, Any]]: +class _ListedSearchTool(TypedDict, total=False): + search_tool_id: ReadOnly[str | None] + search_tool_name: ReadOnly[Required[str]] + litellm_params: ReadOnly[Required[Mapping[str, object]]] + search_tool_info: ReadOnly[dict[object, object] | None] + created_at: ReadOnly[str | None] + updated_at: ReadOnly[str | None] + + +def _with_loaded_tools_where_undecryptable( + db_search_tools: Sequence[_ListedSearchTool], +) -> Sequence[_ListedSearchTool]: from litellm.proxy.proxy_server import llm_router kept_search_tools: Final = keep_loaded_search_tools_that_do_not_decrypt( diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index c8cd1100e08..39510076d22 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -5,7 +5,7 @@ Search Tool Registry for managing search tool configurations. import os from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Final, Protocol +from typing import Final, Protocol, TypeVar from pydantic import TypeAdapter, ValidationError @@ -172,9 +172,13 @@ def _did_not_decrypt(search_tool: Mapping[str, object]) -> bool: ) +_DbSearchToolT: Final = TypeVar("_DbSearchToolT", bound=Mapping[str, object]) +_LoadedSearchToolT: Final = TypeVar("_LoadedSearchToolT", bound=Mapping[str, object]) + + def keep_loaded_search_tools_that_do_not_decrypt( - db_search_tools: Sequence[Mapping[str, object]], loaded_search_tools: Sequence[Mapping[str, object]] -) -> Sequence[Mapping[str, object]]: + db_search_tools: Sequence[_DbSearchToolT], loaded_search_tools: Sequence[_LoadedSearchToolT] +) -> Sequence[_DbSearchToolT | _LoadedSearchToolT]: """Replace each DB search tool whose params do not decrypt with the current key by its loaded version.""" loaded_by_id: Final = {tool.get("search_tool_id"): tool for tool in loaded_search_tools} kept: Final = tuple( diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index b3b5e41328d..f48f853ba2a 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,10 +1,11 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet +from contextlib import AbstractAsyncContextManager from dataclasses import dataclass from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, Literal, TypeVar +from typing import Final, Literal, Protocol, TypeVar from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -168,6 +169,16 @@ async def _db_or_empty( return None +class _SpendLogTransaction(Protocol): + async def execute_raw(self, query: str, /) -> int: ... + + async def query_raw(self, query: str, /, *args: object) -> object: ... + + +def _spend_log_transaction(prisma_client: PrismaClient) -> AbstractAsyncContextManager[_SpendLogTransaction]: + return prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) + + async def _rows_within_the_statement_timeout( prisma_client: PrismaClient, sql: str, @@ -177,7 +188,7 @@ async def _rows_within_the_statement_timeout( ) -> Sequence[Mapping[str, object]]: async with ( db_span("recover_key_metadata", table), - prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction, + _spend_log_transaction(prisma_client) as transaction, ): await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) for setting in planner_settings: diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index e4a98619bfd..86e8ef31cf4 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -18,7 +18,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone from itertools import chain, groupby -from types import MappingProxyType +from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter @@ -71,13 +71,20 @@ class _SentinelBatchTable(Protocol): class _SentinelBatch(Protocol): - litellm_dailyteamspend: _SentinelBatchTable + @property + def litellm_dailyteamspend(self) -> _SentinelBatchTable: ... class _SentinelBatchManager(Protocol): async def __aenter__(self) -> _SentinelBatch: ... - async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + /, + ) -> bool | None: ... class _SupportsGet(Protocol): diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 5b961f1320b..2de589564ac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -3,7 +3,7 @@ import asyncio import json import os from collections import Counter -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import ItemsView, Mapping, MutableMapping, Sequence from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -1306,6 +1306,16 @@ async def _load_environment_variables_param( return param_value +def _submitted_sso_values(submitted: Mapping[str, object]) -> ItemsView[str, object]: + return submitted.items() + + +def _environment_value(value: object) -> str: + if isinstance(value, str): + return value + raise TypeError(f"str expected, not {type(value).__name__}") + + def _decrypted_sso_fields(environment_variables: Mapping[str, object]) -> Mapping[str, str | None]: return MappingProxyType( { @@ -1398,11 +1408,11 @@ async def update_sso_settings( field_name: submitted[field_name] if field_name in submitted else fallback_values.get(field_name) for field_name in SSOConfig.model_fields } - for field_name, value in submitted.items(): + for field_name, value in _submitted_sso_values(submitted): if field_name in SSO_FIELD_ENV_VARS: env_var_name = SSO_FIELD_ENV_VARS[field_name] if value: - os.environ[env_var_name] = value + os.environ[env_var_name] = _environment_value(value) else: # Clear environment variable if value is null/empty os.environ.pop(env_var_name, None) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eac2700551f..874fcbc8d09 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -169,6 +169,7 @@ from litellm.litellm_core_utils.core_helpers import ( attribute_or_none, coerce_token_limit, get_or_create_metadata_bucket, + has_attributes, independent_snapshot, is_expected_client_error, require_attributes, @@ -2373,10 +2374,12 @@ class ProxyLogging: modified_args = self._extract_modified_arguments_from_content(new_content, request_obj) if modified_args is not None: # Return the masked/redacted arguments for the MCP call to use - return MCPPreCallResponseObject( - should_proceed=True, - error_message=None, - modified_arguments=modified_args, + return MCPPreCallResponseObject.model_validate( + { + "should_proceed": True, + "error_message": None, + "modified_arguments": modified_args, + } ) else: # Could not parse modified arguments, allow original call but warn @@ -2395,7 +2398,7 @@ class ProxyLogging: return None - def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> dict | None: + def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> object: """ Extract modified/masked arguments from the guardrail response content. """ @@ -2416,7 +2419,7 @@ class ProxyLogging: # Try to parse as JSON first try: - modified_args = json.loads(args_text) + modified_args = _DECODED_JSON.validate_python(json.loads(args_text)) verbose_proxy_logger.debug("Successfully parsed JSON args: %s", modified_args) return modified_args except json.JSONDecodeError as e: @@ -5643,8 +5646,13 @@ _LOCK_TEAMS_BY_ID_SQL: Final = ( ) -def _user_unique_conflict_message(metadata: object, user_id: object) -> str: - match metadata: +class _UniqueConflictMetadata(Protocol): + @property + def meta(self) -> object: ... + + +def _user_unique_conflict_message(error: _UniqueConflictMetadata, user_id: object) -> str: + match error.meta: case {"target": ["user_id"]}: return f"User with id {user_id} already exists" case {"target": ["sso_user_id"]}: @@ -5658,7 +5666,7 @@ def _raise_user_unique_conflict(error: "UniqueViolationError", user_id: object, raise error raise HTTPException( status_code=409, - detail={"error": _user_unique_conflict_message(error.meta, user_id)}, + detail={"error": _user_unique_conflict_message(error, user_id)}, ) from error @@ -8338,6 +8346,10 @@ class PrismaClient: ### HELPER FUNCTIONS ### +class _JsonDumpableRow(Protocol): + def model_dump_json(self) -> object: ... + + async def cache_user_row(user_id: str, cache: DualCache, db: PrismaClient) -> None: """ Check if a user_id exists in cache, @@ -8349,8 +8361,8 @@ async def cache_user_row(user_id: str, cache: DualCache, db: PrismaClient) -> No user_row: Final = await db.get_data(user_id=user_id) if user_row is not None: print_verbose(f"User Row: {user_row}, type = {type(user_row)}") - if hasattr(user_row, "model_dump_json") and callable(getattr(user_row, "model_dump_json")): - cache_value: Final[str] = user_row.model_dump_json() + if has_attributes(user_row, _JsonDumpableRow, "model_dump_json") and callable(user_row.model_dump_json): + cache_value: Final = user_row.model_dump_json() cache.set_cache(key=cache_key, value=cache_value, ttl=600) # store for 10 minutes diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 4b9547300db..62af5663ea9 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,6 +10,7 @@ All /vector_store management endpoints import copy import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from fastapi import APIRouter, Depends, HTTPException @@ -54,8 +55,16 @@ def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorS return ManagedVectorStoresRepository(prisma_client).table +class _VectorStoreRowFields(TypedDict, total=False, extra_items=object): + pass + + +def _dumped_row(row: "_VectorStoreRow") -> Mapping[str, object]: + return row.model_dump() + + def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: - return LiteLLM_ManagedVectorStore(**row.model_dump()) + return LiteLLM_ManagedVectorStore(**_VectorStoreRowFields(**_dumped_row(row))) class _ConfigOwnedDetail(TypedDict): diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 9004d1951b8..4044e50f8c1 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -2,10 +2,11 @@ import json import re from collections.abc import Iterable, Mapping from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, Protocol from fastapi import HTTPException, Request from pydantic import ConfigDict, TypeAdapter +from typing_extensions import TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -27,6 +28,18 @@ from litellm.utils import ProviderConfigManager _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +class _ManagedVectorStoreFields(TypedDict, total=False, extra_items=object): + pass + + +class _DumpableRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +def _managed_vector_store_from_row(row: _DumpableRow) -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(**_ManagedVectorStoreFields(**row.model_dump())) + + def normalize_litellm_params( vector_store: LiteLLM_ManagedVectorStore, ) -> LiteLLM_ManagedVectorStore: @@ -316,7 +329,7 @@ async def get_litellm_managed_vector_store( ) if not rows: return None - return normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) + return normalize_litellm_params(_managed_vector_store_from_row(rows[0])) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from shared cache: %s", diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 24fd443b68a..8cca272b0b3 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,8 +1,9 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import orjson from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import video_access_group from litellm.proxy.openai_files_endpoints.common_utils import authorize_model_for_key @@ -63,7 +64,11 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: return None -def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> object: +class _CarriesId(Protocol): + id: object + + +def encode_character_id_in_response(response: object, custom_llm_provider: str, model_id: str | None) -> object: if isinstance(response, dict) and response.get("id"): response["id"] = encode_character_id_with_provider( character_id=response["id"], @@ -72,7 +77,9 @@ def encode_character_id_in_response(response: Any, custom_llm_provider: str, mod ) return response - character_id: Final = getattr(response, "id", None) + if not has_attributes(response, _CarriesId, "id"): + return response + character_id: Final = response.id if isinstance(character_id, str) and character_id: response.id = encode_character_id_with_provider( character_id=character_id, diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 8b2cf6e0a4e..0d75fe4e531 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -14,12 +14,13 @@ from __future__ import annotations import base64 from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid4 from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE +from litellm.litellm_core_utils.core_helpers import has_attributes, require_attributes from litellm.litellm_core_utils.url_utils import async_safe_get from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -31,6 +32,8 @@ from litellm.rag.text_splitters import RecursiveCharacterTextSplitter from litellm.types.rag import RAGIngestOptions, RAGIngestResponse if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Iterable + from litellm import Router @@ -38,6 +41,21 @@ class RAGIngestBadRequestError(ValueError): status_code: Final = 400 +class _OcrRouter(Protocol): + @property + def aocr(self) -> Callable[..., Awaitable[object]]: ... + + +class _OcrPages(Protocol): + @property + def pages(self) -> Iterable[object]: ... + + +class _OcrPage(Protocol): + @property + def markdown(self) -> str: ... + + class BaseRAGIngestion(ABC): """ Base class for RAG ingestion. @@ -168,7 +186,8 @@ class BaseRAGIngestion(ABC): # Use router if available if self.router is not None: - ocr_response = await self.router.aocr( + ocr_router: Final = require_attributes(self.router, _OcrRouter, "aocr") + ocr_response = await ocr_router.aocr( model=ocr_model, document={"type": doc_type, url_key: data_url}, ) @@ -179,8 +198,10 @@ class BaseRAGIngestion(ABC): ) # Extract text from pages - if hasattr(ocr_response, "pages") and ocr_response.pages: - return "\n\n".join(page.markdown for page in ocr_response.pages if hasattr(page, "markdown")) + if has_attributes(ocr_response, _OcrPages, "pages") and ocr_response.pages: + return "\n\n".join( + page.markdown for page in ocr_response.pages if has_attributes(page, _OcrPage, "markdown") + ) return None diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 85d87554023..9bd8148b286 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -21,7 +21,7 @@ from litellm.constants import ( request_timeout, ) from litellm.exceptions import BadRequestError -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, keyword_arguments from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig @@ -394,7 +394,7 @@ async def acreate_realtime_client_secret( async def acreate_realtime_translation_client_secret( model: str | None = None, session: Mapping[str, Any] | None = None, - expires_after: Mapping[str, Any] | None = None, + expires_after: Mapping[str, object] | None = None, timeout: float | None = None, **kwargs: object, # kwargs-ok: public API preserves provider-specific request options ) -> httpx.Response: @@ -409,7 +409,11 @@ async def acreate_realtime_translation_client_secret( req: Final = RealtimeClientSecretRequest( model=requested_model_name, session=session_config, - expires_after=RealtimeExpiresAfter(**expires_after) if expires_after else None, + expires_after=RealtimeExpiresAfter( + **_ExpiresAfterFields(**keyword_arguments(expires_after, RealtimeExpiresAfter)) + ) + if expires_after + else None, ) litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") if not isinstance(litellm_logging_obj, LiteLLMLogging): @@ -488,10 +492,18 @@ def _deployment_base_model(kwargs: Mapping[str, object], custom_llm_provider: st return None if base_model is None else base_model.removeprefix(f"{custom_llm_provider}/") +class _ExpiresAfterFields(TypedDict, total=False, extra_items=object): + pass + + +class _TranscriptionSessionFields(TypedDict, total=False, extra_items=object): + pass + + @wrapper_client async def acreate_realtime_transcription_session( model: str | None = None, - transcription_session: dict[str, Any] | None = None, + transcription_session: Mapping[str, object] | None = None, timeout: float | None = None, **kwargs, ): @@ -506,7 +518,9 @@ async def acreate_realtime_transcription_session( """ req: Final = RealtimeTranscriptionSessionRequest( model=model, - **(transcription_session or {}), + **_TranscriptionSessionFields( + **keyword_arguments(transcription_session or {}, RealtimeTranscriptionSessionRequest) + ), ) model_name = req.resolved_model() or "gpt-realtime-whisper" litellm_logging_obj: Final[LiteLLMLogging] = kwargs.get("litellm_logging_obj") @@ -626,7 +640,7 @@ async def arealtime_translation_calls( openai_ephemeral_key: str, sdp_body: bytes, model: str | None = None, - session: Mapping[str, Any] | None = None, + session: Mapping[str, object] | None = None, timeout: float | None = None, **kwargs: object, # kwargs-ok: public API preserves provider-specific request options ) -> httpx.Response: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index fbef73a6114..3c3b87c9b8c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -34,6 +34,16 @@ class _PrismaWriterClientView(Protocol): def writer_db(self) -> object: ... +class _LedgerWriterClientView(Protocol): + @property + def writer_db(self) -> LedgerDatabase: ... + + +class _RawQueryWriterClientView(Protocol): + @property + def writer_db(self) -> LedgerTransaction: ... + + class PrismaTableRepository(Generic[RowT_co]): """Base for repositories that expose a single Prisma table.""" @@ -143,7 +153,7 @@ class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs ) if not ledger: return legacy_count - database: Final[LedgerDatabase] = self.prisma_client.writer_db + database: Final = require_attributes(self._connected_client(), _LedgerWriterClientView, "writer_db").writer_db async with database.tx(timeout=timedelta(seconds=60)) as transaction: accepted: Final = await self._insert_ledger(transaction, ledger, write_spend_logs=write_spend_logs) return legacy_count + accepted @@ -211,7 +221,7 @@ class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs return frozenset(str(row["request_id"]) for row in rows) async def has_delivery(self, delivery_id: str) -> bool: - database: Final[LedgerTransaction] = self.prisma_client.writer_db + database: Final = require_attributes(self._connected_client(), _RawQueryWriterClientView, "writer_db").writer_db rows: Final = await database.query_raw( 'SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = $1', delivery_id ) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index e344ce46be7..8c2a0a36212 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -188,14 +188,16 @@ def rerank( model_response: Final = RerankResponse() - rerank_litellm_params: Final = { - "litellm_call_id": litellm_call_id, - "proxy_server_request": proxy_server_request, - "model_info": model_info, - "preset_cache_key": None, - "stream_response": {}, - **optional_params.model_dump(exclude_unset=True), - } + rerank_litellm_params: Final = dict[str, object]( + { + "litellm_call_id": litellm_call_id, + "proxy_server_request": proxy_server_request, + "model_info": model_info, + "preset_cache_key": None, + "stream_response": {}, + **optional_params.model_dump(exclude_unset=True), + } + ) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c3166b26a74..b2c98282da8 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_allowed_tools_tool_choice, tool_calls_the_choice_excludes, ) +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy.mcp_server.utils import ( @@ -1178,7 +1179,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def should_raise_mcp_guardrail_block(error: Exception) -> bool: - return getattr(error, "mcp_block_behavior", None) == "raise" + return attribute_or_none(error, "mcp_block_behavior") == "raise" @staticmethod def create_follow_up_messages_for_chat( diff --git a/litellm/router.py b/litellm/router.py index 98b016f590f..2bf4db967fd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -234,11 +234,29 @@ def _class_descriptor(cls: type, name: str) -> bool: if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Awaitable from typing import ParamSpec, overload + import httpx + from openai.pagination import AsyncCursorPage + from openai.types.beta.assistant import Assistant + from openai.types.beta.assistant_deleted import AssistantDeleted + from openai.types.beta.threads.message import Message as OpenAIMessage + from openai.types.file_deleted import FileDeleted + + from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.types.llms.openai import AllMessageValues, OpenAIFileObject - from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, ModelResponse + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse + from litellm.types.llms.openai import ( + AllMessageValues, + HttpxBinaryResponseContent, + OpenAIFileObject, + ResponsesAPIResponse, + Thread, + ) + from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, LiteLLMFineTuningJob, ModelResponse + from litellm.types.vector_stores import VectorStoreCreateResponse _P: Final = ParamSpec("_P") @@ -281,6 +299,28 @@ if TYPE_CHECKING: async def acreate_file(self, model: str, **kwargs: object) -> OpenAIFileObject: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion async def acreate_batch(self, model: str, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + + aget_assistants: Callable[..., Awaitable[AsyncCursorPage[Assistant]]] + acreate_assistants: Callable[..., Awaitable[Assistant]] + adelete_assistant: Callable[..., Awaitable[AssistantDeleted]] + acreate_thread: Callable[..., Awaitable[Thread]] + aget_thread: Callable[..., Awaitable[Thread]] + a_add_message: Callable[..., Awaitable[OpenAIMessage]] + aget_messages: Callable[..., Awaitable[AsyncCursorPage[OpenAIMessage]]] + arun_thread: Callable[..., Awaitable[object]] + adecisions: Callable[..., Awaitable[OpenAIDecisionResponse | DecisionsResponse]] + acreate_fine_tuning_job: Callable[..., Awaitable[LiteLLMFineTuningJob]] + aretrieve_fine_tuning_job: Callable[..., Awaitable[LiteLLMFineTuningJob]] + acancel_fine_tuning_job: Callable[..., Awaitable[LiteLLMFineTuningJob]] + alist_fine_tuning_jobs: Callable[..., Awaitable[object]] + alist_batches: Callable[..., Awaitable[Mapping[str, object]]] + acancel_batch: Callable[..., Awaitable[LiteLLMBatch]] + aget_responses: Callable[..., Awaitable[ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] + afile_content: Callable[..., Awaitable[HttpxBinaryResponseContent | FileContentStreamingResult]] + afile_delete: Callable[..., Awaitable[FileDeleted]] + avector_store_create: Callable[..., Awaitable[VectorStoreCreateResponse]] + allm_passthrough_route: Callable[..., Awaitable[httpx.Response | AsyncGenerator[bytes, bytes]]] + def __getattr__(self, name: str) -> Any: ... # the endpoints served by name else: diff --git a/litellm/router_api.py b/litellm/router_api.py index a75397ea43d..1b0f970a537 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -60,6 +60,8 @@ if TYPE_CHECKING: from litellm.messages.dispatch import MessagesResult from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.router import Router + from litellm.router_rust.surface.model_list import HostedModelList + from litellm.router_rust.surface.registry import UntypedEntry from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.auto_router.auto_router import AutoRouter from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -102,14 +104,25 @@ if TYPE_CHECKING: # Re-exported for litellm.types, which may not import router_utils. AttemptedFallbackTargets: TypeAlias = _AttemptedFallbackTargets + ModelGroupAliases: TypeAlias = dict[str, str | RouterModelGroupAliasItem] + RoutingGroups: TypeAlias = dict[str, RoutingGroup] + AccessGroups: TypeAlias = dict[str, list[str]] + AutoRouters: TypeAlias = dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] + ComplexityRouters: TypeAlias = dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] + QualityRouters: TypeAlias = dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] + AdaptiveRouters: TypeAlias = dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] + SwitchyardRouters: TypeAlias = dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]] + RoutingPlugins: TypeAlias = list[RoutingPlugin] + ModelNameSet: TypeAlias = set[str] + # What registry().plugins(kind) returns, one member per kind. RegistryPlugins: TypeAlias = ( - dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] - | dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] - | dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] - | dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] - | dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]] - | list[RoutingPlugin] + AutoRouters + | ComplexityRouters + | QualityRouters + | AdaptiveRouters + | SwitchyardRouters + | RoutingPlugins | RouterBudgetLimiting | DualCache | None @@ -351,6 +364,220 @@ def _json_default_stable_id(obj: object) -> str: return stable_id(obj) +class _RegistryHost(Protocol): + """The host registry members ``RouterRegistry`` reads, typed as the view hands them on.""" + + @property + def _registry_version(self) -> int: ... + + @property + def _core_model_list(self) -> "HostedModelList": ... + + @property + def _core_model_group_alias(self) -> "ModelGroupAliases": ... + + @property + def _core_routing_groups(self) -> "RoutingGroups": ... + + @property + def _core_model_names(self) -> AbstractSet[str]: ... + + @property + def _core_deployment_names(self) -> Sequence[str]: ... + + @property + def _core_team_public_model_names(self) -> frozenset[str]: ... + + @property + def _core_model_name_to_deployment_indices(self) -> Mapping[str, Sequence[int]]: ... + + @property + def _core_team_model_to_deployment_indices(self) -> Mapping[tuple[str, str], Sequence[int]]: ... + + @property + def _core_pattern_router(self) -> "PatternMatchRouter": ... + + @property + def _core_team_pattern_routers(self) -> "Mapping[str, PatternMatchRouter]": ... + + @property + def _core_cached_deployment_model_info(self) -> "_lru_cache_wrapper[ModelMapInfo | None]": ... + + @property + def _core_auto_routers(self) -> "AutoRouters": ... + + @property + def _core_complexity_routers(self) -> "ComplexityRouters": ... + + @property + def _core_quality_routers(self) -> "QualityRouters": ... + + @property + def _core_adaptive_routers(self) -> "AdaptiveRouters": ... + + @property + def _core_switchyard_routers(self) -> "SwitchyardRouters": ... + + @property + def _core_routing_plugins(self) -> "RoutingPlugins": ... + + @property + def _core_prompt_caching_pin_cache(self) -> "DualCache": ... + + def _core_get_deployment(self, model_id: str) -> "Deployment | None": ... + + def _core_has_model_id(self, candidate_id: str) -> bool: ... + + def _core_get_model_info(self, id: str) -> "UntypedEntry | None": ... + + def _core_get_deployment_model_info(self, model_id: str, model_name: str) -> "ModelMapInfo | None": ... + + def _core_get_router_model_info( + self, + deployment: "Mapping[str, object] | Deployment | None", + received_model_name: str, + id: str | None = None, + ) -> "ModelMapInfo": ... + + def _core_get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]: ... + + def _core_get_model_list( + self, model_name: str | None = None, team_id: str | None = None + ) -> Sequence[Mapping[str, object]] | None: ... + + def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: ... + + def _core_get_model_list_from_model_alias( + self, model_name: str | None = None + ) -> Sequence[Mapping[str, object]]: ... + + def _core_get_model_list_from_routing_groups( + self, model_name: str | None = None + ) -> Sequence[Mapping[str, object]]: ... + + def _core_get_all_deployments( + self, model_name: str, model_alias: str | None = None, team_id: str | None = None + ) -> Sequence[Mapping[str, object]]: ... + + def _core_deployments_named(self, model_name: str) -> tuple[Mapping[str, object], ...] | None: ... + + def _core_get_model_list_of_routed_group( + self, model_group: str, team_id: str | None = None + ) -> Sequence[Mapping[str, object]]: ... + + def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: ... + + def _core_get_routable_deployments( + self, model: str, request_kwargs: Mapping[str, object], call_type: str | None = None + ) -> tuple[Mapping[str, object], ...]: ... + + def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> "Deployment | None": ... + + def _core_deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[Mapping[str, object]]: ... + + def _core_should_include_deployment( + self, model_name: str, model: Mapping[str, object], team_id: str | None = None + ) -> bool: ... + + def _core_get_model_names(self, team_id: str | None = None) -> Sequence[str]: ... + + def _core_get_model_ids( + self, model_name: str | None = None, exclude_team_models: bool = False + ) -> Sequence[str]: ... + + def team_public_model_names_of(self, team_id: object) -> frozenset[str]: ... + + def _core_resolved_litellm_models( + self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False + ) -> tuple[str, ...]: ... + + def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: ... + + def _core_get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: ... + + def _core_resolve_model_name_from_model_id( + self, model_id: str | None, custom_llm_provider: str | None = None + ) -> str | None: ... + + def _core_map_team_model( + self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None + ) -> str | None: ... + + def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: ... + + def _core_routable_model_group(self, model_name: str) -> str: ... + + def _core_get_model_from_alias(self, model: str) -> str | None: ... + + def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: ... + + def _core_is_recognized_model(self, model: str) -> bool: ... + + def _core_get_model_group(self, id: str) -> Sequence[Mapping[str, object]] | None: ... + + def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: ... + + def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: ... + + def _core_get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: ... + + def _core_get_routable_upstream_provider(self, model_name: str, team_id: str | None = None) -> str | None: ... + + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: ... + + def _core_get_model_group_info( + self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None + ) -> "ModelGroupInfo | None": ... + + def _core_cached_model_group_info(self, model_group: str) -> "ModelGroupInfo | None": ... + + def _core_get_configured_display_name(self, model_name: str) -> str | None: ... + + def _core_get_configured_mode(self, model_name: str) -> str | None: ... + + def _core_get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: ... + + def _core_get_pinned_reasoning_efforts(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: ... + + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: ... + + def _core_get_routable_deployment_model_info( + self, model_name: str, team_id: str | None = None + ) -> "tuple[ModelMapInfo | None, ...]": ... + + def _core_get_configured_token_limits(self, model_name: str) -> tuple[int | None, int | None]: ... + + def _core_get_model_listing_info(self, model_name: str) -> "DeploymentModelListingInfo | None": ... + + def get_routing_group(self, model_name: str) -> "RoutingGroup | None": ... + + def _get_client( + self, deployment: Mapping[str, object], kwargs: Mapping[str, object], client_type: str | None = None + ) -> object: ... + + def _core_get_model_access_groups( + self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None + ) -> "AccessGroups": ... + + def _core_get_deployment_credentials(self, model_id: str) -> Mapping[str, object] | None: ... + + def _core_get_deployment_credentials_with_provider( + self, model_id: str, team_id: str | None = None + ) -> Mapping[str, object] | None: ... + + def _core_get_credential_deployment(self, model_id: str, team_id: str | None = None) -> "Deployment | None": ... + + def _core_get_router_deployment_budget_limiter(self) -> "RouterBudgetLimiting | None": ... + + def _core_decryptable_reasoning( + self, model_group: str, request_input: Sequence[object], targets: "Sequence[Deployment]" + ) -> tuple[object, ...]: ... + + class RouterRegistry: """ Read-only queries over the configured deployments (``registry()``). @@ -361,7 +588,7 @@ class RouterRegistry: __slots__ = ("_router",) - def __init__(self, router: "Router") -> None: + def __init__(self, router: _RegistryHost) -> None: self._router = router @property @@ -407,12 +634,8 @@ class RouterRegistry: id: str | None = None, ) -> "ModelMapInfo": """Model info used for routing decisions (was ``get_router_model_info``).""" - r: Final = self._router - # The core's overloads split (deployment, id) by which one is None; this forwards either form unchanged. - return r._core_get_router_model_info( # pyright: ignore[reportCallIssue] # see above - deployment=deployment, # pyright: ignore[reportArgumentType] # see above - received_model_name=received_model_name, - id=id, # pyright: ignore[reportArgumentType] # see above + return self._router._core_get_router_model_info( + deployment=deployment, received_model_name=received_model_name, id=id ) def discovered_model_info(self, id: str | None) -> Mapping[str, int]: @@ -1316,12 +1539,52 @@ class RouterSettings: return self._owner("_core_deployment_call_timeout")._core_deployment_call_timeout(litellm_params) +class _HealthHost(Protocol): + """The router members ``RouterHealth`` reads, typed as the view hands them on.""" + + @property + def _core_cache(self) -> "DualCache": ... + + @property + def _core_cooldown_cache(self) -> "CooldownCache": ... + + @property + def _core_health_state_cache(self) -> "DeploymentHealthCache": ... + + @property + def _core_deployment_latency_map(self) -> "UntypedDict": ... + + @property + def _core_ctx(self) -> "RoutingContext": ... + + @property + def _core_cooldown_time(self) -> float | None: ... + + def _core_get_fully_blocked_model_names(self) -> "ModelNameSet": ... + + async def _core_async_get_fully_unhealthy_model_names(self) -> "ModelNameSet": ... + + async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: ... + + async def _core_get_remaining_model_group_usage(self, model_group: str) -> Mapping[str, int]: ... + + def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": ... + + def _core_split_request_priority( + self, model: str, kwargs: Mapping[str, object] + ) -> tuple[int | None, Mapping[str, object]]: ... + + def priority_header_update( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: ... + + class RouterHealth: """Cooldowns, blocked deployments, health checks and usage (``health()``).""" __slots__ = ("_router",) - def __init__(self, router: "Router") -> None: + def __init__(self, router: _HealthHost) -> None: self._router = router @property diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py index 9286e9aa06c..b1de96fd91b 100644 --- a/litellm/router_rust/invoke/streaming.py +++ b/litellm/router_rust/invoke/streaming.py @@ -42,6 +42,7 @@ from collections.abc import ( Mapping, Sequence, ) +from collections.abc import Set as AbstractSet from typing import ( TYPE_CHECKING, Final, @@ -629,8 +630,8 @@ def continuation_routing(core: "RustCore", request: Mapping[str, object]) -> Map registry: Final = core.registry() excluded: Final = [ deployment.model_info.id - for record in registry.all() - if (deployment := registry.deployment(str(_record_id(record)))) is not None + for record_id in _record_ids(registry.all()) + if (deployment := registry.deployment(str(record_id))) is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) ] return {"body": {CONTINUATION_EXCLUDED_IDS: excluded}} @@ -641,13 +642,21 @@ def _record_id(record: Mapping[str, object]) -> object: return model_info.get("id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a plain record +def _record_ids(records: Iterable[Mapping[str, object]]) -> Iterator[object]: + return (_record_id(record) for record in records) + + +def _record_model_names(records: Iterable[Mapping[str, object]]) -> AbstractSet[object]: + return {record.get("model_name") for record in records} + + def _names_deployment_that_cannot_continue(core: "RustCore", target: object) -> bool: from litellm.router_utils.fallback_event_handlers import fallback_target_name from litellm.router_utils.pre_call_checks.continuation_prefill_check import deployment_can_continue name: Final = fallback_target_name(target) registry: Final = core.registry() - model_names: Final = {record.get("model_name") for record in registry.all()} + model_names: Final = _record_model_names(registry.all()) if name is None or name in model_names or not registry.has_deployment(name): return False deployment: Final = registry.deployment(name) diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index 046501f38f3..640fdf739c0 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -289,8 +289,8 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): validate_routing_strategy(routing_strategy) # pyright: ignore[reportArgumentType] # the strategy as given # `_reset_custom_routing_strategy`. self._custom_strategy = None - _ = self.__dict__.pop("get_available_deployment", None) - _ = self.__dict__.pop("_core_async_get_available_deployment", None) + self.__dict__.pop("get_available_deployment", None) + self.__dict__.pop("_core_async_get_available_deployment", None) self._sel.init_default(routing_strategy, routing_strategy_args, configured=False) def set_optional_pre_call_checks(self, optional_pre_call_checks: Iterable[object] | None) -> None: @@ -301,7 +301,7 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): """``initialize_assistants_endpoint`` binds the endpoints; this core answers them from ``__getattr__``, so what is left to do is to drop an assigned replacement.""" for name in _ASSISTANTS_ENDPOINTS: - _ = self.__dict__.pop(name, None) + self.__dict__.pop(name, None) def initialize_router_endpoints(self) -> None: """``initialize_router_endpoints``, as ``initialize_assistants_endpoint`` above.""" @@ -310,7 +310,7 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): ) for name in (*_ROUTER_ENDPOINTS, *_generated_endpoints): - _ = self.__dict__.pop(name, None) + self.__dict__.pop(name, None) # -- pre-routing strategies -------------------------------------------------------------------- diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index c3b038b5f26..094c61e0550 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -103,7 +103,7 @@ class QualityRouter(CustomLogger): self._tier_to_models_cache = self._build_tier_index() return self._tier_to_models_cache - def _get_routing_preferences(self, deployment: object) -> dict[str, Any] | None: + def _get_routing_preferences(self, deployment: object) -> dict[str, object] | None: """ Extract litellm_routing_preferences from a deployment, handling both dict-shaped and Pydantic-object-shaped deployments. diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 9e90b72ba1b..09bbeb92836 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -42,8 +42,7 @@ async def router_cooldown_event_callback( temp_litellm_params = dict(temp_litellm_params) _model_name: Final = _deployment.model_name or "" _api_base: Final = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" - model_info: Final = _deployment["model_info"] - model_id: Final = model_info.id + model_id: Final = _deployment.model_info.id litellm_model_name: Final = temp_litellm_params.get("model") or "" llm_provider = "" @@ -68,7 +67,7 @@ async def router_cooldown_event_callback( prometheusLogger.increment_deployment_cooled_down( litellm_model_name=_model_name, - model_id=model_id, + model_id=str(model_id), api_base=_api_base, api_provider=llm_provider, exception_status=str(exception_status), diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 22bec9e1ead..9e3e0189fd0 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -149,7 +149,7 @@ _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( ) -def _first_present(*sources: Mapping[str, Any] | None, key: str) -> int | float | None: +def _first_present(*sources: Mapping[str, object] | None, key: str) -> int | float | None: """Return *key* from the first source mapping where it's set, so callers can support a setting living in more than one deployment config location. Sources are checked in order from most to least specific to that setting. Only a number diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index f1f38c384cc..afed7532d74 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -7,22 +7,13 @@ Handles retrieving secrets from different secret management systems. import base64 import os from collections.abc import Mapping -from typing import Any, Final, Generic, Protocol, TypeVar - -from typing_extensions import ReadOnly, TypedDict +from typing import Final, Protocol import litellm from litellm._logging import print_verbose +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -_ClientT = TypeVar("_ClientT") - - -class _SecretManagerClientView(TypedDict, Generic[_ClientT]): - """Typed read of the untyped secret manager handle configured for this key manager.""" - - client: ReadOnly[_ClientT] - class _AzureKeyVaultSecret(Protocol): @property @@ -74,7 +65,7 @@ def _is_base64(s: str) -> bool: def get_secret_from_manager( - client: Any, + client: object, key_manager: str, secret_name: str, key_management_settings: KeyManagementSettings | None = None, @@ -96,16 +87,14 @@ def get_secret_from_manager( Exception: For other errors during secret retrieval """ secret = None - raw_view: Final[_SecretManagerClientView[object]] = {"client": client} - client_object: Final = raw_view["client"] + client_object: Final = client if ( key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value or type(client_object).__module__ + "." + type(client_object).__name__ == "azure.keyvault.secrets._client.SecretClient" ): # support Azure Secret Client - from azure.keyvault.secrets import SecretClient - azure_view: Final[_SecretManagerClientView[_AzureKeyVaultClient]] = {"client": client} - azure_client: Final = azure_view["client"] + azure_client: Final = require_attributes(client, _AzureKeyVaultClient, "get_secret") secret = azure_client.get_secret(secret_name).value elif ( @@ -122,8 +111,7 @@ def get_secret_from_manager( raise ValueError( "Google KMS requires the encrypted secret to be encoded in base64" ) # fix for this vulnerability https://huntr.com/bounties/ae623c2f-b64b-4245-9ed4-f13a0a5824ce - google_kms_view: Final[_SecretManagerClientView[_GoogleKmsClient]] = {"client": client} - google_kms_client: Final = google_kms_view["client"] + google_kms_client: Final = require_attributes(client, _GoogleKmsClient, "decrypt") google_kms_response: Final = google_kms_client.decrypt( request={ "name": litellm._google_kms_resource_name, @@ -143,8 +131,7 @@ def get_secret_from_manager( ciphertext_blob: Final = base64.b64decode(encrypted_value) # Perform the decryption - aws_kms_view: Final[_SecretManagerClientView[_AwsKmsClient]] = {"client": client} - aws_kms_client: Final = aws_kms_view["client"] + aws_kms_client: Final = require_attributes(client, _AwsKmsClient, "decrypt") aws_kms_response: Final = aws_kms_client.decrypt(CiphertextBlob=ciphertext_blob) # Extract and decode the plaintext @@ -171,8 +158,9 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: - google_secret_manager_view: Final[_SecretManagerClientView[_GoogleSecretManagerClient]] = {"client": client} - google_secret_manager_client: Final = google_secret_manager_view["client"] + google_secret_manager_client: Final = require_attributes( + client, _GoogleSecretManagerClient, "get_secret_from_google_secret_manager" + ) secret = google_secret_manager_client.get_secret_from_google_secret_manager(secret_name) print_verbose(f"secret from google secret manager: [set={secret is not None}]") if secret is None: @@ -183,8 +171,7 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value: try: - hashicorp_view: Final[_SecretManagerClientView[_SyncSecretReader]] = {"client": client} - hashicorp_client: Final = hashicorp_view["client"] + hashicorp_client: Final = require_attributes(client, _SyncSecretReader, "sync_read_secret") secret = hashicorp_client.sync_read_secret(secret_name=secret_name) if secret is None: raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") @@ -194,8 +181,7 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.CYBERARK.value: try: - cyberark_view: Final[_SecretManagerClientView[_SyncSecretReader]] = {"client": client} - cyberark_client: Final = cyberark_view["client"] + cyberark_client: Final = require_attributes(client, _SyncSecretReader, "sync_read_secret") secret = cyberark_client.sync_read_secret(secret_name=secret_name) if secret is None: raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}") @@ -224,8 +210,7 @@ def get_secret_from_manager( secret = os.getenv(secret_name) else: # assume the default is infisicial client - infisical_view: Final[_SecretManagerClientView[_InfisicalClient]] = {"client": client} - infisical_client: Final = infisical_view["client"] + infisical_client: Final = require_attributes(client, _InfisicalClient, "get_secret") secret = infisical_client.get_secret(secret_name).secret_value return secret diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index daf65b77a85..517a499ec8c 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,11 +1,15 @@ -from typing import Final +from typing import Final, Protocol from litellm.constants import VLLM_CHOICE_FIELDS -from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.core_helpers import decoded_json_object, json_contains, require_attributes, subscript from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + class TextCompletionStreamWrapper: def __init__( self, @@ -43,16 +47,22 @@ class TextCompletionStreamWrapper: raise Exception if chunk["choices"]: - chat_choice: Final = chunk["choices"][0] + chat_choice: Final = subscript(subscript(chunk, "choices"), 0) choice_fields: Final = decoded_json_object( - {field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + { + field: subscript(chat_choice, field) + for field in VLLM_CHOICE_FIELDS + if json_contains(chat_choice, field) + } ) text_choices: Final = TextChoices(index=0, **choice_fields) - delta: Final = chat_choice["delta"] - text_choices["text"] = delta["content"] - text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chat_choice["index"] - text_choices["finish_reason"] = chat_choice["finish_reason"] + delta: Final = subscript(chat_choice, "delta") + text_choices["text"] = subscript(delta, "content") + text_choices["reasoning_content"] = require_attributes(delta, _SupportsGet, "get").get( + "reasoning_content" + ) + text_choices["index"] = subscript(chat_choice, "index") + text_choices["finish_reason"] = subscript(chat_choice, "finish_reason") response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/litellm/utils.py b/litellm/utils.py index fb974b24283..7012741148c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -95,7 +95,9 @@ from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, has_attributes, iter_items, + json_contains, normalize_drop_params, + require_attributes, unbind_budget_reservation_from_callbacks, ) from litellm.litellm_core_utils.fallback_generalizations import ( @@ -581,6 +583,7 @@ additional_details: Final[dict[str, str] | None] = {} local_cache: Final[dict[str, str] | None] = {} last_fetched_at: Final = None last_fetched_at_keys: Final = None +_OPAQUE_VALUE: Final = TypeAdapter(object) ######## Model Response ######################### # All liteLLM Model responses will be in this format, Follows the OpenAI Format @@ -765,14 +768,20 @@ def _is_same_registered_custom_logger(existing: object, callback_class: CustomLo ) -def get_request_guardrails(kwargs: dict[str, Any]) -> list[str]: +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +def _metadata_field(metadata: object, key: str) -> object: + return require_attributes(metadata, _SupportsGet, "get").get(key) if metadata else None + + +def get_request_guardrails(kwargs: Mapping[str, object]) -> object: """ Get the request guardrails from the kwargs """ - metadata: Final = kwargs.get("metadata") or {} - requester_metadata: Final = metadata.get("requester_metadata") or {} - applied_guardrails: Final = requester_metadata.get("guardrails") or [] - return applied_guardrails + requester_metadata: Final = _metadata_field(kwargs.get("metadata"), "requester_metadata") + return _metadata_field(requester_metadata, "guardrails") or [] def get_applied_guardrails(kwargs: dict[str, object]) -> list[str]: @@ -787,7 +796,7 @@ def get_applied_guardrails(kwargs: dict[str, object]) -> list[str]: for callback in litellm.callbacks: if callback is not None and isinstance(callback, CustomGuardrail): if callback.guardrail_name is not None: - if callback.default_on is True or callback.guardrail_name in request_guardrails: + if callback.default_on is True or json_contains(request_guardrails, callback.guardrail_name): applied_guardrails.append(callback.guardrail_name) return applied_guardrails @@ -1431,7 +1440,7 @@ def function_setup( level="info", ) if "logger_fn" in kwargs: - user_logger_fn = kwargs["logger_fn"] + user_logger_fn = _OPAQUE_VALUE.validate_python(kwargs["logger_fn"]) # INIT LOGGER - for user-specified integrations model: Final = args[0] if len(args) > 0 else kwargs.get("model", None) call_type: Final = original_function @@ -1515,9 +1524,9 @@ def function_setup( elif call_type == CallTypes.embedding.value or call_type == CallTypes.aembedding.value: messages = args[1] if len(args) > 1 else kwargs.get("input", None) elif call_type == CallTypes.image_generation.value or call_type == CallTypes.aimage_generation.value: - messages = args[0] if len(args) > 0 else kwargs["prompt"] + messages = _OPAQUE_VALUE.validate_python(args[0] if len(args) > 0 else kwargs["prompt"]) elif call_type == CallTypes.moderation.value or call_type == CallTypes.amoderation.value: - messages = args[1] if len(args) > 1 else kwargs["input"] + messages = _OPAQUE_VALUE.validate_python(args[1] if len(args) > 1 else kwargs["input"]) elif call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value: messages = args[0] if len(args) > 0 else kwargs.get("prompt") elif call_type == CallTypes.rerank.value or call_type == CallTypes.arerank.value: @@ -1544,10 +1553,10 @@ def function_setup( kwargs["metadata"] = {"file_checksum": file_checksum} messages = file_checksum elif call_type == CallTypes.aspeech.value or call_type == CallTypes.speech.value: - messages = kwargs.get("input", "speech") + messages = _OPAQUE_VALUE.validate_python(kwargs.get("input", "speech")) elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle both 'input' (standard Responses API) and 'messages' (Cursor chat format) - messages = ( + messages = _OPAQUE_VALUE.validate_python( args[0] if len(args) > 0 else kwargs.get("input") or kwargs.get("messages", "default-message-value") ) elif ( @@ -1727,9 +1736,9 @@ def _schedule_async_success_logging( async def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, - result, - start_time, - end_time, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, is_completion_with_fallbacks: bool, ): if ( @@ -1748,7 +1757,7 @@ async def _client_async_logging_helper( ) -def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]: +def _get_wrapper_num_retries(kwargs: dict[str, object], exception: Exception) -> tuple[object, dict[str, object]]: """ Get the number of retries from the kwargs and the retry policy. Used for the wrapper functions. @@ -1773,6 +1782,17 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu return num_retries, kwargs +@runtime_checkable +class _ModelFallbackLookup(Protocol): + def __getitem__(self, model: str | None, /) -> object: ... + + +def _context_window_fallback_model(fallbacks: object, model: str | None) -> object: + if isinstance(fallbacks, _ModelFallbackLookup): + return fallbacks[model] + raise TypeError(f"'{type(fallbacks).__name__}' object is not subscriptable") + + def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: from litellm.litellm_core_utils.litellm_logging import Logging @@ -2661,7 +2681,7 @@ def client(original_function): retry_call_type: Final = _get_call_type_from_public_name(original_function.__name__) num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) if retry_call_type == CallTypes.acompletion.value: - context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) + context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict") is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs) @@ -2689,13 +2709,13 @@ def client(original_function): elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict - and model in context_window_fallback_dict + and json_contains(context_window_fallback_dict, model) and not is_acompletion_litellm_router_call ): if len(args) > 0: - args[0] = context_window_fallback_dict[model] + args[0] = _context_window_fallback_model(context_window_fallback_dict, model) else: - kwargs["model"] = context_window_fallback_dict[model] + kwargs["model"] = _context_window_fallback_model(context_window_fallback_dict, model) result = await original_function(*args, **kwargs) return result elif retry_call_type == CallTypes.aresponses.value: @@ -4469,7 +4489,7 @@ def get_optional_params_transcription( passed_params.pop("custom_llm_provider") passed_params.pop("drop_params") normalized_drop_params: Final = normalize_drop_params(drop_params) - special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") + special_params: Final = as_str_mapping(_OPAQUE_VALUE.validate_python(passed_params.pop("kwargs"))) or {} from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( transcription_stream_flag, transcription_stream_unset, @@ -4494,9 +4514,10 @@ def get_optional_params_transcription( message="`stream` must be a boolean. To drop it from the call, set `litellm.drop_params = True`.", ) + passed_entries: Final[Mapping[str, object]] = MappingProxyType(passed_params) non_default_params: Final = { k: v - for k, v in passed_params.items() + for k, v in passed_entries.items() if k in default_params and v != default_params[k] and not (k == "stream" and (stream_unset or malformed_stream)) } optional_params = {} @@ -4735,7 +4756,7 @@ def get_optional_params_embeddings( # retrieve all parameters passed to the function passed_params: Final = locals() passed_params.pop("custom_llm_provider", None) - special_params: Final = passed_params.pop("kwargs") + special_params: Final = as_str_mapping(_OPAQUE_VALUE.validate_python(passed_params.pop("kwargs"))) or {} passed_params.pop("drop_params", None) drop_params = normalize_drop_params(drop_params) @@ -5196,7 +5217,7 @@ class PreProcessNonDefaultParams: @staticmethod def base_pre_process_non_default_params( passed_params: dict, - special_params: dict, + special_params: Mapping[str, object], custom_llm_provider: str, additional_drop_params: list[str] | None, default_param_values: dict, @@ -5241,7 +5262,7 @@ class PreProcessNonDefaultParams: @staticmethod def embedding_pre_process_non_default_params( passed_params: dict, - special_params: dict, + special_params: Mapping[str, object], custom_llm_provider: str, additional_drop_params: list[str] | None, model: str, @@ -5262,7 +5283,7 @@ class PreProcessNonDefaultParams: def pre_process_non_default_params( passed_params: dict, - special_params: dict, + special_params: Mapping[str, object], custom_llm_provider: str, additional_drop_params: list[str] | None, model: str, @@ -7585,7 +7606,8 @@ def get_model_info_helper( token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) - for cost_key, cost_value in _model_info.items(): + cost_entries: Final[Mapping[str, object]] = MappingProxyType(_model_info) + for cost_key, cost_value in cost_entries.items(): if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None: returned_model_info[cost_key] = cost_value return returned_model_info @@ -7805,7 +7827,7 @@ def function_to_dict(input_function) -> dict: # may represent a set of acceptable values # translating as enum for function calling try: - param_enum = str(list(literal_eval(param_type))) + param_enum = str(list(iter_items(_OPAQUE_VALUE.validate_python(literal_eval(param_type))))) param_type = "string" except Exception: pass @@ -10774,7 +10796,7 @@ class ProviderConfigManager: @staticmethod def get_provider_passthrough_config( model: str, - provider: LlmProviders, + provider: LlmProviders | str, ) -> BasePassthroughConfig | None: if LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.passthrough.transformation import ( diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index b83c9ef48b1..2adaf17ac44 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -1,6 +1,6 @@ # litellm/proxy/vector_stores/vector_store_registry.py import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone from itertools import chain from typing import ( @@ -12,6 +12,7 @@ from typing import ( ) from pydantic import ConfigDict, TypeAdapter, ValidationError +from typing_extensions import TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices @@ -39,6 +40,14 @@ _TOOL: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) _STORE_ID_ITEMS: Final = TypeAdapter(list[object] | tuple[object, ...], config=ConfigDict(strict=True)) +class _ManagedVectorStoreRowFields(TypedDict, total=False, extra_items=object): + pass + + +def _row_fields(row: Iterable[tuple[str, object]]) -> Mapping[str, object]: + return dict(row) + + def _named_store_ids(store_ids: object) -> tuple[str, ...]: if isinstance(store_ids, str): return (store_ids,) @@ -538,8 +547,10 @@ class VectorStoreRegistry: order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: - _dict_vector_store = dict(vector_store) - _litellm_managed_vector_store = LiteLLM_ManagedVectorStore(**_dict_vector_store) + _dict_vector_store = _row_fields(vector_store) + _litellm_managed_vector_store = LiteLLM_ManagedVectorStore( + **_ManagedVectorStoreRowFields(**_dict_vector_store) + ) vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 66eb10e30db..f2681b03cb2 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -6,10 +6,13 @@ from functools import partial from typing import Final, Literal, overload from httpx._types import FileContent +from pydantic import TypeAdapter +from typing_extensions import TypedDict import litellm from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT +from litellm.litellm_core_utils.core_helpers import iter_items, keyword_arguments from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.videos.transformation import BaseVideoConfig @@ -26,6 +29,16 @@ from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager, client from litellm.videos.utils import VideoGenerationRequestUtils +_DECODED_JSON: Final = TypeAdapter(object) + + +class _VideoObjectFields(TypedDict, total=False, extra_items=object): + pass + + +class _CharacterObjectFields(TypedDict, total=False, extra_items=object): + pass + def _unsupported_video_call( call_name: str, custom_llm_provider: str, model: str | None = None @@ -204,9 +217,11 @@ def video_generation( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) - response: Final = VideoObject(**mock_response) + response: Final = VideoObject.model_validate( + _VideoObjectFields(**keyword_arguments(mock_response, VideoObject)) + ) return response # get llm provider logic @@ -593,9 +608,11 @@ def video_remix( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) - response: Final = VideoObject(**mock_response) + response: Final = VideoObject.model_validate( + _VideoObjectFields(**keyword_arguments(mock_response, VideoObject)) + ) return response # Try to decode provider from video_id if not explicitly provided @@ -808,8 +825,11 @@ def video_list( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) - return [VideoObject(**item) for item in mock_response] + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) + return [ + VideoObject.model_validate(_VideoObjectFields(**keyword_arguments(item, VideoObject))) + for item in iter_items(mock_response) + ] # Ensure custom_llm_provider is not None - default to openai if not provided if custom_llm_provider is None: @@ -1026,9 +1046,11 @@ def video_status( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) - response: Final = VideoObject(**mock_response) + response: Final = VideoObject.model_validate( + _VideoObjectFields(**keyword_arguments(mock_response, VideoObject)) + ) return response # Try to decode provider from video_id if not explicitly provided @@ -1173,8 +1195,10 @@ def video_create_character( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) - return CharacterObject(**mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) + return CharacterObject.model_validate( + _CharacterObjectFields(**keyword_arguments(mock_response, CharacterObject)) + ) if custom_llm_provider is None: custom_llm_provider = "openai" @@ -1298,8 +1322,10 @@ def video_get_character( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) - return CharacterObject(**mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) + return CharacterObject.model_validate( + _CharacterObjectFields(**keyword_arguments(mock_response, CharacterObject)) + ) if custom_llm_provider is None: custom_llm_provider = "openai" @@ -1430,8 +1456,8 @@ def video_edit( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) - return VideoObject(**mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) + return VideoObject.model_validate(_VideoObjectFields(**keyword_arguments(mock_response, VideoObject))) if custom_llm_provider is None: decoded: Final = decode_video_id_with_provider(video_id) @@ -1564,8 +1590,8 @@ def video_extension( mock_response = kwargs.get("mock_response", None) if mock_response is not None: if isinstance(mock_response, str): - mock_response = json.loads(mock_response) - return VideoObject(**mock_response) + mock_response = _DECODED_JSON.validate_python(json.loads(mock_response)) + return VideoObject.model_validate(_VideoObjectFields(**keyword_arguments(mock_response, VideoObject))) if custom_llm_provider is None: decoded: Final = decode_video_id_with_provider(video_id) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index beb1b58ec27..793c40bcf83 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -9935,3 +9935,134 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( assert usage.prompt_tokens_details.cached_tokens == 80 assert usage.completion_tokens_details is not None assert usage.completion_tokens_details.reasoning_tokens == 5 + + +def _stream_of(chunks: Sequence[object], sync_mode: bool) -> Iterator[object] | AsyncIterator[object]: + return iter(chunks) if sync_mode else _as_async(iter(chunks)) + + +def _streamed_text_and_finish_reasons(received: Sequence[ModelResponseStream]) -> tuple[str, list[str]]: + return ( + "".join(chunk.choices[0].delta.content or "" for chunk in received if chunk.choices), + [chunk.choices[0].finish_reason for chunk in received if chunk.choices and chunk.choices[0].finish_reason], + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_predibase_byte_stream_yields_token_text_then_the_reported_finish_reason( + sync_mode: bool, logging_obj: Logging +) -> None: + lines: Final = ( + b'data: {"token": {"id": 1, "text": "Hel"}}', + b'data: {"token": {"id": 2, "text": "lo"}, "details": {"finish_reason": "length"}}', + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=_stream_of(lines, sync_mode), + model="llama-3-1-8b-instruct", + custom_llm_provider="predibase", + logging_obj=logging_obj, + ) + + received: Final = await _collect(wrapper, sync_mode) + + assert _streamed_text_and_finish_reasons(received) == ("Hello", ["length"]) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_predibase_error_event_fails_the_stream_with_the_upstream_message( + sync_mode: bool, logging_obj: Logging +) -> None: + lines: Final = ( + b'data: {"token": {"id": 1, "text": "Hel"}}', + b'data: {"error": "Request failed during generation: CUDA out of memory"}', + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=_stream_of(lines, sync_mode), + model="llama-3-1-8b-instruct", + custom_llm_provider="predibase", + logging_obj=logging_obj, + ) + + with pytest.raises(MidStreamFallbackError) as raised: + await _collect(wrapper, sync_mode) + + assert "CUDA out of memory" in str(raised.value) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_nlp_cloud_stream_drops_the_done_marker_and_finishes_with_stop( + sync_mode: bool, logging_obj: Logging +) -> None: + lines: Final = ('{"generated_text": "Hello"}', '{"generated_text": " world[DONE]"}') + wrapper: Final = CustomStreamWrapper( + completion_stream=_stream_of(lines, sync_mode), + model="finetuned-llama-3-70b", + custom_llm_provider="nlp_cloud", + logging_obj=logging_obj, + ) + + received: Final = await _collect(wrapper, sync_mode) + + assert _streamed_text_and_finish_reasons(received) == ("Hello world", ["stop"]) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_aleph_alpha_stream_yields_the_completion_text_and_stops(sync_mode: bool, logging_obj: Logging) -> None: + lines: Final = (b'{"completions": [{"completion": " Paris", "finish_reason": "maximum_tokens"}]}',) + wrapper: Final = CustomStreamWrapper( + completion_stream=_stream_of(lines, sync_mode), + model="luminous-base", + custom_llm_provider="aleph_alpha", + logging_obj=logging_obj, + ) + + received: Final = await _collect(wrapper, sync_mode) + + assert _streamed_text_and_finish_reasons(received) == (" Paris", ["stop"]) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_generic_tool_use_streams_as_a_function_call_when_the_caller_passed_functions( + sync_mode: bool, logging_obj: Logging +) -> None: + logging_obj.update_environment_variables( + model="bedrock/invoke-endpoint", + user=None, + optional_params={"functions": [{"name": "get_weather", "parameters": {"type": "object", "properties": {}}}]}, + litellm_params={"litellm_call_id": "functions-stream"}, + ) + tool_use: Final = { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "index": 0, + } + chunks: Final = [ + { + "text": "", + "tool_use": tool_use, + "is_finished": True, + "finish_reason": "tool_calls", + "usage": None, + "index": 0, + } + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="bedrock/invoke-endpoint", + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) + + received: Final = await _collect(wrapper, sync_mode) + + deltas: Final = [chunk.choices[0].delta for chunk in received if chunk.choices] + assert [ + (delta.function_call.name, delta.function_call.arguments) for delta in deltas if delta.function_call is not None + ] == [("get_weather", '{"city": "Paris"}')] + assert [delta.tool_calls for delta in deltas if delta.tool_calls] == [] diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index 14b6094747f..04bc8e2bf97 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -1,8 +1,10 @@ import datetime import json +from collections.abc import Mapping from typing import Final import httpx +import openai import pytest import respx @@ -22,7 +24,12 @@ def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") -def _assistants_api(thread_id: str, answer: str, project_path: str = "/api/projects/litellm-test") -> respx.MockRouter: +def _assistants_api( + thread_id: str, + answer: str, + project_path: str = "/api/projects/litellm-test", + run: Mapping[str, object] | None = None, +) -> respx.MockRouter: thread_path: Final = f"{project_path}/threads/{thread_id}" router: Final = respx.mock(assert_all_called=False) router.post(path=f"{project_path}/threads").respond( @@ -42,7 +49,10 @@ def _assistants_api(thread_id: str, answer: str, project_path: str = "/api/proje }, ) router.get(path=f"{thread_path}/runs/run_1").respond( - 200, json={"id": "run_1", "object": "thread.run", "thread_id": thread_id, "status": "completed"} + 200, + json={"id": "run_1", "object": "thread.run", "thread_id": thread_id, "status": "completed"} + if run is None + else run, ) router.get(path=f"{thread_path}/messages").respond( 200, @@ -172,6 +182,50 @@ async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: assert response._hidden_params["thread_id"] == "thread_existing" +async def _complete(sync_mode: bool, request: Mapping[str, object]) -> object: + if sync_mode: + return litellm.completion(**request) + return await litellm.acompletion(**request) + + +@pytest.mark.parametrize("status", ["failed", "cancelled", "expired"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_a_run_that_ends_without_completing_raises_its_last_error_and_reads_no_messages( + sync_mode: bool, status: str +) -> None: + router: Final = _assistants_api( + "thread_new", + "never read", + run={ + "id": "run_1", + "object": "thread.run", + "thread_id": "thread_new", + "status": status, + "last_error": {"code": "server_error", "message": "Sorry, something went wrong."}, + }, + ) + request: Final = { + "model": f"azure_ai/agents/{AGENT_ID}", + "messages": [{"role": "user", "content": "hi"}], + "api_base": API_BASE, + "api_key": "azure-ad-token", + "stream": False, + } + + with router: + with pytest.raises(openai.APIError) as raised: + await _complete(sync_mode, request) + recorded: Final = _recorded(router) + + assert raised.value.status_code == 500 + assert f"Run {status}: Sorry, something went wrong." in raised.value.message + assert [(method, url.split("?")[0]) for method, url, _, _ in recorded][-1] == ( + "GET", + "/api/projects/litellm-test/threads/thread_new/runs/run_1", + ) + + def _sse(event: str, data: dict[str, object] | str) -> str: return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)}\n\n" diff --git a/tests/unit/llms/azure_ai/vector_stores/test_azure_ai_vector_stores_transformation.py b/tests/unit/llms/azure_ai/vector_stores/test_azure_ai_vector_stores_transformation.py index 7a0df7d96f0..5f86c5adc58 100644 --- a/tests/unit/llms/azure_ai/vector_stores/test_azure_ai_vector_stores_transformation.py +++ b/tests/unit/llms/azure_ai/vector_stores/test_azure_ai_vector_stores_transformation.py @@ -1,4 +1,5 @@ import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -180,3 +181,38 @@ def test_transform_requires_embedding_model(): with pytest.raises(ValueError, match="litellm_embedding_model is required"): AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) + + +class _SearchLogging: + def __init__(self, query: str) -> None: + self.model_call_details = {"input": query} + + +def test_search_response_keeps_document_fields_as_attributes_without_search_internals() -> None: + body: Final = { + "value": [ + { + "@search.score": 0.8, + "id": "doc-1", + "content": "Run litellm --config config.yaml", + "contentVector": [0.5, -0.25], + "title": "Proxy setup", + "tags": ["proxy", "setup"], + } + ] + } + + response: Final = AzureAIVectorStoreConfig().transform_search_vector_store_response( + httpx.Response(200, json=body), _SearchLogging("how do I start the proxy") + ) + + assert response["search_query"] == "how do I start the proxy" + assert response["data"] == [ + { + "score": 0.8, + "content": [{"text": "Run litellm --config config.yaml", "type": "text"}], + "file_id": "doc-1", + "filename": "Document doc-1", + "attributes": {"title": "Proxy setup", "tags": ["proxy", "setup"], "document_id": "doc-1"}, + } + ] diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py index 5b3cb2e0c92..fa4780053c2 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping from typing import Final from unittest.mock import Mock import httpx import pytest +from pydantic import BaseModel, ConfigDict from litellm.llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( AmazonTwelveLabsPegasusConfig, @@ -112,3 +114,47 @@ def test_transform_response_accepts_null_finish_reason(finish_reason: str | None assert result.choices[0].message.content == "A cat jumps." assert result.choices[0].finish_reason == expected + + +class _VideoSummary(BaseModel): + model_config = ConfigDict(frozen=True) + + title: str + duration_seconds: int + + +def _request_with_response_format(response_format: object) -> Mapping[str, object]: + config: Final = AmazonTwelveLabsPegasusConfig() + optional_params: Final = config.map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model="twelvelabs.pegasus-1-2-v1:0", + drop_params=False, + ) + return config.transform_request( + model="twelvelabs.pegasus-1-2-v1:0", + messages=_make_messages(), + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + +def test_pydantic_response_format_is_sent_as_the_json_schema_of_its_fields() -> None: + response_format: Final = _request_with_response_format(_VideoSummary)["responseFormat"] + assert isinstance(response_format, Mapping) + json_schema: Final = response_format["jsonSchema"] + assert isinstance(json_schema, Mapping) + + assert set(json_schema["properties"]) == {"title", "duration_seconds"} + + +def test_null_response_format_sends_no_response_format() -> None: + request: Final = _request_with_response_format(None) + + assert "responseFormat" not in request + + +def test_response_format_that_is_neither_a_schema_nor_a_model_raises_type_error() -> None: + with pytest.raises(TypeError): + _request_with_response_format("json") diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 87c57023c17..e69fb34038f 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -6,6 +6,7 @@ import httpx import pytest import respx +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig @@ -375,3 +376,37 @@ async def test_bedrock_search_with_credentials_managed_registry( assert search_response["data"][0]["score"] == 0.9 assert search_response["data"][0]["content"][0]["text"] == "Mock search result" + + +class _SearchLogging: + def __init__(self, query: str) -> None: + self.model_call_details: Final = {"query": query} + + +def test_search_response_keeps_each_text_result_with_its_score_and_metadata() -> None: + metadata: Final = {"x-amz-bedrock-kb-source-uri": "s3://docs/setup.md", "x-amz-bedrock-kb-chunk-id": "chunk-1"} + body: Final = { + "retrievalResults": [ + {"content": {"text": "Run litellm --config config.yaml"}, "metadata": metadata, "score": 0.42}, + {"content": {"type": "IMAGE"}, "score": 0.9}, + ] + } + + response: Final = BedrockVectorStoreConfig().transform_search_vector_store_response( + httpx.Response(200, json=body), _SearchLogging("how do I start the proxy") + ) + + assert response["search_query"] == "how do I start the proxy" + assert [(result["score"], result["content"][0]["text"], result["attributes"]) for result in response["data"]] == [ + (0.42, "Run litellm --config config.yaml", metadata) + ] + + +@pytest.mark.parametrize("body", [[], [{"retrievalResults": []}], "retrievalResults"]) +def test_search_response_that_is_not_a_json_object_raises_a_bedrock_error(body: object) -> None: + with pytest.raises(BedrockError) as raised: + BedrockVectorStoreConfig().transform_search_vector_store_response( + httpx.Response(200, json=body), _SearchLogging("query") + ) + + assert raised.value.status_code == 200 diff --git a/tests/unit/llms/cohere/test_cohere_common_utils.py b/tests/unit/llms/cohere/test_cohere_common_utils.py index 1d17824a181..0fde387b5ea 100644 --- a/tests/unit/llms/cohere/test_cohere_common_utils.py +++ b/tests/unit/llms/cohere/test_cohere_common_utils.py @@ -2,7 +2,7 @@ from typing import Final import pytest -from litellm.llms.cohere.common_utils import ModelResponseIterator +from litellm.llms.cohere.common_utils import CohereV2ModelResponseIterator, ModelResponseIterator def test_stream_parses_object_chunks_and_rejects_a_chunk_that_is_not_an_object() -> None: @@ -16,3 +16,21 @@ def test_stream_parses_object_chunks_and_rejects_a_chunk_that_is_not_an_object() assert (first["text"], first["index"], first["is_finished"]) == ("Hi", 1, False) with pytest.raises(AttributeError, match=r"^'list' object has no attribute 'get'$"): next(stream) + + +def test_v2_stream_parses_object_chunks_and_rejects_a_chunk_that_is_not_an_object() -> None: + stream: Final = CohereV2ModelResponseIterator( + streaming_response=iter( + [ + '{"type": "content-delta", "index": 0, "delta": {"message": {"content": {"text": "Hi"}}}}', + "[1, 2]", + ] + ), + sync_stream=True, + ) + + first: Final = next(stream) + + assert (first["text"], first["index"], first["is_finished"]) == ("Hi", 0, False) + with pytest.raises(RuntimeError): + next(stream) diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_transformation.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_transformation.py new file mode 100644 index 00000000000..37016d69b53 --- /dev/null +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_transformation.py @@ -0,0 +1,106 @@ +from datetime import datetime, timezone +from typing import Final + +import httpx +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.huggingface.embedding.transformation import HuggingFaceEmbeddingConfig +from litellm.types.utils import ModelResponse + + +def _transform( + body: object, *, task: str | None = None, optional_params: dict[str, object] | None = None +) -> ModelResponse: + return HuggingFaceEmbeddingConfig().transform_response( + model="tgi-endpoint", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=Logging( + model="tgi-endpoint", + messages=[{"role": "user", "content": "Capital of France?"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="huggingface-transform", + function_id="huggingface-transform", + ), + request_data={"inputs": "Capital of France?"}, + messages=[{"role": "user", "content": "Capital of France?"}], + optional_params={} if optional_params is None else optional_params, + litellm_params={} if task is None else {"task": task}, + encoding=None, + ) + + +def test_tgi_response_strips_chat_tokens_and_sums_the_token_logprobs() -> None: + response: Final = _transform( + [ + { + "generated_text": "Paris", + "details": { + "finish_reason": "length", + "tokens": [{"logprob": -0.5}, {"logprob": None}, {"logprob": -0.25}], + }, + } + ] + ) + + assert response.choices[0].message.content == "Paris" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message._logprob == -0.75 + + +def test_tgi_response_with_best_of_adds_one_choice_per_extra_sequence() -> None: + response: Final = _transform( + [ + { + "generated_text": "Paris", + "details": { + "finish_reason": "length", + "tokens": [{"logprob": -0.5}], + "best_of_sequences": [ + { + "generated_text": "Lyon", + "finish_reason": "length", + "tokens": [{"logprob": -1.0}, {"logprob": None}, {"logprob": -0.5}], + }, + {"generated_text": "", "finish_reason": "stop", "tokens": []}, + ], + }, + } + ], + optional_params={"best_of": 2}, + ) + + assert [(choice.index, choice.message.content, choice.finish_reason) for choice in response.choices] == [ + (0, "Paris", "length"), + (1, "Lyon", "length"), + (2, None, "stop"), + ] + assert response.choices[1].message.logprobs == -1.5 + + +def test_text_generation_task_reads_the_first_generated_text() -> None: + response: Final = _transform([{"generated_text": "<|assistant|>Paris"}], task="text-generation") + + assert response.choices[0].message.content == "Paris" + + +@pytest.mark.parametrize( + ("body", "error"), + [ + pytest.param([{"generated_text": 42}], TypeError, id="generated-text-without-a-length"), + pytest.param([{"generated_text": ["Paris"]}], AttributeError, id="generated-text-that-is-not-a-string"), + pytest.param( + [{"generated_text": "Paris", "details": {"finish_reason": "length", "tokens": [{"logprob": "-0.5"}]}}], + TypeError, + id="logprob-that-is-not-a-number", + ), + ], +) +def test_tgi_response_with_malformed_fields_raises_the_python_error_for_that_field( + body: object, error: type[Exception] +) -> None: + with pytest.raises(error): + _transform(body) diff --git a/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py index 35f8d933cee..3cf2518921a 100644 --- a/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py +++ b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py @@ -1,6 +1,13 @@ +from datetime import datetime, timezone +from typing import Final + +import httpx import pytest +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.nlp_cloud.chat.transformation import NLPCloudConfig +from litellm.llms.nlp_cloud.common_utils import NLPCloudError +from litellm.types.utils import ModelResponse @pytest.mark.parametrize( @@ -45,3 +52,48 @@ def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None: ScopedNLPCloudConfig(top_k=None, top_p=0.5) assert ScopedNLPCloudConfig.get_config() == {"top_k": 40, "top_p": 0.5} + + +def _transform(body: object) -> ModelResponse: + return NLPCloudConfig().transform_response( + model="finetuned-llama-3-70b", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=Logging( + model="finetuned-llama-3-70b", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="nlp-cloud-transform", + function_id="nlp-cloud-transform", + ), + request_data={"text": "Hi"}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_response_returns_the_generated_text_and_reported_token_counts() -> None: + response: Final = _transform({"generated_text": "Hello there", "nb_input_tokens": 3, "nb_generated_tokens": 2}) + + assert response.choices[0].message.content == "Hello there" + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (3, 2, 5) + + +@pytest.mark.parametrize( + "body", + [ + pytest.param({"nb_input_tokens": 3, "nb_generated_tokens": 0}, id="missing-generated-text"), + pytest.param( + {"generated_text": 42, "nb_input_tokens": 3, "nb_generated_tokens": 1}, id="generated-text-without-a-length" + ), + ], +) +def test_transform_response_without_usable_generated_text_raises_an_nlp_cloud_error(body: dict[str, object]) -> None: + with pytest.raises(NLPCloudError) as raised: + _transform(body) + + assert raised.value.status_code == 200 diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 36df47deaab..a51acd6dd44 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -5,8 +5,10 @@ import logging import os import sys import traceback +from datetime import datetime, timezone from typing import Final, cast +import httpx import pytest import respx from pydantic import BaseModel @@ -28,6 +30,7 @@ from unittest.mock import ANY, MagicMock, patch import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import ( ChatCompletionDeltaToolCall, Choices, @@ -1444,3 +1447,59 @@ def test_reasoning_effort_with_a_matching_think_value_does_not_warn( assert result["think"] is think assert _warnings_naming(caplog, "reasoning_effort") == () + + +def _transform_json_format_function_call(content: object) -> ModelResponse: + return OllamaChatConfig().transform_response( + model="llama3.2", + raw_response=httpx.Response( + 200, + json={ + "model": "llama3.2", + "message": {"role": "assistant", "content": content}, + "done": True, + "prompt_eval_count": 12, + "eval_count": 8, + }, + ), + model_response=ModelResponse(), + logging_obj=Logging( + model="llama3.2", + messages=[{"role": "user", "content": "Weather in Paris?"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="ollama-json-function-call", + function_id="ollama-json-function-call", + ), + request_data={"format": "json"}, + messages=[{"role": "user", "content": "Weather in Paris?"}], + optional_params={}, + litellm_params={"function_name": "lookup_weather"}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("content", "name", "arguments"), + [ + pytest.param( + '{"name": "get_weather", "arguments": {"city": "Paris"}}', "get_weather", {"city": "Paris"}, id="named-call" + ), + pytest.param('{"city": "Paris"}', "lookup_weather", {"city": "Paris"}, id="bare-arguments"), + ], +) +def test_json_format_response_for_a_function_call_becomes_a_tool_call( + content: str, name: str, arguments: dict[str, object] +) -> None: + response: Final = _transform_json_format_function_call(content) + + choice: Final = response.choices[0] + tool_call: Final = choice.message.tool_calls[0] + assert (choice.finish_reason, choice.message.content) == ("tool_calls", None) + assert (tool_call.function.name, json.loads(tool_call.function.arguments)) == (name, arguments) + + +def test_json_format_function_call_response_without_text_content_raises_type_error() -> None: + with pytest.raises(TypeError): + _transform_json_format_function_call(None) diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index ad3036560a3..4e5b6d665e2 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -5,7 +5,7 @@ import time from collections.abc import Awaitable, Callable, Mapping from datetime import datetime, timezone from pathlib import Path -from types import TracebackType +from types import MappingProxyType, TracebackType from typing import ClassVar, Final, Literal, Protocol from unittest.mock import AsyncMock, MagicMock, patch @@ -1135,6 +1135,93 @@ async def test_realtime_session_requests_forward_an_explicit_null_upstream( assert [key for key, value in session_fields.items() if value is None] == [] +def _create_translation_client_secret(expires_after: object, gateway: _GatewayThatRecordsTheBody) -> Awaitable[object]: + return realtime_main.acreate_realtime_translation_client_secret.__wrapped__( + model="gpt-realtime-translate", + expires_after=expires_after, + custom_llm_provider="openai", + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=_realtime_http_logging("translation_client_secrets"), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "expires_after", + [ + pytest.param({"anchor": "created_at", "seconds": 600}, id="dict"), + pytest.param(MappingProxyType({"anchor": "created_at", "seconds": 600}), id="read-only-mapping"), + ], +) +async def test_translation_client_secret_forwards_a_mapping_expires_after(expires_after: Mapping[str, object]) -> None: + gateway: Final = _GatewayThatRecordsTheBody() + + await _create_translation_client_secret(expires_after, gateway) + + assert gateway.body["expires_after"] == {"anchor": "created_at", "seconds": 600} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "expires_after", + [ + pytest.param([("anchor", "created_at"), ("seconds", 600)], id="list-of-pairs"), + pytest.param("seconds=600", id="string"), + ], +) +async def test_translation_client_secret_rejects_an_expires_after_that_is_not_a_mapping(expires_after: object) -> None: + gateway: Final = _GatewayThatRecordsTheBody() + + with pytest.raises(TypeError, match=r"realtime\.RealtimeExpiresAfter\(\) argument after \*\* must be a mapping"): + await _create_translation_client_secret(expires_after, gateway) + + assert gateway.body is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transcription_session", + [ + pytest.param([("input_audio_format", "pcm16")], id="list-of-pairs"), + pytest.param("input_audio_format=pcm16", id="string"), + ], +) +async def test_transcription_session_rejects_a_session_that_is_not_a_mapping(transcription_session: object) -> None: + gateway: Final = _GatewayThatRecordsTheBody() + + with pytest.raises( + TypeError, match=r"realtime\.RealtimeTranscriptionSessionRequest\(\) argument after \*\* must be a mapping" + ): + await realtime_main.acreate_realtime_transcription_session.__wrapped__( + model="openai/gpt-transcribe", + transcription_session=transcription_session, + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + assert gateway.body is None + + +@pytest.mark.asyncio +async def test_transcription_session_accepts_a_read_only_mapping() -> None: + gateway: Final = _GatewayThatRecordsTheBody() + + await realtime_main.acreate_realtime_transcription_session.__wrapped__( + model="openai/gpt-transcribe", + transcription_session=MappingProxyType({"input_audio_format": "pcm16"}), + api_base="https://realtime-gateway.example/v1", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + ) + + assert gateway.body["input_audio_format"] == "pcm16" + + @pytest.mark.asyncio async def test_transcription_client_secret_sends_the_deployment_model_as_the_session_model() -> None: gateway: Final = _GatewayThatRecordsTheSession() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 577d4b4b11d..97df9cdbfbf 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI, OpenAI +from openai import APIError, APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -12016,3 +12016,57 @@ def test_supplied_openai_client_base_url_decides_data_residency(monkeypatch: pyt ) assert captured["logging_obj"].litellm_params.get("data_residency") is None + + +_ALEPH_ALPHA_COMPLETE_URL: Final = "https://api.aleph-alpha.com/complete" + + +def _aleph_alpha_completion() -> object: + return litellm.completion( + model="luminous-base", + custom_llm_provider="aleph_alpha", + messages=[{"role": "user", "content": "Capital of France?"}], + api_key="aleph-test-key", + ) + + +def test_aleph_alpha_completion_returns_each_upstream_completion_as_a_choice(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ALEPH_ALPHA_COMPLETE_URL).respond( + 200, + json={ + "model_version": "2022-04", + "completions": [ + {"completion": " Paris", "finish_reason": "maximum_tokens"}, + {"completion": " Lyon", "finish_reason": "end_of_text"}, + ], + }, + ) + + response: Final = _aleph_alpha_completion() + + assert [choice.message.content for choice in response.choices] == [" Paris", " Lyon"] + + +def test_aleph_alpha_error_in_a_successful_body_fails_the_call_with_the_upstream_message( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ALEPH_ALPHA_COMPLETE_URL).respond( + 200, json={"error": "Prompt rejected by the content filter", "code": "PROMPT_REJECTED"} + ) + + with pytest.raises(APIError) as raised: + _aleph_alpha_completion() + + assert "Prompt rejected by the content filter" in raised.value.message + + +@pytest.mark.parametrize("completion", [42, None]) +def test_aleph_alpha_completion_that_is_not_text_fails_the_call( + respx_mock: respx.MockRouter, completion: object +) -> None: + respx_mock.post(_ALEPH_ALPHA_COMPLETE_URL).respond( + 200, json={"completions": [{"completion": completion, "finish_reason": "end_of_text"}]} + ) + + with pytest.raises(APIError): + _aleph_alpha_completion() diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py index d43141bdea3..26a6c0687d0 100644 --- a/tests/unit/videos/test_main.py +++ b/tests/unit/videos/test_main.py @@ -30,6 +30,7 @@ helper runs for real against genuinely-encoded ids, so the provider assertions reflect production. """ +from collections.abc import Callable, Mapping from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict, Final @@ -38,6 +39,7 @@ from unittest.mock import MagicMock, patch from litellm.types.videos.main import VideoCreateOptionalRequestParams import pytest +from pydantic import BaseModel import litellm @@ -486,3 +488,80 @@ def test_video_generation__custom_provider_without_a_video_config_is_a_400(monke with pytest.raises(litellm.BadRequestError, match="video generation is not supported for image-only"): videos_main.video_generation(model="image-only/seedance-1", prompt="a cat") handler.video_generation_handler.assert_not_called() + + +_VIDEO_JSON: Final = '{"id": "video_1", "object": "video", "status": "completed", "model": "sora-2"}' +_CHARACTER_JSON: Final = '{"id": "char_1", "object": "character", "created_at": 1767225600, "name": "hero"}' + + +def _fields(value: object) -> object: + if isinstance(value, list): + return [_fields(item) for item in value] + return value.model_dump() if isinstance(value, BaseModel) else value + + +@pytest.mark.parametrize( + ("operation", "arguments", "mock_response", "expected"), + [ + pytest.param( + videos_main.video_generation, + {"prompt": "a cat surfing", "model": "sora-2"}, + _VIDEO_JSON, + VideoObject(id="video_1", object="video", status="completed", model="sora-2"), + id="generation", + ), + pytest.param( + videos_main.video_remix, + {"video_id": "video_1", "prompt": "make it night"}, + _VIDEO_JSON, + VideoObject(id="video_1", object="video", status="completed", model="sora-2"), + id="remix", + ), + pytest.param( + videos_main.video_status, + {"video_id": "video_1"}, + _VIDEO_JSON, + VideoObject(id="video_1", object="video", status="completed", model="sora-2"), + id="status", + ), + pytest.param( + videos_main.video_edit, + {"video_id": "video_1", "prompt": "add rain"}, + _VIDEO_JSON, + VideoObject(id="video_1", object="video", status="completed", model="sora-2"), + id="edit", + ), + pytest.param( + videos_main.video_extension, + {"video_id": "video_1", "prompt": "keep going", "seconds": "4"}, + _VIDEO_JSON, + VideoObject(id="video_1", object="video", status="completed", model="sora-2"), + id="extension", + ), + pytest.param( + videos_main.video_list, + {}, + f"[{_VIDEO_JSON}]", + [VideoObject(id="video_1", object="video", status="completed", model="sora-2")], + id="list", + ), + pytest.param( + videos_main.video_create_character, + {"name": "hero", "video": b"\x00\x00\x00\x18ftypmp42"}, + _CHARACTER_JSON, + CharacterObject(id="char_1", object="character", created_at=1767225600, name="hero"), + id="create-character", + ), + pytest.param( + videos_main.video_get_character, + {"character_id": "char_1"}, + _CHARACTER_JSON, + CharacterObject(id="char_1", object="character", created_at=1767225600, name="hero"), + id="get-character", + ), + ], +) +def test_a_json_string_mock_response_is_parsed_into_the_typed_result( + operation: Callable[..., object], arguments: Mapping[str, object], mock_response: str, expected: object +) -> None: + assert _fields(operation(**arguments, mock_response=mock_response)) == _fields(expected)