WML prompt update

executor.py

@@ -9,7 +9,7 @@ from risk_atlas_nexus.blocks.inference import WMLInferenceEngine
 from risk_atlas_nexus.blocks.inference.params import WMLInferenceEngineParams
 from risk_atlas_nexus.library import RiskAtlasNexus
 
-from functools import lru_cache
+from functools import lru_cache, wraps
 from dotenv import load_dotenv
 
 load_dotenv(override=True)
@@ -65,7 +65,26 @@ def generate_subgraph(risk):
     return gr.Markdown(value = diagram_string)
 
 
-@lru_cache
+def custom_lru_cache(maxsize=128, exclude_values=(None,[],[[]])):
+    """
+    Make the LRU cache not cache result when empty result was returned
+    """
+    def decorator(func):
+        cached_func = lru_cache(maxsize=maxsize)(func)
+        
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            result = cached_func(*args, **kwargs)
+            # check for empty df of risks
+            if result[2].constructor_args["samples"] in exclude_values:
+                return func(*args, **kwargs)
+            return result
+        
+        return wrapper
+    return decorator
+
+
+@custom_lru_cache(exclude_values=(None, []))
 def risk_identifier(usecase: str, 
                     model_name_or_path: str = "meta-llama/llama-3-3-70b-instruct", 
                     taxonomy: str = "ibm-risk-atlas"): # -> List[Dict[str, Any]]: #pd.DataFrame:
@@ -79,7 +98,7 @@ def risk_identifier(usecase: str,
             "project_id": os.environ["WML_PROJECT_ID"],
         },
         parameters=WMLInferenceEngineParams(
-            max_new_tokens=150, decoding_method="greedy", repetition_penalty=1
+            max_new_tokens=1000, decoding_method="greedy", repetition_penalty=1
         ),  # type: ignore
     )
 
@@ -121,6 +140,7 @@ def get_controls_and_actions(riskid, taxonomy):
     related_risk_ids = [r.id for r in ran.get_related_risks(id=riskid)]
     action_ids = []
     control_ids =[]
+    intrinsic_ids=[]
 
     if taxonomy == "ibm-risk-atlas":
         # look for actions associated with related risks    
@@ -133,16 +153,22 @@ def get_controls_and_actions(riskid, taxonomy):
                 rac = ran.get_related_risk_controls(id=i)
                 if rac:
                     control_ids += rac
+                
+                ran_intrinsics = ran.get_related_intrinsics(risk_id=i)
+                if ran_intrinsics: 
+                    intrinsic_ids += ran_intrinsics
     
         else:
             action_ids = []
             control_ids = []
+            intrinsic_ids=[]
     else:
         # Use only actions related to primary risks
         action_ids = ran.get_related_actions(id=riskid)
         control_ids = ran.get_related_risk_controls(id=riskid)
+        intrinsic_ids = ran.get_related_intrinsics(risk_id=riskid)
 
-    return [ran.get_action_by_id(i).name for i in action_ids] + [ran.get_risk_control(i.id).name for i in control_ids] #type: ignore
+    return [ran.get_action_by_id(i).name for i in action_ids] + [ran.get_risk_control(i.id).name for i in control_ids] + [ran.get_intrinsic(i.id).name for i in intrinsic_ids]#type: ignore
 
 
 @lru_cache
@@ -169,26 +195,34 @@ def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Markdown, gr.Dataset, gr
 
     action_ids = []
     control_ids =[]
+    intrinsic_ids=[]
 
     if taxonomy == "ibm-risk-atlas":
         # look for actions associated with related risks    
         if related_risk_ids:
             for i in related_risk_ids:
-                rai = ran.get_related_actions(id=i)
-                if rai:
-                    action_ids += rai
+                ran_actions = ran.get_related_actions(id=i)
+                if ran_actions:
+                    action_ids += ran_actions
                 
-                rac = ran.get_related_risk_controls(id=i)
-                if rac:
-                    control_ids += rac
+                ran_controls = ran.get_related_risk_controls(id=i)
+                if ran_controls:
+                    control_ids += ran_controls
+                
+                ran_intrinsics = ran.get_related_intrinsics(risk_id=i)
+                if ran_intrinsics: 
+                    intrinsic_ids += ran_intrinsics
     
         else:
             action_ids = []
             control_ids = []
+            intrinsic_ids=[]
     else:
         # Use only actions related to primary risks
         action_ids = ran.get_related_actions(id=riskid)
         control_ids = ran.get_related_risk_controls(id=riskid)
+        intrinsic_ids = ran.get_related_intrinsics(risk_id=riskid)
+        
 
     # Sanitize outputs
     if not related_risk_ids:
@@ -200,7 +234,7 @@ def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Markdown, gr.Dataset, gr
         samples = related_risk_ids
         sample_labels = [i.name for i in ran.get_related_risks(id=riskid)] #type: ignore
 
-    if not action_ids and not control_ids:
+    if not action_ids and not control_ids and not intrinsic_ids:
         alabel = "No mitigations found."
         asamples = None
         asample_labels = None
@@ -210,9 +244,11 @@ def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Markdown, gr.Dataset, gr
         alabel = f"Mitigation actions and controls related to risk {riskid}."
         asamples = action_ids
         asamples_ctl = control_ids
-        asample_labels = [ran.get_action_by_id(i).description for i in asamples] + [ran.get_risk_control(i.id).name for i in asamples_ctl]# type: ignore
-        asample_name = [ran.get_action_by_id(i).name for i in asamples] + [ran.get_risk_control(i.id).name for i in asamples_ctl] #type: ignore
-        mitdf = pd.DataFrame({"Mitigation": asample_name, "Description": asample_labels})
+        asamples_int = intrinsic_ids
+        asample_labels = [ran.get_action_by_id(i).description for i in asamples] + [ran.get_risk_control(i.id).name for i in asamples_ctl] + [ran.get_intrinsic(i.id).description for i in asamples_int]# type: ignore
+        asample_name = [ran.get_action_by_id(i).name for i in asamples] + [ran.get_risk_control(i.id).name for i in asamples_ctl] + [ran.get_intrinsic(i.id).name for i in asamples_int] #type: ignore
+        asample_types = ["Action" for i in asamples] + ["Control" for i in asamples_ctl] + ["Intrinsic" for i in asamples_int]
+        mitdf = pd.DataFrame({"Type": asample_types, "Mitigation": asample_name, "Description": asample_labels})
     
     if not related_ai_eval_ids:
         blabel = "No related AI evaluations found."