LCORE-1259: Safety Shield configuration for query & streaming query

[JslYoon]

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

• Refactor
• New feature
• Bug fix
• CVE fix
• Optimization
• Documentation Update
• Configuration Update
• Bump-up service version
• Bump-up dependent library
• Bump-up library or tool used for development (does not change the final image)
• CI configuration change
• Konflux configuration change
• Unit tests improvement
• Integration tests improvement
• End to end tests improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

• Assisted-by: Claude Code

Related Tickets & Documents

• Related Issue: LCORE-1259
• Closes #: LCORE-1259

Checklist before requesting a review

• I have performed a self-review of my code.
• PR has passed all pre-merge test jobs.
• If it is a core feature, I have added thorough tests.

Testing

• Please provide detailed steps to perform tests related to this code change.
• How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

Release Notes

• New Features

  • Added option to specify which shields to use for content moderation in queries. When not specified, all configured shields are used; an empty selection skips moderation.
  • Enhanced error handling for invalid shield selections with detailed feedback.

• Tests

  • Added comprehensive tests validating shield selection, filtering behavior, and error handling for invalid shield configurations.

[coderabbitai]

Walkthrough

This PR adds an optional shield_ids parameter to the QueryRequest model and threads it through the query endpoints and shield moderation utility. The parameter enables clients to filter which shields are executed during content moderation, with None running all configured shields and an empty list skipping all shields.

Changes

Cohort / File(s)	Summary
Query Request Model src/models/requests.py	Added optional shield_ids: Optional[list[str]] field to QueryRequest with documentation describing behavior when None (all shields) or empty list (skip shields).
Query Endpoints src/app/endpoints/a2a.py, src/app/endpoints/query_v2.py, src/app/endpoints/streaming_query_v2.py	Updated to extract shield_ids from QueryRequest and pass it to run_shield_moderation() function calls.
Shield Moderation Utility src/utils/shields.py	Added shield_ids parameter with filtering logic: fetches all shields, filters to requested subset, logs missing shields, raises 422 error if no valid shields found, and iterates moderation only over filtered shields.
Shield Tests tests/unit/utils/test_shields.py	Added three new tests validating shield_ids parameter: empty list skips shields, missing shields raises HTTPException, and filtering to specific shields works correctly.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Endpoint as Query Endpoint
    participant ShieldUtil as Shield Utility
    participant ShieldAPI as Shield API
    
    Client->>Endpoint: POST query with shield_ids
    Endpoint->>ShieldUtil: run_shield_moderation(client, input_text, shield_ids)
    ShieldUtil->>ShieldAPI: list_shields()
    ShieldAPI-->>ShieldUtil: all_shields
    ShieldUtil->>ShieldUtil: filter shields by shield_ids
    alt shield_ids specified but no valid shields
        ShieldUtil-->>Endpoint: HTTPException 422
        Endpoint-->>Client: Error response
    else shield_ids empty list
        ShieldUtil->>Endpoint: ShieldModerationResult (not_blocked)
        Endpoint-->>Client: Success response
    else shield_ids with valid shields
        ShieldUtil->>ShieldAPI: create_moderation(model, input_text)
        ShieldAPI-->>ShieldUtil: moderation_result
        ShieldUtil->>Endpoint: ShieldModerationResult
        Endpoint-->>Client: Success response
    end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• LCORE-849: Shields implementation using Moderations API #979: Modifies shields moderation flow and run_shield_moderation API across related endpoints and utilities.

Suggested reviewers

• tisnik
• manstis

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main feature: adding safety shield configuration for query and streaming query endpoints, which aligns with the PR's objective.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Important

Action Needed: IP Allowlist Update

If your organization protects your Git platform with IP whitelisting, please add the new CodeRabbit IP address to your allowlist:

• ✨ 136.113.208.247/32 (new)
• 34.170.211.100/32
• 35.222.179.152/32

Reviews will stop working after February 8, 2026 if the new IP is not added to your allowlist.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[tisnik]

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.