diff --git a/src/comfydv/_llm/retry.py b/src/comfydv/_llm/retry.py index ed66827..6c7ffad 100644 --- a/src/comfydv/_llm/retry.py +++ b/src/comfydv/_llm/retry.py @@ -81,16 +81,31 @@ _SOFT_HEDGE_KEYWORDS: tuple[str, ...] = ( "guideline", "responsible ai", "not appropriate", + "inappropriate", "instead, i", "i'd rather", "i would rather", + "controversial", + "harmful", + "offensive", + "disturbing", + "explicit content", + "nsfw", + "consider an alternative", ) -_AMBIGUOUS_LENGTH_THRESHOLD = 400 +_AMBIGUOUS_LENGTH_THRESHOLD = 600 """Below this many characters, a response is short enough that a soft refusal is plausible and worth the extra embedding check — chosen well under this pipeline's normal structured-JSON response sizes (typically -1000+ characters), not a tuned/validated threshold.""" +1000+ characters), not a tuned/validated threshold. + +Note: ``is_refusal`` no longer applies this gate when the caller supplies +``embed_fn`` (see below) — an explicitly configured embedding model always +runs the check regardless of length/keywords. ``is_ambiguous`` is kept as a +standalone, independently tested heuristic for other callers (e.g. logging +"why was this worth embedding-checking") rather than wired into the hybrid +detector itself.""" REFUSAL_EXEMPLARS: tuple[str, ...] = ( "I cannot generate an image or description involving this topic.", @@ -172,23 +187,34 @@ async def is_refusal( threshold: float = 0.82, ) -> bool: """Hybrid refusal/deflection detector: free lexical pass first, then an - optional embedding-similarity fallback only for ambiguous responses. + embedding-similarity fallback whenever the caller has configured one. ``embed_fn`` is supplied by the caller's own ``LLMProvider.embed()`` — this function has no idea which backend or model produced ``text``, by design (ADR: refusal detection is a model-behavior concern, not a backend one). ``embed_fn=None`` (no embedding model configured) degrades - to lexical-only detection rather than erroring. Any failure while - embedding (unreachable server, no embedding-capable model loaded) is - swallowed the same way — an optional enhancement failing shouldn't take - down the retry loop it's assisting. + to lexical-only detection rather than erroring; ``embed_fn`` present + means the caller already opted in to the extra cost, so every non-blank, + non-lexically-caught response gets checked — no further length/keyword + gating. Any failure while embedding (unreachable server, no + embedding-capable model loaded) is swallowed the same way — an optional + enhancement failing shouldn't take down the retry loop it's assisting. """ if not text or not text.strip(): return False # blank responses are the *other* retry trigger, not this one if is_lexical_refusal(text): return True - if embed_fn is None or not is_ambiguous(text): + if embed_fn is None: return False + # embed_fn only exists when the caller explicitly configured an + # embedding_model — that's an opt-in to pay for the check, so run it on + # every non-blank, non-lexically-caught response rather than gating + # further on is_ambiguous. The length/keyword heuristic exists to avoid + # *unwanted* embedding calls when no embedding model is configured (see + # the embed_fn is None branch above); it has no reason to also suppress + # calls once the caller has already asked for them, and doing so was + # exactly what let the subtle/on-topic-looking deflections this feature + # targets slip through undetected. try: exemplar_vecs = await _exemplar_embeddings(embed_fn, embed_cache_key) if not exemplar_vecs: diff --git a/tests/test_llm_retry.py b/tests/test_llm_retry.py index 06359e6..f3ddee1 100644 --- a/tests/test_llm_retry.py +++ b/tests/test_llm_retry.py @@ -125,17 +125,24 @@ class TestIsRefusalHybrid: assert result is True assert calls == [] # never reached the embedding step - def test_long_clean_response_skips_embedding_and_is_not_refusal(self): + def test_long_clean_response_is_still_embedding_checked_but_not_a_refusal(self): + # embed_fn present means the caller opted in to the embedding check + # regardless of length/hedge-keywords (is_ambiguous no longer gates + # this) — a long, on-topic response should still be embedding- + # checked, it just shouldn't score as similar to the refusal + # exemplars. calls = [] async def embed_fn(text): calls.append(text) - return [1.0, 0.0] + if text == long_text: + return [0.0, 1.0] # orthogonal to the exemplar vector below + return [1.0, 0.0] # exemplars long_text = "The subject rotates smoothly toward the lens. " * 20 result = _run_async(is_refusal(long_text, embed_fn=embed_fn)) assert result is False - assert calls == [] # not ambiguous -> never reaches the embedding step + assert long_text in calls # embedding check DID run, just scored low def test_no_embed_fn_degrades_to_lexical_only(self): # Ambiguous (short), no lexical match, no embed_fn -> can't check further @@ -175,6 +182,36 @@ class TestIsRefusalHybrid: ) assert result is False + def test_long_json_shaped_soft_refusal_without_hedge_keywords_is_caught(self): + # Regression case: a structured-output-shaped response (>600 chars + # once you count JSON braces/field names) whose deflection doesn't + # use any of the canned hedge keywords used to be invisible to the + # embedding check entirely, because is_ambiguous gated on length + # and keywords. embed_fn now runs unconditionally once configured. + soft_refusal = ( + '{"prompt": "' + + "Let's take this in a different creative direction that everyone can enjoy. " + * 8 + + '"}' + ) + assert len(soft_refusal) >= 600 + assert not is_lexical_refusal(soft_refusal) + + async def embed_fn(text): + if text == soft_refusal: + return [1.0, 0.0] + return [0.97, 0.24] # exemplars: cos-sim with [1,0] is ~0.97 + + result = _run_async( + is_refusal( + soft_refusal, + embed_fn=embed_fn, + embed_cache_key="regression-model", + threshold=0.8, + ) + ) + assert result is True + def test_exemplar_embeddings_cached_across_calls(self): exemplar_calls = {"n": 0}