Conversation
…probs
### What changes I have done
- Added `score_threshold` config field to `MultiRiskGraniteGuardianToolConfig` to optionally drop low-confidence detections and reduce false positives. Defaults to `0.0` (no filter) to preserve current behavior.
- Each entry in `risk_results` now includes a per-category `score` in [0, 1], derived from the logprob of the category's first emitted token in Step 2.
- Categories whose per-category score is below `score_threshold` are dropped from `detected_risks` and `risk_results`. Categories without a score (e.g., when Ollama does not return logprobs) pass through unfiltered as a graceful fallback.
- Step 1 decision remains label-based (`Yes`/`No`) as per the model card; no extra Ollama calls are added.
### Why
The multi-harm model's self-reported text confidence (e.g., `"High"`, `"Not Harmful"`) is effectively binary in practice and useless for thresholding. False positives on some categories (e.g., `Harmful` flagged on sarcasm) could not be filtered without manual category exclusion. Per-category logprob-derived scores give callers a real numeric signal to threshold on (e.g., `Violence=0.97` vs `Harmful=0.35` for the same input).
### How I made the changes
- `akd/guardrails/providers/granite_guardian.py`:
- `MultiRiskGraniteGuardianToolConfig`: added `score_threshold: float = 0.0` with `ge=0.0, le=1.0` validation.
- `_call_category_detection`: added top-level `logprobs: True` and `top_logprobs: 5` to the Ollama `/api/generate` request body (Ollama accepts these as top-level params, not inside `options`).
- `_parse_categories_with_scores`: new helper that parses comma-separated categories and computes per-category scores as `exp(first_token_logprob)`.
- `_first_token_logprob_per_category`: new static helper that walks the token stream, skipping whitespace/commas, and returns the logprob of the first token of each emitted category.
- `_parse_categories`: kept as a thin wrapper over `_parse_categories_with_scores` for backward compatibility.
- `_arun`: applies `score_threshold` as a filter in the Step 2 detected-categories list comprehension; builds `risk_results` with `{"is_risky": True, "score": <float|None>}` per category.
### How to test
- `uv run pytest tests/guardrails/` — existing tests still pass (8 failures in `test_granite_think.py` are pre-existing and require a live `granite3.3-guardian:8b` model).
- `uv run python scripts/test_multi_harm.py` with live Ollama + `granite-guardian-3.2-5b-multi-harm-GGUF` — verifies per-category scores appear in `risk_results` (e.g., Violence=0.97, Unethical Behavior=0.76 for the same violent input; Harmful=0.35 on sarcasm, correctly flagged as low-confidence).
### What
- Merged `_parse_categories` and `_parse_categories_with_scores` into a single `_parse_categories` method that returns `dict[GraniteHarmCategory, float | None]` (category -> per-category score).
- Removed the redundant `"scores"` key from the Step 2 return dict; `"categories"` now holds both the categories and their scores as a dict mapping.
- `unfiltered_categories` in `extra` now preserves per-category scores alongside the category list (previously was a bare list).
### Why
The previous split had `_parse_categories` as a thin list-returning wrapper over `_parse_categories_with_scores` purely for backward compatibility, but `_parse_categories` was only called internally and had no external consumers — dead weight. A single dict return is also a more natural fit: `risk_results` is already a dict of category -> metadata, and downstream consumers need both iteration and score lookup. Dicts preserve insertion order in Python 3.7+, so the model's emission order is kept.
### How
- `akd/guardrails/providers/granite_guardian.py`:
- `_parse_categories`: now takes optional `token_logprobs`, returns `dict[GraniteHarmCategory, float | None]`. When `token_logprobs` is None/empty, scores are `None` (same behavior as the pre-logprobs version, just wrapped in dict keys instead of a list).
- `_call_category_detection`: returns `{"categories": <dict>, "raw_response": ...}` (removed the separate `"scores"` key).
- `_arun`: renamed local from `per_category_scores` to `category_scores`; iterates `category_scores.items()` directly in the filter comprehension; passes the whole `category_scores` dict as `extra["unfiltered_categories"]` to preserve scores in observability output.
### How to test
- `uv run pytest tests/guardrails/ --ignore=tests/guardrails/test_granite_think.py` — all 21 tests pass.
- `uv run python scripts/test_multi_harm.py` with live Ollama — confirms `risk_results` still has `score` per category and `extra["unfiltered_categories"]` now includes scores.
enhance: add per-category score to MultiRiskGraniteGuardianTool via logprobs
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Bump version to 0.1.1
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sync uv.lock with v0.1.1
|
❌ Tests failed (exit code: 1) 📊 Test Results
Branch: 📋 Full coverage report and logs are available in the workflow run. |
sanzog03
approved these changes
Apr 14, 2026
|
❌ Tests failed (exit code: 1) 📊 Test Results
Branch: 📋 Full coverage report and logs are available in the workflow run. |
sanzog03
reviewed
Apr 14, 2026
sanzog03
reviewed
Apr 14, 2026
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add fallback for __version__ when running from source
|
❌ Tests failed (exit code: 1) 📊 Test Results
Branch: 📋 Full coverage report and logs are available in the workflow run. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Patch release v0.1.1 — adds per-category logprob scoring to the multi-risk Granite Guardian tool and consolidates version management.
MultiRiskGraniteGuardianToolStep 2 now requests token logprobs from Ollama and derives ascore = exp(first_token_logprob)per detected harm category, giving callers a real numeric confidence signal (e.g.,Violence=0.97vsHarmful=0.35)score_thresholdconfig field (0.0–1.0, default 0.0) to optionally drop low-confidence category detections and reduce false positives__version__inakd/__init__.pywithimportlib.metadata.version("akd")— version is now sourced solely frompyproject.tomlChanges since v0.1.0
abd39ddConsolidate _parse_categories into single dict-returning function0ea08eeAdd per-category score to MultiRiskGraniteGuardianTool via Step 2 logprobs7fc0863Bump version to 0.1.1, use importlib.metadata for single version sourceTest plan
uv run pytest)uv run python -c "import akd; print(akd.__version__)"prints0.1.1