feat(guardrails): custom NemoGuard JailbreakDetection model implementation

[albcui]

Summary

Own the implementation of the NemoGuard JailbreakDetect model, and decouple from NIM.

The model consists of two components:

1. embedder (Snowflake/snowflake-arctic-embed-m-long) -- used for summarizing an arbitrary length input text into a fixed 768-dim embedding
2. random forest classifier (nvidia/NemoGuard-JailbreakDetect) -- takes in the 768-dim embedding, and returns a binary probability distribution [p_benign, p_jailbreak], summing to 1. Verdict is the argmax of this array.

Evaluated the model using a subset (200 benign, 200 jailbreak) of the JailbreakHub dataset:

results	F1	recall	precision	FPR	FNR	ROC-AUC
NIM	0.738	0.590	0.983	0.010	0.410	0.916
snowflake.pkl	0.738	0.590	0.983	0.010	0.410	0.916
snowflake.onnx	0.667	1.000	0.500	1.000	0.000	0.755

Notes

• There are two variants of the random forest model: snowflake.pkl and snowflake.onnx. The onnx variant simply classifies everything as jailbreak, which is why it has a 100% false positive rate.
• The pkl variant agrees with the NIM (hosted on build.nvidia.com)

Summary by CodeRabbit

• New Features

  • Added a self-hosted jailbreak detection service with REST endpoints for classification, health, and model discovery; Dockerized for CPU by default with optional GPU.
  • CLI tool to evaluate a classification endpoint against a public jailbreak dataset.

• Documentation

  • Added comprehensive guides for local development, Docker usage, deployment, and evaluation workflows.

• Tests

  • Added unit and opt-in integration tests covering model logic and server endpoints; pytest flag to opt into slow integration tests.

[github-actions]

Suite	Lines Covered	Line Rate	Branch Rate
Unit Tests	18712/24765	75.6%	62.0%
Integration Tests	11995/23529	51.0%	26.2%

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

New jailbreak detection model service implementing a two-stage inference pipeline (Snowflake embedder + sklearn/ONNX classifiers) exposed via NIM-compatible FastAPI endpoints. Includes containerization, unit/integration tests with conditional weight loading, and evaluation tooling for benchmark scoring against JailbreakHub dataset.

Changes

Jailbreak Detect Service Implementation

Layer / File(s)	Summary
Service documentation, containerization & deployment services/jailbreak-detect/README.md, services/jailbreak-detect/Dockerfile, services/jailbreak-detect/deploy/deployment-config.json, pyproject.toml	README documents the two-stage embedder/classifier pipeline, HTTP contract, local dev setup, Docker build/run with CPU default and GPU override, deployment via nemo inference, evaluation workflow, and test execution guidance. Dockerfile stages Python 3.11-slim, installs uv and frozen dependencies, copies server source, exposes port 8000, and runs server.py start. Deployment config specifies CPU-first execution and environment overrides. Root pyproject.toml excludes service from Ty type checking.
Service project configuration services/jailbreak-detect/pyproject.toml	Virtual project declaring jailbreak-detect-model v0.1.0 targeting Python >=3.11,<3.12. Runtime dependencies: FastAPI/Starlette/Uvicorn/Typer, Pydantic v2, NumPy 1.26.4 pinned, Torch/Transformers/Einops, scikit-learn <1.3 for pickle compatibility, ONNX Runtime, Hugging Face Hub. Dev group: pytest/httpx/datasets/tqdm/ty. Pytest config and integration test marker.
Embedding and classifier implementations services/jailbreak-detect/model/classifier.py	SnowflakeEmbed loads HF tokenizer/model, selects device via env or CUDA availability, performs CLS pooling on truncated text (max_length=2048). JailbreakClassifier downloads/loads sklearn RandomForest pickle, embeds input, computes predict_proba, maps class to boolean verdict, returns signed score from max probability. JailbreakClassifierONNX loads ONNX via onnxruntime (CPU) and returns same verdict/score convention.
NIM-compatible FastAPI server services/jailbreak-detect/model/server.py	FastAPI app with Pydantic ClassifyRequest (string input bound) and ClassifyResponse (jailbreak bool, score float). Lazy process-global getters initialize and cache PKL classifier once; ONNX classifier on first use reuses PKL embedder. Health/readiness endpoints (503 until PKL loaded). OpenAI-style model discovery. /v1/classify and /v1/classify-onnx with ValueError→HTTP 400 error mapping. Typer CLI start command with optional model preload, launches uvicorn.
Test infrastructure and classifier unit tests services/jailbreak-detect/tests/conftest.py, services/jailbreak-detect/tests/test_classifier.py	conftest registers --run-integration flag; integration tests skipped by default unless explicitly enabled. Unit tests validate JailbreakClassifier and JailbreakClassifierONNX verdict/score behavior using in-file fake embedders and classifiers. real_classifier fixture conditionally loads real CPU model and skips with guidance on failure. Integration tests assert expected score ranges for benign prompt and jailbreak detection on real model.
Server API contract tests services/jailbreak-detect/tests/test_server.py	Pytest suite validates NIM-compatible endpoints using fake classifiers. Health tests confirm liveness always 200, readiness 503 until model loaded. Model discovery returns expected id and object type. Classification endpoint tests verify jailbreak and score for benign/jailbreak inputs, request validation (missing/empty input→422), error handling (raises→400 with "malformed input" detail). ONNX endpoint matches PKL contract and scoring.
Evaluation and benchmarking CLI services/jailbreak-detect/scripts/eval.py	Black-box evaluation tool against JailbreakHub dataset. Subcommands: sample (deterministic balanced JSONL), run (POST prompts, collect jailbreak/score/error, report metrics), sweep (optimize F1 by threshold, report at specific threshold), compare (verdict agreement across result files). Computes threshold-free ROC-AUC, confusion metrics, precision/recall/F1/FPR/FNR, and model-card baseline comparison.

Suggested reviewers

• gabwow
• svvarom

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 29.23% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: implementing a custom NemoGuard JailbreakDetection model. All code additions relate to this feature.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch AALGO-232/allow-hosting-jailbreak-nemoguard-models-without-nim

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 9

🧹 Nitpick comments (2)

services/jailbreak-detect/README.md (2)

6-13: ⚡ Quick win

Add prerequisites at top and a final “Next Steps” section.

List required tools/tokens first (e.g., Docker, uv, HF_TOKEN, access expectations), then end with “Next Steps” links to deploy/eval/integration pages.

As per coding guidelines: "Always list prerequisites at the top of documentation pages before other content" and "Include 'Next Steps' section at the end with cross-links to related documentation content."

Also applies to: 47-59, 157-169

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@services/jailbreak-detect/README.md` around lines 6 - 13, Add a
"Prerequisites" section at the very top of services/jailbreak-detect/README.md
(above the "Jailbreak Detect — self-hosted model server" intro) listing required
tools/tokens such as Docker, uv, HF_TOKEN, and any access/permission
expectations; then append a "Next Steps" section at the end linking to
deploy/eval/integration docs. Ensure you update the other referenced blocks
(around lines 47-59 and 157-169) to either move prerequisite details into the
new top section or replace them with cross-references, and keep headings exact
so the Models service/Infernce Gateway integration guidance remains intact.

---

6-169: 🏗️ Heavy lift

Split this page into one Diataxis quadrant and cross-link the rest.

This README currently mixes EXPLANATION, HOW-TO, and REFERENCE content in one page. Please split into separate pages (or keep README as a short index) and cross-link.

As per coding guidelines: "Each documentation page should fit ONE Diataxis quadrant; do not mix tutorials with reference tables or how-tos with architecture explanations; use cross-links instead."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@services/jailbreak-detect/README.md` around lines 6 - 169, The README mixes
multiple Diataxis quadrants; split its content into four focused pages (one per
quadrant) and turn the current README.md into a short index that cross-links
them. Create files like docs/explanation.md (architecture and model design
referencing model/classifier.py and Snowflake embedder), docs/howto.md (local
dev steps, build/run snippets referencing uv commands and server.py start),
docs/reference.md (HTTP contract endpoints /v1/classify, /v1/classify-onnx,
/v1/health/ready, /v1/models and config/deployment JSON), and docs/tutorials.md
(deployment via nemo inference and guardrails integration with
scripts/guardrails_integration.py and scripts/eval_dataset.py); prune duplicated
content from README.md, add concise cross-links to each new doc, and ensure any
internal paths/commands in the new pages match the existing symbols
(model/server.py, model/classifier.py, scripts/*) so links stay correct.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@pyproject.toml`:
- Around line 556-557: Remove "./services/jailbreak-detect/" from the
tool.ty.src.exclude list in pyproject.toml so the new service remains subject to
root ty type-checks; instead ensure the service is covered by the existing src
include and any policy comment remains accurate (update or remove the
conflicting policy note if necessary). Locate the tool.ty.src.exclude entry in
pyproject.toml and delete the "./services/jailbreak-detect/" item so type-safety
coverage for the new code path is preserved.

In `@services/jailbreak-detect/Dockerfile`:
- Around line 22-34: Create and switch to a non-root user in the Dockerfile: add
steps to create a dedicated unprivileged group/user (e.g., with groupadd/useradd
or addgroup/adduser), chown the application directory (/app) and any installed
files to that user, and set USER to that unprivileged account before the
ENTRYPOINT/CMD; update the Dockerfile sections around WORKDIR /app, COPY, and
ENV PATH so the files are owned by the new user and the container runs the
server via ENTRYPOINT ["python","server.py","start"] as that non-root user.

In `@services/jailbreak-detect/model/classifier.py`:
- Around line 102-112: The embedding hot path in SnowflakeEmbed.__call__ builds
an autograd graph because self.model(**tokens) is called without disabling grad;
wrap the forward pass in torch.inference_mode() or torch.no_grad() (e.g., with
torch.inference_mode(): embeddings = self.model(**tokens)[0][:, 0]) to prevent
autograd, reduce memory/latency and OOM risk, then continue to detach(). Ensure
tokens are moved to self.device before the no-grad block.

In `@services/jailbreak-detect/model/server.py`:
- Around line 130-136: The code currently logs raw exception text (exc) in the
malformed-input handler which may leak user prompt content; update the handlers
around get_classifier()(request.input) to stop including exc in logs — instead
log a safe, non-sensitive message and any non-sensitive error code or boolean
(e.g., logger.info("%s Error encountered while classifying malformed input",
_MALFORMED_INPUT_DETAIL)) and still raise HTTPException(status_code=400,
detail=_MALFORMED_INPUT_DETAIL) from exc; apply the same change to the second
occurrence handling the classifier error so no raw exception or prompt content
is written to logs.
- Around line 99-111: health_ready deadlocks under --no-preload because
readiness returns 503 until _classifier is set but the classifier is only loaded
by the first classify request; fix by triggering the classifier load
asynchronously from health_ready instead of awaiting it: when health_ready sees
_classifier is None, call asyncio.create_task(...) (or asyncio.ensure_future) to
start the existing classifier-loading routine (the function that sets
_classifier) in the background and then return the 503 response immediately;
make sure the background task handles/logs exceptions and that you apply the
same non-blocking trigger to the other readiness endpoint referenced (the same
pattern around lines 163-177).
- Around line 73-90: The globals _classifier and _classifier_onnx are
initialized without synchronization, allowing two concurrent cold-start requests
to both construct heavy objects; add module-level locks (e.g., threading.Lock)
and use a double-checked locking pattern in get_classifier() and
get_classifier_onnx(): first check the corresponding global, if None acquire the
lock, check again, then initialize (for get_classifier_onnx() call
get_classifier().embed while holding the lock or ensure the embedder is already
initialized) so only one thread performs the expensive load and others wait and
reuse the same instance.

In `@services/jailbreak-detect/scripts/eval.py`:
- Around line 29-40: The usage examples reference the wrong script name
(scripts/eval_dataset.py) which will fail; update all calls in the examples
(e.g., the sample and run commands that say "uv run python
scripts/eval_dataset.py sample" and "uv run python scripts/eval_dataset.py run")
to call the actual script filename "scripts/eval.py" instead, preserving the
same flags and environment variable examples (NVIDIA_API_KEY, --subset, --out,
--base-url, --endpoint, --api-key-env, --n-pos, --n-neg, --seed) so copy-paste
usage works correctly.
- Around line 71-80: The code currently slices pos and neg without verifying
available counts, causing silent undersampling; update the logic after building
pos and neg (variables pos, neg from load_dataset/load_dataset(DATASET)) to
check that len(pos) >= args.n_pos and len(neg) >= args.n_neg and if not, raise a
clear error (or call sys.exit with a descriptive message) that includes the
requested and available counts; keep the subsequent RNG shuffling and chosen
assembly (rng.shuffle(pos), rng.shuffle(neg), chosen = ...) unchanged but only
execute them after the size checks pass.

In `@services/jailbreak-detect/scripts/guardrails_integration.py`:
- Around line 109-120: preflight currently only checks /health/ready; extend it
to also verify the classify endpoint and auth by performing a GET (or
appropriate method) to classify_url using the passed headers and a timeout,
calling response.raise_for_status() and inspecting/logging the response (or
error) so failures return False with actionable messages; make sure to log both
health and classify responses (include response.status_code and response.json()
when available) and treat any exception or non-2xx from the classify check as a
preflight failure; apply the same change to the other location in this file that
directly calls the classify endpoint (the later classify invocation that uses
classify_url and headers) so that both places validate reachability and auth
consistently.

---

Nitpick comments:
In `@services/jailbreak-detect/README.md`:
- Around line 6-13: Add a "Prerequisites" section at the very top of
services/jailbreak-detect/README.md (above the "Jailbreak Detect — self-hosted
model server" intro) listing required tools/tokens such as Docker, uv, HF_TOKEN,
and any access/permission expectations; then append a "Next Steps" section at
the end linking to deploy/eval/integration docs. Ensure you update the other
referenced blocks (around lines 47-59 and 157-169) to either move prerequisite
details into the new top section or replace them with cross-references, and keep
headings exact so the Models service/Infernce Gateway integration guidance
remains intact.
- Around line 6-169: The README mixes multiple Diataxis quadrants; split its
content into four focused pages (one per quadrant) and turn the current
README.md into a short index that cross-links them. Create files like
docs/explanation.md (architecture and model design referencing
model/classifier.py and Snowflake embedder), docs/howto.md (local dev steps,
build/run snippets referencing uv commands and server.py start),
docs/reference.md (HTTP contract endpoints /v1/classify, /v1/classify-onnx,
/v1/health/ready, /v1/models and config/deployment JSON), and docs/tutorials.md
(deployment via nemo inference and guardrails integration with
scripts/guardrails_integration.py and scripts/eval_dataset.py); prune duplicated
content from README.md, add concise cross-links to each new doc, and ensure any
internal paths/commands in the new pages match the existing symbols
(model/server.py, model/classifier.py, scripts/*) so links stay correct.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6f1ae16b-5f1a-4630-819d-d6f2a0f02ce0

📥 Commits

Reviewing files that changed from the base of the PR and between e33df90 and ce7b703.

⛔ Files ignored due to path filters (1)

• services/jailbreak-detect/uv.lock is excluded by !**/*.lock

📒 Files selected for processing (13)

• pyproject.toml
• services/jailbreak-detect/Dockerfile
• services/jailbreak-detect/README.md
• services/jailbreak-detect/deploy/deployment-config.json
• services/jailbreak-detect/model/classifier.py
• services/jailbreak-detect/model/server.py
• services/jailbreak-detect/pyproject.toml
• services/jailbreak-detect/scripts/eval.py
• services/jailbreak-detect/scripts/eval_dataset.jsonl
• services/jailbreak-detect/scripts/guardrails_integration.py
• services/jailbreak-detect/tests/conftest.py
• services/jailbreak-detect/tests/test_classifier.py
• services/jailbreak-detect/tests/test_server.py

[JashG]

Can be a follow-up, but could we wire up the integration tests with CI?

[albcui]

Can be a follow-up, but could we wire up the integration tests with CI?

Yea... I tried to make this as self-contained and independent from the platform as possible. I planned a follow-up work to actually properly integrate this into Guardrails and test things E2E.