app.py

import os
import json
import random
from threading import Thread


import gradio as gr
from gradio.themes import Base
from gradio.themes.utils import colors

from transformers import pipeline, TextIteratorStreamer, AutoModelForCausalLM, AutoTokenizer

SYSTEM_PROMPT = "You are a compliance assistant. Use the provided risk data to answer user questions. If a single risk object is given, provide a direct answer. If a list of risks is provided, summarize, compare, or analyze the collection as needed. Always base your response on the data provided."
hf_token = os.environ["HF_TOKEN"]

class HfModelWrapper:
    def __init__(
        self,
        model_path="casperhansen/llama-3.3-70b-instruct-awq",
        sys_prompt=SYSTEM_PROMPT,
        adapter_path="artemisiaai/fine-tuned-adapter",
    ):

        self.model = AutoModelForCausalLM.from_pretrained(
            model_path, device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.sys_prompt = sys_prompt
        self.adapter_path = adapter_path
        self.model.load_adapter(self.adapter_path, token=hf_token)
        self.model.enable_adapters()

    def build_prompt(self, user_msg, history):

        inppt = []
        inppt.append({"role": "system", "content": self.sys_prompt})
        inppt += history
        inppt.append({"role": "user", "content": user_msg})

        prompt = self.tokenizer.apply_chat_template(
            inppt,
            tokenize=False,
        )
        return prompt

    def generate(self, user_input, history):
        input_text = self.build_prompt(user_input, history)
        input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to("cuda")

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)

        gen_kwargs = {
            "inputs": input_ids,
            "streamer": streamer,
            "pad_token_id": self.tokenizer.eos_token_id,
            "max_length": 32768,
            "temperature": 0.1,
            "top_p": 0.8,
            "repetition_penalty": 1.1,
        }

        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
        thread.start()

        return streamer

# Custom theme colors based on brand standards
class ArtemisiaTheme(gr.themes.Base):
    def __init__(self):
        super().__init__()
        
        # Define color variables
        self.primary = "#4f008c"
        self.primary_focus = "#3a0066"
        self.secondary = "#6c757d"
        self.accent = "#7c4dff"
        self.neutral = "#2a2a2a"
        self.success = "#28a745"
        self.warning = "#ffc107"
        self.error = "#dc3545"
        self.info = "#17a2b8"
        
        # Text colors
        self.text = "#2d3748"
        self.text_light = "#4a5568"
        self.text_muted = "#718096"
        
        # Background colors
        self.background = "#ffffff"
        self.surface = "#f8f9fa"
        self.surface_highlight = "#e9ecef"
        
        # Border colors
        self.border = "#e2e8f0"
        self.border_hover = "#cbd5e0"
        
        # Set the theme colors using the new theming API
        self.set(
            # Base colors
            color_background_primary=self.background,
            color_background_secondary=self.surface,
            color_background_tertiary=self.surface_highlight,
            
            # Text colors
            body_text_color=self.text,
            body_text_color_light=self.text_light,
            body_text_color_muted=self.text_muted,
            
            # Primary colors
            primary_color=self.primary,
            primary_hover_color=self.primary_focus,
            
            # Secondary colors
            secondary_color=self.secondary,
            secondary_hover_color="#5a6268",
            
            # Status colors
            success_color=self.success,
            warning_color=self.warning,
            error_color=self.error,
            info_color=self.info,
            
            # Border colors
            border_color_primary=self.border,
            border_color_secondary=self.border_hover,
            
            # Button colors
            button_primary_background=self.primary,
            button_primary_hover_background=self.primary_focus,
            button_primary_text_color="#ffffff",
            
            button_secondary_background=self.secondary,
            button_secondary_hover_background="#5a6268",
            button_secondary_text_color="#ffffff",
            
            # Input fields
            input_background=self.background,
            input_border_color=self.border,
            input_border_hover_color=self.border_hover,
            input_text_color=self.text,
            input_placeholder_color=self.text_muted,
            
            # Chat components
            chatbot_background=self.background,
            chatbot_border_color=self.border,
            chatbot_text_color=self.text,
            chatbot_user_background="#f0f4ff",
            chatbot_bot_background=self.surface,
            
            # Sliders
            slider_color=self.primary,
            slider_color_hover=self.primary_focus,
            
            # Checkboxes and radio buttons
            checkbox_background=self.background,
            checkbox_border_color=self.border,
            checkbox_background_selected=self.primary,
            checkbox_border_selected=self.primary,
            
            # Tabs
            tab_background=self.surface,
            tab_background_selected=self.background,
            tab_text_color=self.text_light,
            tab_text_selected_color=self.primary,
            tab_border_color=self.border,
            
            # Dropdowns
            dropdown_background=self.background,
            dropdown_border_color=self.border,
            dropdown_text_color=self.text,
            dropdown_background_hover=self.surface_highlight,
            
            # Tooltips
            tooltip_background=self.neutral,
            tooltip_text_color="#ffffff",
            
            # Shadows
            shadow_spread="0px 2px 8px rgba(0, 0, 0, 0.1)",
            shadow_drop="0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06)",
            
            # Border radius
            radius_lg="0.5rem",
            radius_md="0.375rem",
            radius_sm="0.25rem",
            
            # Transitions
            transition_all="all 0.2s ease-in-out"
        )

# Custom CSS for exact layout and styling
custom_css = """
:root {
    /* Theme colors */
    --primary: #4f008c;
    --primary-hover: #3d006d;
    --text: #ffffff;
    --text-secondary: #e0e0e0;
    --background: #ffffff;
    --error: #ff5252;
    
    /* Spacing */
    --spacing-xs: 0.25rem;
    --spacing-sm: 0.5rem;
    --spacing-md: 1rem;
    --spacing-lg: 1.5rem;
    --spacing-xl: 2rem;
    
    /* Border radius */
    --radius-sm: 4px;
    --radius-md: 8px;
    --radius-lg: 12px;
}

/* Base styles */
body {
    margin: 0;
    padding: 0;
    width: 100%;
    background-color: var(--background);
    color: var(--text);
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif;
    font-weight: 400;
    line-height: 1.5;
    -webkit-font-smoothing: antialiased;
    -moz-osx-font-smoothing: grayscale;
}

/* Container */
.gradio-container {
    background-color: var(--background) !important;
    padding: 0 !important;
    max-width: 100% !important;
}

/* Header */
.header {
    background-color: var(--background);
    padding: var(--spacing-md) var(--spacing-lg);
    display: flex;
    flex-wrap: nowrap;
    align-items: flex-end;
    gap: var(--spacing-lg);
    border-bottom: 1px solid rgba(0, 0, 0, 0.1);
}

.logo-container {
    display: flex;
    gap: 0;
    flex-direction: row;
    width: fit-content !important;
    flex: 0 0 auto !important;
    min-width: 0 !important;
    align-items: flex-end;
}

.nav-divider {
    width: 1px;
    height: 1rem;
    background-color: var(--primary);
    margin: 0 var(--spacing-sm);
    align-self: flex-end;
}

.dashboard-title {
    font-size: 1rem;
    font-weight: 500;
    margin: 0;
    white-space: nowrap;
    flex: 1 1 0%;
    min-width: 0;
    overflow: hidden;
    text-overflow: ellipsis;
    letter-spacing: -0.02em;
    color: var(--primary) !important;
}

.header h1, .header h2, .header h3, .header h4, .header h5, .header h6 {
    color: var(--primary) !important;
    margin: 0;
}

.stc-logo {
    display: flex;
    align-items: flex-end;
    padding: 0;
    margin: 0;
}

.stc-logo svg {
    height: 1rem;
    width: auto;
    display: block;
    fill: var(--primary);
    overflow: hidden;
    scrollbar-width: none;
    overflow-x: hidden;
    fill: #4f008c !important;
}

.artemsia-logo {
    display: flex;
    align-items: flex-end;
    padding: 0;
    margin: 0;
}

.artemsia-logo svg {
    height: 1rem;
    width: auto;
    display: block;
    overflow: hidden;
    scrollbar-width: none;
    overflow-x: hidden;
    fill: #4f008c !important;
}

/* Make logos purple on white header */
.stc-logo svg path {
    fill: #4f008c !important;
}

.artemsia-logo svg path {
    fill: #4f008c !important;
}

.about-section {
    padding: 12px 20px;
    background-color: #000;
    border-bottom: 1px solid #333;
}

.about-section h2 {
    color: #ff5252;
    margin-bottom: 6px;
    font-size: 1rem;
    font-weight: 600;
}

.about-section p {
    font-size: 0.875rem;
    line-height: 1.6;
}

.main-content {
    display: flex;
    padding: 0;
    background-color: #ffffff;
    gap: 8px;
    justify-content: space-between;
}

.chat-container, .results-container {
    flex: 1;
    background-color: transparent;
    border-radius: 15px;
    padding: 0;
    display: flex;
    flex-direction: column;
    min-height: 550px;
    max-height: 550px;
    height: 550px;
    overflow: hidden;
    margin: 0;
    gap: 0 !important;
}

.section-header {
    gap: 0;
    background-color: #120428;
    border-radius: 15px 15px 0 0;
    padding: 12px;
    margin-bottom: 0;
}

.section-header h2 {
    margin: 0;
    padding: 0;
    font-size: 1rem;
    font-weight: 600;
    letter-spacing: -0.02em;
}

.section-header p {
    margin: 0;
    padding: 4px 0 0 0;
    color: #ccc;
    font-size: 0.875rem;
    font-weight: 400;
    line-height: 1.4;
}

.chat-content {
    flex: 1;
    background-color: #0a0013;
    border-radius: 0 0 15px 15px;
    padding: 0;
    margin: 0;
    display: flex;
    flex-direction: column;
    justify-content: space-between;
    overflow: hidden;
}

.chat-messages {
    flex: 1;
    background: #120428 !important;
    padding: 10px !important;
    margin: 0 !important;
    max-width: 100%;
    border: none;
    overflow-y: auto;
}

.results-content {
    flex: 1;
    background-color: #0a0013;
    border-radius: 0 0 15px 15px;
    padding: 16px;
    margin: 0;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    text-align: center;
}

.placeholder-text {
    color: #aaa;
    font-size: 0.875rem;
    font-weight: 400;
    max-width: 400px;
    line-height: 1.6;
}

/* Updated styles to match the image reference */
.gradio-container {
    max-width: 100% !important;
    padding: 0 !important;
}

.chat-container {
    padding: 0;
    margin: 0;
}

.results-container {
    padding: 0;
    margin: 0;
    min-width: min(50%, 100%) !important;
}

#component-0 > .gr-form > div:not(.about-section):not(.header) {
    background-color: #ffffff;
}

/* Chat messages container adjustments */
.chat-messages {
    padding: 0 !important;
    margin: 0 !important;
    border: none !important;
}

/* Remove any gaps between elements, but exclude section headers */
.chat-container > div:not(.section-header), 
.results-container > div:not(.section-header) {
    margin: 0 !important;
    padding: 0 !important;
}

/* Override flex-grow for section headers */
#component-27, 
.column.section-header,
div[id^="component"][class*="column section-header"] {
    flex-grow: 0 !important;
}

/* Add more forceful display for chatbot component */
.gradio-chatbot {
    height: auto !important;
    min-height: 400px !important;
    display: block !important;
    width: 100% !important;
    overflow-y: auto !important;
    visibility: visible !important;
    opacity: 1 !important;
    background-color: #0a0013 !important;
    background: #0a0013 !important;
}

/* Also target the container of the chatbot */
.message-wrap {
    display: block !important;
    visibility: visible !important;
    opacity: 1 !important;
    background-color: #0a0013 !important;
    background: #0a0013 !important;
}

/* Target the messages themselves */
.message {
    display: block !important;
    visibility: visible !important;
    opacity: 1 !important;
    background-color: #0a0013 !important;
}

/* Ensure chat bubbles have proper max-width */
.bubble-wrap {
    background: #0A0013 !important;
    background-color: #0A0013 !important;
    margin: 0 !important;
    padding: 0 !important;
    min-width: 100%;
    border: none !important;
}

/* Fix scrollable area background */
.gradio-chatbot .messages-wrapper {
    background-color: #0a0013 !important;
    background: #0a0013 !important;
}

/* Ensure the entire chat area has consistent background */
.gradio-chatbot .wrapper,
.gradio-chatbot .container,
.gradio-chatbot .messages {
    background-color: #0a0013 !important;
    background: #0a0013 !important;
}

/* Target the chat bubble containers */
.gradio-chatbot .messages-wrapper > div > div {
    max-width: 70% !important; /* Allow bubbles to be a bit wider */
    width: auto !important; /* Let content determine width, up to max-width */
    min-width: 30% !important; /* Ensure a minimum width */
}

/* Style the individual message bubbles */
.gradio-chatbot .messages-wrapper > div > div > div {
    width: 100% !important; /* Text fills the bubble */
    padding: 10px 15px !important; /* Add more padding for better readability */
    border-radius: 12px !important; /* Rounded corners */
}

/* Focus on the text content inside bubbles */
.gradio-chatbot .bubble {
    width: 100% !important; /* Fill the bubble width */
    max-width: 100% !important;
    word-wrap: break-word !important;
    overflow-wrap: break-word !important;
    display: block !important; /* Ensure it takes full width */
    white-space: normal !important; /* Allow proper wrapping */
}

/* Target the actual text paragraphs inside bubbles */
.gradio-chatbot .bubble p,
.gradio-chatbot .bubble span,
.gradio-chatbot .bubble div {
    width: 100% !important;
    display: inline-block !important;
    word-break: break-word !important;
    white-space: normal !important;
    text-align: left !important;
}

/* Target the first lines of text specifically to ensure they fill the width */
.gradio-chatbot .bubble p:first-child,
.gradio-chatbot .bubble span:first-child {
    width: 100% !important;
    display: block !important;
}

/* Specific gradio selector */
.chat-messages [data-testid="chatbot"] .message-wrap > div {
    max-width: 50% !important;
}

.chat-input-container {
    width: 100%;
    position: relative;
    border-top: 1px solid rgba(255, 255, 255, 0.1);
    padding: 16px;
    background-color: #000000;
    display: flex;
    gap: 0 !important;
}

.chat-input {
    flex-grow: 1;
    padding: 0px;
    background-color: #000000;
    border: none;
    border-radius: 8px 0 0 8px;
    color: #fff;
    font-size: 0.9rem;
    font-weight: 400;
    line-height: 1.5;
    width: 100%;
    margin: 0 !important;
}

.send-button {
    width: 50px !important;
    min-width: 40px !important;
    max-width: 50px !important;
    height: 100%;
    background-color: #993333;
    border: none !important;
    border-radius: 0 8px 8px 0;
    color: white;
    display: flex;
    align-items: center;
    justify-content: center;
    cursor: pointer;
    transition: background-color 0.2s ease;
    margin: 0 !important;
}

.send-button:hover {
    background-color: #b33c3c;
}

.container.show_textbox_border {
    border: none !important;
    background: transparent !important;
    box-shadow: none !important;
    margin: 0 !important;
}

.input-container {
    background: transparent !important;
    border: none !important;
    margin: 0 !important;
}

.input-container textarea {
    background: transparent !important;
    color: #fff !important;
    border: none !important;
    box-shadow: none !important;
    border-radius: 0 !important;
    min-width: min(80%, 100%) !important;
}

span[data-testid="block-info"] {
    display: none !important;
}

.chat-input::placeholder {
    color: #aaa;
    background-color: #000000;
    padding: 0px;
}

/* Paper plane icon styling for the send button */
.send-icon {
    width: 20px;
    height: 20px;
}

.send-button svg {
    width: 20px;
    height: 20px;
    display: block;
    margin: auto;
}

.bot, .user {
    max-width: 70% !important;
    width: fit-content !important;
}

/* Enhanced typography for chat messages */
.gradio-chatbot .message {
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif !important;
    font-size: 0.9rem !important;
    font-weight: 400 !important;
    line-height: 1.5 !important;
    -webkit-font-smoothing: antialiased !important;
    color: #ffffff !important;  /* Force white text for messages */
}

/* Improved text styling for all components */
.gradio-container * {
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif !important;
    color: #000000 !important;  /* Default text color */
}

/* Ensure text in chat bubbles is visible */
.gradio-chatbot .bubble,
.gradio-chatbot .bubble p,
.gradio-chatbot .bubble span,
.gradio-chatbot .bubble div,
.gradio-chatbot .message-content,
.gradio-chatbot .markdown {
    color: #ffffff !important;  /* Force white text in chat bubbles */
}

/* Enhanced headings */
h1, h2, h3, h4, h5, h6 {
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif !important;
    font-weight: 600 !important;
    letter-spacing: -0.02em !important;
    line-height: 1.3 !important;
    color: #4f008c !important;  /* Purple color for headings */
}

/* Ensure input text is visible */
.chat-input,
.chat-input::placeholder,
.chat-input textarea,
.chat-input textarea::placeholder {
    color: #ffffff !important;
    background-color: #000000 !important;
}

/* Fix for any other text that might be inheriting wrong colors */
body, html {
    color: #000000 !important;
}

/* Ensure user messages have proper text color */
.user .bubble,
.user .bubble * {
    background-color: #4f008c !important;
    color: #ffffff !important;
}

/* Ensure bot messages have proper text color */
.bot .bubble,
.bot .bubble * {
    background-color: #120428 !important;
    color: #ffffff !important;
}

/* Button text improvements */
button {
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif !important;
    font-weight: 500 !important;
    letter-spacing: -0.01em !important;
}

/* Input field improvements */
input, textarea {
    font-family: 'Helvetica Neue', Helvetica, Arial, system-ui, -apple-system, BlinkMacSystemFont, sans-serif !important;
    font-weight: 400 !important;
    line-height: 1.5 !important;
}
"""

def create_ui():
    # Initialize the model
    model_wrapper = HfModelWrapper()
    
    # Create the Gradio interface with our custom theme and CSS
    with gr.Blocks(
        theme=ArtemisiaTheme(),
        css=custom_css,
        title="Compliance Assistant",
    ) as demo:
        # Header section
        with gr.Row(elem_classes=["header"]):
            with gr.Column(scale=1, elem_classes=["logo-container"]):
                gr.HTML("""
                <div class="stc-logo">
                    <svg xmlns="http://www.w3.org/2000/svg" width="66" height="32" viewBox="0 0 66 32" fill="#4f008c">
                        <path fill-rule="evenodd" clip-rule="evenodd" d="M31.625 6.80851L31.6213 0.00567279L24.7511 0L24.75 6.79666L31.625 6.80851ZM5.88232 23.0513L0 24.2653C0.424529 29.3706 6.18066 32.1214 11.3441 31.9959C14.2329 31.9257 16.8607 30.9901 18.5472 29.6392C22.4356 26.5245 22.4103 20.0758 17.2744 17.3711C15.3452 16.3551 13.2453 15.9901 11.4475 15.6777C9.10574 15.2707 7.27666 14.9528 7.00534 13.4013C6.4686 10.332 12.5096 9.46738 14.5772 12.2286C14.8637 12.6112 14.954 12.903 15.0505 13.2152L15.0505 13.2152C15.114 13.4203 15.1801 13.6342 15.3064 13.8884L20.889 12.6335C20.9094 12.6247 20.9455 12.6077 20.9725 12.5925C20.6349 7.68784 14.5151 5.00625 9.32399 5.5062C-0.27525 6.43076 -2.55039 16.9042 5.58028 20.2095C6.79516 20.7033 8.15129 20.9902 9.48385 21.2722C10.7248 21.5347 11.9453 21.793 13.0123 22.21C15.9184 23.3457 15.0929 26.4741 11.6771 26.7789C8.56398 27.0567 6.23313 25.7203 5.88232 23.0513ZM66 23.5042L59.9776 21.5819C59.8074 21.8761 59.6967 22.1399 59.5911 22.3913C59.4397 22.7522 59.2989 23.0876 59.0079 23.451C58.5671 24.0016 58.1301 24.4609 57.4746 24.8733C56.2011 25.6745 54.3913 26.0052 52.5943 25.6296C49.5135 24.9858 47.5738 22.504 47.4893 18.9828C47.4011 15.3106 49.2053 12.6962 52.1919 11.9434C56.004 10.9825 58.8882 12.9476 59.6362 15.8925L65.8508 13.9081C65.2461 10.9173 62.8376 8.44026 60.8112 7.24677C52.9767 2.63234 40.5366 7.03787 40.5625 18.7666C40.5881 30.3289 53.233 34.8158 60.8587 30.2249C62.1039 29.4752 63.2247 28.4579 64.0145 27.4692C64.723 26.5823 65.7798 24.9372 66 23.5042ZM24.8097 12.9397L31.6484 12.9362C31.6463 13.2625 31.6421 13.7315 31.6371 14.2928L31.637 14.3022L31.637 14.3033L31.637 14.3036L31.637 14.304C31.6122 17.0787 31.5674 22.0894 31.6626 23.2768C31.7957 24.9371 32.571 25.7899 34.2386 26.0043C35.9984 26.2305 36.9321 25.8072 38.115 25.2708L38.1151 25.2708C38.2402 25.2141 38.3681 25.1561 38.5 25.0975L38.4566 30.5261C35.605 32.6136 28.5618 32.8937 25.8963 28.6417C24.7526 26.8176 24.7717 24.551 24.7916 22.1906L24.7916 22.19C24.7942 21.8817 24.7968 21.5719 24.7968 21.2612C24.7969 20.3989 24.7853 19.5126 24.7737 18.6184V18.6183V18.6183V18.6182V18.6182C24.7489 16.7095 24.7236 14.7647 24.8097 12.9397ZM38.5 12.9362L31.625 12.935L31.6257 6.12766L38.4997 6.13078L38.5 12.9362Z" fill="currentColor"></path>
                    </svg>
                </div>
                """)
                
            with gr.Column(scale=1, elem_classes=["dashboard-title"]):
                gr.Markdown("# Compliance Assistant", elem_classes=["dashboard-title"])
        
        # Main chat interface
        chatbot = gr.Chatbot(
            label="Chat",
            elem_id="chatbot",
            elem_classes=["chatbot"],
            height=600,
            show_label=False,
        )
        
        with gr.Row(elem_classes=["input-container"]):
            msg = gr.Textbox(
                show_label=False,
                placeholder="Ask me about compliance...",
                container=False,
                scale=10,
                min_width=0,
                elem_classes=["input-box"],
            )
            submit_btn = gr.Button("Send", variant="primary", scale=1, min_width=0)
            clear_btn = gr.Button("Clear", variant="secondary", scale=1, min_width=0)
        
        # Handle form submission
        def respond(message, chat_history):
            if not message.strip():
                return "", chat_history
                
            # Add user message to chat history
            chat_history.append((message, None))
            
            # Get model response
            response = ""
            for chunk in model_wrapper.generate(message, chat_history):
                if chunk is not None:
                    response += chunk
                    chat_history[-1] = (message, response)
                    yield "", chat_history
        
        # Connect the UI elements
        msg.submit(
            respond, [msg, chatbot], [msg, chatbot], queue=False
        )
        submit_btn.click(
            respond, [msg, chatbot], [msg, chatbot], queue=False
        )
        clear_btn.click(lambda: ([], []), None, [chatbot, msg], queue=False)
    
    return demo

if __name__ == "__main__":
    # Create and launch the UI
    demo = create_ui()
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True,
        show_error=True,
    )
# SVG assets
artemsia_logo_svg = """
<?xml version="1.0" encoding="UTF-8"?>
<svg id="Layer_2" data-name="Layer 2" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 351.26 430.29">
  <g id="Layer_1-2" data-name="Layer 1">
    <path d="M24.35,126.29c-5.54,0-10.14-1.51-13.79-4.52-3.64-3.03-6.31-7.15-8.02-12.38-1.7-5.22-2.54-11.1-2.54-17.65,0-5.92.78-12.11,2.35-18.58,1.57-6.48,3.92-12.78,7.06-18.88,3.15-6.11,6.99-11.65,11.52-16.63,4.54-4.97,9.76-8.94,15.67-11.9,5.92-2.96,12.46-4.44,19.65-4.44,5.15,0,9.23.73,12.25,2.19,3.03,1.45,5.73,3.3,8.13,5.56l10.02-5.67h6.02l-12.83,60.4c-.25,1.39-.63,3.41-1.13,6.04-.5,2.64-.91,5.35-1.23,8.13-.32,2.77-.48,5.03-.48,6.79,0,10.07,2.58,15.1,7.75,15.1,2.01,0,4.21-.57,6.6-1.71,2.39-1.14,4.96-2.71,7.73-4.71l1.15,1.5c-1.77,1.64-4.16,3.34-7.19,5.1-3.02,1.75-6.22,3.22-9.63,4.42-3.39,1.21-6.66,1.81-9.81,1.81-4.53,0-8.11-1.39-10.75-4.15-2.64-2.78-3.96-6.61-3.96-11.52,0-2.53.13-4.82.38-6.9.25-2.08.63-4.69,1.13-7.83h-.19c-4.41,9.7-9.6,17.19-15.58,22.48-5.97,5.28-12.73,7.92-20.27,7.92ZM32.46,121.19c3.4,0,6.67-1.26,9.81-3.77,3.15-2.51,6.07-5.69,8.77-9.52,2.71-3.84,4.97-7.81,6.79-11.9,1.83-4.09,3.07-7.65,3.71-10.67l11.31-55.13c-1.39-2.26-3.41-3.96-6.04-5.08-2.64-1.14-5.41-1.71-8.29-1.71-5.17,0-9.8,1.58-13.9,4.73-4.08,3.14-7.73,7.26-10.94,12.35-3.21,5.1-5.89,10.73-8.02,16.9-2.14,6.17-3.74,12.3-4.81,18.4-1.07,6.1-1.6,11.67-1.6,16.71,0,10.32,1.26,17.68,3.77,22.08,2.51,4.41,5.66,6.6,9.44,6.6Z"/>
    <path d="M113.81,124.4l17.19-81.92c.38-1.76.65-3.46.83-5.08.19-1.64.29-2.78.29-3.42,0-3.78-1.6-6.2-4.81-7.27-3.21-1.07-6.64-1.6-10.29-1.6h-1.5l.38-1.88h34.54l-4.33,27.17h.38c2.39-5.15,5.03-9.93,7.92-14.33,2.9-4.4,6.3-7.96,10.19-10.67,3.9-2.71,8.81-4.06,14.73-4.06,5.29,0,9.19,1.23,11.71,3.69,2.51,2.46,3.77,5.57,3.77,9.33,0,4.41-1.29,7.9-3.88,10.48-2.57,2.58-6.63,3.88-12.17,3.88.13-1.01.19-1.98.19-2.92v-2.75c0-4.65-.66-8.36-1.98-11.13-1.32-2.78-3.81-4.17-7.46-4.17-2.64,0-5.29,1.17-7.94,3.5-2.64,2.33-5.13,5.27-7.46,8.79-2.32,3.52-4.27,7.13-5.83,10.83-1.57,3.71-2.68,6.95-3.31,9.73l-13.58,63.79h-17.56Z"/>
    <path d="M228.01,126.29c-13.72,0-20.58-7.36-20.58-22.08,0-2.26.19-4.88.56-7.83.39-2.96.83-5.63,1.33-8.02l13.4-63.25h-14.33l.38-1.88c6.92,0,12.51-.76,16.79-2.27,4.28-1.51,7.8-3.96,10.56-7.35,2.78-3.4,5.36-7.94,7.75-13.6h1.9l-4.92,23.23h22.83l-.38,1.88h-22.83l-13.4,62.67c-.64,3.16-1.18,6.15-1.63,8.98-.43,2.83-.65,5.26-.65,7.27,0,5.66,1.13,9.71,3.4,12.17,2.26,2.46,5.47,3.69,9.63,3.69,2.39,0,4.94-.41,7.65-1.23,2.71-.82,5.07-1.86,7.08-3.13l.94,1.71c-3.78,2.64-7.81,4.81-12.08,6.52-4.28,1.69-8.74,2.54-13.4,2.54Z"/>
    <path d="M306.72,126.29c-6.67,0-12.81-1.45-18.42-4.33-5.6-2.9-10.06-7.21-13.4-12.94-3.33-5.73-5-13.01-5-21.81,0-6.92,1.19-14.21,3.58-21.88,2.4-7.68,5.99-14.85,10.77-21.52,4.78-6.68,10.66-12.09,17.65-16.25,6.98-4.15,15.07-6.23,24.25-6.23,7.93,0,14.09,1.93,18.5,5.77,4.4,3.83,6.6,9.15,6.6,15.94,0,6.04-2.36,11.68-7.08,16.9-4.71,5.22-11.16,9.47-19.33,12.75-8.18,3.27-17.43,4.9-27.75,4.9h-6.6c-.63,3.78-.94,7.49-.94,11.15,0,11.32,1.6,20.1,4.81,26.33,3.21,6.22,7.9,9.33,14.06,9.33,7.17,0,12.89-2.2,17.17-6.6,4.28-4.4,7.55-10.06,9.81-16.98,2.64.75,3.96,2.64,3.96,5.65,0,2.39-1.17,5.16-3.5,8.31-2.32,3.14-5.87,5.85-10.65,8.13-4.78,2.26-10.95,3.4-18.5,3.4ZM290.86,75.71h5.1c11.07,0,19.81-3.08,26.23-9.25s9.63-14.41,9.63-24.73c0-12.21-3.4-18.31-10.19-18.31-3.78,0-7.31,1.54-10.58,4.63-3.27,3.08-6.22,7.21-8.88,12.38-2.64,5.16-4.91,10.81-6.79,16.98-1.89,6.17-3.4,12.27-4.52,18.31Z"/>
    <path d="M159.7,278.29c-5.54,0-9.79-1.77-12.75-5.29-2.96-3.53-4.44-8.05-4.44-13.58s1.19-14.28,3.58-25.48l3.21-14.73c.5-2.51,1.28-6.13,2.35-10.85,1.07-4.72,1.6-9.09,1.6-13.13,0-2.01-.16-4.21-.48-6.6-.31-2.39-1.28-4.46-2.92-6.23-1.63-1.76-4.33-2.65-8.1-2.65-3.28,0-6.4,1.01-9.35,3.02-2.96,2.02-5.67,4.44-8.13,7.27-2.45,2.83-4.52,5.6-6.23,8.31-1.7,2.7-2.92,4.74-3.67,6.13l-7.35,33.04c-.38,1.75-.89,4.45-1.52,8.1-.63,3.66-1.22,7.59-1.79,11.81-.57,4.21-1.08,8.08-1.52,11.6-.43,3.52-.65,5.97-.65,7.35h-17.94c0-1.51.22-4,.67-7.46.44-3.46.98-7.35,1.6-11.69.63-4.34,1.31-8.59,2.06-12.75.76-4.15,1.46-7.67,2.08-10.56l3.21-14.73c.64-2.51,1.46-6.13,2.46-10.85,1.01-4.72,1.52-9.09,1.52-13.13,0-2.51-.32-4.96-.96-7.35-.63-2.39-1.79-4.34-3.5-5.85-1.7-1.51-4.24-2.27-7.63-2.27s-6.68,1.17-9.83,3.5c-3.14,2.33-6.03,5.14-8.67,8.42-2.64,3.27-4.78,6.41-6.42,9.44-1.64,3.02-2.58,5.21-2.83,6.58l-14.73,68.71h-17.54l17.17-81.92c.38-1.76.66-3.46.85-5.08.19-1.64.29-2.78.29-3.42,0-3.78-1.29-6.2-3.88-7.27-2.58-1.07-5.7-1.6-9.35-1.6h-1.5l.38-1.88h34.92l-5.1,23.77h.4c5.28-8.3,10.71-14.63,16.31-18.96,5.59-4.34,12.11-6.52,19.54-6.52s12.65,2.11,15.67,6.33c3.01,4.21,4.52,9.08,4.52,14.63,0,.75-.04,1.51-.1,2.27-.06.75-.15,1.5-.27,2.25h.38c5.29-8.3,10.76-14.63,16.42-18.96,5.67-4.34,12.15-6.52,19.46-6.52s12.63,2.11,15.65,6.33c3.03,4.21,4.54,9.08,4.54,14.63,0,3.14-.41,6.57-1.23,10.29-.82,3.71-1.54,7.2-2.17,10.48l-4.9,22.27c-.77,3.52-1.37,6.97-1.81,10.38-.43,3.41-.65,6.55-.65,9.44,0,4.41.72,7.8,2.17,10.19,1.44,2.39,3.68,3.58,6.71,3.58,3.9,0,7.99-1.69,12.27-5.08l.94,1.5c-3.03,2.52-6.4,4.94-10.1,7.27s-8.02,3.5-12.92,3.5Z"/>
    <path d="M227.45,156.35c-2.52,0-4.59-.72-6.23-2.17-1.64-1.46-2.46-3.88-2.46-7.27,0-4.15.88-7.23,2.65-9.25,1.76-2.01,4.03-3.02,6.79-3.02,2.53,0,4.6.66,6.23,1.98,1.64,1.32,2.46,3.81,2.46,7.46,0,4.78-1.01,8.02-3.02,9.73-2.02,1.7-4.16,2.54-6.42,2.54ZM214.42,278.29c-5.03,0-8.9-1.83-11.6-5.48-2.71-3.65-4.06-8.31-4.06-13.98,0-2.64.25-5.69.75-9.15.51-3.46,1.08-6.85,1.71-10.19s1.19-6.26,1.69-8.79l8.69-40.75c.38-1.89.56-3.53.56-4.92,0-3.4-1.04-5.57-3.13-6.52-2.07-.94-4.86-1.42-8.38-1.42h-2.46l.38-1.88h33.79l-12.83,60.58c-.25,1.39-.64,3.41-1.15,6.04-.5,2.64-.91,5.35-1.23,8.13-.31,2.77-.46,5.03-.46,6.79,0,9.18,2.7,13.77,8.1,13.77,2.78,0,5.23-.47,7.38-1.42,2.14-.94,4.34-2.17,6.6-3.67l1.13,1.5c-2.77,2.52-6.54,5.04-11.33,7.56-4.78,2.51-9.49,3.77-14.15,3.77Z"/>
    <path d="M278.78,278.29c-9.43,0-16.98-1.89-22.65-5.67-5.67-3.78-8.5-9-8.5-15.67,0-3.4.78-5.98,2.35-7.73,1.58-1.76,3.38-2.99,5.4-3.69,2.01-.69,3.78-1.04,5.29-1.04-.13,1-.22,2.04-.29,3.1-.06,1.07-.08,2.05-.08,2.94,0,7.17,1.51,13.27,4.52,18.31,3.01,5.03,7.92,7.54,14.73,7.54s11.98-1.95,15.94-5.85c3.97-3.9,5.96-9.25,5.96-16.04,0-5.79-1.45-10.38-4.35-13.77-2.89-3.4-7.79-6.67-14.71-9.81-6.93-3.28-12.47-7.08-16.63-11.42-4.16-4.34-6.23-10.3-6.23-17.85,0-8.05,3.27-14.76,9.81-20.1,6.55-5.34,15.43-8.02,26.63-8.02,6.92,0,12.51.89,16.79,2.65,4.28,1.77,7.45,4.04,9.52,6.81,2.08,2.77,3.13,5.59,3.13,8.48,0,4.03-1.35,6.99-4.06,8.88-2.7,1.89-6.06,2.83-10.08,2.83.13-1.01.21-1.96.27-2.83.07-.89.1-1.83.1-2.83,0-6.17-1.29-11.39-3.88-15.67-2.58-4.28-6.96-6.42-13.13-6.42-7.04,0-11.92,2.05-14.63,6.15-2.71,4.08-4.06,8.64-4.06,13.67,0,5.79,1.54,10.27,4.63,13.42,3.08,3.14,7.96,6.22,14.65,9.23,6.92,3.16,12.42,6.99,16.5,11.52,4.09,4.53,6.15,10.45,6.15,17.75,0,9.18-3.49,16.67-10.48,22.46-6.99,5.79-16.52,8.69-28.6,8.69Z"/>
    <path d="M120.34,309.1c-2.41,0-4.39-.78-5.96-2.35-1.57-1.57-2.35-4.3-2.35-8.21s.78-6.8,2.35-8.31c1.57-1.51,3.55-2.27,5.96-2.27,2.26,0,4.21.76,5.83,2.27,1.64,1.52,2.46,4.29,2.46,8.31s-.82,6.64-2.46,8.21c-1.63,1.57-3.57,2.35-5.83,2.35ZM94.46,428.4v-1.9h2.46c5.54,0,9.51-1.13,11.9-3.4,2.39-2.26,3.58-5.97,3.58-11.13v-68.15c0-5.28-1.2-9.05-3.58-11.31-2.39-2.28-6.35-3.42-11.9-3.42h-.56v-1.88h33.79v83.79c0,5.54,1.19,9.51,3.58,11.9,2.39,2.39,6.35,3.58,11.9,3.58h2.44v1.9h-53.6Z"/>
    <path d="M191.11,430.29c-9.7,0-17.31-2.48-22.85-7.46-5.53-4.97-8.29-11.8-8.29-20.48,0-9.69,3.55-16.93,10.67-21.71,7.11-4.78,17.9-7.42,32.38-7.94l15.67-.56v-14.15c0-7.05-.39-12.84-1.15-17.38-.75-4.53-2.39-7.89-4.9-10.08-2.52-2.21-6.35-3.31-11.52-3.31-7.29,0-11.98,2.27-14.06,6.79-2.07,4.53-3.1,11.64-3.1,21.33-5.29,0-9.29-.88-12-2.65-2.7-1.76-4.04-4.66-4.04-8.69,0-4.53,1.53-8.15,4.6-10.85,3.08-2.71,7.17-4.69,12.27-5.94,5.09-1.26,10.79-1.9,17.08-1.9,11.58,0,20.23,2.46,25.96,7.38,5.72,4.91,8.58,13.33,8.58,25.29v56.44c0,4.53.88,7.68,2.65,9.44,1.76,1.75,4.78,2.63,9.06,2.63h.56v1.9h-28.5v-19.81h-.38c-2.02,7.54-5.57,13.05-10.67,16.52-5.1,3.46-11.1,5.19-18.02,5.19ZM194.88,427.65c4.65,0,8.77-1.19,12.35-3.58,3.59-2.4,6.4-5.61,8.42-9.63,2.01-4.03,3.02-8.38,3.02-13.04v-27.17l-12.08.56c-10.7.5-18.09,2.95-22.19,7.35-4.08,4.41-6.13,11.83-6.13,22.27,0,15.49,5.53,23.23,16.6,23.23Z"/>
  </g>
</svg>
"""


# Paper plane SVG icon for send button
paper_plane_svg = """<svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="white" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round">
  <path d="M22 2L11 13"/>
  <path d="M22 2L15 22L11 13L2 9L22 2Z"/>
</svg>"""

wrapper = HfModelWrapper()
with open("./question_mapping.json", "r") as f:
    database = json.load(f)
database = {k.lower(): v for k, v in database.items()}


# Mock data function for chatbot
"""def send_message(message, history):
    #if not message:
    #    return history
    #history.append({"role": "assistant", "content": f"This is a response about: {message}"})
    #return history
    
    history.append({"role": "user", "content": message})
    # Recover extra data
    if message.lower() in database:
        context = random.choice(database[message.lower()])
        message += " \n"+context
    response_generator = wrapper.generate(message, history)
    partial = ""
    idx = 0
    for t in response_generator:
        if idx <= 3:
            idx += 1
            continue
        else:
            partial += t
            yield history + [{"role": "assistant", "content": partial}]"""    
    
def send_message(message, history):
    history.append({"role": "user", "content": message})
    context = ""
    # Recupera contesto extra se esiste
    if message.lower() in database:
        context = random.choice(database[message.lower()])
        message += " \n" + context

    response_generator = wrapper.generate(message, history)
    partial = ""
    idx = 0
    cleared = False

    if len(context) > 1000:
        context = context[:1000] + "[...]"
    
    for t in response_generator:
        if idx <= 3:
            idx += 1
            continue
        else:
            partial += t
            # Quando arriva il primo token utile, svuota la textbox
            if not cleared:
                cleared = True
                yield gr.update(value=""), history + [{"role": "assistant", "content": partial}], gr.update(value=context)
            else:
                yield "", history + [{"role": "assistant", "content": partial}], gr.update(value=context)


# Create the dashboard
with gr.Blocks(theme=ArtemisiaTheme(), css=custom_css) as demo:
        # Header
    with gr.Row(elem_classes="header"):
        with gr.Row(elem_classes="logo-container"):
            gr.HTML(stc_logo_svg, elem_classes="stc-logo")
          #  gr.HTML('<div class="nav-divider"></div>', elem_classes="divider")
           # gr.HTML(artemsia_logo_svg, elem_classes="stc-logo")
        gr.Markdown("# Risk Insights Dashboard", elem_classes="dashboard-title")



    # Main content area
    with gr.Row(elem_classes="main-content"):
        # Chat column
        with gr.Column(scale=1, elem_classes="chat-container"):
            with gr.Column(elem_classes="section-header"):
                gr.Markdown("## Chat Interface")
                gr.Markdown("Query your Archer risk data here or ask questions for valuable insights.")
            with gr.Column(elem_classes="chat-content"):
                chatbot = gr.Chatbot(value=[
                    {"role": "assistant", "content": "Welcome to Risk Insights. How can I help you analyze your risk data today?"}
                ], 
                type='messages',
                elem_classes="chat-messages",
                layout="bubble")                
                with gr.Row(elem_classes="chat-input-container"):
                    chat_input = gr.Textbox(placeholder="Ask about stc Group's risk data", elem_classes="chat-input", label="", scale=1)
                    send_button = gr.Button(value="→", elem_classes="send-button", scale=0)

        # Results column
        with gr.Column(scale=1, elem_classes="results-container"):
            with gr.Column(elem_classes="section-header"):
                gr.Markdown("## Results & Insights")
                gr.Markdown("Visualizations and data from your queries")
            with gr.Column(elem_classes="results-content"):
                results_placeholder = gr.Markdown("", elem_classes="placeholder-text", render=True)

    # Wire up the chat functionality
    #send_button.click(
    #    fn=send_message,
    #    inputs=[chat_input, chatbot],
    #    outputs=[chatbot]
    #).then(lambda: "", None, chat_input)
    
    #chat_input.submit(
    #    fn=send_message,
    #    inputs=[chat_input, chatbot],
    #    outputs=[chatbot]
    #).then(lambda: "", None, chat_input)
    """
    send_button.click(
        lambda msg: (gr.update(value=""), msg),  # clears textbox immediately
        inputs=[chat_input],
        outputs=[chat_input, gr.State()],  # dummy state to capture message
    ).then(
        fn=send_message,
        inputs=[chat_input, chatbot],
        outputs=[chatbot]
    )
    
    chat_input.submit(
        lambda msg: (gr.update(value=""), msg),
        inputs=[chat_input],
        outputs=[chat_input, gr.State()],
    ).then(
        fn=send_message,
        inputs=[chat_input, chatbot],
        outputs=[chatbot]
    )
    """
    send_button.click(
        fn=send_message,
        inputs=[chat_input, chatbot],
        outputs=[chat_input, chatbot, results_placeholder]  
    )
    
    chat_input.submit(
        fn=send_message,
        inputs=[chat_input, chatbot],
        outputs=[chat_input, chatbot, results_placeholder]
    )
    # JavaScript for UI enhancements
    gr.HTML("""
    <script>
    // Replace the send button text with SVG icon
    document.addEventListener('DOMContentLoaded', function() {
        function updateSendButton() {
            const sendButtons = document.querySelectorAll('.send-button');
            if (sendButtons && sendButtons.length > 0) {
                sendButtons.forEach(button => {
                    button.innerHTML = `<svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="white" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"><path d="M22 2L11 13"></path><path d="M22 2L15 22L11 13L2 9L22 2Z"></path></svg>`;
                });
                return true;
            }
            return false;
        }
        
        // Try immediately and then retry a few times to handle Gradio's dynamic loading
        let attempts = 0;
        const maxAttempts = 10;
        const checkInterval = setInterval(function() {
            if (updateSendButton() || attempts >= maxAttempts) {
                clearInterval(checkInterval);
            }
            attempts++;
        }, 500);
    });
    </script>
    """)

# Launch the app
#demo.launch(share=True)
if __name__ == "__main__":
    demo.launch()