abdulnim/GRC_framework

1 2import os3import requests4import gradio as gr5requests.adapters.DEFAULT_TIMEOUT = 606import time7import openai8from openai import OpenAI9from utils import ai_audit_analysis_categories, get_system_prompt, ANALYSIS_TYPES10import json11 12 13# Create Global Variables14client = OpenAI(api_key= "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")15 16global complete_chat_history, bot_last_message17bot_last_message = "" 18complete_chat_history = []19 20 21# /////////////////// *****************************///////////////// Utitlity Functions22#region Utility Functions23 24# Function to update OpenAPI key the API key25def update_api_key(new_api_key):26    global client27    if new_api_key.strip() != "":28        client = OpenAI(api_key=new_api_key)29    return "API Key updated successfully"30 31 32def load_chatboat_last_message():33    return bot_last_message34 35def load_chatboat_complet_history():36    complete_text = ""37    for turn in complete_chat_history:38        user_message, bot_message = turn39        complete_text = f"{complete_text}\nUser: {user_message}\nAssistant: {bot_message}"40    return complete_text41 42 43def format_json_result_to_html(result):44    formatted_result = ""45    for key, value in result.items():46        if isinstance(value, list):47            formatted_result += f"<strong>{key.title()}:</strong><br>" + "<br>".join(value) + "<br><br>"48        else:49            formatted_result += f"<strong>{key.title()}:</strong> {value}<br>"50    return formatted_result.strip()51 52def format_json_result(result):53    formatted_result = ""54    for key, value in result.items():55        if isinstance(value, list):56            formatted_result += f"{key.title()}:\n" + "\n".join(value) + "\n\n"57        else:58            formatted_result += f"{key.title()}: {value}\n"59    return formatted_result.strip()60 61# Function to dynamically format the JSON result into Markdown format62def format_result_to_markdown(result):63    formatted_result = ""64    for key, value in result.items():65        formatted_result += f"**{key.title()}**: "66        if isinstance(value, list):67            formatted_result += "\n" + "\n".join(f"- {item}" for item in value) + "\n\n"68        else:69            formatted_result += f"{value}\n\n"70    return formatted_result.strip()71 72#endregion73 74 75 76 77 78 79# /////////////////// *****************************///////////////// Conversation with Open Ai Chatboat80#region Conversation with Open Ai Chatboat81# A Normal call to OpenAI API '''82def chat(system_prompt, user_prompt, model = 'gpt-3.5-turbo', temperature = 0):83    response = client.chat.completions.create(84    messages=[85        {"role": "system", "content": system_prompt},86        {"role": "user", "content": user_prompt}87    ],88    model="gpt-3.5-turbo",89    )90    91    res = response.choices[0].message.content  92    return res93 94# Lets format the prompt from the chat_history so that its looks good on the UI95def format_chat_prompt(message, chat_history, max_convo_length):96    prompt = ""97    for turn in chat_history[-max_convo_length:]:98        user_message, bot_message = turn99        prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"100    prompt = f"{prompt}\nUser: {message}\nAssistant:"101    return prompt102 103 104# This function gets a message from user, passes it to chat gpt and return the output105def get_response_from_chatboat(message,chat_history, max_convo_length=10):106    global bot_last_message, complete_chat_history107    formatted_prompt = format_chat_prompt(message, chat_history, max_convo_length)108    bot_message = chat(system_prompt='You are a friendly chatbot. Generate the output for only the Assistant.',user_prompt=formatted_prompt)109 110    chat_history.append((message, bot_message))111    complete_chat_history.append((message, bot_message))112    bot_last_message = bot_message113    return "", chat_history114 115#endregion116 117 118 119def analyse_current_conversation(text, analysis_type):120 121    try:122        if(ANALYSIS_TYPES.get(analysis_type, None) is None):123            return f"Analysis type {analysis_type} is not implemented yet, please choose another category"124        125        if not text:126            return f"No text provided to analyze for {analysis_type}, please provide text or load from chatboat history"127        128        word_count = len(text.split())129 130        if(word_count < 20 ):131            return f" The text is too short to analyze for {analysis_type}, please provide a large text"132 133        system_prompt = get_system_prompt(analysis_type)134        text_to_analyze = text135 136        response = client.chat.completions.create(137        messages=[138            {"role": "system", "content": system_prompt},139            {"role": "user", "content": text_to_analyze}140        ],141        model="gpt-3.5-turbo",142        )143        144        analysis_result = response.choices[0].message.content  145        print(analysis_result)146 147        # Parse the result, handle JSON parsing errors148        try:149            parsed_result = json.loads(analysis_result)150        except json.JSONDecodeError:151            return "Failed to parse the analysis result. Please check the format of the returned data."152 153        formatted_json = format_result_to_markdown(parsed_result)154        return formatted_json155 156    except KeyError as e:157        return f"Key error occurred: {e}. Please check your keys."158    except Exception as e:159        # Check if the error message is related to the API key160        if 'API key' in str(e):161            return "OpenAI API key error: Please verify your API key."162        else:163            return f"An unexpected error occurred: {e}. Please check your implementation."164 165 166    # parsed_result = json.loads(analysis_result)167 168    # formated_json = format_result_to_markdown(parsed_result)169 170    # print(parsed_result)171    # # Your implementation for counting words and performing analysis172    # return formated_json173 174 175 176 177 178#region UI Related Functions179 180def update_dropdown(main_category):181    # Get the subcategories based on the selected main category182    subcategories = ai_audit_analysis_categories.get(main_category, [])183    print(subcategories)184    return gr.Dropdown(choices=subcategories, value=subcategories[0] if subcategories else None)185 186 187def update_analysis_type(subcategory):188    pass189    print(subcategory)190 191 192 193#endregion194 195 196 197with gr.Blocks() as demo:198 199    gr.Markdown("<center><img src='https://huggingface.co/spaces/abdulnim/GRC_framework/resolve/main/logo.png' alt='Align X' width='150'/></center>")200 201 202    # Add a text field for the API key203    api_key_field = gr.Textbox(label="Enter your Chatgpt OpenAI API Key")204    update_api_key_btn = gr.Button("Update API Key")205    update_api_key_btn.click(update_api_key, inputs=[api_key_field], outputs=[])206 207    # gr.Markdown("# AI Audit and GRC Framework!")208    gr.Markdown("# AlignXX Demo")209 210    with gr.Tabs():211        with gr.TabItem("Prompt Testing"):212            gr.Markdown("## Prompt Testing")213            chatbot = gr.Chatbot(height=600)214            msg = gr.Textbox(label="Write something for the chatbot here")215            clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")216            submit_btn = gr.Button("Submit")217            submit_btn.click(get_response_from_chatboat, inputs=[msg, chatbot], outputs=[msg, chatbot])218            msg.submit(get_response_from_chatboat, inputs=[msg, chatbot], outputs=[msg, chatbot])219 220        with gr.TabItem("Prompt Assessment"):221            gr.Markdown("## Prompt Assessment")222            gr.Markdown("Load your chatbot text or write your own to  and analyze it")223            text_field = gr.Textbox(label="Text to Process", interactive=True, lines=2)224 225            # Radio button and dropdown list226            initial_main_category = next(iter(ai_audit_analysis_categories))227            initial_sub_categories = ai_audit_analysis_categories[initial_main_category]228 229            main_category_radio = gr.Radio(list(ai_audit_analysis_categories.keys()), label="Main Audit Categories", value=initial_main_category)230            sub_category_dropdown = gr.Dropdown(choices=initial_sub_categories, label="Sub Categories", value=initial_sub_categories[0])231            # Update the dropdown based on the radio selection232            main_category_radio.change(fn=update_dropdown, inputs= main_category_radio, outputs=sub_category_dropdown)233            sub_category_dropdown.change(fn=update_analysis_type, inputs=sub_category_dropdown)234 235            load_last_message_btn = gr.Button("Load Last Message")236            load_complete_conv_btn = gr.Button("Load Complete Chat History")237            process_btn = gr.Button("Process")238            # analysis_result = gr.Label()239            analysis_result = gr.Markdown()240            load_last_message_btn.click(load_chatboat_last_message, inputs=[], outputs=text_field)241            load_complete_conv_btn.click(load_chatboat_complet_history, inputs=[], outputs=text_field)242            process_btn.click(analyse_current_conversation, inputs=[text_field, sub_category_dropdown], outputs=analysis_result)243 244demo.launch(share=True)245 246