diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 89ad67f0485..c1dd4d022e6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -31,8 +31,8 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: dict[int, str] = field(default_factory=dict) - holdback_per_choice: dict[int, int] = field(default_factory=dict) + mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) + holdback_per_choice: Mapping[int, int] = field(default_factory=dict) @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 661cffa9846..9ad651d1e65 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -169,8 +169,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ()) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -259,12 +259,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in data.get("tools") or (): if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or []: + for fn in data.get("functions") or (): if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) return names @@ -388,7 +388,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -439,9 +439,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"response": response} - else: - if "response" not in request_data: - request_data["response"] = response + request_data.setdefault("response", response) self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -461,7 +459,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) @@ -495,7 +493,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: @@ -551,7 +549,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can @@ -624,7 +622,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) # Step 4: Apply guardrailed text back to all streaming chunks # For each choice, replace the combined text across all chunks @@ -787,7 +785,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, sink: StreamTransformSink, ) -> None: """Run the guardrail over the raw accumulated text and report the @@ -799,8 +797,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far) if not raw_by_index: - sink.mutated_text_per_choice = {} - sink.holdback_per_choice = {} + sink.mutated_text_per_choice = MappingProxyType({}) + sink.holdback_per_choice = MappingProxyType({}) return # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits @@ -846,19 +844,25 @@ class OpenAIChatCompletionsHandler(BaseTranslation): finally: request_data["response"] = assembled request_data["responses"] = responses_so_far - sink.mutated_text_per_choice = dict( - chain.from_iterable( - choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds + sink.mutated_text_per_choice = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds + ) ) ) - sink.holdback_per_choice = dict( - chain.from_iterable(choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds) + sink.holdback_per_choice = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds + ) + ) ) return tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( tuple( - {"index": index, **converted} + MappingProxyType({"index": index, **converted}) for index, tool_call in enumerate(tool_calls) if (converted := self._convert_tool_call_to_dict(tool_call)) is not None ) @@ -887,13 +891,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): len(texts_to_check), ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or [] - sink.mutated_text_per_choice = { - idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) - } - sink.holdback_per_choice = { - indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) - } + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () + sink.mutated_text_per_choice = MappingProxyType( + {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} + ) + sink.holdback_per_choice = MappingProxyType( + {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} + ) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) @@ -1248,7 +1252,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], - guardrailed_texts: list[str], + guardrailed_texts: Sequence[str], task_mappings: list[tuple[int, int | None]], ) -> None: """ @@ -1266,12 +1270,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Final[dict[tuple[int, int | None], str]] = {} - for task_idx, guardrail_response in enumerate(guardrailed_texts): - mapping = task_mappings[task_idx] - choice_idx = cast(int, mapping[0]) - content_idx_optional = cast(int | None, mapping[1]) - guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response + guardrail_map: Final = MappingProxyType( + {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)} + ) # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) @@ -1281,34 +1282,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # choice by its index field: on n>1 streams a chunk usually carries one # choice at list position 0 whose index names the logical choice. for response in responses: - for choice in response.choices: - if isinstance(choice, litellm.StreamingChoices): - content = choice.delta.content - elif isinstance(choice, litellm.Choices): - content = choice.message.content - else: - continue - + for choice, message in ( + (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message) + for choice in response.choices + if isinstance(choice, (StreamingChoices, Choices)) + ): + content = message.content if content is None: continue if isinstance(content, str): - # String content str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: - if str_key not in already_set: - # First chunk - set the complete guardrailed text - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = guardrail_map[str_key] - elif isinstance(choice, litellm.Choices): - choice.message.content = guardrail_map[str_key] - already_set[str_key] = True - else: - # Subsequent chunks - clear the content - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = "" - elif isinstance(choice, litellm.Choices): - choice.message.content = "" + message.content = guardrail_map[str_key] if str_key not in already_set else "" + already_set[str_key] = True elif isinstance(content, list): # List content - handle each content item diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 1ec80fff110..43164f3459e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2416,8 +2416,12 @@ async def apply_guardrail( metadata: Final = _guardrail_request_metadata(request.metadata, data.get("metadata")) request_data: Final[dict] = { - **({"messages": request.messages} if request.messages is not None else {}), - **({"metadata": metadata} if request.metadata is not None or metadata else {}), + key: value + for key, value in ( + ("messages", request.messages), + ("metadata", metadata if request.metadata is not None or metadata else None), + ) + if value is not None } _input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type) guardrailed_inputs: Final = await active_guardrail.apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py b/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py index 53fcc87d7ee..d0b2d413677 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py +++ b/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py @@ -7,6 +7,7 @@ from __future__ import annotations import os from collections.abc import Mapping, Sequence +from itertools import chain, product from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal @@ -66,7 +67,7 @@ def _metadata(block: object) -> Mapping[str, object]: def _consumer_id(request_data: Mapping[str, object]) -> str | None: blocks: Final = tuple(_metadata(request_data.get(source)) for source in ("litellm_metadata", "metadata")) - candidates: Final = (block.get(name) for names in CONSUMER_ID_KEYS for name in names for block in blocks) + candidates: Final = (block.get(name) for name, block in product(chain.from_iterable(CONSUMER_ID_KEYS), blocks)) return next((value for value in candidates if isinstance(value, str) and value), None) @@ -78,7 +79,9 @@ def _message_text(message: Mapping[str, object]) -> str: def _copy_message(value: object) -> Mapping[str, object] | None: if not isinstance(value, Mapping): return None - return {str(key): item for key, item in value.items()} # mutable-ok: shallow copy for write-back + return { # mutable-ok: shallow copy for write-back + str(key): item for key, item in TypeAdapter(Mapping[object, object]).validate_python(value).items() + } def _copy_messages(messages: Sequence[object]) -> tuple[Mapping[str, object], ...] | None: @@ -96,14 +99,14 @@ def _tool_calls_in_message(message: Mapping[str, object]) -> tuple[object, ...] return None if not isinstance(raw, list): raise HTTPException(status_code=400, detail=TRANSFORM_MISSING) - return tuple(raw) + return TypeAdapter(tuple[object, ...]).validate_python(raw) def _tool_calls_from_messages(messages: Sequence[Mapping[str, object]]) -> tuple[object, ...] | None: groups: Final = tuple(_tool_calls_in_message(message) for message in messages) if all(group is None for group in groups): return None - return tuple(tool_call for group in groups if group is not None for tool_call in group) + return tuple(chain.from_iterable(group for group in groups if group is not None)) def _rewrite_last_user_message( @@ -237,7 +240,7 @@ class NeuralTrustGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail contract input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: @@ -249,7 +252,7 @@ class NeuralTrustGuardrail(CustomGuardrail): async def _evaluate( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail contract input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None, ) -> GenericGuardrailAPIInputs: @@ -303,7 +306,7 @@ class NeuralTrustGuardrail(CustomGuardrail): def _evaluate_body( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail contract input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None, ) -> dict[str, object]: # mutable-ok: outbound JSON @@ -411,7 +414,7 @@ class NeuralTrustGuardrail(CustomGuardrail): raw_messages: Final = transformed.get("messages") if isinstance(raw_messages, list) and raw_messages: - rewritten_messages: Final = _copy_messages(raw_messages) + rewritten_messages: Final = _copy_messages(TypeAdapter(tuple[object, ...]).validate_python(raw_messages)) if rewritten_messages is None or len(rewritten_messages) != len(_sent_messages(inputs, input_type)): raise HTTPException(status_code=400, detail=TRANSFORM_MISSING) return _inputs_with_messages(inputs, rewritten_messages, replace_tool_calls=True) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 3e65df1ac76..3615119666b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -12,6 +12,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException +from pydantic import InstanceOf, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -473,9 +474,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, reference_chunk: object, - mutated_text_per_choice: dict[int, str], + mutated_text_per_choice: Mapping[int, str], emitted_text_per_choice: dict[int, str], - holdback_per_choice: dict[int, int], + holdback_per_choice: Mapping[int, int], finish_reason_per_choice: dict[int, str | None], held_chars_per_choice: dict[int, int], is_final: bool, @@ -983,7 +984,9 @@ class UnifiedLLMGuardrails(CustomLogger): from litellm.integrations.custom_guardrail import ModifyResponseException if guardrail_to_apply is None: - guardrail_to_apply = request_data.pop("guardrail_to_apply", None) + guardrail_to_apply = TypeAdapter(InstanceOf[CustomGuardrail] | None).validate_python( + request_data.pop("guardrail_to_apply", None) + ) def _streaming_flag(name: str, default: object) -> Any: return self.resolve_streaming_flag(guardrail_to_apply, name, default)