feat(lasso): extend LassoGuardrail to support tool calling (RND-5748) - #54
Closed
vladpolevoi wants to merge 5 commits into
Closed
feat(lasso): extend LassoGuardrail to support tool calling (RND-5748)#54vladpolevoi wants to merge 5 commits into
vladpolevoi wants to merge 5 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📜 Recent review details⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (23)
📝 WalkthroughWalkthroughThe PR adds end-to-end OpenAI-style tool/function calling support to the Lasso guardrail: it expands messages into Lasso content blocks, maps tool definitions into payloads, extracts tool calls from post-call responses, and remaps masked outputs back onto original OpenAI message/tool schemas. ChangesLasso Guardrail Tool Call Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py (1)
258-279:⚠️ Potential issue | 🟠 Major | ⚡ Quick winTool-call masking never reaches the returned model response.
After
/classifixseestool_useblocks,_apply_masking_to_model_response()only updateschoice.message.content. Any masked tool-call payload still stays inchoice.message.tool_calls[*].function.arguments, so sensitive arguments remain unmasked in the actual response object. Please map maskedtool_useblocks back by call id and rewrite the corresponding tool-call arguments as well.Also applies to: 757-775
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py` around lines 258 - 279, The masking flow updates only choice.message.content but leaves tool call arguments unmasked; modify the logic that handles Lasso responses (called from where _call_lasso_api is awaited and _process_lasso_response is invoked) so that after obtaining masked_messages from the lasso response you map masked tool_use blocks back to model tool calls by their call id and rewrite the matching choice.message.tool_calls[*].function.arguments accordingly; specifically, extend _apply_masking_to_model_response (or create a helper it calls) to iterate messages for entries with tool_use/tool_call ids, find the corresponding tool_calls in the response choices, and replace the sensitive arguments with the masked payloads from lasso_response (the same way you replace choice.message.content) so both content and tool_call.function.arguments are updated.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py`:
- Around line 223-248: The code block building response_messages (inside the
loop over response.choices using variables response_messages, choice, msg, call,
func, name, args_str) is misformatted and causing lint failures; run the
project's formatter (black) on
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py and reformat the
affected sections (including the similar blocks around the other ranges noted)
so indentation, line breaks, and trailing commas follow project style; ensure
the tool-call construction (the dict with "role"/"content" and nested
"type":"tool_use") is consistently formatted after running black.
- Around line 393-396: The current flow in _run_lasso_guardrail expands
assistant.tool_calls / role="tool" into Lasso tool_use/tool_result blocks
(stored in messages) then calls _handle_masking and returns its result, which
causes data["messages"] to be overwritten with Lasso-only structured messages;
instead preserve the original OpenAI-compatible schema: call _handle_masking but
do not directly return the expanded messages—either keep a copy of the original
data["messages"] and restore it after masking, or map the masked response back
into the original message schema before assigning to data["messages"]; update
the code paths in _run_lasso_guardrail (and the similar block around lines
437-440) to ensure response["messages"] are converted back to the OpenAI-style
messages or the original data["messages"] is restored before returning.
- Around line 390-394: After calling
self._expand_messages_for_classification(messages) in _run_lasso_guardrail,
check if the returned messages list is now empty and return data early to avoid
proceeding with an empty payload; update the block that assigns messages to use
the expanded result (messages =
self._expand_messages_for_classification(messages)) and immediately return data
if not messages so you don't make an unnecessary API call or trigger Lasso
errors when expansion dropped every message.
---
Outside diff comments:
In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py`:
- Around line 258-279: The masking flow updates only choice.message.content but
leaves tool call arguments unmasked; modify the logic that handles Lasso
responses (called from where _call_lasso_api is awaited and
_process_lasso_response is invoked) so that after obtaining masked_messages from
the lasso response you map masked tool_use blocks back to model tool calls by
their call id and rewrite the matching
choice.message.tool_calls[*].function.arguments accordingly; specifically,
extend _apply_masking_to_model_response (or create a helper it calls) to iterate
messages for entries with tool_use/tool_call ids, find the corresponding
tool_calls in the response choices, and replace the sensitive arguments with the
masked payloads from lasso_response (the same way you replace
choice.message.content) so both content and tool_call.function.arguments are
updated.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 48f5951c-60bd-4008-b5b6-389ea1a2b5af
📒 Files selected for processing (2)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.pytests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (24)
- GitHub Check: Vertex AI / Run tests
- GitHub Check: All Other Providers / Run tests
- GitHub Check: proxy-endpoints / Run tests
- GitHub Check: integrations / Run tests
- GitHub Check: responses-caching-types / Run tests
- GitHub Check: benchmarks
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: core-utils / Run tests
- GitHub Check: test
- GitHub Check: proxy-infra / Run tests
- GitHub Check: build-ui
- GitHub Check: proxy-auth / Run tests
- GitHub Check: test-server-root-path (/llmproxy)
- GitHub Check: test-server-root-path (/api/v1)
- GitHub Check: auth-and-jwt
- GitHub Check: misc / Run tests
- GitHub Check: proxy-response-and-misc
- GitHub Check: enterprise-routing / Run tests
- GitHub Check: proxy-user-auth-and-spend
- GitHub Check: proxy-utils
- GitHub Check: proxy-server
- GitHub Check: proxy-config
- GitHub Check: proxy-token-counter
🧰 Additional context used
🪛 GitHub Actions: LiteLLM Linting / lint
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
[error] 1-1: black would reformat this file.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py`:
- Around line 543-562: The _update_tool_calls_from_masked function currently
assumes call is a dict when updating "function" and "arguments", causing
AttributeError for object-like calls; change the update logic to branch on
isinstance(call, dict): if dict, keep the existing dict-copy flow (dict(call),
func = dict(call.get("function", {})), func["arguments"] =
json.dumps(masked_input), call["function"]=func); otherwise handle object calls
by making a shallow copy or modifying attributes via getattr/setattr (e.g., func
= getattr(call, "function", None) or set a new function object) and set
func.arguments = json.dumps(masked_input) (or assign a new function object with
updated arguments), then append the modified object; follow the same guarded
pattern used in _apply_masking_to_model_response to avoid AttributeError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 4f73bc76-f649-4d84-882e-08393114e1b1
📒 Files selected for processing (1)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (22)
- GitHub Check: proxy-infra / Run tests
- GitHub Check: core-utils / Run tests
- GitHub Check: integrations / Run tests
- GitHub Check: proxy-endpoints / Run tests
- GitHub Check: misc / Run tests
- GitHub Check: All Other Providers / Run tests
- GitHub Check: responses-caching-types / Run tests
- GitHub Check: Vertex AI / Run tests
- GitHub Check: enterprise-routing / Run tests
- GitHub Check: proxy-auth / Run tests
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: proxy-server
- GitHub Check: Analyze (python)
- GitHub Check: proxy-response-and-misc
- GitHub Check: auth-and-jwt
- GitHub Check: proxy-user-auth-and-spend
- GitHub Check: proxy-utils
- GitHub Check: proxy-token-counter
- GitHub Check: proxy-config
- GitHub Check: test
- GitHub Check: test-server-root-path (/llmproxy)
- GitHub Check: test-server-root-path (/api/v1)
🔇 Additional comments (6)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py (6)
223-275: LGTM!The tool_call extraction correctly handles both dict and object representations, with proper null checks and JSON parsing error handling.
417-422: LGTM!Early return guards correctly prevent unnecessary API calls when messages are empty before or after expansion.
480-541: LGTM!The mapping logic correctly indexes masked content by type and applies it back to the original message structure while preserving the OpenAI-compatible schema.
620-709: LGTM!The expansion logic correctly transforms OpenAI-format messages to Lasso content blocks, with proper handling of both dict and object tool_calls, and appropriate logging for malformed inputs.
761-797: LGTM!The tool mapping correctly transforms OpenAI
ChatCompletionToolParamto LassoToolDefinitionformat, handling both dict and object representations and preserving optional fields.
919-969: LGTM!The masking logic correctly handles both dict and object representations of tool_calls, serving as the reference pattern that should be followed in
_update_tool_calls_from_masked.
…_masked (RND-5748)
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py (1)
223-275: ⚡ Quick winExtract the tool-call normalization into one helper.
The dict-vs-object branching, id/name lookup, and JSON-argument parsing are duplicated in both
async_post_call_success_hook()and_expand_messages_for_classification(). Keeping that logic in sync is already brittle in this file; a small shared helper would make future shape fixes land consistently in both paths.Also applies to: 625-714
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py` around lines 223 - 275, Extract the repeated dict-vs-object branching, id/name lookup and JSON-argument parsing into a single helper (e.g., normalize_tool_call or parse_tool_call) and use it from both async_post_call_success_hook and _expand_messages_for_classification; the helper should accept a tool call (dict or object), return a normalized tuple or dict containing call_id, name, and input_data (with JSON parsing and the existing json error handling/logging using verbose_proxy_logger.debug preserved), and then have both callers build the same {"role":"assistant","content":{...}} structure from that normalized output to replace their duplicated code paths.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py`:
- Around line 498-510: The mapping logic in _map_masked_messages_back currently
only restores masked string contents and skips non-string (dict/list) payloads
and their tool-call mappings, causing redacted multimodal/tool messages to be
lost; update _map_masked_messages_back to preserve and replay masked non-string
content too by: when iterating masked_messages handle dict/list content the same
way as masked strings (store original dict/list in a masked_structs map
alongside masked_tool_use and masked_tool_result), and when reconstructing
messages restore from masked_structs and apply
masked_tool_use/masked_tool_result for content dicts (checking
content.get("type") == "tool_use" / "tool_result" and tool_calls) so assistant
messages with dict/list content properly get their tool_calls remapped; this
aligns behavior with _expand_messages_for_classification and ensures tool_calls
remapping runs for non-string message content.
---
Nitpick comments:
In `@litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py`:
- Around line 223-275: Extract the repeated dict-vs-object branching, id/name
lookup and JSON-argument parsing into a single helper (e.g., normalize_tool_call
or parse_tool_call) and use it from both async_post_call_success_hook and
_expand_messages_for_classification; the helper should accept a tool call (dict
or object), return a normalized tuple or dict containing call_id, name, and
input_data (with JSON parsing and the existing json error handling/logging using
verbose_proxy_logger.debug preserved), and then have both callers build the same
{"role":"assistant","content":{...}} structure from that normalized output to
replace their duplicated code paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e4767e01-954d-423b-a083-b24dcda73782
📒 Files selected for processing (1)
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (25)
- GitHub Check: benchmarks
- GitHub Check: misc / Run tests
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: enterprise-routing / Run tests
- GitHub Check: integrations / Run tests
- GitHub Check: proxy-auth / Run tests
- GitHub Check: Vertex AI / Run tests
- GitHub Check: All Other Providers / Run tests
- GitHub Check: Analyze (python)
- GitHub Check: core-utils / Run tests
- GitHub Check: proxy-endpoints / Run tests
- GitHub Check: responses-caching-types / Run tests
- GitHub Check: proxy-infra / Run tests
- GitHub Check: test-server-root-path (/llmproxy)
- GitHub Check: proxy-utils
- GitHub Check: test-server-root-path (/api/v1)
- GitHub Check: auth-and-jwt
- GitHub Check: proxy-user-auth-and-spend
- GitHub Check: test
- GitHub Check: proxy-server
- GitHub Check: key-generation
- GitHub Check: proxy-config
- GitHub Check: proxy-response-and-misc
- GitHub Check: proxy-token-counter
- GitHub Check: build-ui
orgersh92
reviewed
May 11, 2026
Author
|
Closed due to moving to this BerriAI#27648 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RND-5748
Summary
_prepare_payloadnow mapsdata["tools"](OpenAIChatCompletionToolParam) to LassoToolDefinitionshape and includes it in the request body_expand_messages_for_classificationhelper converts OpenAI-format messages to Lasso content blocks: assistanttool_calls→tool_useblocks,role=toolmessages →developerrole +tool_resultblocksasync_post_call_success_hooknow extractschoice.message.tool_callsfrom the LLM response and classifies them astool_useblocksTest plan
test_payload_preparation_with_tools— tools field mapped and includedtest_payload_preparation_no_tools— no tools key when not provided (regression)test_expand_messages_assistant_tool_calls— tool_calls expand to tool_use blockstest_expand_messages_tool_role— role=tool folds into tool_result blocktest_expand_messages_tool_role_missing_tool_call_id— malformed tool message skippedtest_expand_messages_assistant_with_text_and_tool_calls— text + tool_calls produces two messagestest_expand_messages_plain_text_unchanged— plain text regressiontest_post_call_with_tool_calls— tool_calls in response classified correctlytest_post_call_text_only_regression— text-only post-call still worksSummary by CodeRabbit
New Features
Bug Fixes
Tests