feat: add opt-in prompt-level cache to LlmJudge

Author: baranylcn

gateframe/rules/semantic.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import OrderedDict
 from collections.abc import Callable
 from typing import Any
 
@@ -47,19 +48,53 @@ def validate(self, output: Any, **context: Any) -> FailureResult | None:  # noqa
 
 class LlmJudge:
     """EXPENSIVE: Each call invokes the provided LLM function.
-    The caller supplies the actual LLM call -- gateframe does not import any LLM SDK."""
+    The caller supplies the actual LLM call -- gateframe does not import any LLM SDK.
+
+    Set ``cache=True`` to memoize results keyed on the rendered prompt.  This avoids
+    redundant LLM calls when the same output+context is validated more than once within
+    the lifetime of the judge instance (e.g. retry loops, multi-step pipelines).
+    ``max_cache_size`` caps the in-memory LRU store; the oldest entry is evicted when
+    the limit is reached.
+    """
 
     def __init__(
         self,
         prompt_template: str,
         llm_call: Callable[[str], str],
         passing_response: str = "PASS",
+        *,
+        cache: bool = False,
+        max_cache_size: int = 256,
     ) -> None:
         self._prompt_template = prompt_template
         self._llm_call = llm_call
         self._passing_response = passing_response
+        self._cache: OrderedDict[str, bool] | None = OrderedDict() if cache else None
+        self._max_cache_size = max_cache_size
 
     def __call__(self, output: Any, **context: Any) -> bool:  # noqa: ANN401
         prompt = self._prompt_template.format(output=output, **context)
-        response = self._llm_call(prompt)
-        return self._passing_response.lower() in response.strip().lower()
+
+        if self._cache is not None:
+            if prompt in self._cache:
+                self._cache.move_to_end(prompt)
+                return self._cache[prompt]
+
+        result = self._passing_response.lower() in self._llm_call(prompt).strip().lower()
+
+        if self._cache is not None:
+            if len(self._cache) >= self._max_cache_size:
+                self._cache.popitem(last=False)
+            self._cache[prompt] = result
+
+        return result
+
+    def cache_clear(self) -> None:
+        """Evict all cached prompt results."""
+        if self._cache is not None:
+            self._cache.clear()
+
+    @property
+    def cache_size(self) -> int:
+        """Number of cached entries (0 when caching is disabled)."""
+        return len(self._cache) if self._cache is not None else 0

tests/rules/test_semantic.py

@@ -140,3 +140,53 @@ def test_partial_match(self) -> None:
             llm_call=lambda prompt: "The result is PASS based on criteria",
         )
         assert judge({}) is True
+
+    def test_cache_disabled_by_default(self) -> None:
+        judge = LlmJudge(prompt_template="{output}", llm_call=lambda p: "PASS")
+        judge("x")
+        assert judge.cache_size == 0
+
+    def test_cache_stores_result(self) -> None:
+        calls: list[str] = []
+
+        def mock_llm(prompt: str) -> str:
+            calls.append(prompt)
+            return "PASS"
+
+        judge = LlmJudge(prompt_template="{output}", llm_call=mock_llm, cache=True)
+        assert judge("hello") is True
+        assert judge("hello") is True
+        assert len(calls) == 1  # second call served from cache
+        assert judge.cache_size == 1
+
+    def test_cache_different_prompts_not_shared(self) -> None:
+        calls: list[str] = []
+
+        def mock_llm(prompt: str) -> str:
+            calls.append(prompt)
+            return "PASS"
+
+        judge = LlmJudge(prompt_template="{output}", llm_call=mock_llm, cache=True)
+        judge("a")
+        judge("b")
+        assert len(calls) == 2
+        assert judge.cache_size == 2
+
+    def test_cache_evicts_lru_when_full(self) -> None:
+        judge = LlmJudge(
+            prompt_template="{output}",
+            llm_call=lambda p: "PASS",
+            cache=True,
+            max_cache_size=2,
+        )
+        judge("x")
+        judge("y")
+        judge("z")  # evicts "x"
+        assert judge.cache_size == 2
+
+    def test_cache_clear(self) -> None:
+        judge = LlmJudge(prompt_template="{output}", llm_call=lambda p: "PASS", cache=True)
+        judge("a")
+        assert judge.cache_size == 1
+        judge.cache_clear()
+        assert judge.cache_size == 0