Skip to content
50 changes: 48 additions & 2 deletions backend/app/models/llm/request.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@

import sqlalchemy as sa
from pydantic import (
ConfigDict,
Discriminator,
HttpUrl,
Tag,
SerializerFunctionWrapHandler,
model_serializer,
model_validator,
Expand Down Expand Up @@ -464,7 +467,50 @@ class ProxyCompletionConfig(SQLModel):


class Validator(SQLModel):
"""Reference to a validator config saved in the guardrails service."""

validator_config_id: UUID
is_enabled: bool = Field(
default=True,
description="When false, the reference is neither fetched nor run.",
)


class InlineValidatorTypeEnum(str, Enum):
GENDER_ASSUMPTION_BIAS = "gender_assumption_bias"
PII_REMOVER = "pii_remover"
PROFANITY_FREE = "profanity_free"
TOPIC_RELEVANCE_LLM = "topic_relevance_llm"


class InlineValidator(SQLModel):
"""Validator spec stored directly in the config blob.

Per-type params (on_fail, languages, ban_list_id, ...) are extras validated by the
guardrails service, so they are allowed through untouched here.
"""

model_config = ConfigDict(extra="allow") # pyright: ignore[reportAssignmentType]

type: InlineValidatorTypeEnum = Field(..., description="Guardrails validator type")
is_enabled: bool = Field(
default=True,
description="When false, the spec is kept in the blob but not run.",
)


def _guardrail_entry_tag(value: Any) -> str:
if isinstance(value, dict):
return "ref" if "validator_config_id" in value else "inline"
return "ref" if hasattr(value, "validator_config_id") else "inline"


# Explicit tag: InlineValidator allows extras, so a smart union could also
# swallow a reference dict that happens to carry a `type`.
GuardrailEntry = Annotated[
Annotated[Validator, Tag("ref")] | Annotated[InlineValidator, Tag("inline")],
Discriminator(_guardrail_entry_tag),
]


class PromptTemplate(SQLModel):
Expand Down Expand Up @@ -500,12 +546,12 @@ def _default_proxy_provider(cls, data: Any) -> Any:
description="Prompt template with {{input}} placeholder to wrap around the user input",
)

input_guardrails: list[Validator] | None = Field(
input_guardrails: list[GuardrailEntry] | None = Field(
default=None,
description="Guardrails applied to validate/sanitize the input before the LLM call",
)

output_guardrails: list[Validator] | None = Field(
output_guardrails: list[GuardrailEntry] | None = Field(
default=None,
description="Guardrails applied to validate/sanitize the output after the LLM call",
)
Expand Down
67 changes: 45 additions & 22 deletions backend/app/services/llm/guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,18 @@
from opentelemetry import trace

from app.core.config import settings
from app.models.llm.request import Validator
from app.models.llm.request import InlineValidator, Validator

logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)


GUARDRAILS_PROXY_TIMEOUT_SECONDS = 30.0

# The service logs validator_log.stage from the spec and defaults to input.
GUARDRAIL_STAGE_INPUT = "input"
GUARDRAIL_STAGE_OUTPUT = "output"

# Upstream supports more; only these are productised for tenants.
EXPOSED_VALIDATOR_TYPES = frozenset(
{"topic_relevance_llm", "uli_slur_match", "gender_assumption_bias", "pii_remover"}
Expand Down Expand Up @@ -74,7 +78,7 @@ def proxy_guardrails_request(
url = f"{settings.KAAPI_GUARDRAILS_URL}{path}"
headers = _guardrails_headers(organization_id, project_id)
# Unset query params must be omitted, not sent as empty values.
query = {k: v for k, v in (params or {}).items() if v is not None}
query = {key: value for key, value in (params or {}).items() if value is not None}

try:
with (
Expand Down Expand Up @@ -173,7 +177,7 @@ def blocked(self) -> bool:
def apply_guardrails(
*,
text: str,
validators: list[Validator] | None,
validators: Sequence[Validator | InlineValidator] | None,
job_id: UUID,
project_id: int | None,
organization_id: int | None,
Expand All @@ -189,32 +193,50 @@ def apply_guardrails(

Args:
text: Primary text to validate. Sent as `input` to the service.
validators: Validator references (config IDs). When None/empty the
function short-circuits with a no-op outcome.
validators: Config-ID references or inline specs.
Entries with is_enabled=False are skipped; when nothing enabled
remains the function short-circuits with a no-op outcome.
output_text: When provided, sent as `output` to the service for
output-guardrail validators that compare input/output pairs.
Also routes the IDs through the output-config fetch path.
"""
if not validators:
enabled = [validator for validator in validators or [] if validator.is_enabled]
if not enabled:
return GuardrailsOutcome(
safe_text=text, error=None, bypassed=False, rephrase_needed=False, raw={}
)

is_output = output_text is not None
input_cfgs, output_cfgs = list_validators_config(
organization_id=organization_id,
project_id=project_id,
input_validator_configs=None if is_output else validators,
output_validator_configs=validators if is_output else None,
)
resolved = output_cfgs if is_output else input_cfgs
if not resolved:
logger.info(
f"[apply_guardrails] No validator configs resolved upstream; skipping "
f"POST /guardrails. job_id={job_id}, requested_ids="
f"{[str(v.validator_config_id) for v in validators]}, "
f"is_output={is_output}"
stage = GUARDRAIL_STAGE_OUTPUT if is_output else GUARDRAIL_STAGE_INPUT
references = [
validator for validator in enabled if isinstance(validator, Validator)
]
# The service's validator configs are extra="forbid", so is_enabled must go.
inline_specs = [
{**validator.model_dump(mode="json", exclude={"is_enabled"}), "stage": stage}
for validator in enabled
if isinstance(validator, InlineValidator)
]

fetched: list[dict[str, Any]] = []
if references:
input_cfgs, output_cfgs = list_validators_config(
organization_id=organization_id,
project_id=project_id,
input_validator_configs=None if is_output else references,
output_validator_configs=references if is_output else None,
)
fetched = output_cfgs if is_output else input_cfgs
if not fetched:
logger.info(
f"[apply_guardrails] No validator configs resolved upstream | "
f"job_id={job_id}, is_output={is_output}, requested_ids="
f"{[str(v.validator_config_id) for v in references]}"
)

# Blobs are all-refs or all-inline, so no interleaving or ordering is needed.
resolved = fetched + inline_specs
if not resolved:
return GuardrailsOutcome(
safe_text=text, error=None, bypassed=False, rephrase_needed=False, raw={}
)
Expand Down Expand Up @@ -307,10 +329,11 @@ def run_guardrails_validation(
Returns:
JSON response from the guardrails service with validation results.
"""
# is_enabled is unknown to the service's extra="forbid" validator schemas.
validators = [
validator.model_dump(mode="json")
if isinstance(validator, Validator)
else validator
validator
if isinstance(validator, dict)
else validator.model_dump(mode="json", exclude={"is_enabled"})
for validator in guardrail_config
]

Expand Down
96 changes: 95 additions & 1 deletion backend/app/tests/models/llm/test_request.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,15 @@
from app.models.llm.request import build_kaapi_completion_config
from uuid import uuid4

import pytest
from pydantic import ValidationError

from app.models.llm.request import (
ConfigBlob,
InlineValidator,
NativeCompletionConfig,
Validator,
build_kaapi_completion_config,
)
from app.services.llm.mappers import kaapi_params_as_dict


Expand Down Expand Up @@ -67,3 +78,86 @@ def test_unset_temperature_survives_json_round_trip(self) -> None:
assert "temperature" not in kaapi_params_as_dict(
round_tripped.completion.params
)


def _blob(input_guardrails: list) -> ConfigBlob:
return ConfigBlob(
completion=NativeCompletionConfig(
provider="openai-native", type="text", params={"model": "gpt-4o"}
),
input_guardrails=input_guardrails,
)


class TestGuardrailEntryParsing:
def test_dict_with_config_id_parses_to_validator(self) -> None:
config_id = uuid4()
blob = _blob([{"validator_config_id": str(config_id)}])

(entry,) = blob.input_guardrails or []
assert isinstance(entry, Validator)
assert entry.validator_config_id == config_id
assert entry.is_enabled is True

def test_dict_with_type_parses_to_inline_and_keeps_extras(self) -> None:
blob = _blob(
[
{
"type": "topic_relevance_llm",
"on_fail": "fix",
"languages": ["en", "hi"],
}
]
)

(entry,) = blob.input_guardrails or []
assert isinstance(entry, InlineValidator)
assert entry.model_dump() == {
"type": "topic_relevance_llm",
"on_fail": "fix",
"is_enabled": True,
"languages": ["en", "hi"],
}

def test_ref_dict_carrying_type_still_parses_to_validator(self) -> None:
# InlineValidator allows extras, so without the discriminator a smart
# union could claim this ref.
config_id = uuid4()
blob = _blob([{"validator_config_id": str(config_id), "type": "ban_list"}])

(entry,) = blob.input_guardrails or []
assert isinstance(entry, Validator)
assert entry.model_dump() == {
"validator_config_id": config_id,
"is_enabled": True,
}

def test_unsupported_inline_type_is_rejected(self) -> None:
with pytest.raises(ValidationError):
_blob([{"type": "ban_list"}])

def test_revalidating_parsed_entries_keeps_types(self) -> None:
parsed = _blob(
[{"validator_config_id": str(uuid4())}, {"type": "pii_remover"}]
).input_guardrails

rebuilt = _blob(list(parsed or [])).input_guardrails or []

assert [type(e) for e in rebuilt] == [Validator, InlineValidator]

def test_is_enabled_false_survives_json_round_trip(self) -> None:
config_id = uuid4()
dumped = _blob(
[
{"validator_config_id": str(config_id), "is_enabled": False},
{"type": "pii_remover", "is_enabled": False},
]
).model_dump(mode="json")

ref, inline = ConfigBlob.model_validate(dumped).input_guardrails or []

assert isinstance(ref, Validator)
assert ref.validator_config_id == config_id
assert ref.is_enabled is False
assert isinstance(inline, InlineValidator)
assert inline.is_enabled is False
Loading
Loading