Add support output_shields in agents #248
Conversation
WalkthroughThe changes introduce support for distinguishing input and output shields in agent creation, based on shield ID naming conventions. Code and tests are updated to filter shields into input and output categories, pass them separately to agent constructors, and handle them distinctly in both synchronous and asynchronous query endpoints. Documentation is updated to describe the new conventions. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Endpoint
participant Client
participant Agent
User->>Endpoint: Submit query
Endpoint->>Client: Get available shields
Endpoint->>Endpoint: Filter shields (input/output/inout)
Endpoint->>Agent: Create agent with input_shields, output_shields
Agent->>Client: Process user message (with shields)
Agent-->>Endpoint: Return response
Endpoint-->>User: Return response
Assessment against linked issues
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ 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). (1)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
TamiTakamiya
force-pushed
the
TamiTakamiya/support-output-shields-in-agents
branch
from
756eb51 to
28224e5
Compare
|
@TamiTakamiya |
TamiTakamiya
force-pushed
the
TamiTakamiya/support-output-shields-in-agents
branch
from
28224e5 to
ae6da91
Compare
![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
🧹 Nitpick comments (2)
README.md (1)
142-144: Improve readability by varying sentence structure.The static analysis tool correctly identified repetitive sentence beginnings. Consider rewording for better flow.
-1. If the `shield_id` starts with `input_`, it will be used for input only. -1. If the `shield_id` starts with `output_`, it will be used for output only. -1. If the `shield_id` starts with `inout_`, it will be used both for input and output. +1. When the `shield_id` starts with `input_`, it will be used for input only. +1. For `shield_id` starting with `output_`, it will be used for output only. +1. If the `shield_id` starts with `inout_`, it will be used both for input and output.src/app/endpoints/streaming_query.py (1)
174-180: Remove unnecessary finally block in JSON parsing.The empty finally block serves no purpose and should be removed for cleaner code.
try: meta = json.loads(match.replace("'", '"')) metadata_map[meta["document_id"]] = meta except JSONDecodeError: pass - finally: - pass
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
README.md(1 hunks)src/app/endpoints/query.py(6 hunks)src/app/endpoints/streaming_query.py(7 hunks)tests/unit/app/endpoints/test_query.py(15 hunks)tests/unit/app/endpoints/test_streaming_query.py(15 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (3)
src/app/endpoints/streaming_query.py (1)
src/app/endpoints/query.py (3)
is_input_shield(216-218)is_output_shield(211-213)get_agent(75-105)
src/app/endpoints/query.py (1)
src/app/endpoints/streaming_query.py (1)
get_agent(49-77)
tests/unit/app/endpoints/test_query.py (4)
tests/unit/app/endpoints/test_streaming_query.py (4)
test_retrieve_response_four_available_shields(446-507)MockShield(350-360)MockShield(398-408)MockShield(449-459)src/configuration.py (1)
mcp_servers(79-84)src/models/requests.py (1)
QueryRequest(61-146)src/app/endpoints/query.py (1)
retrieve_response(221-293)
🪛 LanguageTool
README.md
[style] ~144-~144: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ..., it will be used for output only. 1. If the shield_idstarts withinout`, i...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
🔇 Additional comments (14)
README.md (1)
136-145: Documentation accurately describes the new shield naming conventions.The Safety Shields section clearly explains how shield IDs are classified based on their prefixes. The logic matches the implementation in the code.
src/app/endpoints/streaming_query.py (4)
49-56: Function signature correctly updated for separate shield handling.The
get_agentfunction now properly accepts separate lists for input and output shields, aligning with the new shield classification system.
69-70: Agent creation properly uses both input and output shields.The AsyncAgent constructor now receives both input and output shields separately, which is the correct implementation for the new feature.
281-288: Shield filtering logic implemented correctly.The code properly filters shields into input and output categories using the imported helper functions that implement the naming convention logic.
6-6: JSON error handling improvement is well implemented.Adding JSONDecodeError import and proper exception handling for malformed JSON prevents processing interruptions.
tests/unit/app/endpoints/test_streaming_query.py (3)
446-508: Comprehensive test coverage for shield classification.The new test function effectively validates the shield filtering logic with all four shield types (default, input_, output_, inout_). The test correctly verifies that:
- Default and input_ shields go to input list
- output_ shields go to output list
- inout_ shields go to both lists
The assertions properly validate the
get_agentfunction is called with the correct separated shield lists.
920-921: Test function signatures correctly updated for separate shields.All test functions that call
get_agenthave been properly updated to use the new signature with separateavailable_input_shieldsandavailable_output_shieldsparameters.Also applies to: 965-966, 1024-1025, 1083-1084, 1141-1142, 1197-1198
732-733: Test assertions updated to match new function signature.The test assertions for
get_agentcalls have been correctly updated to expect separate empty lists for input and output shields when no shields are available.Also applies to: 799-800, 877-878
tests/unit/app/endpoints/test_query.py (2)
472-533: LGTM! Comprehensive test for shield categorization logic.The test correctly verifies the new shield categorization functionality with four shields representing different naming conventions. The assertions properly validate that:
- Input shields include shields with no prefix,
input_prefix, andinout_prefix- Output shields include shields with
output_prefix andinout_prefix- The
get_agentfunction is called with the correctly separated shield lists
679-680: LGTM! Parameter updates align with new function signature.The updates correctly modify all
get_agentcalls to use the new signature with separateavailable_input_shieldsandavailable_output_shieldsparameters, and theAgentconstructor calls to useinput_shieldsandoutput_shieldsparameters. The changes maintain backward compatibility by passing empty lists where no shields were previously used.Also applies to: 743-744, 814-815, 969-970, 1010-1011, 1063-1064, 1116-1117, 1170-1171, 1222-1223
src/app/endpoints/query.py (4)
15-15: LGTM! Import addition supports new shield categorization functions.The
Shieldimport is necessary for the new helper functions that perform shield categorization based on identifier prefixes.
207-219: LGTM! Shield categorization logic correctly implements the naming convention.The helper functions properly implement the shield categorization based on identifier prefixes:
_is_inout_shield: Identifies shields for both input and output (inout_prefix)is_output_shield: Identifies shields for output monitoring (output_orinout_prefix)is_input_shield: Identifies shields for input monitoring (inout_prefix or any other identifier)The logic correctly handles the documented naming convention where shields without specific prefixes are treated as input-only.
75-82: LGTM! Function signature and implementation updated correctly.The
get_agentfunction signature properly accepts separate input and output shield parameters, and theAgentconstructor is correctly updated to pass bothinput_shieldsandoutput_shieldswith appropriate empty list fallbacks.Also applies to: 96-97
229-243: LGTM! Shield filtering and logging updated correctly.The
retrieve_responsefunction properly:
- Filters shields into separate input and output lists using the new helper functions
- Provides clear logging for both input and output shield availability
- Passes the correctly categorized shield lists to the
get_agentfunctionThe implementation maintains the existing functionality while adding the new shield categorization capability.
Also applies to: 258-260
|
@tisnik Would you take a look? |
Description
For #247
This PR adds the support of output_shields in agents used in lightspeed-stack.
Here are the naming convention of shield_id in this PR (also documented in README.md):
shield_idstarts withinput_, it will be used for input only.shield_idstarts withoutput_, it will be used for output only.shield_idstarts withinout_, it will be used both for input and output.Type of change
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
Documentation
New Features
Bug Fixes
Tests