Repository navigation
Add Guardrails AI scorer integration - #20038
Conversation
Implements mlflow.genai.scorers.guardrails module providing: - GuardrailsScorer base class wrapping Guardrails AI validators - Safety validators: ToxicLanguage, NSFWText, DetectJailbreak - PII validators: DetectPII, SecretsPresent - Quality validators: GibberishText - get_scorer() factory function for custom validator access Key design decisions: - Uses on_fail=OnFailAction.NOOP to return validation results instead of raising exceptions - Minimal mocking in tests: only mock get_validator_class (external hub), use real Guard - Follows Phoenix scorer pattern with validator registry - Returns Feedback objects with pass/fail values and rationale Resolves mlflow#20036 Signed-off-by: debu-sinha <debusinha2009@gmail.com>
🛠 DevTools 🛠
Install mlflow from this PRFor Databricks, use the following command: |
|
@debu-sinha Thank you for the contribution! Could you fix the following issue(s)? ⚠ Invalid PR templateThis PR does not appear to have been filed using the MLflow PR template. Please copy the PR template from here and fill it out. |
|
@smoorjani @B-Step62 - This implementation follows the "Simple Pattern" established by Phoenix and TruLens integrations: Architectural choices:
Why no LLM support:
This makes them fast, deterministic, and cost-free to run. Test approach (per feedback on Phoenix PR):
Happy to adjust if you prefer a different pattern. Looking forward to your review! |
Signed-off-by: debu-sinha <debusinha2009@gmail.com>
smoorjani
left a comment
There was a problem hiding this comment.
left a few comments here too, but overall LGTM! do we need to add guardrails into the test dependencies similar to the other 3p integrations?
| _logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| @experimental(version="3.9.0") |
There was a problem hiding this comment.
we missed the 3.9.0 release candidate, let's set this to 3.10.0 for the rc set in early Feb
| def __init__( | ||
| self, | ||
| validator_name: str | None = None, | ||
| **validator_kwargs: Any, |
There was a problem hiding this comment.
any reason we pass these kwargs during inference (L81) instead of when constructing the class?
| trace=trace, | ||
| ) | ||
|
|
||
| # Run validation |
There was a problem hiding this comment.
nit: let's remove these one line comments - they don't add much in terms of readability - same for places like L118, L122, and so on.
|
|
||
| # Extract validation outcome | ||
| passed = result.validation_passed | ||
| value = "pass" if passed else "fail" |
There was a problem hiding this comment.
can we use the constant value in mlflow/genai/judges/utils/__init__.py on L115/116
| rationale=rationale, | ||
| source=assessment_source, | ||
| metadata={ | ||
| FRAMEWORK_METADATA_KEY: "guardrails", |
There was a problem hiding this comment.
let's still include this in the error case as well.
|
|
||
|
|
||
| def get_validator_class(validator_name: str): | ||
| check_guardrails_installed() |
There was a problem hiding this comment.
we've already verified it is installed by this point, we can drop this
| scorer_cls = getattr(guardrails_scorers, scorer_class) | ||
| scorer = scorer_cls() | ||
|
|
||
| # Guard is REAL, validator is mocked (like Phoenix mocks model) |
There was a problem hiding this comment.
| # Guard is REAL, validator is mocked (like Phoenix mocks model) |
| assert isinstance(result, Feedback) | ||
| assert result.name == validator_name | ||
| assert result.value == "pass" | ||
| assert result.metadata["validation_passed"] is True |
There was a problem hiding this comment.
let's aim to do something like
assert result == Feedback(
name=validator_name,
value="pass",
metadata={...},
...
)for this test and the ones below.
|
|
||
| scorer = ToxicLanguage() | ||
|
|
||
| # Guard is REAL (core package), only validator is mocked (hub package) |
There was a problem hiding this comment.
| # Guard is REAL (core package), only validator is mocked (hub package) |
| ({"query": "test"}, None, "test"), | ||
| ], | ||
| ) | ||
| def test_map_scorer_inputs_to_text(inputs, outputs, expected): |
There was a problem hiding this comment.
can we also test with a real trace? similar to what you were doing in the trulens tests.
Changes: - Update @experimental version to 3.10.0 (missed 3.9.0 RC) - Remove one-liner comments that don't add readability value - Remove validation_passed from metadata (already reflected in value) - Add metadata to error case for consistency - Remove redundant check_guardrails_installed() call from registry - Remove unused list_available_validators() function - Add guardrails-ai to CI test dependencies - Update tests to assert full Feedback fields - Add test with real trace object Signed-off-by: debu-sinha <debusinha2009@gmail.com>
|
Thanks for the review! I've addressed all comments: Code changes:
Test changes:
CI:
All 31 tests pass locally. |
- Remove one-liner category comments (Safety/PII/Quality validators) - Add Feedback import and isinstance checks in all tests - Expand test assertions to cover all Feedback fields - Remove unused mock_validator_class parameter from error_handling test Signed-off-by: debu-sinha <debusinha2009@gmail.com>
The registry was incorrectly mapping DetectJailbreak to DetectPromptInjection, which doesn't exist in the Guardrails Hub. The correct validator class is DetectJailbreak (hub://guardrails/detect_jailbreak). Changes: - Fix registry mapping from "DetectPromptInjection" to "DetectJailbreak" - Update test expectations for correct validator class name - Add detailed docstring with Args for threshold and device parameters Signed-off-by: debu-sinha <debusinha2009@gmail.com>
Updated examples that are verified to trigger detection with real validators: - DetectJailbreak: Use complete jailbreak prompt that triggers detection - SecretsPresent: Use full API key pattern that triggers detection Signed-off-by: debu-sinha <debusinha2009@gmail.com>
smoorjani
left a comment
There was a problem hiding this comment.
Mostly looks good! left a few remaining comments to address
| GuardrailsScorer instance that can be called with MLflow's scorer interface | ||
|
|
||
| Examples: | ||
| >>> scorer = get_scorer("ToxicLanguage", threshold=0.7) |
There was a problem hiding this comment.
can we use the .. code-block:: python format for examples in docstring?
| except ImportError: | ||
| raise MlflowException.invalid_parameter_value( | ||
| "Guardrails AI scorers require the 'guardrails-ai' package. " | ||
| "Install it with: pip install guardrails-ai" |
There was a problem hiding this comment.
| "Install it with: pip install guardrails-ai" | |
| "Install it with: pip install `guardrails-ai`" |
|
|
||
|
|
||
| def test_get_validator_class_unknown(): | ||
| from mlflow.genai.scorers.guardrails.registry import get_validator_class |
There was a problem hiding this comment.
can we keep this import at the top of the file? same for the test above
| attributes["mlflow.spanOutputs"] = json.dumps(outputs) | ||
| attributes["mlflow.spanType"] = json.dumps("CHAIN") | ||
|
|
||
| otel_span = OTelReadableSpan( |
There was a problem hiding this comment.
can we use mlflow.start_span instead of directly instantiating trace classes. This is more akin to traces users will use for testing.
|
|
||
| def test_check_guardrails_installed_success(): | ||
| with patch.dict("sys.modules", {"guardrails": object()}): | ||
| from mlflow.genai.scorers.guardrails.utils import check_guardrails_installed |
There was a problem hiding this comment.
can we use top-level imports?
|
|
||
|
|
||
| def test_check_guardrails_installed_success(): | ||
| with patch.dict("sys.modules", {"guardrails": object()}): |
There was a problem hiding this comment.
this shouldn't be necessary if the library is installed in our test suite - let's remove this
|
|
||
|
|
||
| def test_map_scorer_inputs_to_text_with_trace(): | ||
| from mlflow.genai.scorers.guardrails.utils import map_scorer_inputs_to_text |
There was a problem hiding this comment.
same comment to move imports to file-level instead of function-level, unless absolutely necessary.
|
|
||
| # Extract validation outcome | ||
| passed = result.validation_passed | ||
| value = "pass" if passed else "fail" |
| def __init__( | ||
| self, | ||
| validator_name: str | None = None, | ||
| **validator_kwargs: Any, |
- Add _FRAMEWORK_NAME constant and use "guardrails-ai" instead of "guardrails" - Simplify registry to a list since key and value are the same - Add backticks around `guardrails-ai` in error message - Move imports to file-level in test files Signed-off-by: debu-sinha <debusinha2009@gmail.com>
Signed-off-by: debu-sinha <debusinha2009@gmail.com>
|
Addressed all review comments:
Regarding the kwargs question: The All 31 tests pass. Ready for re-review. |
- Use mlflow.start_span() instead of manually instantiating trace classes - Remove unnecessary test that patches guardrails (library is installed in test suite) Signed-off-by: debu-sinha <debusinha2009@gmail.com>
|
Additional changes addressing remaining review feedback: test_utils.py simplifications:
Clarification on kwargs question: self._guard = Guard().use(validator_class, on_fail=OnFailAction.NOOP, **validator_kwargs)This creates the Guard with the validator and all kwargs when the scorer is instantiated. Already implemented:
All 30 tests pass. Ready for re-review. |
|
Regarding the kwargs question at L66/82: The kwargs (like def __init__(self, validator_name: str | None = None, **validator_kwargs: Any):
...
self._guard = Guard().use(validator_class, on_fail=OnFailAction.NOOP, **validator_kwargs)The |
Signed-off-by: debu-sinha <debusinha2009@gmail.com>
|
Added trace-based tests for Guardrails scorers in
These tests use |
Summary of ChangesAll review comments have been addressed: Code Changes
Test Changes
Design ClarificationThe kwargs (like |
smoorjani
left a comment
There was a problem hiding this comment.
LGTM! Thanks for driving this!
|
Documentation preview for 4166d69 is available at: More info
|
AveshCSingh
left a comment
There was a problem hiding this comment.
I took a quick pass and didn't notice anything surprising, but defer to Samraj's detailed review
| except ImportError: | ||
| raise MlflowException.invalid_parameter_value( | ||
| "Guardrails AI scorers require the `guardrails-ai` package. " | ||
| "Install it with: pip install `guardrails-ai`" |
There was a problem hiding this comment.
| "Install it with: pip install `guardrails-ai`" | |
| "Install it with: `pip install guardrails-ai`" |
Signed-off-by: debu-sinha <debusinha2009@gmail.com>
Related Issues/PRs
Resolves #20036
What changes are proposed in this pull request?
Adds integration with Guardrails AI validators as MLflow scorers, enabling LLM safety, PII detection, and content quality evaluation.
New scorers:
ToxicLanguage,NSFWText,DetectJailbreakDetectPII,SecretsPresentGibberishTextget_scorer()for custom validator accessKey Design Decisions:
on_fail=OnFailAction.NOOPto return validation results instead of raising exceptionsget_validator_class(external hub), use realGuardFeedbackobjects with pass/fail values and rationaleHow is this PR tested?
Unit tests (31 passing):
Manual test examples (verified with real validators):
ToxicLanguage - PASS case:
ToxicLanguage - FAIL case:
DetectPII - FAIL case:
DetectJailbreak - FAIL case:
SecretsPresent - FAIL case:
GibberishText - FAIL case:
Does this PR require documentation update?
Release Notes
Is this a user-facing change?
Added Guardrails AI integration for MLflow scorers, providing safety validators (ToxicLanguage, NSFWText, DetectJailbreak), PII detectors (DetectPII, SecretsPresent), and quality validators (GibberishText) for evaluating LLM outputs.
What component(s), interfaces, languages, and integrations does this PR affect?
Components
area/tracking: Tracking Service, tracking client APIs, autologgingarea/models: MLmodel format, model serialization/deserialization, flavorsarea/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registryarea/scoring: MLflow Model server, model deployment tools, Spark UDFsarea/evaluation: MLflow model evaluation features, evaluation metrics, and evaluation workflowsarea/gateway: MLflow AI Gateway client APIs, server, and third-party integrationsarea/prompts: MLflow prompt engineering features, prompt templates, and prompt managementarea/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionalityarea/projects: MLproject format, project running backendsarea/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev serverarea/build: Build and test infrastructure for MLflowarea/docs: MLflow documentation pagesHow should the PR be classified in the release notes? Choose one:
rn/none- No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" sectionrn/breaking-change- The PR will be mentioned in the "Breaking Changes" sectionrn/feature- A new user-facing feature worth mentioning in the release notesrn/bug-fix- A user-facing bug fix worth mentioning in the release notesrn/documentation- A user-facing documentation change worth mentioning in the release notesShould this PR be included in the next patch release?