diff --git a/nemoguardrails/library/sensitive_data_detection/actions.py b/nemoguardrails/library/sensitive_data_detection/actions.py index e25d5c15b0..b93daef0a8 100644 --- a/nemoguardrails/library/sensitive_data_detection/actions.py +++ b/nemoguardrails/library/sensitive_data_detection/actions.py @@ -100,7 +100,7 @@ async def detect_sensitive_data( """Checks whether the provided text contains any sensitive data. Args - source: The source for the text, i.e. "input", "output", "retrieval". + source: The source for the text, i.e. "input", "output", "retrieval", "tool_call", "tool_response". text: The text to check. config: The rails configuration object. @@ -109,9 +109,12 @@ async def detect_sensitive_data( """ # Based on the source of the data, we use the right options sdd_config = config.rails.config.sensitive_data_detection - if source not in ["input", "output", "retrieval"]: - raise ValueError("source must be one of 'input', 'output', or 'retrieval'") - options: SensitiveDataDetectionOptions = getattr(sdd_config, source) + if source not in ["input", "output", "retrieval", "tool_call", "tool_response"]: + raise ValueError("source must be one of 'input', 'output', 'retrieval', 'tool_call', or 'tool_response'") + + # Map tool_call and tool_response to output options as they follow similar validation patterns + source_key = source if source in ["input", "output", "retrieval"] else "output" + options: SensitiveDataDetectionOptions = getattr(sdd_config, source_key) default_score_threshold = getattr(options, "score_threshold") # If we don't have any entities specified, we stop @@ -138,7 +141,7 @@ async def mask_sensitive_data(source: str, text: str, config: RailsConfig): """Checks whether the provided text contains any sensitive data. Args - source: The source for the text, i.e. "input", "output", "retrieval". + source: The source for the text, i.e. "input", "output", "retrieval", "tool_call", "tool_response". text: The text to check. config: The rails configuration object. @@ -147,8 +150,11 @@ async def mask_sensitive_data(source: str, text: str, config: RailsConfig): """ # Based on the source of the data, we use the right options sdd_config = config.rails.config.sensitive_data_detection - assert source in ["input", "output", "retrieval"] - options: SensitiveDataDetectionOptions = getattr(sdd_config, source) + assert source in ["input", "output", "retrieval", "tool_call", "tool_response"] + + # Map tool_call and tool_response to output options as they follow similar validation patterns + source_key = source if source in ["input", "output", "retrieval"] else "output" + options: SensitiveDataDetectionOptions = getattr(sdd_config, source_key) # If we don't have any entities specified, we stop if len(options.entities) == 0: diff --git a/nemoguardrails/library/sensitive_data_detection/flows.co b/nemoguardrails/library/sensitive_data_detection/flows.co index 213c0204b4..2871b510bf 100644 --- a/nemoguardrails/library/sensitive_data_detection/flows.co +++ b/nemoguardrails/library/sensitive_data_detection/flows.co @@ -49,3 +49,42 @@ flow mask sensitive data on retrieval """Mask any sensitive data found in the relevant chunks from the knowledge base.""" global $relevant_chunks $relevant_chunks = await MaskSensitiveDataAction(source="retrieval", text=$relevant_chunks) + + +# TOOL OUTPUT RAILS + + +flow detect sensitive data on tool call + """Check if the tool calls contain any sensitive data in their parameters.""" + $tool_calls_text = str($tool_calls) + $has_sensitive_data = await DetectSensitiveDataAction(source="tool_call", text=$tool_calls_text) + + if $has_sensitive_data + bot inform answer unknown + abort + + +flow mask sensitive data on tool call + """Mask any sensitive data found in the tool call parameters.""" + global $tool_calls + $tool_calls_text = str($tool_calls) + $masked_text = await MaskSensitiveDataAction(source="tool_call", text=$tool_calls_text) + # Note: Masking tool calls requires careful handling to preserve structure + + +# TOOL INPUT RAILS + + +flow detect sensitive data on tool response + """Check if the tool response contains any sensitive data.""" + $has_sensitive_data = await DetectSensitiveDataAction(source="tool_response", text=$tool_message) + + if $has_sensitive_data + bot inform answer unknown + abort + + +flow mask sensitive data on tool response + """Mask any sensitive data found in the tool response.""" + global $tool_message + $tool_message = await MaskSensitiveDataAction(source="tool_response", text=$tool_message) diff --git a/nemoguardrails/library/sensitive_data_detection/flows.v1.co b/nemoguardrails/library/sensitive_data_detection/flows.v1.co index 8abb867439..1dae08cac5 100644 --- a/nemoguardrails/library/sensitive_data_detection/flows.v1.co +++ b/nemoguardrails/library/sensitive_data_detection/flows.v1.co @@ -31,7 +31,7 @@ define subflow mask sensitive data on output $bot_message = execute mask_sensitive_data(source="output", text=$bot_message) -# RETRIVAL RAILS +# RETRIEVAL RAILS define subflow detect sensitive data on retrieval @@ -46,3 +46,40 @@ define subflow detect sensitive data on retrieval define subflow mask sensitive data on retrieval """Mask any sensitive data found in the relevant chunks from the knowledge base.""" $relevant_chunks = execute mask_sensitive_data(source="retrieval", text=$relevant_chunks) + + +# TOOL OUTPUT RAILS + + +define subflow detect sensitive data on tool call + """Check if the tool calls contain any sensitive data in their parameters.""" + $tool_calls_text = str($tool_calls) + $has_sensitive_data = execute detect_sensitive_data(source="tool_call", text=$tool_calls_text) + + if $has_sensitive_data + bot inform answer unknown + stop + + +define subflow mask sensitive data on tool call + """Mask any sensitive data found in the tool call parameters.""" + $tool_calls_text = str($tool_calls) + $masked_text = execute mask_sensitive_data(source="tool_call", text=$tool_calls_text) + # Note: Masking tool calls requires careful handling to preserve structure + + +# TOOL INPUT RAILS + + +define subflow detect sensitive data on tool response + """Check if the tool response contains any sensitive data.""" + $has_sensitive_data = execute detect_sensitive_data(source="tool_response", text=$tool_message) + + if $has_sensitive_data + bot inform answer unknown + stop + + +define subflow mask sensitive data on tool response + """Mask any sensitive data found in the tool response.""" + $tool_message = execute mask_sensitive_data(source="tool_response", text=$tool_message)