mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
refactor(responses): unify guardrail input processing and drop mutable request params
The staging merge tightened the ruff-strict and type-discipline budgets, so the two `data: dict` parameters the write-back helpers introduced (LIT001) and the 17-branch `process_input_messages` (C901) no longer fit. Fold the duplicated string/list guardrail flow into one path: a pure `_extract_guardrail_inputs` builds the guardrail payload, the write-back helpers become pure functions returning `_RequestFields` (patched input items plus the resulting instructions value), and the request dict is only mutated in `process_input_messages` itself. `_apply_guardrail_responses_to_input` takes Sequence views since it only reads. A non-list `structured_messages` payload now falls through to the plain texts write-back, matching the pre-write-back behavior for guardrails that never touch structured messages. Handler file deltas vs the merge base: LIT001 57 -> 53, LIT002 42 -> 40, LIT010 28 -> 18, C901 3 -> 3.
This commit is contained in:
mateo-berri
parent
574010a2ce
commit
3004b12e90
1 changed files with 121 additions and 137 deletions
|
|
@ -31,7 +31,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
|
|||
from collections.abc import Mapping, Sequence
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
|
||||
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
from openai.types.responses.tool_param import FunctionToolParam
|
||||
|
|
@ -203,18 +203,28 @@ def _input_item_provenance(
|
|||
return item_for_message, tainted
|
||||
|
||||
|
||||
def _patch_rewritten_rows_into_input(
|
||||
data: dict,
|
||||
class _RequestFields(NamedTuple):
|
||||
input: tuple[object, ...]
|
||||
instructions: str | None
|
||||
|
||||
|
||||
class _ExtractedInputs(NamedTuple):
|
||||
inputs: GenericGuardrailAPIInputs
|
||||
task_mappings: tuple[tuple[int, int | None], ...]
|
||||
|
||||
|
||||
def _patched_request_fields(
|
||||
raw_input: object,
|
||||
instructions: object,
|
||||
original_messages: Sequence[object],
|
||||
structured_messages: Sequence[object],
|
||||
) -> bool:
|
||||
raw_input: Final = data.get("input")
|
||||
) -> _RequestFields | None:
|
||||
if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages):
|
||||
return False
|
||||
offset: Final = 1 if data.get("instructions") else 0
|
||||
return None
|
||||
offset: Final = 1 if instructions else 0
|
||||
provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:])
|
||||
if provenance is None:
|
||||
return False
|
||||
return None
|
||||
item_for_message, tainted = provenance
|
||||
changed: Final = tuple(
|
||||
(index, rewritten)
|
||||
|
|
@ -225,13 +235,14 @@ def _patch_rewritten_rows_into_input(
|
|||
rewritten_instructions: Final = (
|
||||
instruction_rewrites[0].get("content")
|
||||
if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping)
|
||||
else None
|
||||
else instructions
|
||||
)
|
||||
if instruction_rewrites and not isinstance(rewritten_instructions, str):
|
||||
return False
|
||||
instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None
|
||||
if rewritten_instructions is not None and instructions_value is None:
|
||||
return None
|
||||
body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset)
|
||||
if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes):
|
||||
return False
|
||||
return None
|
||||
replacements: Final = MappingProxyType(
|
||||
{
|
||||
item_for_message[message_index]: _rewritten_input_item(
|
||||
|
|
@ -242,11 +253,28 @@ def _patch_rewritten_rows_into_input(
|
|||
}
|
||||
)
|
||||
if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()):
|
||||
return False
|
||||
data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)] # mutable-ok: JSON body
|
||||
if isinstance(rewritten_instructions, str):
|
||||
data["instructions"] = rewritten_instructions
|
||||
return True
|
||||
return None
|
||||
return _RequestFields(
|
||||
input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)),
|
||||
instructions=instructions_value,
|
||||
)
|
||||
|
||||
|
||||
def _written_back_request_fields(
|
||||
raw_input: object,
|
||||
instructions: object,
|
||||
original_messages: Sequence[object],
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
) -> _RequestFields | None:
|
||||
if not isinstance(structured_messages, list):
|
||||
return None
|
||||
patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages)
|
||||
if patched is not None:
|
||||
return patched
|
||||
input_items, converted_instructions = (
|
||||
LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages)
|
||||
)
|
||||
return _RequestFields(input=tuple(input_items), instructions=converted_instructions)
|
||||
|
||||
|
||||
class OpenAIResponsesHandler(BaseTranslation):
|
||||
|
|
@ -288,136 +316,92 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
Handles both string input and list of message objects.
|
||||
"""
|
||||
input_data: Final[str | ResponseInputParam | None] = data.get("input")
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = []
|
||||
if input_data is None:
|
||||
if not isinstance(input_data, (str, list)):
|
||||
return data
|
||||
|
||||
structured_messages: Final = self.get_structured_messages(data)
|
||||
|
||||
# Handle simple string input
|
||||
if isinstance(input_data, str):
|
||||
inputs = GenericGuardrailAPIInputs(texts=[input_data])
|
||||
original_tools: list[dict[str, object]] = []
|
||||
|
||||
# Extract and transform tools if present
|
||||
if "tools" in data and data["tools"]:
|
||||
original_tools = list(data["tools"])
|
||||
self._extract_and_transform_tools(data["tools"], tools_to_check)
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = structured_messages
|
||||
# Include model information if available
|
||||
model = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
|
||||
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
|
||||
if (
|
||||
guardrailed_structured_messages is not None
|
||||
and guardrailed_structured_messages is not structured_messages
|
||||
):
|
||||
self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
|
||||
extracted: Final = self._extract_guardrail_inputs(data, input_data)
|
||||
if not extracted.inputs.get("texts"):
|
||||
return data
|
||||
if structured_messages:
|
||||
extracted.inputs["structured_messages"] = structured_messages
|
||||
original_tools: Final[list[dict[str, object]]] = list(data.get("tools") or [])
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=extracted.inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
|
||||
written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs)
|
||||
if written_back is not None:
|
||||
data["input"] = list(written_back.input) # mutable-ok: JSON body
|
||||
if written_back.instructions is None:
|
||||
data.pop("instructions", None)
|
||||
else:
|
||||
guardrailed_texts = guardrailed_inputs.get("texts") or ()
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
|
||||
self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
|
||||
return data
|
||||
|
||||
# Handle list input (ResponseInputParam)
|
||||
if not isinstance(input_data, list):
|
||||
return data
|
||||
data["instructions"] = written_back.instructions
|
||||
elif isinstance(input_data, str):
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
|
||||
else:
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=guardrailed_inputs.get("texts") or (),
|
||||
task_mappings=extracted.task_mappings,
|
||||
)
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
|
||||
return data
|
||||
|
||||
def _extract_guardrail_inputs(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
input_data: "str | ResponseInputParam",
|
||||
) -> _ExtractedInputs:
|
||||
texts_to_check: Final[list[str]] = []
|
||||
images_to_check: Final[list[str]] = []
|
||||
task_mappings: Final[list[tuple[int, int | None]]] = []
|
||||
original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or [])
|
||||
|
||||
# Step 1: Extract all text content, images, and tools
|
||||
for msg_idx, message in enumerate(input_data):
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts_to_check,
|
||||
images_to_check=images_to_check,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
|
||||
# Extract and transform tools if present
|
||||
if "tools" in data and data["tools"]:
|
||||
self._extract_and_transform_tools(data["tools"], tools_to_check)
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check:
|
||||
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = structured_messages
|
||||
# Include model information if available
|
||||
model = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
self._apply_guardrailed_tools_to_data(
|
||||
data,
|
||||
original_tools_list,
|
||||
guardrailed_inputs.get("tools"),
|
||||
)
|
||||
|
||||
guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
|
||||
if (
|
||||
guardrailed_structured_messages is not None
|
||||
and guardrailed_structured_messages is not structured_messages
|
||||
):
|
||||
self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages)
|
||||
else:
|
||||
# Step 3: Map guardrail responses back to original input structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=guardrailed_inputs.get("texts", []), # mutable-ok: callee signature takes a list
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = []
|
||||
if isinstance(input_data, str):
|
||||
texts_to_check.append(input_data)
|
||||
else:
|
||||
for msg_idx, message in enumerate(input_data):
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts_to_check,
|
||||
images_to_check=images_to_check,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
|
||||
|
||||
return data
|
||||
tools: Final = data.get("tools")
|
||||
if tools:
|
||||
self._extract_and_transform_tools(
|
||||
cast("list[FunctionToolParam | OpenAIMcpServerTool]", tools), # cast-ok: request body tools
|
||||
tools_to_check,
|
||||
)
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
model: Final = data.get("model")
|
||||
if isinstance(model, str):
|
||||
inputs["model"] = model
|
||||
return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings))
|
||||
|
||||
@staticmethod
|
||||
def _write_back_structured_messages(
|
||||
data: dict,
|
||||
original_messages: Sequence[object],
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
) -> None:
|
||||
if not isinstance(structured_messages, list):
|
||||
return
|
||||
if _patch_rewritten_rows_into_input(data, original_messages, structured_messages):
|
||||
return
|
||||
input_items, instructions = (
|
||||
LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
|
||||
list(structured_messages) # mutable-ok: converter signature takes a list
|
||||
)
|
||||
def _written_back_request_fields(
|
||||
data: Mapping[str, object],
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
guardrailed_inputs: GenericGuardrailAPIInputs,
|
||||
) -> _RequestFields | None:
|
||||
guardrailed: Final = guardrailed_inputs.get("structured_messages")
|
||||
if guardrailed is None or guardrailed is structured_messages:
|
||||
return None
|
||||
return _written_back_request_fields(
|
||||
data.get("input"),
|
||||
data.get("instructions"),
|
||||
structured_messages or (),
|
||||
guardrailed,
|
||||
)
|
||||
data["input"] = input_items
|
||||
if instructions is None:
|
||||
data.pop("instructions", None)
|
||||
return
|
||||
data["instructions"] = instructions
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from Responses API request (tools[].name for function
|
||||
|
|
@ -543,8 +527,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
async def _apply_guardrail_responses_to_input(
|
||||
self,
|
||||
messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam
|
||||
responses: list[str],
|
||||
task_mappings: list[tuple[int, int | None]],
|
||||
responses: Sequence[str],
|
||||
task_mappings: Sequence[tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to input messages.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue