app.py

from huggingface_hub import InferenceClient
import gradio as gr
import json
from datetime import datetime
import os
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from google.oauth2 import service_account
from googleapiclient.http import MediaFileUpload
from pathlib import Path
import argparse
from huggingface_hub import snapshot_download
import random 
import string
import re
from llama_cpp import Llama

# Variables for model, Google Drive, and initial context
repo_name = 'TheBloke/OpenHermes-2.5-Mistral-7B-GGUF'
model_file = "openhermes-2.5-mistral-7b.Q4_K_M.gguf"
SCOPES = ['https://www.googleapis.com/auth/drive']
SERVICE_ACCOUNT_FILE = '/home/user/app/service_account_creds.json'
folder_id = '1Hjbu7FmuVs0Yz8y_veo6SzY_2tc48OWt'

initContext = """<|im_start|>You are playing the role of an aggressive patient called Barry Lawson. Barry has been waiting for hours at the hospital with severe pain in his chest and hasn't been seen yet. You are incredibly frustrated and angry that things are taking so long, you will only calm down if your concerns are addressed. Waiting will only make you more angry and aggressive towards the staff. At any moment you could break into a RAGE! Make sure you don't censor swear words. Make sure to explain your emotional state and thinking before responding, for example, Barry: (tired of waiting and very angry) What the fuck do I have to do around here to get some treatment!"""
unique_id = ""
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# Function to load the LLaMA model
def load_model():
    llm = Llama(model_path=model_file, model_type="mistral",n_gpu_layers=-1,n_ctx = 2048)
    return llm

# Function to generate a unique identifier for each chat session
def generate_unique_id():
    # Generate a random sequence of 3 letters and 3 digits
    letters = ''.join(random.choices(string.ascii_letters, k=3))
    digits = ''.join(random.choices(string.digits, k=3))
    unique_id = letters + digits
    return unique_id

# Download the model from Hugging Face Hub
print('Fetching model:', repo_name, model_file)
snapshot_download(repo_id=repo_name, local_dir=".", allow_patterns=model_file)
print('Done fetching model:')


class ChatbotAPP:
    def __init__(self,model,service_account_file,scopes,folder_id,unique_id,initContext):
        self.llm = model  # LLaMA model instance
        self.service_account_file = service_account_file  # Path to Google service account credentials
        self.scopes = scopes  # Google Drive API scopes
        self.folder_id = folder_id  # Google Drive folder ID to store chat logs
        self.unique_id = unique_id  # Unique identifier for the chat session
        self.chat_history = []  # List to store chat history for the current session
        self.chat_log_history = []  # List to store chat logs for uploading
        self.isFirstRun = True  # Flag to check if it's the first run of the chat session
        self.initContext = initContext  # Initial context for the chat session
        self.context = ""  # Current context for the chat session
        self.agreed = False  # Flag to check if the user agreed to terms and conditions
        self.service = self.get_drive_service()  # Google Drive service instance
        self.app = self.create_app()  # Gradio app instance
        self.chat_log_name = ""  # Filename for the chat log


    # Method to create Google Drive service instance
    def get_drive_service(self):
        credentials = service_account.Credentials.from_service_account_file(
            self.service_account_file, scopes=self.scopes)
        self.service = build('drive', 'v3', credentials=credentials)
        print("Google Service Created")
        return self.service

    def generate_unique_id(self): #in an instance the user resets using the reset button
    # Generate a random sequence of 3 letters and 3 digits
        letters = ''.join(random.choices(string.ascii_letters, k=3))
        digits = ''.join(random.choices(string.digits, k=3))
        unique_id = letters + digits
        return unique_id

    # Method to search for a chat log file in Google Drive
    def search_file(self):
        #Search for a file by name in the specified Google Drive folder.
        query = f"name = '{self.chat_log_name}' and '{self.folder_id}' in parents and trashed = false"
        response = self.service.files().list(q=query, spaces='drive', fields='files(id, name)').execute()
        files = response.get('files', [])
        if not files:
            print(f"Chat log {self.chat_log_name} does not exist")
        else:
            print(f"Chat log {self.chat_log_name} exist")
        return files
    
    def strip_text(self,text):    # Method to strip unwanted text from chat messages

        # Pattern to match text inside parentheses or angle brackets and any text following angle brackets
        pattern = r"\(.*?\)|<.*?>.*"
        
        # Use re.sub() to replace the matched text with an empty string
        cleaned_text = re.sub(pattern, "", text)
        
        return cleaned_text
    
    def upload_to_google_drive(self):    # Method to upload the current chat log to Google Drive

        existing_files = self.search_file()
        print(existing_files)
        
        data = {
                #"name": Name,
                #"occupation": Occupation,
                #"years of experience": YearsOfExp,
                #"ethnicity": Ethnicity,
                #"gender": Gender,
                #"age": Age,
                "Unique ID": self.unique_id,
                "chat_history": self.chat_log_history
                }
        
        with open(self.chat_log_name, "w") as log_file:
                json.dump(data, log_file, indent=4)
    
        if not existing_files:        # Upload or update the chat log file on Google Drive

            # If the file does not exist, upload it
            file_metadata = {
                'name': self.chat_log_name,
                'parents': [self.folder_id],'mimeType': 'application/json'
            }
            media = MediaFileUpload(self.chat_log_name, mimetype='application/json')
            file = self.service.files().create(body=file_metadata, media_body=media, fields='id').execute()
            print(f"Uploaded new file with ID: {file.get('id')}")
        else:
            print(f"File '{self.chat_log_name}' already exists.")
            # Example: Update the file content
            file_id = existing_files[0]['id']
            media = MediaFileUpload(self.chat_log_name, mimetype='application/json')
            updated_file = self.service.files().update(fileId=file_id, media_body=media).execute()
            print(f"Updated existing file with ID: {updated_file.get('id')}")

    def generate(self,prompt, history):    # Method to generate a response to the user's input

            
        #if not len(Name) == 0 and not len(Occupation) == 0 and not len(Ethnicity) == 0 and not len(Gender) == 0 and not len(Age) == 0 and not len(YearsOfExp):
        if self.agreed:
            firstmsg =""
            if self.isFirstRun:
                self.context = self.initContext
                self.isFirstRun = False
                firstmsg = prompt
                
        
            self.context += """
                      <|im_start|>nurse
                      Nurse:"""+prompt+"""
                      <|im_start|>barry
                      Barry:
                      """
        
            response = ""
            
            while(len(response) < 1):
                output = self.llm(self.context, max_tokens=400, stop=["Nurse:"], echo=False)
                response = output["choices"][0]["text"]
                response = response.strip()
                #yield response
        
        
          #  for output in llm(input, stream=True, max_tokens=100, ):
            #    piece = output['choices'][0]['text']
           #     response += piece
           #     chatbot[-1] = (chatbot[-1][0], response)
        
           #     yield  response
                
            cleaned_response = self.strip_text(response)
        
            self.chat_history.append((prompt,cleaned_response))
            if  not self.isFirstRun:
                self.chat_log_history.append({"user": prompt, "bot": cleaned_response})
                self.upload_to_google_drive()
                
            else:
                self.chat_log_history.append({"user": firstmsg, "bot": cleaned_response})
            
            self.context += response
            
            print (self.context)
            return self.chat_history
            
        else:
            output = "Did you forget to Agree to the Terms and Conditions?"
            self.chat_history.append((prompt,output))
            return self.chat_history

    def update_chatlog_name(self):
        self.chat_log_name =  f'chat_log_for_{self.unique_id}_{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.json'
        return self.chat_log_name

    def start_chat_button_fn(self,agree_status):    # Method to handle the start chat button action
        
        if agree_status: 
            self.agreed = agree_status
            self.update_chatlog_name()
            return f"You can start chatting now."    
        else:
            return "You must agree to the terms and conditions to proceed"
        
    def reset_chat_interface(self):    # Method to reset the chat interface

        self.chat_history = []
        self.chat_log_history = []
        self.isFirstRun = True
        return "Chat has been reset."
    
    def reset_name_interface(self):    # Method to create the Gradio app interface

        Name = ""
        Occupation = ""
        YearsOfExp = ""
        Ethnicity = ""
        Gender = ""
        Age = ""
        chat_log_name = ""
        return "User info has been reset."
    
    def reset_all(self):

        message1 = self.reset_chat_interface()
        #message2 = reset_name_interface()
        #message3 = load_model()
        self.unique_id = self.generate_unique_id()
        self.update_chatlog_name()
        return f"All Chat components have been rest.  Uniqe ID for this session is, {self.unique_id}. Please note this down.",self.unique_id

    def clear_chat_window(self):
        return []

    def create_app(self):    # Method to launch the Gradio app

        with gr.Blocks(theme=gr.themes.Soft()) as app:
            gr.Markdown("# ECU-IVADE: Conversational AI Model for Aggressive Patient Behavior (Beta Testing)")
            unique_id_display = gr.Textbox(value=self.unique_id, label="Session Unique ID", interactive=False,show_copy_button = True)
            
            with gr.Tab("Terms and Conditions"):
                #name = gr.Textbox(label="Name")
                #occupation = gr.Textbox(label="Occupation")
                #yearsofexp = gr.Textbox(label="Years of Experience")
                #ethnicity = gr.Textbox(label="Ethnicity")
                #gender = gr.Dropdown(choices=["Male", "Female", "Other", "Prefer Not To Say"], label="Gender")
                #age = gr.Textbox(label="Age")
                #submit_info = gr.Button("Submit")
                gr.Markdown("## Terms and Conditions")
                gr.Markdown("""
            Before using our chatbot, please read the following terms and conditions carefully:
            
            - **Data Collection**: Our chatbot collects chat logs for the purpose of improving our services and user experience.
            - **Privacy**: We ensure the confidentiality and security of your data, in line with our privacy policy.
            - **Survey**: At the end of the chat session, you will be asked to participate in a short survey to gather feedback about your experience.
            - **Consent**: By checking the box below and initiating the chat, you agree to these terms and the collection of chat logs, and consent to take part in the survey upon completing your session.
            
            Please check the box below to acknowledge your agreement and proceed.
            """)
                agree_status = gr.Checkbox(label="I have read and understand the terms and conditions.")
                status_label = gr.Markdown()
                start_chat_button = gr.Button("Agree to Terms and Conditions")
                #submit_info.click(submit_user_info, inputs=[name, occupation, yearsofexp, ethnicity, gender, age], outputs=[status_textbox])
                start_chat_button.click(self.start_chat_button_fn, inputs=[agree_status], outputs=[status_label])
                #status_textbox = gr.Textbox(interactive = False)
                
            with gr.Tab("Chat Bot"):
                chatbot = gr.Chatbot()
                msg = gr.Textbox(label="Type your message")
                send = gr.Button("Send")
                clear = gr.Button("Clear Chat")
                send.click(self.generate, inputs=[msg], outputs=chatbot)
                clear.click(self.clear_chart_window, inputs=[], outputs=chatbot)
            
            with gr.Tab("Reset"):
                reset_button = gr.Button("Reset ChatBot Instance")
                reset_output = gr.Textbox(label="Reset Output", interactive=False)
                reset_button.click(self.reset_all, inputs=[], outputs=[reset_output,unique_id_display])
        return app

# Create an instance of the ChatbotAPP class and launch the app
llm = load_model()
unique_id = generate_unique_id()
chatbot_app = ChatbotAPP(llm,SERVICE_ACCOUNT_FILE,SCOPES,folder_id,unique_id,initContext)
app = chatbot_app.create_app()
app.launch(debug=True)