app.py

import time
import openai
import os
import pandas as pd
import numpy as np
import gradio as gr
from llama_index import GPTVectorStoreIndex, StorageContext, load_index_from_storage
from openai.embeddings_utils import get_embedding, cosine_similarity

openai.api_key = os.getenv("OPENAI_API_KEY")
passwd = os.getenv("PASSWD_SECRET")

title = "Confidential forensics tool with ChatGPT"
examples = ["Who is Phillip Allen?", "What the project in Austin is about?", "Give me more details about the real estate project"]

file_metadata = lambda x: {"filename": x}
storage_context = StorageContext.from_defaults(persist_dir="./storage")
index = load_index_from_storage(storage_context)
query_engine = index.as_query_engine(
 similarity_top_k=3,
)

default_mail_name = "no title"

df = pd.read_csv("metadata.csv", delimiter=";")
df["subject"] = df["subject"].replace([None], default_mail_name)
df2 = df["subject"].head(10).to_frame()
df2 = df2.rename({'subject': 'Emails titles'}, axis=1)

def get_email_subject(response):
 emails = []
 for node in response.source_nodes:
 email = node.node.extra_info["filename"].split("\\")[-1].split(".")[0]
 emails.append(email)

 mask = df.email_name.apply(lambda x: x in emails)
 return df.loc[mask]

# def get_email_subject(response):
# podcasts = []
# for node in response.source_nodes:
# podcast = node.node.extra_info["filename"].split("/")[-1].split(".")[0]
# podcasts.append(podcast)

def search_emails(opt, message, n=3):
 "Outputs the top n emails that match the most the pattern"
 if len(message.strip()) < 1:
 message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
 else:
 try:
 embedding = get_embedding(message)
 message = ""
 df['similarities'] = df.embedding.apply(func=(lambda x: cosine_similarity(x, embedding)))

 message_tmp = df.sort_values('similarities', ascending=False).head(n)
 message_tmp = [(row.file, row.body, row.similarities) for index, row in message_tmp.iterrows()]
 for msg in message_tmp:
 message += f"Mail ID: {msg[0]}\nContent: {msg[1].strip()}\nSimilarity score: {msg[2]}\n\n"
 except Exception as e:
 message = "An error occured when handling your query, please try again."
 print(e)
 return message, ""

def respond_upload(btn_upload, message, chat_history):
 time.sleep(2)
 message = "***File uploaded***"
 bot_message = "Your document has been uploaded and will be accounted for your queries."
 chat_history.append((message, bot_message))
 return btn_upload, "", chat_history

def respond2(message, chat_history, box, btn):
 message, chat_history = respond_common(message, chat_history, box, btn)
 return message, chat_history, box

def respond(message, chat_history):
 message, chat_history = respond_common(message, chat_history)
 return "", chat_history

def respond_common(message, chat_history, box=None, btn=None):
 if len(message.strip()) < 1:
 message = "***Empty***"
 bot_message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
 else:
 try:
 resp = query_engine.query(message)
 bot_message = str(resp).strip()
 get_email_subject(resp)
 bot_message += "\n\nSource(s):\n"
 for i, row in get_email_subject(resp).iterrows():
 bot_message += f"Email ID: **{row.email_name}**\n**Subject: {row.subject}**\n"
 except Exception as e:
 bot_message = "An error occured when handling your query, please try again."
 print(e)
 chat_history.append((message, bot_message))
 return message, chat_history


with gr.Blocks(title=title) as demo:
 gr.Markdown(
 """

 # """ + title + """
 """)
 dat = gr.Dataframe(
 value=df2,
 max_cols=1,
 max_rows=4,
 overflow_row_behaviour="paginate",
 )
 btn_upload = gr.UploadButton("Upload a new document...", file_types=["text"])
 gr.Markdown(
 """
 ## Chatbot
 """)
 chatbot = gr.Chatbot().style(height=400)
 with gr.Row():
 with gr.Column(scale=0.85):
 msg = gr.Textbox(
 show_label=False,
 placeholder="Enter text and press enter, or click on Send.",
 ).style(container=False)
 with gr.Column(scale=0.15, min_width=0):
 btn_send = gr.Button("Send your query")
 with gr.Row():
 gr.Markdown(
 """
 Example of queries
 """)
 for ex in examples:
 btn = gr.Button(ex)
 btn.click(respond2, [btn, chatbot, msg], [btn, chatbot, msg])

 msg.submit(respond, [msg, chatbot], [msg, chatbot])
 btn_send.click(respond, [msg, chatbot], [msg, chatbot])
 btn_upload.upload(respond_upload, [btn_upload, msg, chatbot], [btn_upload, msg, chatbot])
 # gr.Markdown(
 # """

 # ## Search the matching document
 # """)
 # opt = gr.Textbox(
 # show_label=False,
 # placeholder="The document matching with your query will be shown here.",
 # interactive=False,
 # lines=8
 # )
 # with gr.Row():
 # with gr.Column(scale=0.85):
 # msg2 = gr.Textbox(
 # show_label=False,
 # placeholder="Enter text and press enter, or click on Send.",
 # ).style(container=False)
 # with gr.Column(scale=0.15, min_width=0):
 # btn_send2 = gr.Button("Send your query")

 # btn_send2.click(search_emails, [opt, msg2], [opt, msg2])
 
if __name__ == "__main__":
 demo.launch(auth=("mithril", passwd))