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