app.py

import langchain
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.document_loaders import UnstructuredPDFLoader,UnstructuredWordDocumentLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain.vectorstores import FAISS
from zipfile import ZipFile
import gradio as gr
import openpyxl
import os
import shutil
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
import tiktoken
import secrets
import time
import requests
import tempfile

from groq import Groq



tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")

# create the length function
def tiktoken_len(text):
    tokens = tokenizer.encode(
        text,
        disallowed_special=()
    )
    return len(tokens)

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=750,
    chunk_overlap=350,
    length_function=tiktoken_len,
    separators=["\n\n", "\n", " ", ""]
)

embeddings = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
foo = Document(page_content='foo is fou!',metadata={"source":'foo source'})


def reset_database(ui_session_id):
  session_id = f"PDFAISS-{ui_session_id}"
  if 'drive' in session_id:
    print("RESET DATABASE: session_id contains 'drive' !!")
    return None

  try:
    shutil.rmtree(session_id)
  except:
    print(f'no {session_id} directory present')
  
  try:
    os.remove(f"{session_id}.zip")
  except:
    print("no {session_id}.zip present")

  return None

def is_duplicate(split_docs,db):
  epsilon=0.0
  print(f"DUPLICATE: Treating: {split_docs[0].metadata['source'].split('/')[-1]}")
  for i in range(min(3,len(split_docs))):
    query = split_docs[i].page_content
    docs = db.similarity_search_with_score(query,k=1)
    _ , score = docs[0]
    epsilon += score
  print(f"DUPLICATE: epsilon: {epsilon}")
  return epsilon < 0.1

def merge_split_docs_to_db(split_docs,db,progress,progress_step=0.1):
  progress(progress_step,desc="merging docs")
  if len(split_docs)==0:
    print("MERGE to db: NO docs!!")
    return

  filename = split_docs[0].metadata['source']
  if is_duplicate(split_docs,db):
    print(f"MERGE: Document is duplicated: {filename}")
    return
  print(f"MERGE: number of split docs: {len(split_docs)}")
  batch = 10
  for i in range(0, len(split_docs), batch):
    progress(i/len(split_docs),desc=f"added {i} chunks of {len(split_docs)} chunks")
    db1 = FAISS.from_documents(split_docs[i:i+batch], embeddings)
    db.merge_from(db1)
  return db

def merge_pdf_to_db(filename,db,progress,progress_step=0.1):
  progress_step+=0.05
  progress(progress_step,'unpacking pdf')
  doc = UnstructuredPDFLoader(filename).load()
  doc[0].metadata['source'] = filename.split('/')[-1]
  split_docs = text_splitter.split_documents(doc)
  progress_step+=0.3
  progress(progress_step,'docx unpacked')
  return merge_split_docs_to_db(split_docs,db,progress,progress_step)

def merge_docx_to_db(filename,db,progress,progress_step=0.1):
  progress_step+=0.05
  progress(progress_step,'unpacking docx')
  doc = UnstructuredWordDocumentLoader(filename).load()
  doc[0].metadata['source'] = filename.split('/')[-1]
  split_docs = text_splitter.split_documents(doc)
  progress_step+=0.3
  progress(progress_step,'docx unpacked')
  return merge_split_docs_to_db(split_docs,db,progress,progress_step)

def merge_txt_to_db(filename,db,progress,progress_step=0.1):
  progress_step+=0.05
  progress(progress_step,'unpacking txt')
  with open(filename) as f:
      docs = text_splitter.split_text(f.read())
      split_docs = [Document(page_content=doc,metadata={'source':filename.split('/')[-1]}) for doc in docs]
  progress_step+=0.3
  progress(progress_step,'txt unpacked')
  return merge_split_docs_to_db(split_docs,db,progress,progress_step)

def unpack_zip_file(filename,db,progress):
    with ZipFile(filename, 'r') as zipObj:
        contents = zipObj.namelist()
        print(f"unpack zip: contents: {contents}")
        tmp_directory = filename.split('/')[-1].split('.')[-2]
        shutil.unpack_archive(filename, tmp_directory)

        if 'index.faiss' in [item.lower() for item in contents]:
            db2 = FAISS.load_local(tmp_directory, embeddings, allow_dangerous_deserialization=True)
            db.merge_from(db2)
            return db
        
        for file in contents:
            if file.lower().endswith('.docx'):
              db = merge_docx_to_db(f"{tmp_directory}/{file}",db,progress)
            if file.lower().endswith('.pdf'):
              db = merge_pdf_to_db(f"{tmp_directory}/{file}",db,progress)
            if file.lower().endswith('.txt'):
              db = merge_txt_to_db(f"{tmp_directory}/{file}",db,progress)
        return db

def add_files_to_zip(session_id):
    zip_file_name = f"{session_id}.zip"
    with ZipFile(zip_file_name, "w") as zipObj:
        for root, dirs, files in os.walk(session_id):
            for file_name in files:
                file_path = os.path.join(root, file_name)
                arcname = os.path.relpath(file_path, session_id)
                zipObj.write(file_path, arcname)

#### UI Functions ####

def embed_files(files,ui_session_id,progress=gr.Progress(),progress_step=0.05):
    if ui_session_id not in os.environ['users'].split(', '):
        return "README.md", ""
    print(files)
    progress(progress_step,desc="Starting...")
    split_docs=[]
    if len(ui_session_id)==0:
      ui_session_id = secrets.token_urlsafe(16)
    session_id = f"PDFAISS-{ui_session_id}"

    try:
      db = FAISS.load_local(session_id,embeddings, allow_dangerous_deserialization=True)
    except:
      print(f"SESSION: {session_id} database does not exist, create a FAISS db")
      db =  FAISS.from_documents([foo], embeddings)
      db.save_local(session_id)
      print(f"SESSION: {session_id} database created")
    
    print("EMBEDDED, before embeddeding: ",session_id,len(db.index_to_docstore_id))
    for file_id,file in enumerate(files):
        print("ID : ", file_id, "FILE : ", file)
        file_type = file.name.split('.')[-1].lower()
        source = file.name.split('/')[-1]
        print(f"current file: {source}")
        progress(file_id/len(files),desc=f"Treating {source}")

        if file_type == 'pdf':
            db2 = merge_pdf_to_db(file.name,db,progress)
        
        if file_type == 'txt':
            db2 = merge_txt_to_db(file.name,db,progress)
        
        if file_type == 'docx':
            db2 = merge_docx_to_db(file.name,db,progress)

        if file_type == 'zip':
            db2 = unpack_zip_file(file.name,db,progress)

        if db2 != None:
            db = db2
            db.save_local(session_id)
            ### move file to store ###
            progress(progress_step, desc = 'moving file to store')
            directory_path = f"{session_id}/store/"
            if not os.path.exists(directory_path):
                os.makedirs(directory_path)
            try:
                shutil.move(file.name, directory_path)
            except:
                pass

    ### load the updated db and zip it ###
    progress(progress_step, desc = 'loading db')
    db = FAISS.load_local(session_id,embeddings, allow_dangerous_deserialization=True)
    print("EMBEDDED, after embeddeding: ",session_id,len(db.index_to_docstore_id))
    progress(progress_step, desc = 'zipping db for download')
    add_files_to_zip(session_id)
    print(f"EMBEDDED: db zipped")
    progress(progress_step, desc = 'db zipped')
    return f"{session_id}.zip", ui_session_id, ""



def add_to_db(references,ui_session_id):
    files = store_files(references)
    return embed_files(files,ui_session_id)

def export_files(references):
    files = store_files(references, ret_names=True)
    #paths = [file.name for file in files]
    return files
    

def display_docs(docs):
  output_str = ''
  for i, doc in enumerate(docs):
      source = doc.metadata['source'].split('/')[-1]
      output_str += f"Ref: {i+1}\n{repr(doc.page_content)}\nSource: {source}\n*§*§*\n"
  return output_str


# def display_docs_modal(docs):
#   output_list = []
#   for i, doc in enumerate(docs):
#       source = doc.metadata['source'].split('/')[-1]
#       output_str.append(f"Ref: {i+1}\n{repr(doc.page_content)}\nSource: {source}\n*§*§*\n")
#   return output_list


def hide_source():
    return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)


def ask_llm(system, user_input):
    messages = [
        {
            "role": "system",
            "content": system
        },
        {
            "role": "user",
            "content": user_input,
        }
    ]
    client = Groq(api_key=os.environ["GROQ_KEY"])
    chat_completion = client.chat.completions.create(
        messages=messages,
        model="llama3-70b-8192",#'mixtral-8x7b-32768',
    )
    return chat_completion.choices[0].message.content

def ask_llm_stream(system, user_input):
    llm_response = ""
    client = Groq(api_key=os.environ["GROQ_KEY"])
    if user_input is None or user_input == "":
        user_input = "What is the introduction of the document about?"
    messages = [
        {
            "role": "system",
            "content": system
        },
        {
            "role": "user",
            "content": user_input,
        }
    ]

    stream = client.chat.completions.create(
        messages=messages,
        model="mixtral-8x7b-32768",
        temperature=0.5,
        max_tokens=1024,
        top_p=1,
        stop=None,
        stream=True,
    )

    for chunk in stream:
        llm_response += str(chunk.choices[0].delta.content) if chunk.choices[0].delta.content is not None else ""
        yield llm_response
    

def ask_gpt(query, ui_session_id, history):
    print(f"before: {os.environ['prompts']}")
    os.environ['prompts'] += ', ' + query
    print(f"after: {os.environ['prompts']}")

    if ui_session_id not in os.environ['users'].split(', '):
        return "Please Login", "", ""
    session_id = f"PDFAISS-{ui_session_id}"
    try:
      db = FAISS.load_local(session_id,embeddings, allow_dangerous_deserialization=True)
      print("ASKGPT after loading",session_id,len(db.index_to_docstore_id))
    except:
      print(f"SESSION: {session_id} database does not exist")
      return f"SESSION: {session_id} database does not exist","",""

    docs = db.similarity_search(query, k=4)

    documents = "\n\n*-*-*-*-*-*\n\n".join(f"Content: {doc.page_content}\n" for doc in docs)
    system = f"# Instructions\nTake a deep breath and resonate step by step.\nYou are a helpful standard assistant. Your have only one mission and that consists in answering to the user input based on the **provided documents**. If the answer to the question that is asked by the user isn't contained in the **provided documents**, say so but **don't make up an answer**. I chose you because you can say 'I don't know' so please don't do like the other LLMs and don't define acronyms that aren\'t present in the following **PROVIDED DOCUMENTS** double check if it is present before answering. If some of the information can be useful for the user you can tell him.\nFinish your response by **ONE** follow up question that the provided documents could answer.\n\nThe documents are separated by the string \'*-*-*-*-*-*\'. Do not provide any explanations or details.\n\n# **Provided documents**: {documents}."    
    gen = ask_llm_stream(system, query)
    last_value=""
    displayable_docs = display_docs(docs)
    yn_display = len(docs)*[True]+(5-len(docs))*[False]
        
    while True:
        try:
            last_value = next(gen)
            yield last_value, displayable_docs, history + f"[query]\n{query}\n[answer]\n{last_value}\n[references]\n{displayable_docs}\n\n", gr.Button("Ref 1", visible=yn_display[0]), gr.Button("Ref 2", visible=yn_display[1]), gr.Button("Ref 3", visible=yn_display[2]), gr.Button("Ref 4", visible=yn_display[3]), gr.Button("Ref 5", visible=yn_display[4])
        except StopIteration as e:
            break
    history += f"[query]\n{query}\n[answer]\n{last_value}\n[references]\n{displayable_docs}\n\n"
    return last_value, displayable_docs, history, gr.Button("Ref 1", visible=yn_display[0]), gr.Button("Ref 2", visible=yn_display[1]), gr.Button("Ref 3", visible=yn_display[2]), gr.Button("Ref 4", visible=yn_display[3]), gr.Button("Ref 5", visible=yn_display[4])


def auth_user(ui_session_id):
    if ui_session_id in os.environ['users'].split(', '):
        return gr.Textbox(label='Username', visible=False), gr.File(file_count="multiple", file_types=[".txt", ".pdf",".zip",".docx"], visible=True), gr.Button("Reset AI Knowledge", visible=True), gr.Markdown(label='AI Answer', visible=True), gr.Textbox(placeholder="Type your question", label="Question ❔", scale=9, visible=True), gr.Button("▶", scale=1, visible=True), gr.Textbox(label='Sources', show_copy_button=True, visible=True), gr.File(label="Zipped database", visible=True), gr.Textbox(label='History', show_copy_button=True, visible=True)
    else:
        return gr.Textbox(label='Username', visible=True), gr.File(file_count="multiple", file_types=[".txt", ".pdf",".zip",".docx"], visible=False), gr.Button("Reset AI Knowledge", visible=False), gr.Markdown(label='AI Answer', visible=False), gr.Textbox(placeholder="Type your question", label="Question ❔", scale=9, visible=False), gr.Button("▶", scale=1, visible=False), gr.Textbox(label='Sources', show_copy_button=True, visible=False), gr.File(label="Zipped database", visible=False), gr.Textbox(label='History', show_copy_button=True, visible=False)

def display_info0(documents):
    try:
        return gr.Markdown(value=documents.split("\n*§*§*\n")[0], label='Source', visible=True), gr.Button('Hide', visible=True)
    except Exception as e:
        gr.Info("No Document")
        return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)

def display_info1(documents):
    try:
        return gr.Markdown(value=documents.split("\n*§*§*\n")[1], label='Source', visible=True), gr.Button('Hide', visible=True)
    except Exception as e:
        gr.Info("No Document")
        return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)
        
def display_info2(documents):
    try:
        return gr.Markdown(value=documents.split("\n*§*§*\n")[2], label='Source', visible=True), gr.Button('Hide', visible=True)
    except Exception as e:
        gr.Info("No Document")
        return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)
    
def display_info3(documents):
    try:
        return gr.Markdown(value=documents.split("\n*§*§*\n")[3], label='Source', visible=True), gr.Button('Hide', visible=True)
    except Exception as e:
        gr.Info("No Document")
        return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)
    
def display_info4(documents):
    try:
        return gr.Markdown(value=documents.split("\n*§*§*\n")[4], label='Source', visible=True), gr.Button('Hide', visible=True)
    except Exception as e:
        gr.Info("No Document")
        return gr.Markdown(label='Source', visible=False), gr.Button('Hide', visible=False)

with gr.Blocks() as demo:
    gr.Markdown("# Enrich an LLM knowledge with your own documents 🧠🤖")
        
    with gr.Column():        
        tb_session_id = gr.Textbox(label='Username')
        docs_input = gr.File(file_count="multiple", file_types=[".txt", ".pdf",".zip",".docx"], visible=False)
        btn_reset_db = gr.Button("Reset AI Knowledge", visible=False)


    with gr.Column():
        answer_output = gr.Markdown(label='AI Answer', visible=False)
        btn_hide_source = gr.Button('Hide', visible=False)
        md_ref = gr.Markdown(label='Source', visible=False)
        with gr.Row():
            query_input = gr.Textbox(placeholder="Type your question", label="Question ❔", scale=9, visible=False)
            btn_askGPT = gr.Button("▶", scale=1, visible=False)
        with gr.Row():
            btn1 = gr.Button("Ref 1", visible=False)
            btn2 = gr.Button("Ref 2", visible=False)
            btn3 = gr.Button("Ref 3", visible=False)
            btn4 = gr.Button("Ref 4", visible=False)
            btn5 = gr.Button("Ref 5", visible=False)
        
        
        tb_sources = gr.Textbox(label='Sources', show_copy_button=True, visible=False)
        

    with gr.Accordion("Download your knowledge base and see your conversation history", open=False):
        db_output = gr.File(label="Zipped database", visible=False)
        tb_history = gr.Textbox(label='History', show_copy_button=True, visible=False, interactive=False)
        

    tb_session_id.submit(auth_user, inputs=tb_session_id, outputs=[tb_session_id, docs_input, btn_reset_db, answer_output, query_input, btn_askGPT, tb_sources, db_output, tb_history])

    docs_input.upload(embed_files, inputs=[docs_input,tb_session_id], outputs=[db_output,tb_session_id, query_input])
    btn_reset_db.click(reset_database,inputs=[tb_session_id],outputs=[db_output])
    btn_askGPT.click(ask_gpt, inputs=[query_input, tb_session_id, tb_history], outputs=[answer_output, tb_sources, tb_history, btn1, btn2, btn3, btn4, btn5])
    query_input.submit(ask_gpt, inputs=[query_input, tb_session_id, tb_history], outputs=[answer_output, tb_sources, tb_history, btn1, btn2, btn3, btn4, btn5])

    btn1.click(display_info0, inputs=tb_sources, outputs=[md_ref, btn_hide_source])
    btn2.click(display_info1, inputs=tb_sources, outputs=[md_ref, btn_hide_source])
    btn3.click(display_info2, inputs=tb_sources, outputs=[md_ref, btn_hide_source])
    btn4.click(display_info3, inputs=tb_sources, outputs=[md_ref, btn_hide_source])
    btn5.click(display_info4, inputs=tb_sources, outputs=[md_ref, btn_hide_source])
    btn_hide_source.click(hide_source, inputs=None, outputs=[md_ref, btn_hide_source])
    


demo.launch(debug=False,share=False)