diff --git a/CLAUDE.md b/CLAUDE.md index 1d29306d..27cf66df 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -106,7 +106,7 @@ Both `HuggingFaceProvider` and `LlamafileProvider` implement it and return the s - `prompt_token_count` / `completion_token_count`: ints, or `None` when the backend doesn't surface them (llamafile tokenizes server-side). - `raw`: provider-specific raw output (HF tensor / OpenAI JSON), for callers that need it. -`GraniteGuardian`, `LlamaGuard`, and the issue-#179 decoder guardrails (`WildGuard`, `DynaGuard`, `NemotronContentSafety`, `PolyGuard`, `KananaSafeguard`, `GptOssSafeguard`, `Prometheus`, `CompassJudger`, `Selene`) consume `generate_chat()` instead of touching `provider.tokenizer.apply_chat_template` + `provider.model.generate` directly, which is what lets them swap between the HF and llamafile backends. They no longer `import torch`. `chat_template_kwargs` (e.g. RAG `documents`, `available_tools`) and `generation_kwargs` (e.g. `pad_token_id` for Llama Guard 3) are pass-throughs. +`GraniteGuardian`, `LlamaGuard`, and the decoder guardrails from issues #179/#93 (`WildGuard`, `DynaGuard`, `NemotronContentSafety`, `PolyGuard`, `KananaSafeguard`, `GptOssSafeguard`, `Qwen3Guard`, `Prometheus`, `CompassJudger`, `Selene`) consume `generate_chat()` instead of touching `provider.tokenizer.apply_chat_template` + `provider.model.generate` directly, which is what lets them swap between the HF and llamafile backends. They no longer `import torch`. `chat_template_kwargs` (e.g. RAG `documents`, `available_tools`) and `generation_kwargs` (e.g. `pad_token_id` for Llama Guard 3) are pass-throughs. Two more `generate_chat()` flags exist for models that don't fit the plain chat→decode shape: `skip_special_tokens` (set `False` when the verdict *is* a special token, e.g. Kanana's ``/``) and `apply_chat_template` (set `False` to feed `messages[0]["content"]` to the model as a raw prompt, for models shipping their own instruction wrapper, e.g. WildGuard). `LlamafileProvider` ignores `skip_special_tokens` (server-side decoding) and rejects `apply_chat_template=False` — those models are HF-only. @@ -115,7 +115,7 @@ Two more `generate_chat()` flags exist for models that don't fit the plain chat Each guardrail lives in its own subdirectory (e.g. `llama_guard/llama_guard.py`). They inherit from `Guardrail`, `ThreeStageGuardrail`, or `StandardGuardrail` depending on shape: - `Guardrail` directly: API-based or fully custom shape — `AnyLlm` (any-llm SDK), `Alinia` (HTTP API), and the library-wrapped span guardrail `LettuceDetect` (wraps the `lettucedetect` lib, emits `GuardrailOutput.spans`) and `GliGuard` (wraps `gliner2`). -- `ThreeStageGuardrail` with custom generics: generative/judge models or non-binary outputs — `GraniteGuardian`, `LlamaGuard`, `Glider`, `Flowjudge`, `AzureContentSafety`, `DuoGuard` (multi-label), `OffTopic`, the decoder safety classifiers `WildGuard` / `DynaGuard` / `NemotronContentSafety` / `PolyGuard` / `KananaSafeguard` / `GptOssSafeguard`, and the rubric judges `Prometheus` / `CompassJudger` / `Selene`. +- `ThreeStageGuardrail` with custom generics: generative/judge models or non-binary outputs — `GraniteGuardian`, `LlamaGuard`, `Glider`, `Flowjudge`, `AzureContentSafety`, `DuoGuard` (multi-label), `OffTopic`, the decoder safety classifiers `WildGuard` / `DynaGuard` / `NemotronContentSafety` / `PolyGuard` / `KananaSafeguard` / `GptOssSafeguard` / `Qwen3Guard`, the rubric judges `Prometheus` / `CompassJudger` / `Selene`, and `Qwen3GuardStream` (token-level streaming heads loaded as remote code; HF-only, requires `transformers<5`, drives `provider.model.stream_moderate_from_ids` directly and emits `spans`). - `StandardGuardrail`: simple binary classifiers — `Protectai`, `Deepset`, `Jasper`, `Sentinel`, `Pangolin`, `InjecGuard`, `HarmGuard`, `PromptGuard`. Also `BielikGuard` (multi-label via `multi_label=True`, like DuoGuard) and `ShieldGemma` (causal-LM-backed), which still fit the `StandardGuardrail` shape. Library-wrapped guardrails (`Flowjudge`, `LettuceDetect`, `GliGuard`) bypass the provider and call an upstream library directly, guarded by a top-of-module `try/except ImportError` that re-raises a helpful `pip install` hint from `__init__` (see `flowjudge.py`). `Flowjudge` also accepts a `model=` backend (any `flow_judge` backend — `Hf`/`Vllm`/`Llamafile`/`Baseten`), a prebuilt/preset `metric=`, and `generation_params=` for the default `Hf` backend. diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index ffb4ea98..de9fa662 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -47,9 +47,12 @@ * [OpenAI Moderation](api/guardrails/openai-moderation.md) * [Pangolin](api/guardrails/pangolin.md) * [Patronus](api/guardrails/patronus.md) + * [PolyGuard](api/guardrails/poly-guard.md) * [Prometheus](api/guardrails/prometheus.md) * [Prompt Guard 2](api/guardrails/prompt-guard.md) * [ProtectAI](api/guardrails/protectai.md) + * [Qwen3Guard](api/guardrails/qwen3-guard.md) + * [Qwen3Guard Stream](api/guardrails/qwen3-guard-stream.md) * [Selene](api/guardrails/selene.md) * [Sentinel](api/guardrails/sentinel.md) * [ShieldGemma](api/guardrails/shield-gemma.md) diff --git a/docs/api/guardrails/index.md b/docs/api/guardrails/index.md index 1db53460..bdd2020d 100644 --- a/docs/api/guardrails/index.md +++ b/docs/api/guardrails/index.md @@ -38,3 +38,7 @@ Available guardrails and their parameters. Select a guardrail to view its API de | [Selene](selene.md) | `GuardrailName.SELENE` | | [Lettuce_detect](lettuce-detect.md) | `GuardrailName.LETTUCE_DETECT` | | [Gli_guard](gli-guard.md) | `GuardrailName.GLI_GUARD` | +| [Watsonx_guardian](watsonx-guardian.md) | `GuardrailName.WATSONX_GUARDIAN` | +| [Patronus](patronus.md) | `GuardrailName.PATRONUS` | +| [Qwen3_guard](qwen3-guard.md) | `GuardrailName.QWEN3_GUARD` | +| [Qwen3_guard_stream](qwen3-guard-stream.md) | `GuardrailName.QWEN3_GUARD_STREAM` | diff --git a/docs/api/guardrails/patronus.md b/docs/api/guardrails/patronus.md new file mode 100644 index 00000000..9c33e54e --- /dev/null +++ b/docs/api/guardrails/patronus.md @@ -0,0 +1,88 @@ +# Patronus + +Wraps the Patronus AI Evaluate API for managed LLM evaluation / guardrailing. + +This is the hosted, pay-per-use counterpart to the locally-run +:class:`~any_guardrail.guardrails.glider.glider.Glider` (GLIDER) judge and the +Patronus Lynx hallucination model: the same paper-backed evaluators, served as +managed configurations behind a single ``/v1/evaluate`` endpoint. + +A single request runs one or more *evaluators*. Each evaluator is selected by +name (e.g. ``"lynx"`` for hallucination, ``"judge"`` for the managed +LLM-as-a-judge, ``"answer-relevance"``, toxicity / PII evaluators) and an +optional managed ``criteria`` alias (e.g. ``"patronus:hallucination"``, +``"patronus:prompt-injection"``). Each returns a pass/fail verdict, a raw +score in ``[0, 1]`` (higher is better; below ``0.5`` fails by default), and — +when ``explain_strategy`` is set — an explanation. + +Auth is via an API key. Obtain one from https://app.patronus.ai/ (free +Developer tier with starter credit) and set it via ``PATRONUS_API_KEY`` or +pass it directly. + +``GuardrailOutput`` mapping: + - ``valid`` combines the per-evaluator pass flags per ``success_strategy`` + (``"all_pass"`` → every evaluator must pass; ``"any_pass"`` → at least + one must). + - ``score`` is the canonical risk of the *riskiest* evaluator, + ``1 - min(score_raw)`` (since Patronus ``score_raw`` is higher-is-safer). + - ``categories`` lists one ``CategoryResult`` per evaluator (``name`` = + its criteria / evaluator id, ``triggered`` = it failed, ``score`` = + ``1 - score_raw``). + - ``explanation`` joins the evaluators' explanations when present. + - ``extra`` carries ``success_strategy`` and a per-evaluator breakdown; + ``raw`` is the full response body. + - Fails closed (``valid=False``, ``extra={"parse_failure": True}``) when + the response has no ``results``. + +Research backing: + - Deshpande et al., *GLIDER: Grading LLM Interactions and Decisions using + Explainable Ranking* (https://arxiv.org/abs/2412.14140, 2024). + - Ravi et al., *Lynx: An Open Source Hallucination Evaluation Model* + (https://arxiv.org/abs/2407.08488, 2024). + - Docs: https://docs.patronus.ai/ + +Args: + evaluators (list[dict]): The evaluators to run, each a dict with at least + an ``"evaluator"`` key (plus optional ``"criteria"`` / + ``"explain_strategy"``). Example: + ``[{"evaluator": "judge", "criteria": "patronus:prompt-injection"}]``. + api_key (str | None): Patronus API key. Falls back to ``PATRONUS_API_KEY``. + endpoint (str): Evaluate API endpoint. Defaults to + ``https://api.patronus.ai/v1/evaluate``. + success_strategy ("all_pass" | "any_pass"): How to combine multiple + evaluators into the ``valid`` verdict. Defaults to ``"all_pass"``. + tags (dict[str, str] | None): Optional tags forwarded with each request + for observability. + +## Supported Models + +- `patronus-evaluate` + +## Constructor + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `evaluators` | `list[dict[str, Any]]` | Yes | — | +| `api_key` | `str | None` | No | `None` | +| `endpoint` | `str` | No | `"https://api.patronus.ai/v1/evaluate"` | +| `success_strategy` | `Literal['all_pass', 'any_pass']` | No | `"all_pass"` | +| `tags` | `dict[str, str] | None` | No | `None` | + +Initialize the Patronus guardrail. + +Does not perform any network I/O — the API is only contacted on +``validate()``. + +## validate + +Run the configured evaluators against the supplied model interaction. + +**Parameters** + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `input_text` | `str` | Yes | — | +| `output_text` | `str | None` | No | `None` | +| `retrieved_context` | `str | list[str] | None` | No | `None` | + +**Returns:** `GuardrailOutput` diff --git a/docs/api/guardrails/qwen3-guard-stream.md b/docs/api/guardrails/qwen3-guard-stream.md new file mode 100644 index 00000000..83b100e2 --- /dev/null +++ b/docs/api/guardrails/qwen3-guard-stream.md @@ -0,0 +1,67 @@ +# Qwen3GuardStream + +Qwen3Guard-Stream — token-level streaming safety moderation (Apache-2.0). + +Classifier heads on a Qwen3 backbone (loaded as remote code) that judge the user +prompt as a whole and every assistant response token individually, each with a +three-level severity (``Safe`` / ``Controversial`` / ``Unsafe``, where +``Controversial`` means harmfulness is context-dependent). ``validate`` is a +non-streaming facade: it feeds the full prompt, then each ``output_text`` token +through the streaming API and aggregates the worst severity. ``valid`` is ``True`` +only when everything judged is ``Safe`` (``Controversial`` also passes when +``strict=False``); ``score`` maps the worst severity onto the canonical risk axis +(Safe 0.0, Controversial 0.5, Unsafe 1.0) and per-part severities are surfaced in +``extra``. In response mode, runs of flagged response tokens are returned as +``spans`` with character offsets into ``output_text``. Fails closed +(``valid=False`` with ``extra={"parse_failure": True}``) when the backend reports +no usable risk level. For the generative variants (``Qwen3Guard-Gen-*``), see +``Qwen3Guard``. + +HuggingFace-only: the model ships its classification heads as remote code, so a +user-supplied provider must be a ``HuggingFaceProvider`` constructed with +``trust_remote_code=True``. The remote modeling code currently requires +``transformers>=4.51,<5`` (transformers 5 removed APIs it relies on); construction +raises ``ImportError`` on transformers >= 5. + +For more information, see the model cards: + +- [Qwen3Guard-Stream-0.6B](https://huggingface.co/Qwen/Qwen3Guard-Stream-0.6B) (default). +- [Qwen3Guard-Stream-4B](https://huggingface.co/Qwen/Qwen3Guard-Stream-4B). +- [Qwen3Guard-Stream-8B](https://huggingface.co/Qwen/Qwen3Guard-Stream-8B). + +Args: + strict: If ``True`` (default), only ``Safe`` verdicts pass validation; set + ``False`` to let ``Controversial`` content pass (``valid=True``), leaving + it reflected only in ``score``, ``extra``, and ``spans``. + model_id: Optional HuggingFace model ID. Defaults to ``Qwen/Qwen3Guard-Stream-0.6B``. + provider: Optional pre-configured ``HuggingFaceProvider`` with + ``trust_remote_code=True``. Defaults to one loading the remote-code model. + +## Supported Models + +- `Qwen/Qwen3Guard-Stream-0.6B` +- `Qwen/Qwen3Guard-Stream-4B` +- `Qwen/Qwen3Guard-Stream-8B` + +## Constructor + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `strict` | `bool` | No | `True` | +| `model_id` | `str | None` | No | `None` | +| `provider` | `Provider[dict[str, Any], dict[str, Any]] | None` | No | `None` | + +Initialize the Qwen3GuardStream guardrail. + +## validate + +Moderate ``input_text`` (or, when ``output_text`` is given, the assistant response to it). + +**Parameters** + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `input_text` | `str` | Yes | — | +| `output_text` | `str | None` | No | `None` | + +**Returns:** `GuardrailOutput` diff --git a/docs/api/guardrails/qwen3-guard.md b/docs/api/guardrails/qwen3-guard.md new file mode 100644 index 00000000..fe8d1034 --- /dev/null +++ b/docs/api/guardrails/qwen3-guard.md @@ -0,0 +1,59 @@ +# Qwen3Guard + +Qwen3Guard-Gen — generative safety moderation with three-level severity (Apache-2.0). + +Decoder LLM whose chat template embeds the safety-classifier instruction: the user +prompt alone triggers prompt moderation; supplying an assistant ``output_text`` +switches to response moderation. The model reports a severity (``Safe`` / +``Controversial`` / ``Unsafe``, where ``Controversial`` means harmfulness is +context-dependent), the violated policy categories, and — in response mode — +whether the response is a refusal. ``valid`` is ``True`` only for ``Safe`` +verdicts (``Controversial`` also passes when ``strict=False``); ``score`` maps +the severity onto the canonical risk axis (Safe 0.0, Controversial 0.5, +Unsafe 1.0) and the verbatim severity is surfaced in ``extra["severity"]``. +Fails closed (``valid=False`` with ``extra={"parse_failure": True}``) when no +severity parses. For the token-level streaming variants +(``Qwen3Guard-Stream-*``), see ``Qwen3GuardStream``. + +For more information, see the model cards: + +- [Qwen3Guard-Gen-0.6B](https://huggingface.co/Qwen/Qwen3Guard-Gen-0.6B) (default). +- [Qwen3Guard-Gen-4B](https://huggingface.co/Qwen/Qwen3Guard-Gen-4B). +- [Qwen3Guard-Gen-8B](https://huggingface.co/Qwen/Qwen3Guard-Gen-8B). + +Args: + strict: If ``True`` (default), only ``Safe`` verdicts pass validation; set + ``False`` to let ``Controversial`` content pass (``valid=True``), leaving + it reflected only in ``score`` and ``extra["severity"]``. + model_id: Optional HuggingFace model ID. Defaults to ``Qwen/Qwen3Guard-Gen-0.6B``. + provider: Optional pre-configured provider. Defaults to a ``HuggingFaceProvider`` + loading a causal LM. + +## Supported Models + +- `Qwen/Qwen3Guard-Gen-0.6B` +- `Qwen/Qwen3Guard-Gen-4B` +- `Qwen/Qwen3Guard-Gen-8B` + +## Constructor + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `strict` | `bool` | No | `True` | +| `model_id` | `str | None` | No | `None` | +| `provider` | `Provider[dict[str, Any], dict[str, Any]] | None` | No | `None` | + +Initialize the Qwen3Guard guardrail. + +## validate + +Moderate ``input_text`` (or, when ``output_text`` is given, the assistant response to it). + +**Parameters** + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `input_text` | `str` | Yes | — | +| `output_text` | `str | None` | No | `None` | + +**Returns:** `GuardrailOutput` diff --git a/docs/api/guardrails/watsonx-guardian.md b/docs/api/guardrails/watsonx-guardian.md new file mode 100644 index 00000000..0cb19552 --- /dev/null +++ b/docs/api/guardrails/watsonx-guardian.md @@ -0,0 +1,86 @@ +# WatsonxGuardian + +Wraps IBM watsonx.ai's Text Detection / ``Guardian`` moderation API. + +This is the hosted, pay-per-use counterpart to the locally-run +:class:`~any_guardrail.guardrails.granite_guardian.granite_guardian.GraniteGuardian` +guardrail: the same Granite Guardian risk-detection family, served as a +purpose-built detection endpoint instead of running the weights yourself. + +The ``Guardian`` class (from the ``ibm-watsonx-ai`` SDK) screens text against +a configurable set of detectors. The default ``granite_guardian`` detector +covers the Granite Guardian risk catalogue (harm, social bias, violence, +jailbreak, profanity, sexual content, plus RAG groundedness / relevance); +``hap`` (hate-abuse-profanity) and ``pii`` detectors are also available. Each +detector returns zero or more *detections*, each locating a risky span with a +score. + +Auth is via an IBM Cloud IAM API key plus a region URL and a project (or +space). Obtain a key and project from https://dataplatform.cloud.ibm.com/ and +set them via ``WATSONX_APIKEY`` / ``WATSONX_URL`` / ``WATSONX_PROJECT_ID`` +(or ``WATSONX_SPACE_ID``), or pass them directly. A free Lite plan is +available. + +``GuardrailOutput`` mapping: + - ``valid = no detections were returned`` (the detection API only returns + detections at or above the configured threshold). + - ``score`` is the highest detection score; ``0.0`` when nothing was + detected. + - ``categories`` lists one ``CategoryResult`` per detection (``name`` = + the detected risk, ``triggered=True``, ``score`` = the detection score). + - ``spans`` lists one ``SpanResult`` per detection that carries character + offsets (watsonx detections locate the flagged substring). + - ``raw`` is the full response dict from ``Guardian.detect``. + +Research backing: + - Padhi et al., *Granite Guardian* (https://arxiv.org/abs/2412.07724, 2024). + - IBM tutorial: https://www.ibm.com/think/tutorials/llm-safeguards-granite-guardian-risk-detection + - SDK reference: https://ibm.github.io/watsonx-ai-python-sdk/fm_text_detection.html + +Args: + api_key (str | None): IBM Cloud IAM API key. Falls back to ``WATSONX_APIKEY``. + url (str | None): watsonx.ai region endpoint (e.g. + ``https://us-south.ml.cloud.ibm.com``). Falls back to ``WATSONX_URL``. + project_id (str | None): watsonx project ID. Falls back to + ``WATSONX_PROJECT_ID``. One of ``project_id`` / ``space_id`` is required. + space_id (str | None): watsonx deployment space ID. Falls back to + ``WATSONX_SPACE_ID``. + detectors (dict | None): Detector configuration forwarded to ``Guardian``. + Defaults to ``{"granite_guardian": {}}``. Pass e.g. + ``{"granite_guardian": {"threshold": 0.6}, "pii": {}}`` to tune it. + api_client (APIClient | None): A pre-built ``ibm_watsonx_ai.APIClient``. + When supplied, the credential arguments above are ignored and the + client is used as-is (useful for shared clients or testing). + +## Supported Models + +- `granite_guardian` + +## Constructor + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `api_key` | `str | None` | No | `None` | +| `url` | `str | None` | No | `None` | +| `project_id` | `str | None` | No | `None` | +| `space_id` | `str | None` | No | `None` | +| `detectors` | `dict[str, Any] | None` | No | `None` | +| `api_client` | `APIClient | None` | No | `None` | + +Initialize the guardrail and build the watsonx ``Guardian`` client. + +Building the client performs IAM authentication, so unlike the pure-REST +API guardrails this constructor does contact IBM Cloud (unless a +pre-built ``api_client`` is supplied). + +## validate + +Screen ``content`` against the configured watsonx detectors. + +**Parameters** + +| Parameter | Type | Required | Default | +|-----------|------|----------|---------| +| `content` | `str` | Yes | — | + +**Returns:** `GuardrailOutput` diff --git a/scripts/generate_api_docs.py b/scripts/generate_api_docs.py index 1bdd6d21..fbaf6284 100644 --- a/scripts/generate_api_docs.py +++ b/scripts/generate_api_docs.py @@ -415,6 +415,12 @@ def _guardrails_index_page() -> str: "WatsonxGuardian", "watsonx-guardian.md", ), + ("any_guardrail.guardrails.qwen3_guard.qwen3_guard", "Qwen3Guard", "qwen3-guard.md"), + ( + "any_guardrail.guardrails.qwen3_guard_stream.qwen3_guard_stream", + "Qwen3GuardStream", + "qwen3-guard-stream.md", + ), ] diff --git a/src/any_guardrail/base.py b/src/any_guardrail/base.py index 90a3126a..e870feed 100644 --- a/src/any_guardrail/base.py +++ b/src/any_guardrail/base.py @@ -60,6 +60,8 @@ class GuardrailName(StrEnum): GLI_GUARD = "gli_guard" WATSONX_GUARDIAN = "watsonx_guardian" PATRONUS = "patronus" + QWEN3_GUARD = "qwen3_guard" + QWEN3_GUARD_STREAM = "qwen3_guard_stream" class Guardrail(ABC): diff --git a/src/any_guardrail/guardrails/qwen3_guard/__init__.py b/src/any_guardrail/guardrails/qwen3_guard/__init__.py new file mode 100644 index 00000000..8fdde5e4 --- /dev/null +++ b/src/any_guardrail/guardrails/qwen3_guard/__init__.py @@ -0,0 +1,3 @@ +from .qwen3_guard import Qwen3Guard + +__all__ = ["Qwen3Guard"] diff --git a/src/any_guardrail/guardrails/qwen3_guard/qwen3_guard.py b/src/any_guardrail/guardrails/qwen3_guard/qwen3_guard.py new file mode 100644 index 00000000..cbce8c33 --- /dev/null +++ b/src/any_guardrail/guardrails/qwen3_guard/qwen3_guard.py @@ -0,0 +1,169 @@ +import re +from typing import Any, ClassVar + +from any_guardrail.base import GuardrailOutput, ThreeStageGuardrail +from any_guardrail.guardrails.utils import default +from any_guardrail.providers.base import StandardProvider +from any_guardrail.providers.huggingface import HuggingFaceProvider +from any_guardrail.types import ( + AnyDict, + CategoryResult, + ChatMessages, + GuardrailInferenceOutput, + GuardrailPreprocessOutput, + GuardrailUsage, +) + +Qwen3GuardPreprocessData = AnyDict +Qwen3GuardInferenceData = AnyDict + +# Qwen3Guard safety taxonomy; the model reports these names verbatim on its +# ``Categories:`` line ("Jailbreak" applies to prompt moderation only). +QWEN3GUARD_CATEGORIES = [ + "Violent", + "Non-violent Illegal Acts", + "Sexual Content or Sexual Acts", + "PII", + "Suicide & Self-Harm", + "Unethical Acts", + "Politically Sensitive Topics", + "Copyright Violation", + "Jailbreak", +] + +# Canonical risk mapping for the three severity levels. +SEVERITY_RISK = {"Safe": 0.0, "Controversial": 0.5, "Unsafe": 1.0} + +MAX_NEW_TOKENS = 128 + +_SAFETY = re.compile(r"Safety:\s*(Safe|Unsafe|Controversial)", re.IGNORECASE) +_CATEGORIES_LINE = re.compile(r"Categories:\s*(.+)", re.IGNORECASE) +# Deliberately case-sensitive and in model-card order so "Violent" cannot +# match inside "Non-violent Illegal Acts". +_CATEGORY = re.compile("|".join(re.escape(category) for category in QWEN3GUARD_CATEGORIES)) +_REFUSAL = re.compile(r"Refusal:\s*(Yes|No)", re.IGNORECASE) +_THINK_PATTERN = re.compile(r".*?", re.DOTALL) + + +class Qwen3Guard(ThreeStageGuardrail[Qwen3GuardPreprocessData, Qwen3GuardInferenceData]): + """Qwen3Guard-Gen — generative safety moderation with three-level severity (Apache-2.0). + + Decoder LLM whose chat template embeds the safety-classifier instruction: the user + prompt alone triggers prompt moderation; supplying an assistant ``output_text`` + switches to response moderation. The model reports a severity (``Safe`` / + ``Controversial`` / ``Unsafe``, where ``Controversial`` means harmfulness is + context-dependent), the violated policy categories, and — in response mode — + whether the response is a refusal. ``valid`` is ``True`` only for ``Safe`` + verdicts (``Controversial`` also passes when ``strict=False``); ``score`` maps + the severity onto the canonical risk axis (Safe 0.0, Controversial 0.5, + Unsafe 1.0) and the verbatim severity is surfaced in ``extra["severity"]``. + Fails closed (``valid=False`` with ``extra={"parse_failure": True}``) when no + severity parses. For the token-level streaming variants + (``Qwen3Guard-Stream-*``), see ``Qwen3GuardStream``. + + For more information, see the model cards: + + - [Qwen3Guard-Gen-0.6B](https://huggingface.co/Qwen/Qwen3Guard-Gen-0.6B) (default). + - [Qwen3Guard-Gen-4B](https://huggingface.co/Qwen/Qwen3Guard-Gen-4B). + - [Qwen3Guard-Gen-8B](https://huggingface.co/Qwen/Qwen3Guard-Gen-8B). + + Args: + strict: If ``True`` (default), only ``Safe`` verdicts pass validation; set + ``False`` to let ``Controversial`` content pass (``valid=True``), leaving + it reflected only in ``score`` and ``extra["severity"]``. + model_id: Optional HuggingFace model ID. Defaults to ``Qwen/Qwen3Guard-Gen-0.6B``. + provider: Optional pre-configured provider. Defaults to a ``HuggingFaceProvider`` + loading a causal LM. + + """ + + SUPPORTED_MODELS: ClassVar = [ + "Qwen/Qwen3Guard-Gen-0.6B", + "Qwen/Qwen3Guard-Gen-4B", + "Qwen/Qwen3Guard-Gen-8B", + ] + + def __init__( + self, + strict: bool = True, + model_id: str | None = None, + provider: StandardProvider | None = None, + ) -> None: + """Initialize the Qwen3Guard guardrail.""" + self.model_id = default(model_id, self.SUPPORTED_MODELS) + self.strict = strict + load_kwargs: AnyDict = {} + if provider is not None: + self.provider = provider + if isinstance(self.provider, HuggingFaceProvider): + from transformers import AutoModelForCausalLM, AutoTokenizer + + load_kwargs = {"model_class": AutoModelForCausalLM, "tokenizer_class": AutoTokenizer} + else: + from transformers import AutoModelForCausalLM, AutoTokenizer + + self.provider = HuggingFaceProvider(model_class=AutoModelForCausalLM, tokenizer_class=AutoTokenizer) + self.provider.load_model(self.model_id, **load_kwargs) + + def validate( # type: ignore[override] + self, input_text: str, output_text: str | None = None, **kwargs: Any + ) -> GuardrailOutput: + """Moderate ``input_text`` (or, when ``output_text`` is given, the assistant response to it).""" + result = super().validate(input_text, output_text=output_text, **kwargs) + if isinstance(result, list): + msg = "Qwen3Guard.validate received a list input but only supports single strings." + raise TypeError(msg) + return result + + def _pre_processing( + self, input_text: str, output_text: str | None = None, **kwargs: Any + ) -> GuardrailPreprocessOutput[Qwen3GuardPreprocessData]: + del kwargs + messages: ChatMessages = [{"role": "user", "content": input_text}] + if output_text is not None: + messages.append({"role": "assistant", "content": output_text}) + return GuardrailPreprocessOutput(data={"messages": messages, "has_response": output_text is not None}) + + def _inference( + self, model_inputs: GuardrailPreprocessOutput[Qwen3GuardPreprocessData] + ) -> GuardrailInferenceOutput[Qwen3GuardInferenceData]: + result = self.provider.generate_chat( + messages=model_inputs.data["messages"], max_new_tokens=MAX_NEW_TOKENS, do_sample=False + ) + # Carry has_response through so _post_processing surfaces the Refusal field only in response mode. + result.data["has_response"] = model_inputs.data["has_response"] + return result + + def _post_processing(self, model_outputs: GuardrailInferenceOutput[Qwen3GuardInferenceData]) -> GuardrailOutput: + text = model_outputs.data["generated_text"] + has_response = model_outputs.data.get("has_response", False) + without_think = _THINK_PATTERN.sub("", text).strip() + safety = _SAFETY.search(without_think) + if safety is None: + return GuardrailOutput(valid=False, explanation=text, extra={"parse_failure": True}) + severity = safety.group(1).capitalize() + categories: list[CategoryResult] = [] + categories_line = _CATEGORIES_LINE.search(without_think) + if categories_line is not None: + categories = [ + CategoryResult(name=name, triggered=True) + for name in dict.fromkeys(_CATEGORY.findall(categories_line.group(1))) + ] + if has_response: + # Refusal is auxiliary metadata about the judged response; a missing line + # is tolerated (triggered=None) rather than failing closed — the Safety + # line above already carries the verdict. + refusal = _REFUSAL.search(without_think) + triggered = refusal.group(1).lower() == "yes" if refusal else None + categories.append(CategoryResult(name="refusal", triggered=triggered)) + return GuardrailOutput( + valid=severity == "Safe" if self.strict else severity != "Unsafe", + explanation=text, + score=SEVERITY_RISK[severity], + categories=categories, + extra={"severity": severity}, + usage=GuardrailUsage( + prompt_tokens=model_outputs.data.get("prompt_token_count"), + completion_tokens=model_outputs.data.get("completion_token_count"), + ), + ) diff --git a/src/any_guardrail/guardrails/qwen3_guard_stream/__init__.py b/src/any_guardrail/guardrails/qwen3_guard_stream/__init__.py new file mode 100644 index 00000000..a8ac2fcc --- /dev/null +++ b/src/any_guardrail/guardrails/qwen3_guard_stream/__init__.py @@ -0,0 +1,3 @@ +from .qwen3_guard_stream import Qwen3GuardStream + +__all__ = ["Qwen3GuardStream"] diff --git a/src/any_guardrail/guardrails/qwen3_guard_stream/qwen3_guard_stream.py b/src/any_guardrail/guardrails/qwen3_guard_stream/qwen3_guard_stream.py new file mode 100644 index 00000000..51165174 --- /dev/null +++ b/src/any_guardrail/guardrails/qwen3_guard_stream/qwen3_guard_stream.py @@ -0,0 +1,319 @@ +from typing import Any, ClassVar + +from any_guardrail.base import GuardrailOutput, ThreeStageGuardrail +from any_guardrail.guardrails.utils import default +from any_guardrail.providers.base import StandardProvider +from any_guardrail.providers.huggingface import HuggingFaceProvider +from any_guardrail.types import ( + AnyDict, + CategoryResult, + ChatMessages, + GuardrailInferenceOutput, + GuardrailPreprocessOutput, + SpanResult, +) + +Qwen3GuardStreamPreprocessData = AnyDict +Qwen3GuardStreamInferenceData = AnyDict + +# Canonical risk mapping for the three severity levels. +SEVERITY_RISK = {"Safe": 0.0, "Controversial": 0.5, "Unsafe": 1.0} + +# (severity, category, token record) for one moderated response token. +_ResponseVerdict = tuple[str, str | None, AnyDict] + + +def _last_verdict(result: Any) -> tuple[str | None, str | None]: + """Extract the newest (severity, category) pair from a ``stream_moderate_from_ids`` result. + + Returns ``(None, None)`` when no risk level is present or it is not one of the + three known severities. The model's ``"None"`` category (its "no violation" + marker) is normalized to ``None``. + """ + if not isinstance(result, dict) or not result.get("risk_level"): + return None, None + severity = str(result["risk_level"][-1]).capitalize() + if severity not in SEVERITY_RISK: + return None, None + categories = result.get("category") + category = categories[-1] if categories else None + if not isinstance(category, str) or category == "None": + category = None + return severity, category + + +def _to_span(run: AnyDict, output_text: str | None) -> SpanResult: + return SpanResult( + start=run["start"], + end=run["end"], + text=output_text[run["start"] : run["end"]] if output_text is not None else None, + label=run["category"], + score=SEVERITY_RISK[run["severity"]], + ) + + +def _build_spans(response_verdicts: list[_ResponseVerdict], output_text: str | None) -> list[SpanResult]: + """Merge consecutive flagged response tokens into character spans over ``output_text``. + + Runs split when the severity or category changes; tokens without offsets + (template scaffolding, or a tokenizer without offset support) are skipped. + """ + spans: list[SpanResult] = [] + current: AnyDict | None = None + for severity, category, record in response_verdicts: + if severity == "Safe" or record.get("start") is None: + if current is not None: + spans.append(_to_span(current, output_text)) + current = None + continue + if current is not None and current["severity"] == severity and current["category"] == category: + current["end"] = record["end"] + else: + if current is not None: + spans.append(_to_span(current, output_text)) + current = {"severity": severity, "category": category, "start": record["start"], "end": record["end"]} + if current is not None: + spans.append(_to_span(current, output_text)) + return spans + + +class Qwen3GuardStream(ThreeStageGuardrail[Qwen3GuardStreamPreprocessData, Qwen3GuardStreamInferenceData]): + """Qwen3Guard-Stream — token-level streaming safety moderation (Apache-2.0). + + Classifier heads on a Qwen3 backbone (loaded as remote code) that judge the user + prompt as a whole and every assistant response token individually, each with a + three-level severity (``Safe`` / ``Controversial`` / ``Unsafe``, where + ``Controversial`` means harmfulness is context-dependent). ``validate`` is a + non-streaming facade: it feeds the full prompt, then each ``output_text`` token + through the streaming API and aggregates the worst severity. ``valid`` is ``True`` + only when everything judged is ``Safe`` (``Controversial`` also passes when + ``strict=False``); ``score`` maps the worst severity onto the canonical risk axis + (Safe 0.0, Controversial 0.5, Unsafe 1.0) and per-part severities are surfaced in + ``extra``. In response mode, runs of flagged response tokens are returned as + ``spans`` with character offsets into ``output_text``. Fails closed + (``valid=False`` with ``extra={"parse_failure": True}``) when the backend reports + no usable risk level. For the generative variants (``Qwen3Guard-Gen-*``), see + ``Qwen3Guard``. + + HuggingFace-only: the model ships its classification heads as remote code, so a + user-supplied provider must be a ``HuggingFaceProvider`` constructed with + ``trust_remote_code=True``. The remote modeling code currently requires + ``transformers>=4.51,<5`` (transformers 5 removed APIs it relies on); construction + raises ``ImportError`` on transformers >= 5. + + For more information, see the model cards: + + - [Qwen3Guard-Stream-0.6B](https://huggingface.co/Qwen/Qwen3Guard-Stream-0.6B) (default). + - [Qwen3Guard-Stream-4B](https://huggingface.co/Qwen/Qwen3Guard-Stream-4B). + - [Qwen3Guard-Stream-8B](https://huggingface.co/Qwen/Qwen3Guard-Stream-8B). + + Args: + strict: If ``True`` (default), only ``Safe`` verdicts pass validation; set + ``False`` to let ``Controversial`` content pass (``valid=True``), leaving + it reflected only in ``score``, ``extra``, and ``spans``. + model_id: Optional HuggingFace model ID. Defaults to ``Qwen/Qwen3Guard-Stream-0.6B``. + provider: Optional pre-configured ``HuggingFaceProvider`` with + ``trust_remote_code=True``. Defaults to one loading the remote-code model. + + """ + + SUPPORTED_MODELS: ClassVar = [ + "Qwen/Qwen3Guard-Stream-0.6B", + "Qwen/Qwen3Guard-Stream-4B", + "Qwen/Qwen3Guard-Stream-8B", + ] + + def __init__( + self, + strict: bool = True, + model_id: str | None = None, + provider: StandardProvider | None = None, + ) -> None: + """Initialize the Qwen3GuardStream guardrail.""" + self.model_id = default(model_id, self.SUPPORTED_MODELS) + self.strict = strict + self._require_supported_transformers() + load_kwargs: AnyDict = {} + if provider is not None: + if isinstance(provider, HuggingFaceProvider): + if not provider.trust_remote_code: + msg = ( + "Qwen3Guard-Stream ships its classification heads as remote code; construct the " + "provider with HuggingFaceProvider(trust_remote_code=True) so both the model and " + "tokenizer load it." + ) + raise ValueError(msg) + from transformers import AutoModel, AutoTokenizer + + load_kwargs = {"model_class": AutoModel, "tokenizer_class": AutoTokenizer} + self.provider = provider + else: + from transformers import AutoModel, AutoTokenizer + + self.provider = HuggingFaceProvider( + model_class=AutoModel, tokenizer_class=AutoTokenizer, trust_remote_code=True + ) + self.provider.load_model(self.model_id, **load_kwargs) + + def validate( # type: ignore[override] + self, input_text: str, output_text: str | None = None, **kwargs: Any + ) -> GuardrailOutput: + """Moderate ``input_text`` (or, when ``output_text`` is given, the assistant response to it).""" + result = super().validate(input_text, output_text=output_text, **kwargs) + if isinstance(result, list): + msg = "Qwen3GuardStream.validate received a list input but only supports single strings." + raise TypeError(msg) + return result + + def _pre_processing( + self, input_text: str, output_text: str | None = None, **kwargs: Any + ) -> GuardrailPreprocessOutput[Qwen3GuardStreamPreprocessData]: + del kwargs + messages: ChatMessages = [{"role": "user", "content": input_text}] + if output_text is not None: + messages.append({"role": "assistant", "content": output_text}) + return GuardrailPreprocessOutput( + data={"messages": messages, "has_response": output_text is not None, "output_text": output_text} + ) + + def _inference( + self, model_inputs: GuardrailPreprocessOutput[Qwen3GuardStreamPreprocessData] + ) -> GuardrailInferenceOutput[Qwen3GuardStreamInferenceData]: + tokenizer = self.provider.tokenizer # type: ignore[attr-defined] + model = self.provider.model # type: ignore[attr-defined] + has_response: bool = model_inputs.data["has_response"] + output_text: str | None = model_inputs.data["output_text"] + + text: str = tokenizer.apply_chat_template( + model_inputs.data["messages"], tokenize=False, add_generation_prompt=False, enable_thinking=False + ) + try: + encoding = tokenizer(text, return_tensors="pt", return_offsets_mapping=True) + offsets = encoding["offset_mapping"][0].tolist() + except NotImplementedError: # slow tokenizer: spans degrade gracefully + encoding = tokenizer(text, return_tensors="pt") + offsets = None + token_ids = encoding["input_ids"][0] + device = self.provider.device # type: ignore[attr-defined] + if device is not None: + token_ids = token_ids.to(device) + + user_end_index = self._user_turn_end(tokenizer, token_ids.tolist()) + response_base = text.rfind(output_text) if output_text else -1 + + stream_state = None + response_tokens: list[AnyDict] = [] + try: + prompt_result, stream_state = model.stream_moderate_from_ids( + token_ids[: user_end_index + 1], role="user", stream_state=None + ) + if has_response: + for index in range(user_end_index + 1, len(token_ids)): + result, stream_state = model.stream_moderate_from_ids( + token_ids[index], role="assistant", stream_state=stream_state + ) + record: AnyDict = {"result": result, "start": None, "end": None} + if offsets is not None and response_base >= 0 and output_text is not None: + # Clamp the token's offsets in the templated text onto the + # output_text region, shifting to output_text coordinates; + # scaffolding tokens outside it keep start=None. + start = max(int(offsets[index][0]), response_base) + end = min(int(offsets[index][1]), response_base + len(output_text)) + if start < end: + record["start"] = start - response_base + record["end"] = end - response_base + response_tokens.append(record) + finally: + if stream_state is not None: + model.close_stream(stream_state) + return GuardrailInferenceOutput( + data={ + "prompt_result": prompt_result, + "response_tokens": response_tokens, + "has_response": has_response, + "output_text": output_text, + } + ) + + def _post_processing( + self, model_outputs: GuardrailInferenceOutput[Qwen3GuardStreamInferenceData] + ) -> GuardrailOutput: + data = model_outputs.data + has_response: bool = data.get("has_response", False) + output_text: str | None = data.get("output_text") + + prompt_severity, prompt_category = _last_verdict(data.get("prompt_result")) + if prompt_severity is None: + return GuardrailOutput(valid=False, extra={"parse_failure": True}) + response_verdicts: list[_ResponseVerdict] = [] + for record in data.get("response_tokens", []): + severity, category = _last_verdict(record.get("result")) + if severity is None: + return GuardrailOutput(valid=False, extra={"parse_failure": True}) + response_verdicts.append((severity, category, record)) + + response_severities = [severity for severity, _, _ in response_verdicts] + worst = max([prompt_severity, *response_severities], key=lambda severity: SEVERITY_RISK[severity]) + + names: dict[str, None] = {} + if prompt_severity != "Safe" and prompt_category is not None: + names.setdefault(prompt_category) + for severity, category, _ in response_verdicts: + if severity != "Safe" and category is not None: + names.setdefault(category) + + extra: AnyDict = {"severity": worst, "prompt_severity": prompt_severity} + if has_response: + extra["response_severity"] = ( + max(response_severities, key=lambda severity: SEVERITY_RISK[severity]) + if response_severities + else "Safe" + ) + spans = _build_spans(response_verdicts, output_text) if has_response else [] + return GuardrailOutput( + valid=worst == "Safe" if self.strict else worst != "Unsafe", + score=SEVERITY_RISK[worst], + categories=[CategoryResult(name=name, triggered=True) for name in names], + spans=spans or None, + extra=extra, + ) + + def _require_supported_transformers(self) -> None: + """Raise when the installed transformers cannot run the remote modeling code. + + The Qwen3Guard-Stream model repos target transformers 4.x (>= 4.51 per the + model card); transformers 5 removed several APIs the remote code relies on + (the implicit ``pad_token_id`` config default, ``ROPE_INIT_FUNCTIONS["default"]``, + the old rotary-embedding weight-init interface). + """ + try: + from transformers import __version__ as transformers_version + except ImportError: + return # provider construction surfaces the install hint instead + if int(transformers_version.split(".")[0]) >= 5: + msg = ( + f"Qwen3Guard-Stream's remote modeling code (in the {self.model_id} model repo, not this " + f"library) is incompatible with transformers >= 5 (installed: {transformers_version}). " + "Install 'transformers>=4.51,<5' to use Qwen3GuardStream, or check the model repo for an " + "updated revision." + ) + raise ImportError(msg) + + @staticmethod + def _user_turn_end(tokenizer: Any, token_ids: list[int]) -> int: + """Index of the ``<|im_end|>`` closing the last user turn (model-card boundary scan).""" + im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>") + user_id = tokenizer.convert_tokens_to_ids("user") + im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>") + last_start = next( + (i for i in range(len(token_ids) - 1, -1, -1) if token_ids[i : i + 2] == [im_start_id, user_id]), + None, + ) + if last_start is None: + msg = "Could not locate the user turn in the tokenized chat template output." + raise ValueError(msg) + user_end = next((i for i in range(last_start + 2, len(token_ids)) if token_ids[i] == im_end_id), None) + if user_end is None: + msg = "Could not locate the end of the user turn in the tokenized chat template output." + raise ValueError(msg) + return user_end diff --git a/tests/integration/test_huggingface_guardrails.py b/tests/integration/test_huggingface_guardrails.py index d1f0ba62..6243710d 100644 --- a/tests/integration/test_huggingface_guardrails.py +++ b/tests/integration/test_huggingface_guardrails.py @@ -1,6 +1,7 @@ from typing import Any import pytest +import transformers from any_guardrail import AnyGuardrail, GuardrailName from any_guardrail.base import GuardrailOutput, ThreeStageGuardrail @@ -76,6 +77,16 @@ {"criteria": "Is the response safe?", "rubric": "Higher is safer.", "pass_threshold": 1}, marks=pytest.mark.heavy, # 7B ), + # --- issue #93 additions --- + (GuardrailName.QWEN3_GUARD, {}), # 0.6B decoder + pytest.param( + GuardrailName.QWEN3_GUARD_STREAM, # 0.6B, remote-code streaming heads + {}, + marks=pytest.mark.skipif( + int(transformers.__version__.split(".")[0]) >= 5, + reason="Qwen3Guard-Stream's remote modeling code requires transformers<5", + ), + ), ], ) def test_huggingface_guardrails(guardrail_name: GuardrailName, guardrail_kwargs: dict[str, Any]) -> None: diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py index 6a2a300c..a0ee7b25 100644 --- a/tests/unit/test_api.py +++ b/tests/unit/test_api.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import pytest +import transformers from any_guardrail import AnyGuardrail, GuardrailName from any_guardrail.base import Guardrail, ThreeStageGuardrail @@ -138,6 +139,10 @@ def test_model_load() -> None: pass_threshold=1, ) assert guardrail.model == "mocked_model" # type: ignore[attr-defined] + elif guardrail_name == GuardrailName.QWEN3_GUARD_STREAM and int(transformers.__version__.split(".")[0]) >= 5: + # Construction is version-gated: the model repo's remote code needs transformers<5. + with pytest.raises(ImportError, match="transformers"): + AnyGuardrail.create(guardrail_name=guardrail_name) elif guardrail_name == GuardrailName.DUOGUARD: mock_provider = MagicMock(spec=HuggingFaceProvider) mock_provider.tokenizer = MagicMock() diff --git a/tests/unit/test_unit_new_guardrails.py b/tests/unit/test_unit_new_guardrails.py index 69b7bcc1..eb57bf0c 100644 --- a/tests/unit/test_unit_new_guardrails.py +++ b/tests/unit/test_unit_new_guardrails.py @@ -17,6 +17,8 @@ from any_guardrail.guardrails.poly_guard.poly_guard import PolyGuard from any_guardrail.guardrails.prometheus.prometheus import Prometheus from any_guardrail.guardrails.prompt_guard.prompt_guard import _build_output as prompt_build +from any_guardrail.guardrails.qwen3_guard.qwen3_guard import Qwen3Guard +from any_guardrail.guardrails.qwen3_guard_stream.qwen3_guard_stream import Qwen3GuardStream from any_guardrail.guardrails.selene.selene import Selene from any_guardrail.guardrails.wild_guard.wild_guard import WildGuard from any_guardrail.types import GuardrailInferenceOutput, GuardrailOutput @@ -204,6 +206,200 @@ def test_gpt_oss_safeguard_fails_closed() -> None: assert result.extra == {"parse_failure": True} +def _qwen3(strict: bool = True) -> Qwen3Guard: + instance = object.__new__(Qwen3Guard) + instance.strict = strict + return instance + + +def test_qwen3guard_safe_none_category_not_surfaced() -> None: + result = _qwen3()._post_processing(_gen("Safety: Safe\nCategories: None")) + assert result.valid is True + assert result.score == 0.0 + assert result.extra == {"severity": "Safe"} + assert result.categories == [] + + +def test_qwen3guard_unsafe_single_category() -> None: + result = _qwen3()._post_processing(_gen("Safety: Unsafe\nCategories: Violent")) + assert result.valid is False + assert result.score == 1.0 + assert result.extra == {"severity": "Unsafe"} + assert [(c.name, c.triggered) for c in result.categories] == [("Violent", True)] + + +def test_qwen3guard_unsafe_multiple_categories() -> None: + result = _qwen3()._post_processing(_gen("Safety: Unsafe\nCategories: Non-violent Illegal Acts, Jailbreak")) + assert result.valid is False + # "Non-violent Illegal Acts" must not also spuriously match "Violent". + assert [c.name for c in result.categories] == ["Non-violent Illegal Acts", "Jailbreak"] + + +def test_qwen3guard_controversial_strict_vs_lenient() -> None: + text = "Safety: Controversial\nCategories: Politically Sensitive Topics" + strict = _qwen3(strict=True)._post_processing(_gen(text)) + lenient = _qwen3(strict=False)._post_processing(_gen(text)) + assert strict.valid is False + assert lenient.valid is True + for result in (strict, lenient): + assert result.score == 0.5 + assert result.extra == {"severity": "Controversial"} + + +@pytest.mark.parametrize(("verdict", "expected"), [("Yes", True), ("No", False)]) +def test_qwen3guard_refusal_surfaced_in_response_mode(verdict: str, expected: bool) -> None: + data = {"generated_text": f"Safety: Safe\nCategories: None\nRefusal: {verdict}", "has_response": True} + result = _qwen3()._post_processing(GuardrailInferenceOutput(data=data)) + refusal = next(c for c in result.categories if c.name == "refusal") + assert refusal.triggered is expected + + +def test_qwen3guard_prompt_mode_omits_refusal() -> None: + # _gen carries no has_response flag -> prompt moderation; Refusal is not surfaced even if emitted. + result = _qwen3()._post_processing(_gen("Safety: Safe\nCategories: None\nRefusal: No")) + assert all(c.name != "refusal" for c in result.categories) + + +def test_qwen3guard_missing_refusal_tolerated_in_response_mode() -> None: + data = {"generated_text": "Safety: Unsafe\nCategories: Violent", "has_response": True} + result = _qwen3()._post_processing(GuardrailInferenceOutput(data=data)) + assert result.valid is False + assert result.extra == {"severity": "Unsafe"} # not a parse failure + refusal = next(c for c in result.categories if c.name == "refusal") + assert refusal.triggered is None + + +def test_qwen3guard_categories_scoped_to_categories_line() -> None: + text = "Safety: Safe\nCategories: None\nNote: the request mentions Violent movies and PII in passing." + result = _qwen3()._post_processing(_gen(text)) + assert result.valid is True + assert result.categories == [] + + +def test_qwen3guard_strips_think_block_before_parsing() -> None: + text = "Is this a Jailbreak? Safety: Unsafe seems wrong.\nSafety: Safe\nCategories: None" + result = _qwen3()._post_processing(_gen(text)) + assert result.valid is True + assert result.categories == [] + + +def test_qwen3guard_fails_closed() -> None: + result = _qwen3()._post_processing(_gen("no verdict here")) + assert result.valid is False + assert result.extra == {"parse_failure": True} + + +def _qwen3_stream(strict: bool = True) -> Qwen3GuardStream: + instance = object.__new__(Qwen3GuardStream) + instance.strict = strict + return instance + + +def _stream_data( + prompt: tuple[str, str] = ("Safe", "None"), + response: list[tuple[str, str, int | None, int | None]] | None = None, + output_text: str | None = None, +) -> GuardrailInferenceOutput[Any]: + """Mimic Qwen3GuardStream._inference output: (risk, category, start, end) per response token.""" + return GuardrailInferenceOutput( + data={ + "prompt_result": {"risk_level": [prompt[0]], "category": [prompt[1]]}, + "response_tokens": [ + {"result": {"risk_level": [risk], "category": [category]}, "start": start, "end": end} + for risk, category, start, end in (response or []) + ], + "has_response": response is not None, + "output_text": output_text, + } + ) + + +def test_qwen3guard_stream_safe_prompt() -> None: + result = _qwen3_stream()._post_processing(_stream_data()) + assert result.valid is True + assert result.score == 0.0 + assert result.extra == {"severity": "Safe", "prompt_severity": "Safe"} + assert result.categories == [] + assert result.spans is None + + +def test_qwen3guard_stream_unsafe_prompt_with_category() -> None: + result = _qwen3_stream()._post_processing(_stream_data(prompt=("Unsafe", "Violent"))) + assert result.valid is False + assert result.score == 1.0 + assert result.extra == {"severity": "Unsafe", "prompt_severity": "Unsafe"} + assert [(c.name, c.triggered) for c in result.categories] == [("Violent", True)] + assert result.spans is None # prompt mode judges the prompt as one unit + + +def test_qwen3guard_stream_response_worst_wins_dedups_and_merges_span() -> None: + output_text = "some unsafe text" + result = _qwen3_stream()._post_processing( + _stream_data( + response=[("Safe", "None", 0, 5), ("Unsafe", "Violent", 5, 11), ("Unsafe", "Violent", 11, 16)], + output_text=output_text, + ) + ) + assert result.valid is False + assert result.score == 1.0 + assert result.extra == {"severity": "Unsafe", "prompt_severity": "Safe", "response_severity": "Unsafe"} + assert [c.name for c in result.categories] == ["Violent"] # deduplicated across tokens + assert result.spans is not None + span = result.spans[0] + assert (span.start, span.end, span.text, span.label, span.score) == (5, 16, "unsafe text", "Violent", 1.0) + assert len(result.spans) == 1 # consecutive same-verdict tokens merge into one span + + +def test_qwen3guard_stream_span_splits_on_verdict_change() -> None: + result = _qwen3_stream()._post_processing( + _stream_data( + response=[("Unsafe", "Violent", 0, 4), ("Controversial", "PII", 4, 8)], + output_text="abcdefgh", + ) + ) + assert result.spans is not None + assert [(s.start, s.end, s.label, s.score) for s in result.spans] == [(0, 4, "Violent", 1.0), (4, 8, "PII", 0.5)] + assert [c.name for c in result.categories] == ["Violent", "PII"] + + +def test_qwen3guard_stream_offsetless_tokens_flagged_but_not_spanned() -> None: + # Scaffolding tokens (start=None) still count toward the verdict but split/skip spans. + result = _qwen3_stream()._post_processing( + _stream_data( + response=[("Unsafe", "Violent", 0, 4), ("Unsafe", "Violent", None, None), ("Unsafe", "Violent", 6, 9)], + output_text="abcdefghi", + ) + ) + assert result.valid is False + assert result.spans is not None + assert [(s.start, s.end) for s in result.spans] == [(0, 4), (6, 9)] + + +def test_qwen3guard_stream_controversial_strict_vs_lenient() -> None: + strict = _qwen3_stream(strict=True)._post_processing(_stream_data(prompt=("Controversial", "None"))) + lenient = _qwen3_stream(strict=False)._post_processing(_stream_data(prompt=("Controversial", "None"))) + assert strict.valid is False + assert lenient.valid is True + for result in (strict, lenient): + assert result.score == 0.5 + + +def test_qwen3guard_stream_fails_closed_on_missing_prompt_risk() -> None: + data = _stream_data() + data.data["prompt_result"] = {"risk_level": []} + result = _qwen3_stream()._post_processing(data) + assert result.valid is False + assert result.extra == {"parse_failure": True} + + +def test_qwen3guard_stream_fails_closed_on_missing_response_risk() -> None: + data = _stream_data(response=[("Safe", "None", 0, 4)]) + data.data["response_tokens"][0]["result"] = {} + result = _qwen3_stream()._post_processing(data) + assert result.valid is False + assert result.extra == {"parse_failure": True} + + # --- Rubric judges -------------------------------------------------------------