LCORE-1259: Safety Shield configuration for query & streaming query#1100
LCORE-1259: Safety Shield configuration for query & streaming query#1100JslYoon wants to merge 4 commits intolightspeed-core:mainfrom
Conversation
Signed-off-by: Lucas <lyoon@redhat.com>
Signed-off-by: Lucas <lyoon@redhat.com>
WalkthroughThis PR adds support for overriding which safety shields are applied to query requests. It introduces a new optional Changes
Sequence DiagramsequenceDiagram
participant Client
participant Handler as Endpoint Handler
participant Validator as validate_shield_ids_override
participant Moderation as run_shield_moderation
participant ShieldAPI as Shield Service
Client->>Handler: POST /query with shield_ids
Handler->>Validator: validate_shield_ids_override(request, config)
alt Override disabled & shield_ids provided
Validator-->>Handler: HTTPException 422
Handler-->>Client: Error Response
else Override allowed or shield_ids not provided
Validator-->>Handler: Valid
end
Handler->>Moderation: run_shield_moderation(..., shield_ids)
alt shield_ids is empty list
Moderation-->>Handler: Return (blocked=false)
else shield_ids specified
Moderation->>Moderation: Filter shields_to_run by shield_ids
alt Invalid shield_ids
Moderation-->>Handler: HTTPException 422
else Valid shields found
Moderation->>ShieldAPI: Apply filtered shields
ShieldAPI-->>Moderation: Moderation results
Moderation-->>Handler: ShieldModerationResult
end
else shield_ids is None
Moderation->>ShieldAPI: Apply all configured shields
ShieldAPI-->>Moderation: Moderation results
Moderation-->>Handler: ShieldModerationResult
end
Handler-->>Client: Response with moderation applied
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
No actionable comments were generated in the recent review. 🎉 🧹 Recent nitpick comments
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
tisnik
left a comment
tisnik
left a comment
There was a problem hiding this comment.
Thank you for your PR. It seems ok, but we'd also need to introduce new configuration parameter that will enable or disable this behavior. Because let's say this is ok for your purposes, but other systems want not to allow any overrides like this. Please look at existing configuration option disable_query_system_prompt. Would you be so kind to add something similar to our PR? Thank you in advance.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/utils/shields.py`:
- Line 12: Remove the trailing whitespace after the import symbol
NotFoundResponse in src/utils/shields.py (the line containing
"NotFoundResponse,") so the line ends immediately after the comma; then re-run
formatting/linting (Black and pylint) to ensure C0303 is resolved and the file
is correctly formatted.
🧹 Nitpick comments (3)
src/utils/shields.py (3)
89-95: Avoid unnecessary API call whenshield_idsis an empty list.
client.shields.list()is called on every invocation, even whenshield_ids == []leads to an immediate early return. Move the empty-list guard before the network call.♻️ Suggested reorder
+ # Early exit: empty list means skip all shields + if shield_ids is not None and len(shield_ids) == 0: + logger.info("shield_ids=[] provided, skipping all shields") + return ShieldModerationResult(blocked=False) + all_shields = await client.shields.list() # Filter shields based on shield_ids parameter if shield_ids is not None: - if len(shield_ids) == 0: - logger.info("shield_ids=[] provided, skipping all shields") - return ShieldModerationResult(blocked=False) - shields_to_run = [s for s in all_shields if s.identifier in shield_ids]
97-97: Consider convertingshield_idsto a set for O(1) membership lookups.
s.identifier in shield_idsperforms a linear scan on a list. While the list is likely small, using a set is both idiomatic and future-proof.- shields_to_run = [s for s in all_shields if s.identifier in shield_ids] + shield_id_set = set(shield_ids) + shields_to_run = [s for s in all_shields if s.identifier in shield_id_set]
86-88: DocstringRaisessection is incomplete — missing the 422 case.The new code path on Lines 107-112 raises
HTTPException(422) when none of the requestedshield_idsare found. Document this alongside the existing 404 case.
… config Signed-off-by: Lucas <lyoon@redhat.com>
Thank you for the great feedback @tisnik! I just made the changes, please let me know what you think. |
2 participants
Description
As a developer using Lightspeed Core Stack, I want to specify which safety shields to apply per query, so that I can use different shield configurations for different use cases without restarting the service.
Updated query v1, query v2, and streaming query to take in an extra parameter to identify which shields to apply. If the request parameter is not provided, it will apply all safety shields (same as before).
Type of change
Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
shield_idsparameter to query endpoints, enabling users to selectively apply specific safety shields or disable all shields for individual requests.Documentation