executor.py

from ast import Attribute
from dotenv import load_dotenv

load_dotenv(override=True)

import re
import os
import pandas as pd
import json
from typing import List, Dict, Any
import pandas as pd
import gradio as gr
import datetime
from pathlib import Path
import json

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

# Load the taxonomies
ran = RiskAtlasNexus() # type: ignore


def clear_previous_risks():
    return gr.Markdown("""<h2> Potential Risks </h2> """), [], gr.Dataset(samples=[], 
                                     sample_labels=[], 
                                     samples_per_page=50, visible=False), gr.DownloadButton("Download JSON", visible=False, ), "", gr.Dataset(samples=[], sample_labels=[], visible=False), gr.DataFrame([], wrap=True, show_copy_button=True, show_search="search", visible=False), gr.DataFrame([], wrap=True, show_copy_button=True, show_search="search", visible=False), gr.Markdown(" ") 

def clear_previous_mitigations():
     return "", gr.Dataset(samples=[], sample_labels=[], visible=False), gr.DataFrame([], wrap=True, show_copy_button=True, show_search="search", visible=False), gr.DataFrame([], wrap=True, show_copy_button=True, show_search="search", visible=False), gr.Markdown(" ") 

@lru_cache
def risk_identifier(usecase: str, 
                    model_name_or_path: str = "ibm/granite-3-3-8b-instruct", 
                    taxonomy: str = "ibm-risk-atlas"): # -> List[Dict[str, Any]]: #pd.DataFrame:

    downloadable = False
    inference_engine = WMLInferenceEngine(
        model_name_or_path= model_name_or_path,
        credentials={
            "api_key": os.environ["WML_API_KEY"],
            "api_url": os.environ["WML_API_URL"],
            "project_id": os.environ["WML_PROJECT_ID"],
        },
        parameters=WMLInferenceEngineParams(
            max_new_tokens=150, decoding_method="greedy", repetition_penalty=1
        ),  # type: ignore
    )

    risks = ran.identify_risks_from_usecases( # type: ignore
        usecases=[usecase],
        inference_engine=inference_engine,
        taxonomy=taxonomy,
        max_risk=5
    )[0]
    

    sample_labels = [r.name if r else r.id for r in risks]

    out_sec = gr.Markdown("""<h2> Potential Risks </h2> """)

    # write out a JSON
    data = {'time': str(datetime.datetime.now(datetime.timezone.utc)),
                'intent': usecase,
                'model': model_name_or_path,
                'taxonomy': taxonomy,
                'risks': [json.loads(r.json()) for r in risks]
        }
    file_path = Path("static/download.json")
    with open(file_path, mode='w') as f:
        f.write(json.dumps(data, indent=4))
        downloadable = True

   
    #return out_df
    return out_sec, gr.State(risks), gr.Dataset(samples=[r.id for r in risks], 
                                     sample_labels=sample_labels, 
                                     samples_per_page=50, visible=True, label="Estimated by an LLM."), gr.DownloadButton("Download JSON", "static/download.json", visible=(downloadable and len(risks) > 0))
    

@lru_cache
def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Markdown, gr.Dataset, gr.DataFrame, gr.DataFrame, gr.Markdown]:
    """
    For a specific risk (riskid), returns
    (a) a risk description
    (b) related risks - as a dataset
    (c) mitigations
    (d) related ai evaluations

    """
    
    try:
        risk_desc = ran.get_risk(id=riskid).description # type: ignore
        risk_sec = f"<h3>Description: </h3> {risk_desc}"
    except AttributeError:
        risk_sec = ""

    related_risk_ids = [r.id for r in ran.get_related_risks(id=riskid)]
    related_ai_eval_ids = [ai_eval.id for ai_eval in ran.get_related_evaluations(risk_id=riskid)]

    action_ids = []
    control_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
                
                rac = ran.get_related_risk_controls(id=i)
                if rac:
                    control_ids += rac
    
        else:
            action_ids = []
            control_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)
       
    
    # Sanitize outputs
    if not related_risk_ids:
        label = "No related risks found."
        samples = None
        sample_labels = None
    else:
        label = f"Risks from other taxonomies related to {riskid}"
        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:
        alabel = "No mitigations found."
        asamples = None
        asample_labels = None
        mitdf = pd.DataFrame()
        
    else:
        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})
    
    if not related_ai_eval_ids:
        blabel = "No related AI evaluations found."
        bsamples = None
        bsample_labels = None
        aievalsdf = pd.DataFrame()
    else:
        blabel = f"AI Evaluations related to {riskid}"
        bsamples = related_ai_eval_ids
        bsample_labels = [ran.get_evaluation(i).description for i in bsamples] # type: ignore
        bsample_name = [ran.get_evaluation(i).name for i in bsamples] #type: ignore
        aievalsdf = pd.DataFrame({"AI Evaluation": bsample_name, "Description": bsample_labels})
    
    status = gr.Markdown(" ") if len(mitdf) > 0 else gr.Markdown("No mitigations found.")

    return (gr.Markdown(risk_sec), 
            gr.Dataset(samples=samples, label=label, sample_labels=sample_labels, visible=True),
            gr.DataFrame(mitdf, wrap=True, show_copy_button=True, show_search="search", label=alabel, visible=True),
            gr.DataFrame(aievalsdf, wrap=True, show_copy_button=True, show_search="search", label=blabel, visible=True),
            status)