Skip to content

LCORE-3202: Wire shield into response - #2234

Open
Jazzcort wants to merge 1 commit into
lightspeed-core:feature/upgrade-ogx-to-1.0.2from
Jazzcort:wire-shield-into-response
Open

LCORE-3202: Wire shield into response#2234
Jazzcort wants to merge 1 commit into
lightspeed-core:feature/upgrade-ogx-to-1.0.2from
Jazzcort:wire-shield-into-response

Conversation

@Jazzcort

Copy link
Copy Markdown
Contributor

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

  • 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
  • Benchmarks improvement

Tools used to create PR

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

  • Assisted-by: (e.g., Claude, CodeRabbit, Ollama, etc., N/A if not used)
  • Generated by: (e.g., tool name and version; N/A if not used)

Related Tickets & Documents

  • Related Issue #
  • Closes #

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.

@coderabbitai

coderabbitai Bot commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 98fa1d9f-2fdc-4123-b34b-ce7516f74c97

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
✨ Simplify code
  • Create PR with simplified 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@Jazzcort

Copy link
Copy Markdown
Contributor Author

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. OpenAIResponseObject uses both moderation_id and refusal_message. Will it be safe to just remove id field and use moderation_result.message as output? Or, we will have a bigger refactor there for how moderation result is used. Just want to have more info before I move forward. 😁 @asimurka @jrobertboos

@Jazzcort Jazzcort changed the title Wire shield into response LCORE-3202: Wire shield into response Jul 23, 2026

@asimurka asimurka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, finally we can delete old run_shield_moderation and replace it with v2 but it's ok for now

Comment thread src/utils/shields.py Outdated
return ShieldModerationPassedV2()


def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability[Any]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use object as type var here as well

@Jazzcort
Jazzcort force-pushed the wire-shield-into-response branch 2 times, most recently from 0237198 to b448aac Compare July 24, 2026 20:37
Comment thread src/utils/shields.py Outdated
HTTPException: 503 if the shield model is unreachable or returns an
HTTP error, 500 if the model returns an unexpected response.
"""
selected_shield_configs = (

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Replace for get_shields_for_request function from utils.shields module - the function has more detailed verification steps.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will change! Thanks for pointing that out! 😁

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can even extract that to the main handler and pass here already filtered shields... but that's really up to you ;)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll keep it as is to make run_shield_moderation_v2 more centralized, but no strong opinion here. 😁

Comment thread src/utils/shields.py Outdated

if shield_result.decision == "blocked":
return shield_result
except (ModelHTTPError, ModelAPIError) as e:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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. 😁

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 err

which 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. 😁

@Jazzcort
Jazzcort force-pushed the wire-shield-into-response branch 6 times, most recently from dd54c99 to 27d8af9 Compare July 28, 2026 18:52
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.
@Jazzcort
Jazzcort force-pushed the wire-shield-into-response branch from 27d8af9 to 9f2ffe5 Compare July 28, 2026 19:04
@Jazzcort
Jazzcort marked this pull request as ready for review July 28, 2026 20:32
@Jazzcort

Copy link
Copy Markdown
Contributor Author

I moved map_agent_inference_error and map_pydantic_agent_run_error from utils/agents/query.py to utils/agents/error_handler.py because I was hitting a circular import error. Let me know if you guys think of a better solution. 😁

@Jazzcort
Jazzcort requested a review from asimurka July 28, 2026 20:35

@asimurka asimurka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants