fix _get_action_details_from_flow

Author: Pouyanpi

nemoguardrails/rails/llm/llmrails.py

@@ -1309,22 +1309,56 @@ def _update_explain_info():
 
 
 def _get_action_details_from_flow_id(
-    flow_id: str, flows: List[Union[Dict, Any]]
+    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."""
+    """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",
+        "content safety check",
+        "topic safety check",
+    ]
+    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 (
-            flow_id == flow["id"]
-            or flow_id.startswith("content safety check")
-            or flow_id.startswith("topic safety check")
+            element["_type"] == "run_action"
+            and element["_source_mapping"]["filename"] == "flows.v1.co"
+            and "execute" in element["_source_mapping"]["line_text"]
+            and "action_name" in element
         ):
-            for element in flow["elements"]:
-                if (
-                    element["_type"] == "run_action"
-                    and element["_source_mapping"]["filename"] == "flows.v1.co"
-                    and "execute" in element["_source_mapping"]["line_text"]
-                    and "action_name" in element
-                ):
-                    return element["action_name"], element["action_params"]
+            return element["action_name"], element["action_params"]
 
-    raise ValueError(f"No action found for flow_id: {flow_id}")
+    raise ValueError(f"No run_action element found for flow_id: {flow_id}")