From 07c122f8e4723ae46fb80daa7c90d32f9862b130 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 20:38:21 +0000 Subject: [PATCH] refactor(typing): clear 143 basedpyright Any errors (wave 4) --- .../opik/opik_payload_builder/extractors.py | 19 +- .../prompt_templates/factory.py | 34 ++-- litellm/llms/azure_ai/agents/handler.py | 64 +++---- .../azure_ai/vector_stores/transformation.py | 15 +- .../llms/base_llm/videos/transformation.py | 2 +- .../guardrail_translation/handler.py | 89 ++++++---- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- litellm/llms/edenai/videos/transformation.py | 2 +- litellm/llms/gigachat/chat/transformation.py | 30 ++-- litellm/llms/minimax/videos/transformation.py | 2 +- litellm/llms/ollama/chat/transformation.py | 12 +- .../guardrail_translation/handler.py | 47 +++-- litellm/llms/openai/videos/transformation.py | 99 +++++++---- .../audio_transcription/transformation.py | 27 ++- litellm/llms/vertex_ai/common_utils.py | 48 ++--- .../llms/volcengine/videos/transformation.py | 2 +- .../llms/watsonx/completion/transformation.py | 14 +- litellm/proxy/common_request_processing.py | 58 ++++--- .../proxy/db/budget_window_spend_writer.py | 17 +- litellm/proxy/db/db_spend_update_writer.py | 64 +++++-- litellm/proxy/db/model_usage_rollup.py | 9 +- litellm/proxy/db/prisma_client.py | 164 ++++++++++++++++++ litellm/proxy/db/routing_prisma_wrapper.py | 10 ++ litellm/proxy/db/spend_log_tool_index.py | 12 +- .../access_group_endpoints.py | 24 ++- .../key_management_endpoints.py | 13 +- .../organization_endpoints.py | 17 +- litellm/proxy/mcp_server/db.py | 16 +- litellm/proxy/proxy_server.py | 3 +- .../spend_tracking/ptu_flat_cost_rollup.py | 13 +- litellm/rerank_api/main.py | 18 +- litellm/utils.py | 30 ++-- 32 files changed, 703 insertions(+), 273 deletions(-) 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/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/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/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/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/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index eb91744127e..db20065071c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1580,7 +1580,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( 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/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/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/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 0bdc3399737..c9a0028afa2 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -47,6 +47,12 @@ _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__}") + + class OllamaChatConfig(BaseConfig): """ Reference: https://github.com/ollama/ollama/blob/main/docs/api.md#parameters @@ -354,7 +360,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, @@ -370,7 +376,9 @@ class OllamaChatConfig(BaseConfig): and response_json_message is not None ): function_call: Final = decoded_json_object( - _DECODED_JSON.validate_python(json.loads(response_json_message["content"])) + _DECODED_JSON.validate_python( + json.loads(_json_document(_DECODED_JSON.validate_python(response_json_message["content"]))) + ) ) message: Final = litellm.Message( content=None, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index f81957401ff..e9626977f59 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, @@ -675,6 +675,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 @@ -1075,7 +1093,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: @@ -1090,26 +1108,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, @@ -1240,7 +1259,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/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 2223653b2a2..4d3efc29fa7 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -2,14 +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 +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 @@ -48,6 +55,42 @@ 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: @@ -348,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/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/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/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/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/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 55f2ee94d0e..bf909fec26e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -168,7 +168,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") @@ -1324,7 +1324,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, ...]: @@ -1332,8 +1342,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, ...]: @@ -1352,10 +1365,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)) @@ -1579,7 +1596,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 @@ -1602,7 +1619,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 @@ -3175,7 +3192,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 """ @@ -3436,20 +3453,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, @@ -3691,9 +3709,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/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 ff57596232f..00213c1d989 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 @@ -246,21 +246,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[ @@ -285,7 +307,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): @@ -297,7 +325,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/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 6c4e2bbfdcf..dec225cd348 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 typing import Self + from prisma import Prisma from prisma import models as prisma_models from prisma.types import HttpConfig 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__ = ( @@ -911,6 +1069,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/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/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d4a0a35d9eb..fa2162d2341 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 @@ -238,13 +238,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/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 15d05633679..964fb5742ff 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -200,7 +200,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): @@ -211,7 +223,7 @@ class _TransactionTables(Protocol): def litellm_budgettable(self) -> "_BudgetTxClient": ... @property - def litellm_organizationtable(self) -> "_OrganizationTableClient": ... + def litellm_organizationtable(self) -> "_OrganizationTxClient": ... class _TransactionManager(Protocol): @@ -222,6 +234,7 @@ class _TransactionManager(Protocol): exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: "TracebackType | None", + /, ) -> bool | None: ... diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index 255dfe37c2a..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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 594c3c410ec..be03fad47af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2506,6 +2506,7 @@ class _InvitationAcceptTxManager(Protocol): exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, + /, ) -> bool | None: ... @@ -20006,7 +20007,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/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/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/utils.py b/litellm/utils.py index d797493f559..2962f76e050 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_chat_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, has_attributes, iter_items, @@ -567,6 +568,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 @@ -1423,7 +1425,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 @@ -1507,9 +1509,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: @@ -1536,10 +1538,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 ( @@ -4530,7 +4532,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, @@ -4555,9 +4557,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 = {} @@ -4796,7 +4799,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) @@ -5257,7 +5260,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, @@ -5302,7 +5305,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, @@ -5323,7 +5326,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, @@ -7690,7 +7693,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 @@ -7910,7 +7914,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