LCORE-3202: Wire shield into response - #2234
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 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 |
|
Should only check the last commit. This PR should land after those two previous commits. This is just a draft PR. I removed refusal_message and moderation_id from the moderation result since I think we don't need those anymore. But then, we need to dig into all the functions that use moderation result and safely modify them so the missing fields won't cause errors. i.e. |
asimurka
left a comment
There was a problem hiding this comment.
LGTM, finally we can delete old run_shield_moderation and replace it with v2 but it's ok for now
| return ShieldModerationPassedV2() | ||
|
|
||
|
|
||
| def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability[Any]: |
There was a problem hiding this comment.
use object as type var here as well
0237198 to
b448aac
Compare
| HTTPException: 503 if the shield model is unreachable or returns an | ||
| HTTP error, 500 if the model returns an unexpected response. | ||
| """ | ||
| selected_shield_configs = ( |
There was a problem hiding this comment.
Replace for get_shields_for_request function from utils.shields module - the function has more detailed verification steps.
There was a problem hiding this comment.
Will change! Thanks for pointing that out! 😁
There was a problem hiding this comment.
You can even extract that to the main handler and pass here already filtered shields... but that's really up to you ;)
There was a problem hiding this comment.
I'll keep it as is to make run_shield_moderation_v2 more centralized, but no strong opinion here. 😁
|
|
||
| if shield_result.decision == "blocked": | ||
| return shield_result | ||
| except (ModelHTTPError, ModelAPIError) as e: |
There was a problem hiding this comment.
AFAIK OGX as a provider will propagate its own errors. Therefore, I recommend catching OGX's ApiConnectionError and raise 503, and also ApiConnectionError and raising corresponding erro response. See map_agent_inference_error function in utils.agents.query.py (I think you can reuse it here)
There was a problem hiding this comment.
I was checking _map_api_errors function of the OpenAIResponsesModel in pydantic_ai/models/openai.py which transfers these OpenAI's errors (APIConnectionError and APIStatusError) into pydantic ai errors (ModelHTTPError and ModelAPIError). But there might be some place where raw OpenAI errors can be raise.
Didn't know we have map_agent_inference_error which maps the errors pretty well, I'll use that instead. 😁
There was a problem hiding this comment.
Yeah, that's true but the problem is that we're not using raw openai provider but OGX (acting as openai provider). OGX does not raise native openai errors, so all the errors actually pass the _map_api_errors untouched.
There was a problem hiding this comment.
I spent too much time digging into this. The ogx exception that can be raised through OgxServerTransport/OgxLibraryTransport will be caught in request() of AsyncAPIClient at openai/_base_client.py and get coverted into openai's own APIConnectionError.
except Exception as err:
log.debug("Encountered Exception", exc_info=True)
if remaining_retries > 0:
await self._sleep_for_retry(
retries_taken=retries_taken,
max_retries=max_retries,
options=input_options,
response=None,
)
continue
log.debug("Raising connection error")
raise APIConnectionError(request=request) from errwhich means model_request() and agent.run() won't raise ogx's APIConnectionError and APIStatusError.
These ogx exceptions will only be raised if we're directly using ogx client.
I'm like 80% positive on this, but let me know if I'm wrong.
After all, it does not matter that much. It's just that we can eliminate some exceptions from the try block if we're using the inference through pydantic ai. 😁
dd54c99 to
27d8af9
Compare
Add run_shield_moderation_v2 and build_shield to utils/shields.py to run shield moderation through AbstractSafetyCapability instances instead of the Llama Stack client. Update the responses endpoint to call the new function with shield configs directly. Include unit tests covering pass, block, filtering, and error handling paths.
27d8af9 to
9f2ffe5
Compare
|
I moved |
Description
Wire V2 shield moderation into the Responses API endpoint
Add run_shield_moderation_v2 and build_shield to utils/shields.py to run safety capabilities (question validity, PII redaction) from the new ShieldConfiguration model. Call the new function in the responses endpoint handler alongside the existing Llama Stack shield path.
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