refactor(llm): move get_action_details_from_flow_id from llmrails.py to utils.py (#1341)

Pouyanpi · web-flow · commit 57a3abcf1c3f · 2025-08-19T18:12:56.000+02:00
diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py
@@ -73,7 +73,10 @@
     GenerationOptions,
     GenerationResponse,
 )
-from nemoguardrails.rails.llm.utils import get_history_cache_key
+from nemoguardrails.rails.llm.utils import (
+    get_action_details_from_flow_id,
+    get_history_cache_key,
+)
 from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler
 from nemoguardrails.utils import (
     extract_error_json,
@@ -1590,7 +1593,7 @@ def _prepare_params(
         output_rails_flows_id = self.config.rails.output.flows
         stream_first = stream_first or output_rails_streaming_config.stream_first
         get_action_details = partial(
-            _get_action_details_from_flow_id, flows=self.config.flows
+            get_action_details_from_flow_id, flows=self.config.flows
         )
 
         parallel_mode = getattr(self.config.rails.output, "parallel", False)
@@ -1746,58 +1749,3 @@ def _prepare_params(
                 # yield the individual chunks directly from the buffer strategy
                 for chunk in user_output_chunks:
                     yield chunk
-
-
-def _get_action_details_from_flow_id(
-    flow_id: str,
-    flows: List[Union[Dict, Any]],
-    prefixes: Optional[List[str]] = None,
-) -> Tuple[str, Any]:
-    """Get the action name and parameters from the flow id.
-
-    First, try to find an exact match.
-    If not found, then if the provided flow_id starts with one of the special prefixes,
-    return the first flow whose id starts with that same prefix.
-    """
-
-    supported_prefixes = [
-        "content safety check output",
-        "topic safety check output",
-    ]
-    if prefixes:
-        supported_prefixes.extend(prefixes)
-
-    candidate_flow = None
-
-    for flow in flows:
-        # If exact match, use it
-        if flow["id"] == flow_id:
-            candidate_flow = flow
-            break
-
-        # If no exact match, check if both the provided flow_id and this flow's id share a special prefix
-        for prefix in supported_prefixes:
-            if flow_id.startswith(prefix) and flow["id"].startswith(prefix):
-                candidate_flow = flow
-                # We don't break immediately here because an exact match would have been preferred,
-                # but since we're in the else branch it's fine to choose the first matching candidate.
-                # TODO:we should avoid having multiple matchin prefixes
-                break
-
-        if candidate_flow is not None:
-            break
-
-    if candidate_flow is None:
-        raise ValueError(f"No action found for flow_id: {flow_id}")
-
-    # we have identified a candidate, look for the run_action element.
-    for element in candidate_flow["elements"]:
-        if (
-            element["_type"] == "run_action"
-            and element["_source_mapping"]["filename"].endswith(".co")
-            and "execute" in element["_source_mapping"]["line_text"]
-            and "action_name" in element
-        ):
-            return element["action_name"], element["action_params"]
-
-    raise ValueError(f"No run_action element found for flow_id: {flow_id}")
diff --git a/nemoguardrails/rails/llm/utils.py b/nemoguardrails/rails/llm/utils.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
-from typing import List
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 
 def get_history_cache_key(messages: List[dict]) -> str:
@@ -56,3 +56,57 @@ def get_history_cache_key(messages: List[dict]) -> str:
     history_cache_key = ":".join(key_items)
 
     return history_cache_key
+
+
+def get_action_details_from_flow_id(
+    flow_id: str,
+    flows: List[Union[Dict, Any]],
+    prefixes: Optional[List[str]] = None,
+) -> Tuple[str, Any]:
+    """Get the action name and parameters from the flow id.
+
+    First, try to find an exact match.
+    If not found, then if the provided flow_id starts with one of the special prefixes,
+    return the first flow whose id starts with that same prefix.
+    """
+    supported_prefixes = [
+        "content safety check output",
+        "topic safety check output",
+    ]
+    if prefixes:
+        supported_prefixes.extend(prefixes)
+
+    candidate_flow = None
+
+    for flow in flows:
+        # If exact match, use it
+        if flow["id"] == flow_id:
+            candidate_flow = flow
+            break
+
+        # If no exact match, check if both the provided flow_id and this flow's id share a special prefix
+        for prefix in supported_prefixes:
+            if flow_id.startswith(prefix) and flow["id"].startswith(prefix):
+                candidate_flow = flow
+                # We don't break immediately here because an exact match would have been preferred,
+                # but since we're in the else branch it's fine to choose the first matching candidate.
+                # TODO:we should avoid having multiple matchin prefixes
+                break
+
+        if candidate_flow is not None:
+            break
+
+    if candidate_flow is None:
+        raise ValueError(f"No action found for flow_id: {flow_id}")
+
+    # we have identified a candidate, look for the run_action element.
+    for element in candidate_flow["elements"]:
+        if (
+            element["_type"] == "run_action"
+            and element["_source_mapping"]["filename"].endswith(".co")
+            and "execute" in element["_source_mapping"]["line_text"]
+            and "action_name" in element
+        ):
+            return element["action_name"], element["action_params"]
+
+    raise ValueError(f"No run_action element found for flow_id: {flow_id}")
diff --git a/tests/test_llmrails.py b/tests/test_llmrails.py
@@ -21,7 +21,7 @@
 
 from nemoguardrails import LLMRails, RailsConfig
 from nemoguardrails.rails.llm.config import Model
-from nemoguardrails.rails.llm.llmrails import _get_action_details_from_flow_id
+from nemoguardrails.rails.llm.llmrails import get_action_details_from_flow_id
 from tests.utils import FakeLLM, clean_events, event_sequence_conforms
 
 
@@ -627,124 +627,6 @@ async def compute(what: Optional[str] = "2 + 3"):
     }
 
 
-# get_action_details_from_flow_id used in llmrails.py
-
-
-@pytest.fixture
-def dummy_flows() -> List[Union[Dict, Any]]:
-    return [
-        {
-            "id": "test_flow",
-            "elements": [
-                {
-                    "_type": "run_action",
-                    "_source_mapping": {
-                        "filename": "flows.v1.co",
-                        "line_text": "execute something",
-                    },
-                    "action_name": "test_action",
-                    "action_params": {"param1": "value1"},
-                }
-            ],
-        },
-        # Additional flow that should match on a prefix
-        {
-            "id": "other_flow is prefix",
-            "elements": [
-                {
-                    "_type": "run_action",
-                    "_source_mapping": {
-                        "filename": "flows.v1.co",
-                        "line_text": "execute something else",
-                    },
-                    "action_name": "other_action",
-                    "action_params": {"param2": "value2"},
-                }
-            ],
-        },
-        {
-            "id": "test_rails_co",
-            "elements": [
-                {
-                    "_type": "run_action",
-                    "_source_mapping": {
-                        "filename": "rails.co",
-                        "line_text": "execute something",
-                    },
-                    "action_name": "test_action_supported",
-                    "action_params": {"param1": "value1"},
-                }
-            ],
-        },
-        {
-            "id": "test_rails_co_v2",
-            "elements": [
-                {
-                    "_type": "run_action",
-                    "_source_mapping": {
-                        "filename": "rails.co",
-                        "line_text": "await something",  # in colang 2 we use await
-                    },
-                    "action_name": "test_action_not_supported",
-                    "action_params": {"param1": "value1"},
-                }
-            ],
-        },
-    ]
-
-
-def test_get_action_details_exact_match(dummy_flows):
-    action_name, action_params = _get_action_details_from_flow_id(
-        "test_flow", dummy_flows
-    )
-    assert action_name == "test_action"
-    assert action_params == {"param1": "value1"}
-
-
-def test_get_action_details_exact_match_any_co_file(dummy_flows):
-    action_name, action_params = _get_action_details_from_flow_id(
-        "test_rails_co", dummy_flows
-    )
-    assert action_name == "test_action_supported"
-    assert action_params == {"param1": "value1"}
-
-
-def test_get_action_details_exact_match_not_colang_2(dummy_flows):
-    with pytest.raises(ValueError) as exc_info:
-        _get_action_details_from_flow_id("test_rails_co_v2", dummy_flows)
-
-    assert "No run_action element found for flow_id" in str(exc_info.value)
-
-
-def test_get_action_details_prefix_match(dummy_flows):
-    # For a flow_id that starts with the prefix "other_flow",
-    # we expect to retrieve the action details from the flow whose id starts with that prefix.
-    # we expect a result since we are passing the prefixes argument.
-    action_name, action_params = _get_action_details_from_flow_id(
-        "other_flow", dummy_flows, prefixes=["other_flow"]
-    )
-    assert action_name == "other_action"
-    assert action_params == {"param2": "value2"}
-
-
-def test_get_action_details_prefix_match_unsupported_prefix(dummy_flows):
-    # For a flow_id that starts with the prefix "other_flow",
-    # we expect to retrieve the action details from the flow whose id starts with that prefix.
-    # but as the prefix is not supported, we expect a ValueError.
-
-    with pytest.raises(ValueError) as exc_info:
-        _get_action_details_from_flow_id("other_flow", dummy_flows)
-
-    assert "No action found for flow_id" in str(exc_info.value)
-
-
-def test_get_action_details_no_match(dummy_flows):
-    # Tests that a non matching flow_id raises a ValueError
-    with pytest.raises(ValueError) as exc_info:
-        _get_action_details_from_flow_id("non_existing_flow", dummy_flows)
-    assert "No action found for flow_id" in str(exc_info.value)
-
-
 @pytest.fixture
 def llm_config_with_main():
     """Fixture providing a basic config with a main LLM."""
diff --git a/tests/test_llmrails_reasoning.py b/tests/test_llmrails_reasoning.py
@@ -13,13 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Optional
 
 import pytest
 
 from nemoguardrails import LLMRails, RailsConfig
-from nemoguardrails.rails.llm.llmrails import _get_action_details_from_flow_id
-from tests.utils import FakeLLM, clean_events, event_sequence_conforms
+from tests.utils import FakeLLM
 
 
 @pytest.fixture
diff --git a/tests/test_rails_llm_utils.py b/tests/test_rails_llm_utils.py
