LCORE-3201: Expand Shield Interface for direct running#2220
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
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 |
|
Need to add a few unit tests for this. Will do it later |
Jazzcort
changed the title
|
|
||
| @dataclass | ||
| class QuestionValidity(AbstractCapability[None]): | ||
| class QuestionValidity(AbstractSafetyCapability[Any]): |
There was a problem hiding this comment.
object is the most suitable typevar here
| ] | ||
|
|
||
|
|
||
| class ShieldModerationPassedV2(BaseModel): |
There was a problem hiding this comment.
I think you can reuse old model instead of creating new ones.
You can generate moderation_id like: f"modr-{uuid.uuid4()}", it will be later useful when integrating with openai responses and conversations APIs. You can set refusal_response to optional for now.
Eliminates redundant stored state by deriving refusal_response from message at access time, removing the now-unnecessary create_refusal_response helper and all constructor-site arguments.
a5980fd to
71f4641
Compare
Introduce AbstractSafetyCapability(AbstractCapability[T]) in capabilities/base.py, requiring subclasses to implement a run() method that accepts raw text and returns ShieldModerationResult. This enables safety capabilities to be invoked outside the pydantic-ai agent lifecycle. Migrate QuestionValidity and PiiRedactionCapability to extend the new base class and implement run(): - QuestionValidity.run() delegates to model_request and maps ALLOWED/REJECTED to ShieldModerationPassed/Blocked - PiiRedactionCapability.run() applies regex redaction and returns Blocked when PII is detected Add unit tests for both run() implementations.
71f4641 to
6c7bac7
Compare
| moderation_id: str | ||
| refusal_response: ResponseMessage | ||
|
|
||
| @property |
There was a problem hiding this comment.
Perfect, this is even better than I suggested
2 participants
Description
Introduce simplified V2 shield moderation models (ShieldModerationPassedV2, ShieldModerationBlockedV2, ShieldModerationResultV2) alongside the existing types to avoid breaking the responses pipeline. Add AbstractSafetyCapability in capabilities/base.py, consolidating the old safety module interface with pydantic-ai's AbstractCapability. Migrate QuestionValidity and PiiRedactionCapability to use V2 types and the new base class.
Follow-up PR will clean up the old moderation logic and rename V2 to the standard names.
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