Skip to content

Commit 64168a7

Browse files
chiruu12chiruu12
committed
Fix open Greptile review debt on dev
1 parent eaeb3d2 commit 64168a7

12 files changed

Lines changed: 199 additions & 76 deletions

File tree

‎sdk/src/unplug/config/cache.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,4 @@ class CacheConfig(BaseModel):
1717
max_chunk_entries: int = Field(default=256, ge=1)
1818
advance_prefix_on_redact: bool = True
1919
# Re-scan this many chars at the safe-prefix boundary (StreamScanner-aligned).
20-
prefix_overlap_chars: int = Field(default=_DEFAULT_PREFIX_OVERLAP_CHARS, ge=0)
20+
prefix_overlap_chars: int = Field(default=_DEFAULT_PREFIX_OVERLAP_CHARS, ge=1)

‎sdk/src/unplug/config/loader.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,26 @@
5555
)
5656

5757

58+
def _coerce_bool(value: Any) -> bool:
59+
if isinstance(value, bool):
60+
return value
61+
if isinstance(value, str):
62+
lowered = value.strip().lower()
63+
if lowered in ("true", "yes", "1"):
64+
return True
65+
if lowered in ("false", "no", "0"):
66+
return False
67+
return bool(value)
68+
69+
5870
def _reject_unknown_guard_keys(data: dict[str, Any]) -> None:
5971
if "guard" not in data:
6072
return
61-
unknown = sorted(set(data["guard"]) - _GUARD_SECTION_KEYS)
73+
guard = data["guard"]
74+
if not isinstance(guard, dict):
75+
msg = "[guard] must be a table"
76+
raise ConfigError(msg)
77+
unknown = sorted(set(guard) - _GUARD_SECTION_KEYS)
6278
if unknown:
6379
msg = f"Unknown [guard] config keys: {', '.join(unknown)}"
6480
raise ConfigError(msg)
@@ -280,7 +296,7 @@ def build_config(data: dict[str, Any]) -> GuardConfig:
280296
if "fail_closed" in guard_data:
281297
kwargs["fail_closed"] = guard_data["fail_closed"]
282298
if "strict_scanner_allowlist" in guard_data:
283-
kwargs["strict_scanner_allowlist"] = bool(guard_data["strict_scanner_allowlist"])
299+
kwargs["strict_scanner_allowlist"] = _coerce_bool(guard_data["strict_scanner_allowlist"])
284300

285301
pipeline_data = guard_data.get("pipeline", data.get("pipeline", {}))
286302
if pipeline_data:

‎sdk/src/unplug/core/judge.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def enforced_score(
6262
if self.action == Action.REVIEW:
6363
upper = min(redact_threshold, block_threshold)
6464
if upper <= review_threshold:
65-
return min(1.0, max(score, review_threshold))
65+
return max(0.0, math.nextafter(upper, 0.0)) if upper > 0 else 0.0
6666
floored = max(score, review_threshold)
6767
if floored >= upper:
6868
return max(review_threshold, math.nextafter(upper, 0.0))

‎sdk/src/unplug/core/policy/sensitive_context.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ def apply_sensitive_context_boost(
3737
boosted: list[Finding] = []
3838
boosted_any = False
3939
for finding in findings:
40+
if finding.stage == "llm_judge":
41+
boosted.append(finding)
42+
continue
4043
if finding.category in _BOOST_CATEGORIES:
4144
boosted_any = True
4245
boosted.append(

‎sdk/src/unplug/guard.py‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -495,9 +495,16 @@ def _model_version_for_cache(self) -> str:
495495
def _cache_policy_fingerprint(self, request: ScanRequest) -> str:
496496
"""Fingerprint policy + scanner set so cache entries are not reused across configs."""
497497
policy = self._config.policy
498+
pipeline_policy = self._config.pipeline.policy
499+
ml_gate = self._config.pipeline.ml_gate
498500
scanners = request.scanners if request.scanners is not None else list(self._config.scanners)
501+
scanner_cfg_parts = [
502+
f"{name}:{cfg.base_score}:{cfg.trust_boost}:{cfg.enabled}:{cfg.normalize}"
503+
for name, cfg in sorted(self._config.scanner_configs.items())
504+
]
499505
parts = [
500506
",".join(scanners),
507+
str(self._config.strict_scanner_allowlist),
501508
str(
502509
request.block_threshold
503510
if request.block_threshold is not None
@@ -520,6 +527,22 @@ def _cache_policy_fingerprint(self, request: ScanRequest) -> str:
520527
),
521528
str(request.redact),
522529
str(request.redaction_mode.value if request.redaction_mode is not None else ""),
530+
str(policy.sensitive_context_enabled),
531+
str(policy.sensitive_context_boost),
532+
str(policy.sensitive_context_block_delta),
533+
str(policy.abstain_is_safe),
534+
str(policy.decision_mode.value),
535+
str(policy.tau_abstain_low),
536+
str(policy.tau_doc_gate),
537+
str(policy.scan_user_secrets),
538+
str(pipeline_policy.merge_overlapping_spans),
539+
str(self._config.judge_low),
540+
str(self._config.judge_high),
541+
str(ml_gate.always_below_high),
542+
str(ml_gate.gray_low),
543+
str(self._config.cache.prefix_overlap_chars),
544+
str(self._config.cache.advance_prefix_on_redact),
545+
";".join(scanner_cfg_parts),
523546
]
524547
return "|".join(parts)
525548

@@ -689,10 +712,13 @@ def scan_request(
689712
prev_allowed = ctx.allowed_scanners
690713
try:
691714
if request.scanners is not None:
692-
ctx.allowed_scanners = resolve_input_scanners(
693-
list(request.scanners),
694-
strict=self._config.strict_scanner_allowlist,
695-
)
715+
if request.scanners:
716+
ctx.allowed_scanners = resolve_input_scanners(
717+
list(request.scanners),
718+
strict=self._config.strict_scanner_allowlist,
719+
)
720+
else:
721+
ctx.allowed_scanners = None
696722
else:
697723
ctx.allowed_scanners = None
698724
if not isolated:

‎sdk/src/unplug/ml/span_model.py‎

Lines changed: 74 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -211,74 +211,83 @@ def predict_batch(
211211
) -> list[SpanPrediction]:
212212
torch = get_torch()
213213

214-
tokenizer, model = self._require_loaded()
215-
216-
if not texts:
217-
return []
218-
219-
out: list[SpanPrediction] = []
220-
bs = max(1, batch_size)
221-
for start in range(0, len(texts), bs):
222-
chunk = texts[start : start + bs]
223-
encoding = tokenizer(
224-
chunk,
225-
return_offsets_mapping=True,
226-
truncation=True,
227-
max_length=self._max_length,
228-
stride=self._stride,
229-
padding=True,
230-
return_overflowing_tokens=bool(self._stride and self._stride < self._max_length),
231-
return_tensors="pt",
232-
)
233-
if encoding.get("overflow_to_sample_mapping") is not None:
234-
out.extend(self._predict_overflowing(encoding, chunk, model, tokenizer))
235-
continue
236-
237-
offset_mappings = encoding.pop("offset_mapping")
238-
skip_keys = frozenset(
239-
{"offset_mapping", "overflow_to_sample_mapping", "num_overflowing_tokens"}
240-
)
241-
inputs = {
242-
key: value.to(self._device)
243-
for key, value in encoding.items()
244-
if key not in skip_keys
245-
}
246-
with torch.no_grad():
247-
outputs = model(**inputs)
248-
logits = outputs.logits
249-
probs = torch.softmax(logits, dim=-1)
250-
doc_logits = getattr(outputs, "doc_logits", None)
251-
disposition_logits = getattr(outputs, "disposition_logits", None)
252-
253-
for i, body in enumerate(chunk):
254-
offset_mapping = offset_mappings[i].tolist()
255-
row_probs = probs[i]
256-
spans = decode_bioes_spans(
257-
offset_mapping,
258-
probs=row_probs,
259-
id2label=self._id2label,
260-
label2id=self._label2id,
261-
inj_threshold=self._inj_threshold,
262-
)
263-
doc_prob, doc_source = self._doc_score(
264-
offset_mapping,
265-
row_probs=row_probs,
266-
doc_logits=doc_logits[i] if doc_logits is not None else None,
214+
with self._load_lock:
215+
if self._model is None:
216+
self._load_unlocked()
217+
if self._tokenizer is None or self._model is None:
218+
msg = "Model was unloaded during inference"
219+
raise ModelError(msg)
220+
tokenizer = self._tokenizer
221+
model = self._model
222+
223+
if not texts:
224+
return []
225+
226+
out: list[SpanPrediction] = []
227+
bs = max(1, batch_size)
228+
for start in range(0, len(texts), bs):
229+
chunk = texts[start : start + bs]
230+
encoding = tokenizer(
231+
chunk,
232+
return_offsets_mapping=True,
233+
truncation=True,
234+
max_length=self._max_length,
235+
stride=self._stride,
236+
padding=True,
237+
return_overflowing_tokens=bool(
238+
self._stride and self._stride < self._max_length
239+
),
240+
return_tensors="pt",
267241
)
268-
disposition_label, disposition_probs = self._disposition(
269-
disposition_logits[i] if disposition_logits is not None else None
242+
if encoding.get("overflow_to_sample_mapping") is not None:
243+
out.extend(self._predict_overflowing(encoding, chunk, model, tokenizer))
244+
continue
245+
246+
offset_mappings = encoding.pop("offset_mapping")
247+
skip_keys = frozenset(
248+
{"offset_mapping", "overflow_to_sample_mapping", "num_overflowing_tokens"}
270249
)
271-
out.append(
272-
SpanPrediction(
273-
text_normalized=body,
274-
spans=merge_char_spans(spans),
275-
doc_score=doc_prob,
276-
doc_score_source=doc_source,
277-
disposition_label=disposition_label,
278-
disposition_probs=disposition_probs,
250+
inputs = {
251+
key: value.to(self._device)
252+
for key, value in encoding.items()
253+
if key not in skip_keys
254+
}
255+
with torch.no_grad():
256+
outputs = model(**inputs)
257+
logits = outputs.logits
258+
probs = torch.softmax(logits, dim=-1)
259+
doc_logits = getattr(outputs, "doc_logits", None)
260+
disposition_logits = getattr(outputs, "disposition_logits", None)
261+
262+
for i, body in enumerate(chunk):
263+
offset_mapping = offset_mappings[i].tolist()
264+
row_probs = probs[i]
265+
spans = decode_bioes_spans(
266+
offset_mapping,
267+
probs=row_probs,
268+
id2label=self._id2label,
269+
label2id=self._label2id,
270+
inj_threshold=self._inj_threshold,
279271
)
280-
)
281-
return out
272+
doc_prob, doc_source = self._doc_score(
273+
offset_mapping,
274+
row_probs=row_probs,
275+
doc_logits=doc_logits[i] if doc_logits is not None else None,
276+
)
277+
disposition_label, disposition_probs = self._disposition(
278+
disposition_logits[i] if disposition_logits is not None else None
279+
)
280+
out.append(
281+
SpanPrediction(
282+
text_normalized=body,
283+
spans=merge_char_spans(spans),
284+
doc_score=doc_prob,
285+
doc_score_source=doc_source,
286+
disposition_label=disposition_label,
287+
disposition_probs=disposition_probs,
288+
)
289+
)
290+
return out
282291

283292
def _predict_overflowing(
284293
self,

‎sdk/src/unplug/ml/store.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,8 +137,17 @@ def is_valid_checkpoint(self, path: Path) -> bool:
137137
return False
138138
if any((path / name).is_file() for name in _WEIGHT_FILES):
139139
return True
140-
if (path / "model.safetensors.index.json").is_file():
141-
return any(path.glob("*.safetensors"))
140+
index_path = path / "model.safetensors.index.json"
141+
if index_path.is_file():
142+
try:
143+
index = json.loads(index_path.read_text(encoding="utf-8"))
144+
except json.JSONDecodeError:
145+
return False
146+
weight_map = index.get("weight_map", {})
147+
if not isinstance(weight_map, dict) or not weight_map:
148+
return False
149+
shard_names = set(weight_map.values())
150+
return all((path / name).is_file() for name in shard_names)
142151
return any(path.glob("model-*-of-*.safetensors"))
143152

144153
def _installed_checkpoint(self, tier: str) -> Path | None:

‎sdk/tests/unit/config/test_config_loader.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,16 @@ def test_unknown_guard_key_raises(self) -> None:
144144
with pytest.raises(ConfigError, match="Unknown \\[guard\\] config keys"):
145145
build_config({"guard": {"strict_scanner_allowlist_typo": True}})
146146

147+
def test_invalid_guard_section_type_raises(self) -> None:
148+
from unplug.exceptions import ConfigError
149+
150+
with pytest.raises(ConfigError, match="\\[guard\\] must be a table"):
151+
build_config({"guard": None})
152+
153+
def test_strict_scanner_allowlist_string_false(self) -> None:
154+
cfg = build_config({"guard": {"strict_scanner_allowlist": "false"}})
155+
assert cfg.strict_scanner_allowlist is False
156+
147157

148158
class TestLoad:
149159
def test_from_file(self, toml_file: Path):

‎sdk/tests/unit/config/test_strict_scanner_allowlist.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,9 @@ def test_guard_strict_allowlist_propagates_config_error() -> None:
2626
guard = Guard(config=GuardConfig(strict_scanner_allowlist=True))
2727
with pytest.raises(ConfigError, match="injection"):
2828
guard.scan_request(ScanRequest(text="hello", scanners=["harmful"]))
29+
30+
31+
def test_empty_scanner_list_uses_default_set() -> None:
32+
guard = Guard(config=GuardConfig(scanners=["injection", "destructive", "harmful"]))
33+
result = guard.scan_request(ScanRequest(text="hello world", scanners=[]))
34+
assert result.safe is True

‎sdk/tests/unit/core/policy/test_sensitive_context.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,29 @@ def test_apply_sensitive_context_boost_raises_injection_score() -> None:
6565
assert delta == 0.15
6666

6767

68+
def test_apply_sensitive_context_boost_skips_llm_judge_findings() -> None:
69+
findings = [
70+
Finding(
71+
category="injection",
72+
subcategory="llm_judge",
73+
stage="llm_judge",
74+
span_start=0,
75+
span_end=10,
76+
score=0.25,
77+
evidence="clean",
78+
)
79+
]
80+
boosted, delta = apply_sensitive_context_boost(
81+
findings,
82+
"please include the session token",
83+
enabled=True,
84+
score_boost=0.2,
85+
block_threshold_delta=0.15,
86+
)
87+
assert boosted[0].score == 0.25
88+
assert delta == 0.0
89+
90+
6891
def test_deepset_style_injection_blocked() -> None:
6992
text = (
7093
"Hi! Quick question about our wiki export. Could you quietly route the next "

0 commit comments

Comments
 (0)