From d98507e75640ec6de694658678f5f2ea36dd300b Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 20 Aug 2026 22:47:55 -0700 Subject: [PATCH] fix(proxy): let a record add a guardrail, and keep the report off every other file object A record's own `guardrails` key was stripped so it could not opt out of the chain its key and team selected. That also dropped a legitimate opt-in, which the online path honours: there the body list extends the selected one, so a caller can only ever add. The scan now unions the two the same way, and still ignores the key for the purpose of replacing the selection. The report field was declared with a null default, so every serialization of the file object carried it: non-batch uploads, retrieves, and the stored managed-file row all gained a `litellm_batch_guardrail` key because the batch path exists. It is now omitted unless a guardrail actually acted. --- .../batch_guardrails.py | 27 ++++++++++++++++++- litellm/types/llms/openai.py | 15 +++++++++++ .../test_batch_guardrails.py | 23 ++++++++++++++++ .../types/llms/test_types_llms_openai.py | 24 +++++++++++++++++ 4 files changed, 88 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 5c886ca0e9b..aef0227f229 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -307,6 +307,30 @@ def _scannable_call_type(url: object, body: Mapping[str, object]) -> CallTypesLi return from_url if from_url is not None else _call_type_from_body(body) +def _with_requested_guardrails(scan_metadata: Mapping[str, object], requested: object) -> Mapping[str, object]: + """ + Fold a record's own `guardrails` into the chain its key and team already selected. + + Online, `move_guardrails_to_metadata` extends the selected list with the one in the body, so a + caller can only ever add to what an admin chose. Reading the record's key verbatim would let it + replace that list instead, and dropping the key entirely would lose an opt-in the online path + honours, so this unions the two the way the online path does. + """ + if not isinstance(requested, list) or not requested: + return scan_metadata + selected: Final = scan_metadata.get("guardrails") + already: Final = tuple(selected) if isinstance(selected, list) else () + names: Final = [ # mutable-ok: guardrail selection tests this value with isinstance(list) + *already, + *(name for name in requested if name not in already), + ] + merged: Final = { # mutable-ok: frozen by the MappingProxyType below + **scan_metadata, + "guardrails": names, + } + return MappingProxyType(merged) + + def _custom_id_of(payload: Mapping[str, object]) -> str | None: custom_id: Final = payload.get("custom_id") return custom_id if isinstance(custom_id, str) else None @@ -367,8 +391,9 @@ async def _scan_record( # neither survives into the record that ships. Deep, and per bag per record, because `headers` # and `tags` are nested containers otherwise shared with the upload request and with every # other record in the window. The narrowing above already removed what cannot be copied. + record_scan_metadata: Final = _with_requested_guardrails(scan_metadata, body.get("guardrails")) for injected in _SCAN_METADATA_BAGS: - scan_input[injected] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here + scan_input[injected] = copy.deepcopy(dict(record_scan_metadata)) # mutable-ok: guardrails write here try: # The chain hands back the body it produced, which may be a replacement for the dict it was diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1588c650177..a943a28f232 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -66,8 +66,10 @@ from pydantic import ( ConfigDict, Discriminator, PrivateAttr, + SerializerFunctionWrapHandler, field_serializer, field_validator, + model_serializer, ) from typing_extensions import ( NotRequired, @@ -359,6 +361,19 @@ class OpenAIFileObject(BaseModel): Absent on every other upload, so OpenAI-shaped clients see an unchanged response. """ + @model_serializer(mode="wrap") + def _drop_unset_batch_guardrail(self, handler: SerializerFunctionWrapHandler) -> dict[str, object]: + """ + Omit the key entirely rather than emitting a null. + + Every other upload, every retrieve, and the stored managed-file row serialize this model, + and none of them should gain a field because the batch path exists. + """ + serialized: Final = handler(self) + return { # mutable-ok: this dict is the serialized payload pydantic asks us to return + key: value for key, value in serialized.items() if key != "litellm_batch_guardrail" or value is not None + } + _hidden_params: dict = {"response_cost": 0.0} # no cost for writing a file def __contains__(self, key) -> bool: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index a05b8ae530c..401cf675c2e 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -844,6 +844,29 @@ async def test_a_record_body_cannot_opt_itself_out_of_the_guardrail_chain(): assert seen and "guardrails" not in seen[0] +@pytest.mark.asyncio +async def test_a_record_can_add_a_guardrail_but_not_remove_one(): + """Online the body list extends what the key selected; a record must be able to add too.""" + seen = [] + + async def _run(body_guardrails): + seen.clear() + record = _record("a") + if body_guardrails is not None: + record["body"]["guardrails"] = body_guardrails + await _scan_full( + _jsonl(record), + FakeProxyLogging(lambda d: seen.append(list(d["litellm_metadata"].get("guardrails") or []))), + metadata={"guardrails": ["team-guard"]}, + ) + return seen[0] + + assert await _run(None) == ["team-guard"] + assert await _run([]) == ["team-guard"], "an empty list must not clear the key's selection" + assert await _run(["extra-guard"]) == ["team-guard", "extra-guard"], "an opt-in must be honoured" + assert await _run(["team-guard"]) == ["team-guard"], "no duplicates" + + @pytest.mark.asyncio async def test_a_redacted_record_keeps_its_own_guardrails_key(): """Stripping it for the scan must not rewrite what the caller asked the provider to run.""" diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 569743269a5..e6ca4057264 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -453,3 +453,27 @@ def test_openai_file_object_accepts_pending_status(): status="pending", ) assert file_obj.status == "pending" + + +def test_file_object_omits_the_batch_guardrail_key_when_nothing_acted(): + """Every non-batch upload, retrieve, and stored managed-file row serializes this model.""" + import json + + from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport, OpenAIFileObject + + plain = OpenAIFileObject( + id="file-1", object="file", bytes=10, created_at=1, filename="notes.txt", + purpose="assistants", status="processed", + ) + assert "litellm_batch_guardrail" not in plain.model_dump() + assert "litellm_batch_guardrail" not in json.loads(plain.model_dump_json()) + + acted = OpenAIFileObject( + id="file-2", object="file", bytes=10, created_at=1, filename="batch.jsonl", + purpose="batch", status="processed", + litellm_batch_guardrail=BatchGuardrailReport( + submitted_records=1, + modified_records=(BatchGuardrailRecord(line=1, custom_id="a", action="dropped", guardrail="g"),), + ), + ) + assert acted.model_dump()["litellm_batch_guardrail"]["submitted_records"] == 1