fix: trigger pre-call guardrail hooks for batch file#21933
fix: trigger pre-call guardrail hooks for batch file#21933Harshit28j wants to merge 3 commits intomainfrom
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR adds pre-call guardrail enforcement for batch file uploads by intercepting JSONL files at
Confidence Score: 1/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/batches_endpoints/batch_guardrail_utils.py | New utility for running pre-call guardrails on batch JSONL files. Skips string-registered callbacks (diverges from ProxyLogging.pre_call_hook behavior), and the exception wrapping loses HTTP status codes. |
| litellm/proxy/batches_endpoints/endpoints.py | Minor documentation addition explaining guardrails run at file upload time. No functional changes. |
| litellm/proxy/openai_files_endpoints/files_endpoints.py | Adds guardrail hook invocation for batch file uploads. Most changes are whitespace reformatting. The functional change (lines 465-475) correctly integrates the guardrail call but uses an inline import. |
| tests/test_litellm/proxy/batches_endpoints/test_batch_guardrail_utils.py | Unit tests pass wrong parameter name (proxy_logging_obj instead of user_api_key_cache) and assert on a method the implementation never calls. All tests will fail with TypeError at runtime. |
| tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py | Integration test mocks proxy_logging_obj.pre_call_hook, but the implementation iterates litellm.callbacks directly. The mock is never invoked, so assertions will fail. |
Sequence Diagram
sequenceDiagram
participant Client
participant FilesEndpoint as /v1/files endpoint
participant GuardrailUtils as batch_guardrail_utils
participant Callbacks as litellm.callbacks
participant Provider as LLM Provider
Client->>FilesEndpoint: POST /v1/files (purpose=batch, file=batch.jsonl)
FilesEndpoint->>FilesEndpoint: Read file content
FilesEndpoint->>FilesEndpoint: Check purpose == "batch"
alt purpose is "batch"
FilesEndpoint->>GuardrailUtils: run_pre_call_guardrails_on_batch_file(file_content, cache, key_dict)
loop For each JSONL line
GuardrailUtils->>GuardrailUtils: Parse JSON line, extract body
alt body has "messages"
GuardrailUtils->>Callbacks: Iterate litellm.callbacks
Callbacks-->>GuardrailUtils: CustomGuardrail.async_pre_call_hook()
alt Guardrail rejects
GuardrailUtils-->>FilesEndpoint: Raise Exception
FilesEndpoint-->>Client: 500 Error (guardrail rejection)
end
end
end
GuardrailUtils-->>FilesEndpoint: All items passed
end
FilesEndpoint->>Provider: route_create_file()
Provider-->>FilesEndpoint: File created
FilesEndpoint-->>Client: 200 OK (OpenAIFileObject)
Last reviewed commit: 208ac5a
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| try: | ||
| json_obj = json.loads(line) | ||
| except json.JSONDecodeError: | ||
| continue |
There was a problem hiding this comment.
Silently skipping malformed JSON is a guardrail bypass vector
Malformed or invalid JSON lines are silently skipped without running guardrails on them. If the provider's batch parser is more lenient than json.loads (or handles encoding differently), a malicious user could craft lines that fail Python's JSON parsing but succeed on the provider side, thereby bypassing guardrails entirely.
Consider logging a warning for malformed lines, or rejecting the entire file if any line fails to parse, to prevent guardrail evasion.
| if not body or "messages" not in body: | ||
| continue |
There was a problem hiding this comment.
Embedding requests bypass guardrails
Lines without a messages key are skipped, but batch JSONL files can also contain embedding requests that use input instead of messages (e.g., {"url": "/v1/embeddings", "body": {"input": "sensitive text", "model": "text-embedding-ada-002"}}). The _get_call_type_from_endpoint function already maps embeddings to "aembedding", but those lines will never reach the guardrail check because they're filtered out here.
Consider also checking for "input" in body, or passing all lines with a non-empty body through the guardrail hook.
…m_fix_pre-call-hook
|
@greptile can you please review this PR ? |
| await run_pre_call_guardrails_on_batch_file( | ||
| file_content=file_content, | ||
| proxy_logging_obj=mock_proxy_logging_obj, | ||
| user_api_key_dict=mock_user_api_key_dict | ||
| ) |
There was a problem hiding this comment.
Tests pass wrong parameter name — will crash at runtime
All unit tests call run_pre_call_guardrails_on_batch_file with proxy_logging_obj=..., but the actual function signature accepts user_api_key_cache (a DualCache), not proxy_logging_obj. Every test in this file will fail with TypeError: run_pre_call_guardrails_on_batch_file() got an unexpected keyword argument 'proxy_logging_obj'.
Additionally, the tests assert on mock_proxy_logging_obj.pre_call_hook.call_count, but the implementation never calls proxy_logging_obj.pre_call_hook — it directly iterates litellm.callbacks. This means even if the parameter issue were fixed, the mock would never be invoked and the assertions would fail.
The tests need to be rewritten to match the actual implementation: either mock litellm.callbacks with CustomGuardrail/CustomLogger instances, or change the implementation to use proxy_logging_obj.pre_call_hook() and pass it as a parameter.
| mock_pre_call_hook = mocker.patch.object( | ||
| proxy_logging_obj, | ||
| "pre_call_hook", | ||
| side_effect=Exception("Guardrail blocked this batch item") | ||
| ) |
There was a problem hiding this comment.
Mock targets the wrong object — assertion will fail
This test mocks proxy_logging_obj.pre_call_hook, but run_pre_call_guardrails_on_batch_file never calls that method. Instead, it directly iterates litellm.callbacks and invokes callback.async_pre_call_hook() on each CustomGuardrail/CustomLogger instance.
The call at files_endpoints.py:471 passes user_api_key_cache, not the proxy_logging_obj itself. So the assertion on line 1175 that mock_pre_call_hook.call_count == 1 will always be 0, and the test won't validate the guardrail behavior.
To properly test this, mock litellm.callbacks with a CustomGuardrail instance whose async_pre_call_hook raises the exception.
| try: | ||
| for callback in litellm.callbacks: |
There was a problem hiding this comment.
String-registered callbacks are silently skipped
The original pre_call_hook in ProxyLogging (see litellm/proxy/utils.py:1311-1316) resolves string callbacks via get_custom_logger_compatible_class() before checking their type. This implementation only checks isinstance(callback, CustomGuardrail) and isinstance(callback, CustomLogger), so guardrails registered as strings in litellm.callbacks (which is a common pattern) will be silently skipped.
Consider adding the same string-to-object resolution logic:
for callback in litellm.callbacks:
_callback = callback
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(callback)
if _callback is None:
continue
if isinstance(_callback, CustomGuardrail):
...1 participant
Relevant issues
Fixes LIT-2026 [Pre-call guardrail hook not triggered for batch requests]
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Link:
Link:
Links:
Type
🐛 Bug Fix
Changes
Issue:
Pre-call guardrail hooks (
async_pre_call_hook) were not being triggered for individual items inside batch requests. When a user uploads a JSONL file via/v1/filesand creates a batch via/v1/batches, the guardrail never inspected the actual messages content of the individual requests.Fix:
This PR intercepts the batch file at upload time (
POST /v1/fileswherepurpose="batch") to run pre-call guardrails on every line inside the JSONL before the upload is completed and sent to the provider.proxy_logging_obj.pre_call_hook()for each item.purpose="batch". If any guardrail raises an exception (e.g., content violation), the request is aborted and the file is rejected with the custom_id and line number reported./v1/batchesbatch creation time.tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py.