scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 77

refactor(guardrails): keep the rescanned inputs instead of rebuilding the texts list

Browse source
This commit is contained in:
mateo-berri 2026-09-08 14:52:38 -07:00
parent fe836e830d
commit e700e79cce
1 changed files with 10 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -128,7 +128,7 @@ class _StreamRewriteObserver(CustomGuardrail):
class _ScannedTextRecorder(CustomGuardrail):
def __init__(self, guardrail_name: str) -> None:
super().__init__(guardrail_name=guardrail_name)
self.texts: tuple[str, ...] | None = None
self.inputs: GenericGuardrailAPIInputs | None = None
@_logged_by_inner_guardrail
async def apply_guardrail(
@ -138,7 +138,7 @@ class _ScannedTextRecorder(CustomGuardrail):
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
self.texts = _text_snapshot(inputs.get("texts"))
self.inputs = inputs
return inputs
@ -183,12 +183,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
if replacement is None:
return inputs
scanned: Final = _text_snapshot(inputs.get("texts"))
rewritten: Final = await self._scanned_texts(replacement, logging_obj)
rescanned: Final = await self._rescan(replacement, logging_obj)
rewritten: Final = None if rescanned is None else rescanned.get("texts")
if scanned is None or rewritten is None or len(rewritten) != len(scanned):
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
return {**inputs, "texts": list(rewritten)}
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
return rewritten_inputs
async def _scanned_texts(self, response: object, logging_obj: "LiteLLMLoggingObj | None") -> tuple[str, ...] | None:
async def _rescan(
self, response: object, logging_obj: "LiteLLMLoggingObj | None"
) -> GenericGuardrailAPIInputs | None:
recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown")
await self.endpoint_translation.process_output_response(
response=response,
@ -196,7 +200,7 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
litellm_logging_obj=logging_obj,
user_api_key_dict=self.user_api_key_dict,
)
return recorder.texts
return recorder.inputs
def _prepare_hook_input(