114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
import json
|
|
from typing import Any, Dict, List, Optional
|
|
from jsonschema import Draft202012Validator, ValidationError
|
|
from referencing import Registry, jsonschema as jsonschema_ref
|
|
|
|
from guardrails.actions.reask import SkeletonReAsk
|
|
from guardrails_ai.types import FailResult
|
|
|
|
|
|
class SchemaValidationError(Exception):
|
|
fields: Dict[str, List[str]] = {}
|
|
|
|
def __init__(self, *args: object, fields: Dict[str, List[str]]):
|
|
self.fields = fields
|
|
super().__init__(*args)
|
|
|
|
|
|
def validate_against_schema(
|
|
payload: Any,
|
|
validator: Draft202012Validator,
|
|
*,
|
|
validate_subschema: Optional[bool] = False,
|
|
):
|
|
fields: Dict[str, List[str]] = {}
|
|
error: ValidationError
|
|
for error in validator.iter_errors(payload):
|
|
if validate_subschema is True and error.message.endswith(
|
|
"is a required property"
|
|
):
|
|
continue
|
|
fields[error.json_path] = fields.get(error.json_path, [])
|
|
fields[error.json_path].append(error.message)
|
|
|
|
if fields:
|
|
error_message = (
|
|
"The provided payload is not compliant with the provided schema!"
|
|
)
|
|
raise SchemaValidationError(error_message, fields=fields)
|
|
|
|
|
|
def validate_json_schema(json_schema: Dict[str, Any]):
|
|
"""Validates a json_schema, against the JSON Meta Schema Draft 2020-12.
|
|
|
|
Raises a SchemaValidationError if invalid.
|
|
"""
|
|
json_schema_validator = Draft202012Validator(
|
|
{
|
|
"$ref": "https://json-schema.org/draft/2020-12/schema",
|
|
}
|
|
)
|
|
try:
|
|
validate_against_schema(json_schema, json_schema_validator)
|
|
except SchemaValidationError as e:
|
|
schema_name = json_schema.get("title", json_schema.get("$id"))
|
|
error_message = (
|
|
f"Schema {schema_name} is not compliant with JSON Schema Draft 2020-12!"
|
|
)
|
|
raise SchemaValidationError(error_message, fields=e.fields)
|
|
|
|
|
|
def validate_payload(
|
|
payload: Any,
|
|
json_schema: Dict[str, Any],
|
|
*,
|
|
validate_subschema: Optional[bool] = False,
|
|
):
|
|
"""Validates a payload, against the provided JSON Schema.
|
|
|
|
Raises a SchemaValidationError if invalid.
|
|
"""
|
|
schema_id = json_schema.get("$id", "temp-schema")
|
|
registry = Registry().with_resources(
|
|
[
|
|
(
|
|
f"urn:{schema_id}",
|
|
jsonschema_ref.DRAFT202012.create_resource(json_schema),
|
|
)
|
|
]
|
|
)
|
|
validator = Draft202012Validator(
|
|
{
|
|
"$ref": f"urn:{schema_id}",
|
|
},
|
|
registry=registry,
|
|
# TODO: Add custom checks for date: format,
|
|
# time: format, date-time: format, etc.
|
|
# format_checker=draft202012_format_checker
|
|
)
|
|
validate_against_schema(payload, validator, validate_subschema=validate_subschema)
|
|
|
|
|
|
def schema_validation(llm_output: Any, output_schema: Dict[str, Any], **kwargs):
|
|
validate_subschema = kwargs.get("validate_subschema", False)
|
|
|
|
schema_error = None
|
|
try:
|
|
validate_payload(
|
|
llm_output, output_schema, validate_subschema=validate_subschema
|
|
)
|
|
except SchemaValidationError as sve:
|
|
formatted_error_fields = json.dumps(sve.fields, indent=2)
|
|
schema_error = f"JSON does not match schema:\n{formatted_error_fields}"
|
|
|
|
if schema_error:
|
|
return SkeletonReAsk(
|
|
incorrectValue=llm_output,
|
|
failResults=[
|
|
FailResult(
|
|
fixValue=None,
|
|
errorMessage=schema_error,
|
|
)
|
|
],
|
|
)
|