From 0b49f391372794df328751050cce8eafcfb8114e Mon Sep 17 00:00:00 2001 From: ekam-walia <164983405+ekam-walia@users.noreply.github.com> Date: Wed, 22 Jul 2026 13:49:36 +0530 Subject: [PATCH 1/4] bug: missing "global $relevant_chunks" declaration in hf_classifier retrieval flows --- nemoguardrails/library/hf_classifier/flows.co | 4 +- .../library/hf_classifier/flows.v1.co | 4 +- tests/test_hf_classifier_retrieval_global.py | 425 ++++++++++++++++++ 3 files changed, 431 insertions(+), 2 deletions(-) create mode 100644 tests/test_hf_classifier_retrieval_global.py diff --git a/nemoguardrails/library/hf_classifier/flows.co b/nemoguardrails/library/hf_classifier/flows.co index 34cc4dcd78..187791855e 100644 --- a/nemoguardrails/library/hf_classifier/flows.co +++ b/nemoguardrails/library/hf_classifier/flows.co @@ -22,4 +22,6 @@ flow hf classifier check retrieval $classifier $response = await HfClassifierCheckRetrievalAction(classifier=$classifier) if $response.is_transform - $relevant_chunks = $response.transform_text["relevant_chunks"] + # FIX: Added global declaration to propagate the change to calling context + global $relevant_chunks + $relevant_chunks = $response.transform_text["relevant_chunks"] \ No newline at end of file diff --git a/nemoguardrails/library/hf_classifier/flows.v1.co b/nemoguardrails/library/hf_classifier/flows.v1.co index 015ca8530b..f128ace9c1 100644 --- a/nemoguardrails/library/hf_classifier/flows.v1.co +++ b/nemoguardrails/library/hf_classifier/flows.v1.co @@ -25,4 +25,6 @@ define subflow hf classifier check retrieval $response = execute hf_classifier_check_retrieval(classifier=$classifier) if $response.is_transform - $relevant_chunks = $response.transform_text["relevant_chunks"] + # FIX: Added global declaration to propagate the change to calling context + global $relevant_chunks + $relevant_chunks = $response.transform_text["relevant_chunks"] \ No newline at end of file diff --git a/tests/test_hf_classifier_retrieval_global.py b/tests/test_hf_classifier_retrieval_global.py new file mode 100644 index 0000000000..482e3d173d --- /dev/null +++ b/tests/test_hf_classifier_retrieval_global.py @@ -0,0 +1,425 @@ +# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest.mock import AsyncMock, MagicMock + +from nemoguardrails import LLMRails, RailsConfig + + +class MockHfClassifierResponse: + """Mock response from HF classifier""" + + def __init__(self, is_transform=False, transform_text=None, is_blocked=False): + self.is_transform = is_transform + self.transform_text = transform_text or {} + self.is_blocked = is_blocked + + +class TestHfClassifierRetrievalGlobal(unittest.TestCase): + """Test cases for HF classifier retrieval global variable propagation""" + + def test_hf_classifier_v1_retrieval_transform_propagates_global(self): + """Test that transformed relevant_chunks propagates to global context in v1""" + + # Create mock action that returns transform response + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["new_chunk_1", "new_chunk_2"], + "response": "Transformed response" + } + ) + + # Create config with the v1 flow + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_retrieval + $user_input = "test query" + $classifier = "test_classifier" + execute hf classifier check retrieval + $test_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + # Initialize LLMRails + app = LLMRails(config) + + # Override the action with mock + app.register_action("hf_classifier_check_retrieval", mock_action) + + # Run the flow + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify that $relevant_chunks was propagated to global context + self.assertEqual(app.context.get("relevant_chunks"), ["new_chunk_1", "new_chunk_2"]) + self.assertEqual(app.context.get("test_output"), ["new_chunk_1", "new_chunk_2"]) + + def test_hf_classifier_v1_retrieval_transform_propagates_to_caller_flow(self): + """Test that transformed relevant_chunks is accessible in caller flow for v1""" + + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["caller_chunk_1", "caller_chunk_2"], + "response": "Transformed response" + } + ) + + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow parent_flow + $user_input = "test query" + $classifier = "test_classifier" + execute hf classifier check retrieval + $parent_output = $relevant_chunks + + define flow test_parent + execute parent_flow + $test_result = $parent_output + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app = LLMRails(config) + app.register_action("hf_classifier_check_retrieval", mock_action) + + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify propagation to parent flow and beyond + self.assertEqual(app.context.get("relevant_chunks"), ["caller_chunk_1", "caller_chunk_2"]) + self.assertEqual(app.context.get("parent_output"), ["caller_chunk_1", "caller_chunk_2"]) + self.assertEqual(app.context.get("test_result"), ["caller_chunk_1", "caller_chunk_2"]) + + def test_hf_classifier_v1_retrieval_no_transform_doesnt_modify_chunks(self): + """Test that when is_transform is False, relevant_chunks is not modified in v1""" + + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=False, + is_blocked=False + ) + + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_no_transform + $user_input = "test query" + $original_chunks = ["original_chunk_1", "original_chunk_2"] + $relevant_chunks = $original_chunks + $classifier = "test_classifier" + execute hf classifier check retrieval + $test_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app = LLMRails(config) + app.register_action("hf_classifier_check_retrieval", mock_action) + + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify relevant_chunks remains unchanged + self.assertEqual(app.context.get("relevant_chunks"), ["original_chunk_1", "original_chunk_2"]) + self.assertEqual(app.context.get("test_output"), ["original_chunk_1", "original_chunk_2"]) + + def test_hf_classifier_v1_retrieval_with_empty_transform_text(self): + """Test transform with empty transform_text in v1""" + + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={} + ) + + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_empty_transform + $user_input = "test query" + $original_chunks = ["original"] + $relevant_chunks = $original_chunks + $classifier = "test_classifier" + execute hf classifier check retrieval + $test_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app = LLMRails(config) + app.register_action("hf_classifier_check_retrieval", mock_action) + + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Should set relevant_chunks to whatever transform_text returns + # In this case, it should be an empty dict or the value from transform_text + self.assertEqual(app.context.get("relevant_chunks"), {}) + + def test_hf_classifier_v1_retrieval_blocked_with_exceptions(self): + """Test blocked flow with exceptions enabled in v1""" + + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=False, + is_blocked=True + ) + + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_blocked + $user_input = "test query" + $original_chunks = ["original"] + $relevant_chunks = $original_chunks + $classifier = "test_classifier" + $config.enable_rails_exceptions = True + execute hf classifier check input + $test_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app = LLMRails(config) + app.register_action("hf_classifier_check_input", mock_action) + + # Should create event and stop, but relevant_chunks should remain unchanged + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # relevant_chunks should still be the original value since the flow stopped + self.assertEqual(app.context.get("relevant_chunks"), ["original"]) + + def test_hf_classifier_v1_multiple_calls_keep_transformed_value(self): + """Test that transformed relevant_chunks persists across multiple flow calls""" + + mock_action = AsyncMock() + # First call returns transform + mock_action.side_effect = [ + MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["first_transform_chunk"], + "response": "First transform" + } + ), + MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["second_transform_chunk"], + "response": "Second transform" + } + ) + ] + + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_multiple + $user_input = "test query" + $classifier = "test_classifier" + + # First call + $original_chunks = ["original"] + $relevant_chunks = $original_chunks + execute hf classifier check retrieval + $first_output = $relevant_chunks + + # Second call + execute hf classifier check retrieval + $second_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app = LLMRails(config) + app.register_action("hf_classifier_check_retrieval", mock_action) + + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify the transformations persist correctly + self.assertEqual(app.context.get("relevant_chunks"), ["second_transform_chunk"]) + self.assertEqual(app.context.get("first_output"), ["first_transform_chunk"]) + self.assertEqual(app.context.get("second_output"), ["second_transform_chunk"]) + + def test_original_hf_classifier_retrieval_transform_propagates_global(self): + """Test that transformed relevant_chunks propagates to global context in original flow""" + + # Create mock action that returns transform response + mock_action = AsyncMock() + mock_action.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["original_chunk_1", "original_chunk_2"], + "response": "Original transformed response" + } + ) + + # Create config with the original flow + config = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_retrieval + $user_input = "test query" + $classifier = "test_classifier" + execute hf classifier check retrieval $classifier + $test_output = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + # Initialize LLMRails + app = LLMRails(config) + + # Override the action with mock + app.register_action("HfClassifierCheckRetrievalAction", mock_action) + + # Run the flow + result = app.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify that $relevant_chunks was propagated to global context + self.assertEqual(app.context.get("relevant_chunks"), ["original_chunk_1", "original_chunk_2"]) + self.assertEqual(app.context.get("test_output"), ["original_chunk_1", "original_chunk_2"]) + + def test_both_versions_with_classifier_parameter(self): + """Test both versions with classifier parameter""" + + # Test v1 version + mock_action_v1 = AsyncMock() + mock_action_v1.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["v1_test_chunk"], + "response": "V1 test response" + } + ) + + config_v1 = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_v1 + $user_input = "test query" + $classifier = "test_classifier" + execute hf classifier check retrieval + $result_v1 = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app_v1 = LLMRails(config_v1) + app_v1.register_action("hf_classifier_check_retrieval", mock_action_v1) + result_v1 = app_v1.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Test original version + mock_action_original = AsyncMock() + mock_action_original.return_value = MockHfClassifierResponse( + is_transform=True, + transform_text={ + "relevant_chunks": ["original_test_chunk"], + "response": "Original test response" + } + ) + + config_original = RailsConfig.from_content( + colang_content=""" + define user express greeting + "Hello" + + define flow test_original + $user_input = "test query" + $classifier = "test_classifier" + execute hf classifier check retrieval $classifier + $result_original = $relevant_chunks + """, + yaml_content=""" + models: + - type: main + engine: openai + model: gpt-3.5-turbo + """ + ) + + app_original = LLMRails(config_original) + app_original.register_action("HfClassifierCheckRetrievalAction", mock_action_original) + result_original = app_original.generate(messages=[{"role": "user", "content": "Hello"}]) + + # Verify the fix works for both versions + self.assertEqual(app_v1.context.get("relevant_chunks"), ["v1_test_chunk"]) + self.assertEqual(app_v1.context.get("result_v1"), ["v1_test_chunk"]) + + self.assertEqual(app_original.context.get("relevant_chunks"), ["original_test_chunk"]) + self.assertEqual(app_original.context.get("result_original"), ["original_test_chunk"]) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From c10d0b5526b85ff2704f4df1c3b222c941edecd1 Mon Sep 17 00:00:00 2001 From: ekam-walia <164983405+ekam-walia@users.noreply.github.com> Date: Wed, 22 Jul 2026 13:53:35 +0530 Subject: [PATCH 2/4] line --- nemoguardrails/library/hf_classifier/flows.co | 2 +- nemoguardrails/library/hf_classifier/flows.v1.co | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/library/hf_classifier/flows.co b/nemoguardrails/library/hf_classifier/flows.co index 187791855e..3d1f2de8b4 100644 --- a/nemoguardrails/library/hf_classifier/flows.co +++ b/nemoguardrails/library/hf_classifier/flows.co @@ -24,4 +24,4 @@ flow hf classifier check retrieval $classifier if $response.is_transform # FIX: Added global declaration to propagate the change to calling context global $relevant_chunks - $relevant_chunks = $response.transform_text["relevant_chunks"] \ No newline at end of file + $relevant_chunks = $response.transform_text["relevant_chunks"] diff --git a/nemoguardrails/library/hf_classifier/flows.v1.co b/nemoguardrails/library/hf_classifier/flows.v1.co index f128ace9c1..df9d2d9a27 100644 --- a/nemoguardrails/library/hf_classifier/flows.v1.co +++ b/nemoguardrails/library/hf_classifier/flows.v1.co @@ -27,4 +27,4 @@ define subflow hf classifier check retrieval if $response.is_transform # FIX: Added global declaration to propagate the change to calling context global $relevant_chunks - $relevant_chunks = $response.transform_text["relevant_chunks"] \ No newline at end of file + $relevant_chunks = $response.transform_text["relevant_chunks"] From bf122760cafd6d64e444d0a51a70fed3731eac3d Mon Sep 17 00:00:00 2001 From: ekam-walia <164983405+ekam-walia@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:43:22 +0530 Subject: [PATCH 3/4] .. --- .../library/hf_classifier/flows.v1.co | 8 +- tests/test_hf_classifier_retrieval_global.py | 282 ++---------------- 2 files changed, 31 insertions(+), 259 deletions(-) diff --git a/nemoguardrails/library/hf_classifier/flows.v1.co b/nemoguardrails/library/hf_classifier/flows.v1.co index df9d2d9a27..b2e24b6c51 100644 --- a/nemoguardrails/library/hf_classifier/flows.v1.co +++ b/nemoguardrails/library/hf_classifier/flows.v1.co @@ -25,6 +25,8 @@ define subflow hf classifier check retrieval $response = execute hf_classifier_check_retrieval(classifier=$classifier) if $response.is_transform - # FIX: Added global declaration to propagate the change to calling context - global $relevant_chunks - $relevant_chunks = $response.transform_text["relevant_chunks"] + # Return the transformed chunks as the subflow result + return $response.transform_text["relevant_chunks"] + else + # If no transform, return whatever was passed in (or None) + return $relevant_chunks \ No newline at end of file diff --git a/tests/test_hf_classifier_retrieval_global.py b/tests/test_hf_classifier_retrieval_global.py index 482e3d173d..3ce742bf5f 100644 --- a/tests/test_hf_classifier_retrieval_global.py +++ b/tests/test_hf_classifier_retrieval_global.py @@ -1,26 +1,10 @@ -# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - import unittest -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock from nemoguardrails import LLMRails, RailsConfig class MockHfClassifierResponse: - """Mock response from HF classifier""" - def __init__(self, is_transform=False, transform_text=None, is_blocked=False): self.is_transform = is_transform self.transform_text = transform_text or {} @@ -28,12 +12,8 @@ def __init__(self, is_transform=False, transform_text=None, is_blocked=False): class TestHfClassifierRetrievalGlobal(unittest.TestCase): - """Test cases for HF classifier retrieval global variable propagation""" - - def test_hf_classifier_v1_retrieval_transform_propagates_global(self): - """Test that transformed relevant_chunks propagates to global context in v1""" - # Create mock action that returns transform response + def test_v1_retrieval_transform_propagates_global(self): mock_action = AsyncMock() mock_action.return_value = MockHfClassifierResponse( is_transform=True, @@ -43,7 +23,6 @@ def test_hf_classifier_v1_retrieval_transform_propagates_global(self): } ) - # Create config with the v1 flow config = RailsConfig.from_content( colang_content=""" define user express greeting @@ -52,7 +31,7 @@ def test_hf_classifier_v1_retrieval_transform_propagates_global(self): define flow test_retrieval $user_input = "test query" $classifier = "test_classifier" - execute hf classifier check retrieval + $relevant_chunks = execute hf classifier check retrieval $test_output = $relevant_chunks """, yaml_content=""" @@ -63,28 +42,20 @@ def test_hf_classifier_v1_retrieval_transform_propagates_global(self): """ ) - # Initialize LLMRails app = LLMRails(config) - - # Override the action with mock app.register_action("hf_classifier_check_retrieval", mock_action) + app.generate(messages=[{"role": "user", "content": "Hello"}]) - # Run the flow - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify that $relevant_chunks was propagated to global context self.assertEqual(app.context.get("relevant_chunks"), ["new_chunk_1", "new_chunk_2"]) self.assertEqual(app.context.get("test_output"), ["new_chunk_1", "new_chunk_2"]) - def test_hf_classifier_v1_retrieval_transform_propagates_to_caller_flow(self): - """Test that transformed relevant_chunks is accessible in caller flow for v1""" - + def test_original_retrieval_transform_propagates_global(self): mock_action = AsyncMock() mock_action.return_value = MockHfClassifierResponse( is_transform=True, transform_text={ - "relevant_chunks": ["caller_chunk_1", "caller_chunk_2"], - "response": "Transformed response" + "relevant_chunks": ["original_chunk_1", "original_chunk_2"], + "response": "Original transformed response" } ) @@ -93,15 +64,11 @@ def test_hf_classifier_v1_retrieval_transform_propagates_to_caller_flow(self): define user express greeting "Hello" - define flow parent_flow + define flow test_retrieval $user_input = "test query" $classifier = "test_classifier" - execute hf classifier check retrieval - $parent_output = $relevant_chunks - - define flow test_parent - execute parent_flow - $test_result = $parent_output + execute hf classifier check retrieval $classifier + $test_output = $relevant_chunks """, yaml_content=""" models: @@ -112,18 +79,13 @@ def test_hf_classifier_v1_retrieval_transform_propagates_to_caller_flow(self): ) app = LLMRails(config) - app.register_action("hf_classifier_check_retrieval", mock_action) - - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify propagation to parent flow and beyond - self.assertEqual(app.context.get("relevant_chunks"), ["caller_chunk_1", "caller_chunk_2"]) - self.assertEqual(app.context.get("parent_output"), ["caller_chunk_1", "caller_chunk_2"]) - self.assertEqual(app.context.get("test_result"), ["caller_chunk_1", "caller_chunk_2"]) + app.register_action("HfClassifierCheckRetrievalAction", mock_action) + app.generate(messages=[{"role": "user", "content": "Hello"}]) - def test_hf_classifier_v1_retrieval_no_transform_doesnt_modify_chunks(self): - """Test that when is_transform is False, relevant_chunks is not modified in v1""" + self.assertEqual(app.context.get("relevant_chunks"), ["original_chunk_1", "original_chunk_2"]) + self.assertEqual(app.context.get("test_output"), ["original_chunk_1", "original_chunk_2"]) + def test_no_transform_doesnt_modify_chunks(self): mock_action = AsyncMock() mock_action.return_value = MockHfClassifierResponse( is_transform=False, @@ -153,20 +115,16 @@ def test_hf_classifier_v1_retrieval_no_transform_doesnt_modify_chunks(self): app = LLMRails(config) app.register_action("hf_classifier_check_retrieval", mock_action) + app.generate(messages=[{"role": "user", "content": "Hello"}]) - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify relevant_chunks remains unchanged self.assertEqual(app.context.get("relevant_chunks"), ["original_chunk_1", "original_chunk_2"]) self.assertEqual(app.context.get("test_output"), ["original_chunk_1", "original_chunk_2"]) - def test_hf_classifier_v1_retrieval_with_empty_transform_text(self): - """Test transform with empty transform_text in v1""" - + def test_v1_retrieval_with_empty_transform_text(self): mock_action = AsyncMock() mock_action.return_value = MockHfClassifierResponse( is_transform=True, - transform_text={} + transform_text={"relevant_chunks": []} ) config = RailsConfig.from_content( @@ -179,7 +137,7 @@ def test_hf_classifier_v1_retrieval_with_empty_transform_text(self): $original_chunks = ["original"] $relevant_chunks = $original_chunks $classifier = "test_classifier" - execute hf classifier check retrieval + $relevant_chunks = execute hf classifier check retrieval $test_output = $relevant_chunks """, yaml_content=""" @@ -192,16 +150,12 @@ def test_hf_classifier_v1_retrieval_with_empty_transform_text(self): app = LLMRails(config) app.register_action("hf_classifier_check_retrieval", mock_action) + app.generate(messages=[{"role": "user", "content": "Hello"}]) - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Should set relevant_chunks to whatever transform_text returns - # In this case, it should be an empty dict or the value from transform_text - self.assertEqual(app.context.get("relevant_chunks"), {}) - - def test_hf_classifier_v1_retrieval_blocked_with_exceptions(self): - """Test blocked flow with exceptions enabled in v1""" + self.assertEqual(app.context.get("relevant_chunks"), []) + self.assertEqual(app.context.get("test_output"), []) + def test_v1_retrieval_blocked_flow(self): mock_action = AsyncMock() mock_action.return_value = MockHfClassifierResponse( is_transform=False, @@ -219,108 +173,7 @@ def test_hf_classifier_v1_retrieval_blocked_with_exceptions(self): $relevant_chunks = $original_chunks $classifier = "test_classifier" $config.enable_rails_exceptions = True - execute hf classifier check input - $test_output = $relevant_chunks - """, - yaml_content=""" - models: - - type: main - engine: openai - model: gpt-3.5-turbo - """ - ) - - app = LLMRails(config) - app.register_action("hf_classifier_check_input", mock_action) - - # Should create event and stop, but relevant_chunks should remain unchanged - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # relevant_chunks should still be the original value since the flow stopped - self.assertEqual(app.context.get("relevant_chunks"), ["original"]) - - def test_hf_classifier_v1_multiple_calls_keep_transformed_value(self): - """Test that transformed relevant_chunks persists across multiple flow calls""" - - mock_action = AsyncMock() - # First call returns transform - mock_action.side_effect = [ - MockHfClassifierResponse( - is_transform=True, - transform_text={ - "relevant_chunks": ["first_transform_chunk"], - "response": "First transform" - } - ), - MockHfClassifierResponse( - is_transform=True, - transform_text={ - "relevant_chunks": ["second_transform_chunk"], - "response": "Second transform" - } - ) - ] - - config = RailsConfig.from_content( - colang_content=""" - define user express greeting - "Hello" - - define flow test_multiple - $user_input = "test query" - $classifier = "test_classifier" - - # First call - $original_chunks = ["original"] - $relevant_chunks = $original_chunks execute hf classifier check retrieval - $first_output = $relevant_chunks - - # Second call - execute hf classifier check retrieval - $second_output = $relevant_chunks - """, - yaml_content=""" - models: - - type: main - engine: openai - model: gpt-3.5-turbo - """ - ) - - app = LLMRails(config) - app.register_action("hf_classifier_check_retrieval", mock_action) - - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify the transformations persist correctly - self.assertEqual(app.context.get("relevant_chunks"), ["second_transform_chunk"]) - self.assertEqual(app.context.get("first_output"), ["first_transform_chunk"]) - self.assertEqual(app.context.get("second_output"), ["second_transform_chunk"]) - - def test_original_hf_classifier_retrieval_transform_propagates_global(self): - """Test that transformed relevant_chunks propagates to global context in original flow""" - - # Create mock action that returns transform response - mock_action = AsyncMock() - mock_action.return_value = MockHfClassifierResponse( - is_transform=True, - transform_text={ - "relevant_chunks": ["original_chunk_1", "original_chunk_2"], - "response": "Original transformed response" - } - ) - - # Create config with the original flow - config = RailsConfig.from_content( - colang_content=""" - define user express greeting - "Hello" - - define flow test_retrieval - $user_input = "test query" - $classifier = "test_classifier" - execute hf classifier check retrieval $classifier $test_output = $relevant_chunks """, yaml_content=""" @@ -331,94 +184,11 @@ def test_original_hf_classifier_retrieval_transform_propagates_global(self): """ ) - # Initialize LLMRails app = LLMRails(config) + app.register_action("hf_classifier_check_retrieval", mock_action) + app.generate(messages=[{"role": "user", "content": "Hello"}]) - # Override the action with mock - app.register_action("HfClassifierCheckRetrievalAction", mock_action) - - # Run the flow - result = app.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify that $relevant_chunks was propagated to global context - self.assertEqual(app.context.get("relevant_chunks"), ["original_chunk_1", "original_chunk_2"]) - self.assertEqual(app.context.get("test_output"), ["original_chunk_1", "original_chunk_2"]) - - def test_both_versions_with_classifier_parameter(self): - """Test both versions with classifier parameter""" - - # Test v1 version - mock_action_v1 = AsyncMock() - mock_action_v1.return_value = MockHfClassifierResponse( - is_transform=True, - transform_text={ - "relevant_chunks": ["v1_test_chunk"], - "response": "V1 test response" - } - ) - - config_v1 = RailsConfig.from_content( - colang_content=""" - define user express greeting - "Hello" - - define flow test_v1 - $user_input = "test query" - $classifier = "test_classifier" - execute hf classifier check retrieval - $result_v1 = $relevant_chunks - """, - yaml_content=""" - models: - - type: main - engine: openai - model: gpt-3.5-turbo - """ - ) - - app_v1 = LLMRails(config_v1) - app_v1.register_action("hf_classifier_check_retrieval", mock_action_v1) - result_v1 = app_v1.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Test original version - mock_action_original = AsyncMock() - mock_action_original.return_value = MockHfClassifierResponse( - is_transform=True, - transform_text={ - "relevant_chunks": ["original_test_chunk"], - "response": "Original test response" - } - ) - - config_original = RailsConfig.from_content( - colang_content=""" - define user express greeting - "Hello" - - define flow test_original - $user_input = "test query" - $classifier = "test_classifier" - execute hf classifier check retrieval $classifier - $result_original = $relevant_chunks - """, - yaml_content=""" - models: - - type: main - engine: openai - model: gpt-3.5-turbo - """ - ) - - app_original = LLMRails(config_original) - app_original.register_action("HfClassifierCheckRetrievalAction", mock_action_original) - result_original = app_original.generate(messages=[{"role": "user", "content": "Hello"}]) - - # Verify the fix works for both versions - self.assertEqual(app_v1.context.get("relevant_chunks"), ["v1_test_chunk"]) - self.assertEqual(app_v1.context.get("result_v1"), ["v1_test_chunk"]) - - self.assertEqual(app_original.context.get("relevant_chunks"), ["original_test_chunk"]) - self.assertEqual(app_original.context.get("result_original"), ["original_test_chunk"]) + self.assertEqual(app.context.get("relevant_chunks"), ["original"]) if __name__ == "__main__": From c55221042495b3bee79eb8d262ae8879d6127259 Mon Sep 17 00:00:00 2001 From: ekam-walia <164983405+ekam-walia@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:47:36 +0530 Subject: [PATCH 4/4] . --- nemoguardrails/library/hf_classifier/flows.v1.co | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/library/hf_classifier/flows.v1.co b/nemoguardrails/library/hf_classifier/flows.v1.co index b2e24b6c51..91052a11e7 100644 --- a/nemoguardrails/library/hf_classifier/flows.v1.co +++ b/nemoguardrails/library/hf_classifier/flows.v1.co @@ -25,8 +25,6 @@ define subflow hf classifier check retrieval $response = execute hf_classifier_check_retrieval(classifier=$classifier) if $response.is_transform - # Return the transformed chunks as the subflow result - return $response.transform_text["relevant_chunks"] - else - # If no transform, return whatever was passed in (or None) - return $relevant_chunks \ No newline at end of file + # In Colang 1.0, subflows share global scope. + # The caller must initialize $relevant_chunks before calling this subflow. + $relevant_chunks = $response.transform_text["relevant_chunks"] \ No newline at end of file