Repository navigation
feat(guardrails): add Airia guardrail integration (internal copy of #40784) - #43657
Draft
devin-ai-integration[bot] wants to merge 6 commits into
Draft
devin-ai-integration[bot] wants to merge 6 commits into
devin-ai-integration[bot] wants to merge 6 commits into
Conversation
Adds `airia` as a built-in guardrail provider. The hook sends the prompt under `pre_call` and the model's response under `post_call` to the Airia AI Gateway, which answers allow, block, or redacted content; the only configuration is `api_base`, `api_key`, and an optional `timeout`. `during_call` is not offered: it runs concurrently with the model call, so a block could land after the prompt has already reached the provider. A BLOCKED verdict, or any action this version does not recognise, raises with `blocked_content=True`; a transport error, a non-2xx, or a response body that is not a JSON object raises with `blocked_content=False`, so callers can tell "could not evaluate" from "evaluated and blocked" while both still fail closed. On GUARDRAIL_INTERVENED a copy of the inputs is returned with every rewritten field substituted; an intervention carrying no applicable rewrite, a rewrite of the wrong shape, or a rewrite whose item count does not match what was sent (which would misalign the positional write-back downstream) blocks instead of letting the original through. Streamed responses are moderated whole and then emitted redacted: the hook opts into the unified hook's `incremental_diff` mode (the default `block_only` drops rewrites) with `streaming_end_of_stream_only`, so a redaction can never span transform rounds and underflow. Covered by a test that drives the real streaming iterator hook end to end and checks the assembled client output, not just the guardrail's own config attributes. The HTTP client is dependency-injected (an `async_handler` constructor parameter, defaulting to a real client) rather than mutated post-construction, so tests inject a fake client instead of monkeypatching an instance method. Known limitation, not fixable from this guardrail alone: LiteLLM's unified guardrail framework yields a tool-call chunk to the client immediately during `incremental_diff` streaming, inspecting the assembled tool calls against the guardrail's verdict only once the stream ends (`_run_incremental_transform_stream` in `unified_guardrail.py`). A blocked tool call can therefore reach an executing client before the terminal error arrives. Switching to `block_only` with `streaming_buffer_until_moderated` does not close this without a regression: that flag has no effect under `incremental_diff` (the mode branches away before it is read), and `block_only` itself replays the original, unredacted chunks on release, silently dropping any REDACT verdict on streamed text. Relatedly, LiteLLM's OpenAI chat completions guardrail translation skips `post_call` guardrails entirely for a response with no text content, so a tool-call-only response bypasses this guardrail (and every other post-call guardrail) even outside streaming. Both are shared-framework behavior, not specific to Airia; fixing either is out of scope for an isolated single-guardrail contribution.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Airia Open Source <opensource@airia.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
Author
|
I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".
|
Contributor
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
|
Hi @airia-open-source, can you send me the relevant longstanding credentials e.g api key for airia guardrail through email so that I can verify this PR e2e? Also we require them for maintaining the integration going forward. |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TLDR
Problem this solves:
How it solves it:
airiaguardrail from feat(guardrails): add Airia guardrail integration #40784 unchanged, merged with mainUser Flow
Before: a proxy admin who configures
guardrail: airiagets no enforcementguardrail: airiawithapi_baseandapi_keyunderguardrails:and start the proxySkipping guardrail 'airia' ... Unsupported guardrail: airiaada@example.testin the prompt reaches the model unchangedAfter: the same config calls Airia before and after the model
ada@example.testreaches the model as[EmailAddress1]when Airia redacts itRelevant issues
Internal copy of #40784 by @airia-open-source, credited with a Co-authored-by trailer. The contributor branch could not be pushed to from our side
Pre-Submission checklist
uv run pytest tests/unit/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
This is the deterministic audit run: a real two-worker proxy, real Postgres and Redis, a scripted upstream and a local sink that speaks the Airia
/v1/guardrails/litellmprotocol. A live run against the real Airia service is still pending an Airia credential and will replace nothing below, only add a caseBase is
eae8ed7f3c448d9dde3ecc3291e0f92ddc70a2e4(merge base with main). Head is1d0fbcbb2feb5de551a00e6ee64c10cb8d8b7567; the PR tip adds an empty credit commit and 75d2ded, which only drops a stale# mutable-okcomment, so the code under test is unchanged. Runs usedtests/integration/run.py extensionswith seed 4106601 and order seed 0. The unit mutation check removedget_supported_event_hooksand the new unit test went red, then green on restoreBefore (eae8ed7)
1d0fbcb-baseover the two Airia filesairiais not a known guardrail and traffic is unchecked, 3 skipped as BUG markersAfter (1d0fbcb)
1d0fbcb-headAand1d0fbcb-headBover the same filestest_airia_guardrail.py::test_r1_pre_call_none_chat_clients[httpx]test_airia_guardrail.py::test_r1_pre_call_none_chat_clients[openai_async]test_airia_guardrail.py::test_r1_pre_call_none_chat_clients[openai_sync]test_airia_guardrail.py::test_r2_blocked_pre_call_never_reaches_provider[httpx]test_airia_guardrail.py::test_r2_blocked_pre_call_never_reaches_provider[openai_async]test_airia_guardrail.py::test_r2_blocked_pre_call_never_reaches_provider[openai_sync]test_airia_guardrail.py::test_r3_blocked_without_reason_uses_default_messagetest_airia_guardrail.py::test_r4_pre_call_redaction_reaches_providertest_airia_guardrail.py::test_r5_structured_message_redaction_is_positionaltest_airia_guardrail.py::test_r6_multiple_texts_are_rewritten_without_reorderingtest_airia_guardrail.py::test_r7_multi_choice_response_redaction_rewrites_each_choicetest_airia_guardrail.py::test_r8_streamed_chat_post_call_redaction_has_no_leaktest_airia_guardrail.py::test_r9_streamed_chat_post_call_block_has_no_provider_error_leaktest_airia_guardrail.py::test_r10_messages_pre_call_verdict_matrix[anthropic_async-REDACT]test_airia_guardrail.py::test_r10_messages_pre_call_verdict_matrix[anthropic_sync-BLOCK]test_airia_guardrail.py::test_r10_messages_pre_call_verdict_matrix[httpx-NONE]test_airia_guardrail.py::test_r11_messages_post_call_redaction[anthropic_async]test_airia_guardrail.py::test_r11_messages_post_call_redaction[anthropic_sync]test_airia_guardrail.py::test_r11_messages_post_call_redaction[httpx]test_airia_guardrail.py::test_r12_streamed_messages_post_call_redaction_has_no_leaktest_airia_guardrail.py::test_r13_responses_pre_call_verdict_matrix[httpx-NONE]test_airia_guardrail.py::test_r13_responses_pre_call_verdict_matrix[openai_async-REDACT]test_airia_guardrail.py::test_r13_responses_pre_call_verdict_matrix[openai_sync-BLOCK]test_airia_guardrail.py::test_r14_nonstream_responses_post_call_redactiontest_airia_guardrail.py::test_r14_streamed_responses_post_call_redaction_has_no_leaktest_airia_guardrail.py::test_r15_default_on_and_default_off_are_distincttest_airia_guardrail.py::test_r16_key_attached_guardrail_runs_without_request_fieldtest_airia_guardrail.py::test_r17_both_mode_runs_exactly_once_per_phasetest_airia_guardrail.py::test_r18_during_call_is_skipped_and_request_is_uncheckedtest_airia_guardrail.py::test_r19_logging_only_records_block_but_does_not_rejecttest_airia_guardrail.py::test_r20_missing_credentials_skip_and_environment_fallback_workstest_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-401]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-403]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-404]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-500]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-INVALIDJSON]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NO-ACTION]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NOFIELD]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NONLIST-TEXTS]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NONOBJECT-LIST]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NONOBJECT-NULL]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-NONOBJECT-STRING]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-UNKNOWN-ACTION]test_airia_guardrail.py::test_r21_malformed_or_failed_verdicts_fail_closed_and_proxy_recovers[AIRIA-WRONGLEN]test_airia_guardrail.py::test_r22_transport_and_timeout_fail_closedtest_airia_guardrail.py::test_r23_hostile_inputs_are_forwarded_without_proxy_crash[5kb]test_airia_guardrail.py::test_r23_hostile_inputs_are_forwarded_without_proxy_crash[empty]test_airia_guardrail.py::test_r23_hostile_inputs_are_forwarded_without_proxy_crash[image_part]test_airia_guardrail.py::test_r23_hostile_inputs_are_forwarded_without_proxy_crash[text_part]test_airia_guardrail.py::test_r23_repeated_identical_requests_each_reach_airia_oncetest_airia_guardrail.py::test_r24_caller_master_key_never_reaches_airiatest_airia_guardrail.py::test_r25_tool_only_post_call_is_still_guardedtest_airia_guardrail.py::test_r26_streamed_tool_call_deltas_do_not_leak_before_post_verdicttest_airia_guardrail.py::test_r27_management_and_ui_surfaces_include_airiatest_airia_guardrail.py::test_r28_existing_generic_guardrail_control_still_firestest_airia_guardrail_chaos.py::test_c1_sink_outage_mid_mixed_endpoint_burst_fails_closed_and_recoverstest_airia_guardrail_chaos.py::test_c2_slow_sink_is_bounded_and_proxy_recoverstest_airia_guardrail_chaos.py::test_c3_surviving_worker_keeps_enforcing_airiatest_airia_guardrail_chaos.py::test_c4_owned_proxy_restart_resumes_airia_enforcementType
🆕 New Feature
✅ Test
Caveats (if any)
High
/v1/messagespost_call redaction never reaches the client (R12)model: null/v1/responsespost_call leaks the original text in deltas (R14)response.completedevent is redactedMedium
mode: during_callboots the proxy without Airiaunreachable_fallback: fail_openis ignored; Airia always fails closedLow
Final Attestation
Link to Devin session: https://app.devin.ai/sessions/7749cfdda367473a8133a29ac0415803
Open in Devin Desktop: https://app.devin.ai/desktop/session/7749cfdda367473a8133a29ac0415803?variant=devin
Requested by: @yucheng-berri