scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 128

refactor(guardrails): satisfy upstream type and collection gates

Browse source
This commit is contained in:
albertbausili 2026-10-01 10:23:01 +02:00
parent 79ed687f9d
commit 8fc40af932
5 changed files with 70 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -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)

101
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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

8
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -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(

19
litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py
View file

@ -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)

9
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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)