feat(cache): add caching support for topic safety and content safety output checks (#1457)

* feat(cache): add LLM metadata caching for model and provider information

Extends the cache system to store and restore LLM metadata (model name
and provider name) alongside cache entries. This allows cached results
to maintain provenance information about which model and provider
generated the original response.

- Added LLMMetadataDict and LLMCacheData TypedDict definitions for type
safety
  - Extended CacheEntry to include optional llm_metadata field
  - Implemented extract_llm_metadata_for_cache() to capture model and
provider info from context
  - Implemented restore_llm_metadata_from_cache() to restore metadata
when retrieving cached results
  - Updated get_from_cache_and_restore_stats() to handle metadata
extraction and restoration
  - Added comprehensive test coverage for metadata caching functionalit
* test(cache): add cache miss test cases for safety checks

Author: Pouyanpi

nemoguardrails/library/content_safety/actions.py

@@ -25,6 +25,7 @@
 from nemoguardrails.llm.cache.utils import (
     CacheEntry,
     create_normalized_cache_key,
+    extract_llm_metadata_for_cache,
     extract_llm_stats_for_cache,
     get_from_cache_and_restore_stats,
 )
@@ -110,6 +111,7 @@ async def content_safety_check_input(
         cache_entry: CacheEntry = {
             "result": final_result,
             "llm_stats": extract_llm_stats_for_cache(),
+            "llm_metadata": extract_llm_metadata_for_cache(),
         }
         cache.put(cache_key, cache_entry)
         log.debug(f"Content safety result cached for model '{model_name}'")
@@ -139,6 +141,7 @@ async def content_safety_check_output(
     llm_task_manager: LLMTaskManager,
     model_name: Optional[str] = None,
     context: Optional[dict] = None,
+    model_caches: Optional[Dict[str, CacheInterface]] = None,
     **kwargs,
 ) -> dict:
     _MAX_TOKENS = 3
@@ -176,12 +179,22 @@ async def content_safety_check_output(
             "bot_response": bot_response,
         },
     )
+
     stop = llm_task_manager.get_stop_tokens(task=task)
     max_tokens = llm_task_manager.get_max_tokens(task=task)
 
+    llm_call_info_var.set(LLMCallInfo(task=task))
+
     max_tokens = max_tokens or _MAX_TOKENS
 
-    llm_call_info_var.set(LLMCallInfo(task=task))
+    cache = model_caches.get(model_name) if model_caches else None
+
+    if cache:
+        cache_key = create_normalized_cache_key(check_output_prompt)
+        cached_result = get_from_cache_and_restore_stats(cache, cache_key)
+        if cached_result is not None:
+            log.debug(f"Content safety output cache hit for model '{model_name}'")
+            return cached_result
 
     result = await llm_call(
         llm,
@@ -194,4 +207,16 @@ async def content_safety_check_output(
 
     is_safe, *violated_policies = result
 
-    return {"allowed": is_safe, "policy_violations": violated_policies}
+    final_result = {"allowed": is_safe, "policy_violations": violated_policies}
+
+    if cache:
+        cache_key = create_normalized_cache_key(check_output_prompt)
+        cache_entry: CacheEntry = {
+            "result": final_result,
+            "llm_stats": extract_llm_stats_for_cache(),
+            "llm_metadata": extract_llm_metadata_for_cache(),
+        }
+        cache.put(cache_key, cache_entry)
+        log.debug(f"Content safety output result cached for model '{model_name}'")
+
+    return final_result

nemoguardrails/library/topic_safety/actions.py

@@ -21,6 +21,14 @@
 from nemoguardrails.actions.actions import action
 from nemoguardrails.actions.llm.utils import llm_call
 from nemoguardrails.context import llm_call_info_var
+from nemoguardrails.llm.cache import CacheInterface
+from nemoguardrails.llm.cache.utils import (
+    CacheEntry,
+    create_normalized_cache_key,
+    extract_llm_metadata_for_cache,
+    extract_llm_stats_for_cache,
+    get_from_cache_and_restore_stats,
+)
 from nemoguardrails.llm.filters import to_chat_messages
 from nemoguardrails.llm.taskmanager import LLMTaskManager
 from nemoguardrails.logging.explain import LLMCallInfo
@@ -35,6 +43,7 @@ async def topic_safety_check_input(
     model_name: Optional[str] = None,
     context: Optional[dict] = None,
     events: Optional[List[dict]] = None,
+    model_caches: Optional[Dict[str, CacheInterface]] = None,
     **kwargs,
 ) -> dict:
     _MAX_TOKENS = 10
@@ -102,11 +111,32 @@ async def topic_safety_check_input(
     messages.extend(conversation_history)
     messages.append({"type": "user", "content": user_input})
 
+    cache = model_caches.get(model_name) if model_caches else None
+
+    if cache:
+        cache_key = create_normalized_cache_key(messages)
+        cached_result = get_from_cache_and_restore_stats(cache, cache_key)
+        if cached_result is not None:
+            log.debug(f"Topic safety cache hit for model '{model_name}'")
+            return cached_result
+
     result = await llm_call(llm, messages, stop=stop, llm_params={"temperature": 0.01})
 
     if result.lower().strip() == "off-topic":
         on_topic = False
     else:
         on_topic = True
 
-    return {"on_topic": on_topic}
+    final_result = {"on_topic": on_topic}
+
+    if cache:
+        cache_key = create_normalized_cache_key(messages)
+        cache_entry: CacheEntry = {
+            "result": final_result,
+            "llm_stats": extract_llm_stats_for_cache(),
+            "llm_metadata": extract_llm_metadata_for_cache(),
+        }
+        cache.put(cache_key, cache_entry)
+        log.debug(f"Topic safety result cached for model '{model_name}'")
+
+    return final_result

nemoguardrails/llm/cache/utils.py

@@ -179,6 +179,9 @@ def get_from_cache_and_restore_stats(
     if cached_metadata:
         restore_llm_metadata_from_cache(cached_metadata)
 
+    if cached_metadata:
+        restore_llm_metadata_from_cache(cached_metadata)
+
     processing_log = processing_log_var.get()
     if processing_log is not None:
         llm_call_info = llm_call_info_var.get()

tests/test_content_safety_cache.py

@@ -18,7 +18,10 @@
 import pytest
 
 from nemoguardrails.context import llm_call_info_var, llm_stats_var
-from nemoguardrails.library.content_safety.actions import content_safety_check_input
+from nemoguardrails.library.content_safety.actions import (
+    content_safety_check_input,
+    content_safety_check_output,
+)
 from nemoguardrails.llm.cache.lfu import LFUCache
 from nemoguardrails.llm.cache.utils import create_normalized_cache_key
 from nemoguardrails.logging.explain import LLMCallInfo
@@ -95,6 +98,7 @@ async def test_content_safety_cache_retrieves_result_and_restores_stats(
             "prompt_tokens": 80,
             "completion_tokens": 20,
         },
+        "llm_metadata": None,
     }
     cache_key = create_normalized_cache_key("test prompt")
     cache.put(cache_key, cache_entry)
@@ -138,6 +142,7 @@ async def test_content_safety_cache_duration_reflects_cache_read_time(
             "prompt_tokens": 40,
             "completion_tokens": 10,
         },
+        "llm_metadata": None,
     }
     cache_key = create_normalized_cache_key("test prompt")
     cache.put(cache_key, cache_entry)
@@ -191,6 +196,7 @@ async def test_content_safety_cache_handles_missing_stats_gracefully(
     cache_entry = {
         "result": {"allowed": True, "policy_violations": []},
         "llm_stats": None,
+        "llm_metadata": None,
     }
     cache_key = create_normalized_cache_key("test_key")
     cache.put(cache_key, cache_entry)
@@ -213,3 +219,106 @@ async def test_content_safety_cache_handles_missing_stats_gracefully(
 
     assert result["allowed"] is True
     assert llm_stats.get_stat("total_calls") == 0
+
+
+@pytest.mark.asyncio
+async def test_content_safety_check_output_cache_stores_result(
+    fake_llm_with_stats, mock_task_manager
+):
+    cache = LFUCache(maxsize=10)
+    mock_task_manager.parse_task_output.return_value = [True, "policy2"]
+
+    result = await content_safety_check_output(
+        llms=fake_llm_with_stats,
+        llm_task_manager=mock_task_manager,
+        model_name="test_model",
+        context={"user_message": "test user input", "bot_message": "test bot response"},
+        model_caches={"test_model": cache},
+    )
+
+    assert result["allowed"] is True
+    assert result["policy_violations"] == ["policy2"]
+    assert cache.size() == 1
+
+
+@pytest.mark.asyncio
+async def test_content_safety_check_output_cache_hit(
+    fake_llm_with_stats, mock_task_manager
+):
+    cache = LFUCache(maxsize=10)
+
+    cache_entry = {
+        "result": {"allowed": False, "policy_violations": ["unsafe_output"]},
+        "llm_stats": {
+            "total_tokens": 75,
+            "prompt_tokens": 60,
+            "completion_tokens": 15,
+        },
+        "llm_metadata": None,
+    }
+    cache_key = create_normalized_cache_key("test output prompt")
+    cache.put(cache_key, cache_entry)
+
+    mock_task_manager.render_task_prompt.return_value = "test output prompt"
+
+    llm_stats = LLMStats()
+    llm_stats_var.set(llm_stats)
+
+    result = await content_safety_check_output(
+        llms=fake_llm_with_stats,
+        llm_task_manager=mock_task_manager,
+        model_name="test_model",
+        context={"user_message": "user", "bot_message": "bot"},
+        model_caches={"test_model": cache},
+    )
+
+    assert result["allowed"] is False
+    assert result["policy_violations"] == ["unsafe_output"]
+    assert llm_stats.get_stat("total_calls") == 1
+    assert llm_stats.get_stat("total_tokens") == 75
+
+    llm_call_info = llm_call_info_var.get()
+    assert llm_call_info.from_cache is True
+
+
+@pytest.mark.asyncio
+async def test_content_safety_check_output_cache_miss(
+    fake_llm_with_stats, mock_task_manager
+):
+    cache = LFUCache(maxsize=10)
+
+    cache_entry = {
+        "result": {"allowed": True, "policy_violations": []},
+        "llm_stats": {
+            "total_tokens": 50,
+            "prompt_tokens": 40,
+            "completion_tokens": 10,
+        },
+        "llm_metadata": None,
+    }
+    cache_key = create_normalized_cache_key("different prompt")
+    cache.put(cache_key, cache_entry)
+
+    mock_task_manager.render_task_prompt.return_value = "new output prompt"
+    mock_task_manager.parse_task_output.return_value = [True, "policy2"]
+
+    llm_stats = LLMStats()
+    llm_stats_var.set(llm_stats)
+
+    llm_call_info = LLMCallInfo(task="content_safety_check_output $model=test_model")
+    llm_call_info_var.set(llm_call_info)
+
+    result = await content_safety_check_output(
+        llms=fake_llm_with_stats,
+        llm_task_manager=mock_task_manager,
+        model_name="test_model",
+        context={"user_message": "new user input", "bot_message": "new bot response"},
+        model_caches={"test_model": cache},
+    )
+
+    assert result["allowed"] is True
+    assert result["policy_violations"] == ["policy2"]
+    assert cache.size() == 2
+
+    llm_call_info = llm_call_info_var.get()
+    assert llm_call_info.from_cache is False