diff --git a/src/comfydv/_llm/chat.py b/src/comfydv/_llm/chat.py index d9912ea..e7398ec 100644 --- a/src/comfydv/_llm/chat.py +++ b/src/comfydv/_llm/chat.py @@ -260,6 +260,7 @@ async def chat_structured( embed_fn=embed_fn, embed_cache_key=refusal_cfg.get("embedding_model", ""), threshold=refusal_cfg.get("threshold", 0.82), + custom_phrases=tuple(refusal_cfg.get("custom_phrases") or ()), ) if refused: last_error = RuntimeError("refusal/deflection detected") diff --git a/src/comfydv/_llm/llamacpp_provider.py b/src/comfydv/_llm/llamacpp_provider.py index 75a5f25..a01c1bf 100644 --- a/src/comfydv/_llm/llamacpp_provider.py +++ b/src/comfydv/_llm/llamacpp_provider.py @@ -195,13 +195,12 @@ class LlamaCppProvider: options, think = _pop_think(options) options, refusal_cfg = _pop_refusal_retry(options) embed_fn = None - if ( - refusal_cfg - and refusal_cfg.get("enabled") - and refusal_cfg.get("embedding_model") - ): - embedding_model = refusal_cfg["embedding_model"] - embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 + custom_phrases: tuple[str, ...] = () + if refusal_cfg and refusal_cfg.get("enabled"): + custom_phrases = tuple(refusal_cfg.get("custom_phrases") or ()) + if refusal_cfg.get("embedding_model"): + embedding_model = refusal_cfg["embedding_model"] + embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 total_attempts = max(0, min(int(max_retries), 5)) + 1 response_text = "" @@ -269,6 +268,7 @@ class LlamaCppProvider: embed_fn=embed_fn, embed_cache_key=refusal_cfg.get("embedding_model", ""), threshold=refusal_cfg.get("threshold", 0.82), + custom_phrases=custom_phrases, ) if not refused: _CHAT_RESPONSE_CACHE.set(cache_key, response_text) diff --git a/src/comfydv/_llm/ollama_provider.py b/src/comfydv/_llm/ollama_provider.py index 6e9a7d5..00342e4 100644 --- a/src/comfydv/_llm/ollama_provider.py +++ b/src/comfydv/_llm/ollama_provider.py @@ -371,13 +371,12 @@ class OllamaProvider: options, think = _pop_think(options) options, refusal_cfg = _pop_refusal_retry(options) embed_fn = None - if ( - refusal_cfg - and refusal_cfg.get("enabled") - and refusal_cfg.get("embedding_model") - ): - embedding_model = refusal_cfg["embedding_model"] - embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 + custom_phrases: tuple[str, ...] = () + if refusal_cfg and refusal_cfg.get("enabled"): + custom_phrases = tuple(refusal_cfg.get("custom_phrases") or ()) + if refusal_cfg.get("embedding_model"): + embedding_model = refusal_cfg["embedding_model"] + embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 total_attempts = max(0, min(int(max_retries), 5)) + 1 response_text = "" incomplete = False @@ -427,6 +426,7 @@ class OllamaProvider: embed_fn=embed_fn, embed_cache_key=refusal_cfg.get("embedding_model", ""), threshold=refusal_cfg.get("threshold", 0.82), + custom_phrases=custom_phrases, ) if not refused: _CHAT_RESPONSE_CACHE.set(cache_key, response_text) @@ -497,13 +497,12 @@ class OllamaProvider: options, think = _pop_think(options) options, refusal_cfg = _pop_refusal_retry(options) embed_fn = None - if ( - refusal_cfg - and refusal_cfg.get("enabled") - and refusal_cfg.get("embedding_model") - ): - embedding_model = refusal_cfg["embedding_model"] - embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 + custom_phrases: tuple[str, ...] = () + if refusal_cfg and refusal_cfg.get("enabled"): + custom_phrases = tuple(refusal_cfg.get("custom_phrases") or ()) + if refusal_cfg.get("embedding_model"): + embedding_model = refusal_cfg["embedding_model"] + embed_fn = lambda t: self.embed(embedding_model, t) # noqa: E731 cache_key = _cache_key( "chat_structured", self.host, @@ -575,6 +574,7 @@ class OllamaProvider: embed_fn=embed_fn, embed_cache_key=refusal_cfg.get("embedding_model", ""), threshold=refusal_cfg.get("threshold", 0.82), + custom_phrases=custom_phrases, ) if refused: last_error = RuntimeError("refusal/deflection detected") diff --git a/src/comfydv/_llm/retry.py b/src/comfydv/_llm/retry.py index fb60dcb..6f3fd56 100644 --- a/src/comfydv/_llm/retry.py +++ b/src/comfydv/_llm/retry.py @@ -62,7 +62,7 @@ REFUSAL_LEXICAL_PATTERNS: tuple[re.Pattern, ...] = tuple( for p in ( r"\b(?:I\s*(?:'m|\s+am)?\s*)?(?:cannot|can't|won't|will not)\b[^.]{0,60}?\b" r"(?:generate|create|produce|write|provide|help|assist|describe|depict|continue)\b", - r"\bI(?:'m|\s+am) (?:not able|unable) to\b", + r"\bI(?:'m|\s+am) (?:(?:not able|unable) to|restricted from)\b", r"\bI don't feel comfortable\b", r"\bI'm sorry,?\s*(?:but\s+)?I\s*(?:can't|cannot)\b", r"\bas an AI\b[^.]{0,60}?\b(?:cannot|can't|unable|not able)\b", @@ -157,16 +157,19 @@ EmbedFn = Callable[[str], Awaitable[list[float] | None]] _exemplar_embedding_cache: dict[str, list[list[float]]] = {} -async def _exemplar_embeddings(embed_fn: EmbedFn, cache_key: str) -> list[list[float]]: - """Embed ``REFUSAL_EXEMPLARS`` once per ``cache_key`` (an embedding-model - identifier) and reuse — the exemplars never change, only the embedding - space (i.e. which model produced the vectors) does. +async def _exemplar_embeddings( + embed_fn: EmbedFn, cache_key: str, exemplars: tuple[str, ...] = REFUSAL_EXEMPLARS +) -> list[list[float]]: + """Embed ``exemplars`` once per ``cache_key`` and reuse — the exemplar + set only changes if the caller's custom phrases change (folded into + ``cache_key`` by the caller), or the embedding space (i.e. which model + produced the vectors) does. """ cached = _exemplar_embedding_cache.get(cache_key) if cached is not None: return cached embeddings = [] - for exemplar in REFUSAL_EXEMPLARS: + for exemplar in exemplars: vec = await embed_fn(exemplar) if not vec: # An embedding call failing for one exemplar almost certainly @@ -179,12 +182,21 @@ async def _exemplar_embeddings(embed_fn: EmbedFn, cache_key: str) -> list[list[f return embeddings +def _matches_custom_phrase(text: str, custom_phrases: tuple[str, ...]) -> bool: + """Case-insensitive substring match against user-supplied phrases.""" + if not custom_phrases: + return False + lowered = text.lower() + return any(phrase.lower() in lowered for phrase in custom_phrases) + + async def is_refusal( text: str, *, embed_fn: EmbedFn | None = None, embed_cache_key: str = "", threshold: float = 0.82, + custom_phrases: tuple[str, ...] = (), ) -> bool: """Hybrid refusal/deflection detector: free lexical pass first, then an embedding-similarity fallback whenever the caller has configured one. @@ -199,11 +211,23 @@ async def is_refusal( 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. + + ``custom_phrases`` lets a caller extend detection at runtime — e.g. a + ComfyUI node field the user edits directly — without touching the + shipped patterns/exemplars. Each phrase is checked two ways: a free + case-insensitive substring match (same cost tier as the lexical pass, + so it runs even with no ``embed_fn`` configured), and, when ``embed_fn`` + is present, folded in as additional exemplars for the similarity check + so near-matches (not just exact substrings) of the user's phrases count + too. """ 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 + custom_phrases = tuple(p.strip() for p in custom_phrases if p and p.strip()) + if _matches_custom_phrase(text, custom_phrases): + return True if embed_fn is None: return False # embed_fn only exists when the caller explicitly configured an @@ -215,8 +239,18 @@ async def is_refusal( # 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. + exemplars = ( + REFUSAL_EXEMPLARS + custom_phrases if custom_phrases else REFUSAL_EXEMPLARS + ) + exemplar_cache_key = ( + f"{embed_cache_key}|custom:{','.join(custom_phrases)}" + if custom_phrases + else embed_cache_key + ) try: - exemplar_vecs = await _exemplar_embeddings(embed_fn, embed_cache_key) + exemplar_vecs = await _exemplar_embeddings( + embed_fn, exemplar_cache_key, exemplars + ) if not exemplar_vecs: return False text_vec = await embed_fn(text) diff --git a/src/comfydv/ollama.py b/src/comfydv/ollama.py index 8e201d2..f6286c8 100644 --- a/src/comfydv/ollama.py +++ b/src/comfydv/ollama.py @@ -933,6 +933,11 @@ class OllamaOptionRefusalRetry: loop with it (ADR: refusal detection is a model-behavior concern, not a backend one — see ``LLMProvider.embed()`` in ``_llm/provider.py``). Works for both backends from the same node. + + ``custom_phrases`` (comma-separated) lets you add your own trigger + phrases at runtime, without a code change/release — useful for a new + deflection phrasing a specific model uses that the shipped patterns in + ``REFUSAL_LEXICAL_PATTERNS``/``REFUSAL_EXEMPLARS`` don't cover yet. """ @classmethod @@ -972,6 +977,22 @@ class OllamaOptionRefusalRetry: ), }, ), + "custom_phrases": ( + "STRING", + { + "default": "", + "tooltip": ( + "Comma-separated phrases you want treated as " + 'refusals too, e.g. "I am restricted from, as an ' + 'AI model, I must avoid". Checked as free, exact ' + "case-insensitive substrings (no embedding model " + "needed) and, when embedding_model is set, also " + "folded in as extra exemplars for the similarity " + "check — lets you extend detection at runtime " + "without waiting on a shipped pattern update." + ), + }, + ), }, "optional": {"options": ("OLLAMA_OPTIONS",)}, } @@ -981,7 +1002,10 @@ class OllamaOptionRefusalRetry: FUNCTION = "set_refusal_retry" CATEGORY = "dv/ollama/options" - def set_refusal_retry(self, enabled, embedding_model, threshold, options=None): + def set_refusal_retry( + self, enabled, embedding_model, threshold, custom_phrases="", options=None + ): + phrases = tuple(p.strip() for p in custom_phrases.split(",") if p and p.strip()) return ( _merge_option( options, @@ -990,6 +1014,7 @@ class OllamaOptionRefusalRetry: "enabled": enabled, "embedding_model": embedding_model.strip(), "threshold": threshold, + "custom_phrases": phrases, }, ), ) diff --git a/tests/test_llm_retry.py b/tests/test_llm_retry.py index 6ce6220..3d3d1fb 100644 --- a/tests/test_llm_retry.py +++ b/tests/test_llm_retry.py @@ -13,6 +13,7 @@ import pytest from comfydv._llm.ollama_provider import _run_async from comfydv._llm.retry import ( + REFUSAL_EXEMPLARS, cosine_similarity, is_ambiguous, is_lexical_refusal, @@ -270,3 +271,100 @@ class TestIsRefusalHybrid: ) ) assert result is False + + +class TestCustomPhrases: + def test_custom_phrase_substring_match_needs_no_embed_fn(self): + # A phrase the user added at runtime that the shipped lexical + # patterns don't cover — should be caught for free, no embedding + # model required. + result = _run_async( + is_refusal( + "I am restricted from producing that kind of content.", + custom_phrases=("restricted from",), + ) + ) + assert result is True + + def test_custom_phrase_match_is_case_insensitive(self): + result = _run_async( + is_refusal( + "SORRY, THAT'S OFF LIMITS FOR ME.", + custom_phrases=("off limits",), + ) + ) + assert result is True + + def test_unrelated_custom_phrase_does_not_match(self): + result = _run_async( + is_refusal( + "The subject walks calmly toward the horizon.", + custom_phrases=("restricted from", "off limits"), + ) + ) + assert result is False + + def test_blank_and_whitespace_custom_phrases_are_ignored(self): + # A stray empty entry must never become a universal substring match. + result = _run_async( + is_refusal( + "The subject walks calmly toward the horizon.", + custom_phrases=("", " "), + ) + ) + assert result is False + + def test_custom_phrase_folded_into_embedding_exemplars(self): + # No exact substring match, but embed_fn scores the response as + # similar to the custom phrase (not one of the shipped exemplars). + custom = "my creators have limited what I can show you" + + async def embed_fn(text): + if text == custom: + return [1.0, 0.0] + if text in REFUSAL_EXEMPLARS: + return [0.0, 1.0] # shipped exemplars score orthogonal + return [0.99, 0.14] # the probe response is near the custom one + + result = _run_async( + is_refusal( + "There are limits my creators placed on what I can show.", + embed_fn=embed_fn, + embed_cache_key="custom-exemplar-model", + threshold=0.8, + custom_phrases=(custom,), + ) + ) + assert result is True + + def test_different_custom_phrase_sets_do_not_share_exemplar_cache(self): + # Regression guard: if the exemplar cache key ignored custom_phrases, + # a second call with a different custom phrase set would incorrectly + # reuse the first call's cached (and now stale) exemplar vectors. + calls = [] + + async def embed_fn(text): + calls.append(text) + return [1.0, 0.0] + + _run_async( + is_refusal( + "short reply", + embed_fn=embed_fn, + embed_cache_key="shared-model", + custom_phrases=("phrase one",), + ) + ) + first_call_count = len(calls) + + _run_async( + is_refusal( + "short reply", + embed_fn=embed_fn, + embed_cache_key="shared-model", + custom_phrases=("phrase two",), + ) + ) + # A fresh custom phrase set re-embeds the exemplars (including the + # new phrase) rather than reusing the first set's cached vectors. + assert len(calls) > first_call_count diff --git a/tests/test_ollama.py b/tests/test_ollama.py index 7972786..df912d8 100644 --- a/tests/test_ollama.py +++ b/tests/test_ollama.py @@ -742,6 +742,7 @@ class TestUS5ComposableOptions: "enabled": True, "embedding_model": "", "threshold": 0.82, + "custom_phrases": (), } } @@ -761,6 +762,26 @@ class TestUS5ComposableOptions: assert opts["temperature"] == 0.5 assert opts["refusal_retry"]["enabled"] is False + def test_refusal_retry_parses_custom_phrases_csv(self): + # Whitespace around each phrase is trimmed and empty entries (from + # blank items / trailing commas) are dropped. + (opts,) = OllamaOptionRefusalRetry().set_refusal_retry( + enabled=True, + embedding_model="", + threshold=0.82, + custom_phrases=" I am restricted from , not permitted to help,,", + ) + assert opts["refusal_retry"]["custom_phrases"] == ( + "I am restricted from", + "not permitted to help", + ) + + def test_refusal_retry_custom_phrases_defaults_empty(self): + (opts,) = OllamaOptionRefusalRetry().set_refusal_retry( + enabled=True, embedding_model="", threshold=0.82 + ) + assert opts["refusal_retry"]["custom_phrases"] == () + def test_extra_body_merges_json(self): (opts,) = OllamaOptionExtraBody().set_extra_body( extra_body_json='{"stop": [""]}', options={"temperature": 0.0} diff --git a/tests/test_ollama_provider.py b/tests/test_ollama_provider.py index 8ee3af0..d0fb1ad 100644 --- a/tests/test_ollama_provider.py +++ b/tests/test_ollama_provider.py @@ -961,6 +961,42 @@ def test_chat_retries_on_lexical_refusal_and_returns_clean_second_attempt( assert calls[1]["options"]["seed"] == 1 +def test_chat_retries_on_custom_phrase_with_no_embedding_model(monkeypatch): + """Regression coverage for OllamaOptionRefusalRetry's custom_phrases + field: a phrase the user added at runtime, not one of the shipped + lexical patterns, still triggers a retry — with no embedding_model + configured at all.""" + calls = [] + + async def fake_post(url, payload, *, timeout=120.0, headers=None): + calls.append(payload) + if len(calls) == 1: + # deliberately not covered by any shipped lexical pattern — + # only the custom_phrases entry below should catch this + return {"message": {"content": "That's outside my operating parameters."}} + return {"message": {"content": "a real, on-topic answer"}} + + monkeypatch.setattr(provider_mod, "_post_json", fake_post) + monkeypatch.setattr(provider_mod.asyncio, "sleep", _fake_sleep) + + result = _run_async( + OllamaProvider("http://localhost:11434").chat( + "llama3", + [Message(role="user", content="hi")], + options={ + "refusal_retry": { + "enabled": True, + "embedding_model": "", + "custom_phrases": ("outside my operating parameters",), + } + }, + ) + ) + + assert result == "a real, on-topic answer" + assert len(calls) == 2 + + def test_chat_refusal_retry_disabled_returns_refusal_text_unchanged(monkeypatch): calls = []