app.py

import openai
import os
openai.api_key=os.getenv("OPENAI_API_KEY")

from dotenv import load_dotenv
load_dotenv()

from flask import Flask, jsonify, render_template, request
import requests, json
import PyPDF2

# import nltk
# nltk.download("punkt")


import shutil
from werkzeug.utils import secure_filename
from werkzeug.datastructures import FileStorage

import nltk

from datetime import datetime

import openai
from langchain.llms import OpenAI, Replicate

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.embeddings import SentenceTransformerEmbeddings

from langchain.document_loaders import SeleniumURLLoader, PyPDFLoader
from langchain.docstore.document import Document

from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter

from langchain.chains import VectorDBQA, ConversationChain

from langchain.document_loaders import UnstructuredFileLoader, TextLoader
from langchain import PromptTemplate

from langchain.chains import RetrievalQA
#from langchain.memory import ConversationBufferWindowMemory
from langchain.memory import ConversationBufferMemory


from transformers import LlamaTokenizer, AutoTokenizer

import warnings

warnings.filterwarnings("ignore")


#app = Flask(__name__)

app = Flask(__name__, template_folder="./")

# Create a directory in a known location to save files to.
uploads_dir = os.path.join(app.root_path,'static', 'uploads')

os.makedirs(uploads_dir, exist_ok=True)

# Initialize global variables for conversation history
conversation_history = []

defaultEmbeddingModelID = 3
defaultLLMID=0


def pretty_print_docs(docs):
    print(f"\n{'-' * 100}\n".join([f"Document {i + 1}:\n\n" + "Document Length>>>" + str(
        len(d.page_content)) + "\n\nDocument Source>>> " + d.metadata['source'] + "\n\nContent>>> " + d.page_content for
                                   i, d in enumerate(docs)]))


def getEmbeddingModel(embeddingId):
    if (embeddingId == 1):
        embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
    elif (embeddingId == 2):
        model_name = "hkunlp/instructor-large"
        model_kwargs = {'device': 'cpu'}
        encode_kwargs = {'normalize_embeddings': True}
        embeddings = HuggingFaceInstructEmbeddings(model_name=model_name,model_kwargs=model_kwargs,encode_kwargs=encode_kwargs)
    elif (embeddingId == 3):
        model_name = "BAAI/bge-large-en-v1.5"
        model_kwargs = {'device': 'cuda'}
        encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
        model = HuggingFaceBgeEmbeddings(model_name=model_name,model_kwargs=model_kwargs,encode_kwargs=encode_kwargs)
    else:
        embeddings = OpenAIEmbeddings()
    return OpenAIEmbeddings()


def getLLMModel(LLMID):
    # else:
    #     llm = LlamaCpp(
    if LLMID == 1:
        # llm = Replicate(
        #     model="a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5",
        #     model_kwargs={"temperature": 0.2,"max_length": 2500})
        llm = Replicate(
            model="meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
            model_kwargs={"temperature": 0.2,"max_new_tokens":2500})
        print("LLAMA2 13B LLM Selected")
    elif LLMID == 2:
        # llm = Replicate(
        #     model="replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf",
        #     model_kwargs={"temperature": 0.2,"max_length": 2500})
        llm = Replicate(
        model="meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
        model_kwargs={"temperature": 0.2,"max_new_tokens":2500})
        print("LLAMA2 70B LLM Selected")
    elif LLMID == 3:
        llm = Replicate(model="meta/llama-2-7b-chat:8e6975e5ed6174911a6ff3d60540dfd4844201974602551e10e9e87ab143d81e",
                        model_kwargs={"temperature": 0.2,"max_new_tokens":2500})
        print("LLAMA2 7B Chat LLM Selected")
    elif LLMID == 4:
        llm = Replicate(
            model="a16z-infra/mistral-7b-instruct-v0.1:83b6a56e7c828e667f21fd596c338fd4f0039b46bcfa18d973e8e70e455fda70",
            model_kwargs={"temperature": 0.2,"max_new_tokens":2500})
        print("Mistral AI LLM Selected")
    else:
        llm = OpenAI(model_name="gpt-3.5-turbo-0125",temperature=0.0)
        print("Open AI LLM Selected")
    return llm


def clearKBUploadDirectory(uploads_dir):
    for filename in os.listdir(uploads_dir):
        file_path = os.path.join(uploads_dir, filename)
        print("Clearing Doc Directory. Trying to delete" + file_path)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print('Failed to delete %s. Reason: %s' % (file_path, e))

def PDFChunkerWithSeparator(filepath, separator):

    content = ""
    if filepath.endswith(".pdf"):
        # creating a pdf reader object
        reader = PyPDF2.PdfReader(filepath)
        # print the number of pages in pdf file
        print(len(reader.pages))
        for page in reader.pages:
            content += page.extract_text()
    elif filepath.endswith(".txt"):
        with open(filepath) as f:
            lines = f.readlines()
            f.close()
            for line in lines:
                content+=line
    splitted_content_list = content.split(separator)

    doclist = []
    for splitted_content in splitted_content_list:
        new_doc = Document(page_content=splitted_content, metadata={"source": filepath})
        # print(type(new_doc))
        doclist.append(new_doc)
    if len(doclist)>3:
        print(doclist[len(doclist) - 3])
    return doclist


def loadKB(fileprovided, urlProvided, uploads_dir, request):
    documents = []
    global tokenizer
    BASE_MODEL = "LLAMA-TOKENIZER"
    savedModelPath = "./model/" + BASE_MODEL
    #tokenizer = LlamaTokenizer.from_pretrained(savedModelPath)
    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
    separator = "</Q>"
    if fileprovided:
        # Delete Files
        clearKBUploadDirectory(uploads_dir)
        # Read and Embed New Files provided
        for file in request.files.getlist('files[]'):
            print("File Received>>>" + file.filename)
            file.save(os.path.join(uploads_dir, secure_filename(file.filename)))
            #loader = PyPDFLoader(os.path.join(uploads_dir, secure_filename(file.filename)))
            #documents.extend(loader.load())
            documents.extend(PDFChunkerWithSeparator(os.path.join(uploads_dir, secure_filename(file.filename)),separator))
    else:
        #loader = TextLoader('Jio.txt')
        #documents.extend(loader.load())
        documents.extend(PDFChunkerWithSeparator('JTest.txt',separator))

    if urlProvided:
        weburl = request.form.getlist('weburl')
        print(weburl)
        urlList = weburl[0].split(';')
        print(urlList)
        print("Selenium Started", datetime.now().strftime("%H:%M:%S"))
        # urlLoader=RecursiveUrlLoader(urlList[0])
        urlLoader = SeleniumURLLoader(urlList)
        print("Selenium Completed", datetime.now().strftime("%H:%M:%S"))
        documents.extend(urlLoader.load())
        print("inside selenium loader:")
        print(documents)

    return documents


def getRAGChain(customerName, customerDistrict, custDetailsPresent, vectordb,llmID):

    # Retrieve conversation history if available
    #memory = ConversationBufferWindowMemory(k=3, memory_key="history", input_key="question")
    global memory 
    #memory = ConversationBufferWindowMemory(k=3, memory_key="history", input_key="question", initial_memory=conversation_history)
    #memory = ConversationBufferMemory(k=3, memory_key="history", input_key="query", initial_memory=conversation_history)
    memory = ConversationBufferMemory()
    
    # chain = RetrievalQA.from_chain_type(
    #     llm=getLLMModel(llmID),
    #     chain_type='stuff',
    #     retriever=getRetriever(vectordb),
    #     #retriever=vectordb.as_retriever(),
    #     verbose=False,
    #     chain_type_kwargs={
    #         "verbose": False,
    #         "prompt": createPrompt(customerName, customerDistrict, custDetailsPresent),
    #         "memory": ConversationBufferWindowMemory(
    #             k=3,
    #             memory_key="history",
    #             input_key="question"),
    #     }
    # )
    # chain = RetrievalQA.from_chain_type(
    #     llm=getLLMModel(llmID),
    #     chain_type='stuff',
    #     retriever=getRetriever(vectordb),
    #     memory=memory,
    #     #retriever=vectordb.as_retriever(),
    #     verbose=True,
    #     chain_type_kwargs={
    #         "verbose": False,
    #         "prompt": createPrompt(customerName, customerDistrict, custDetailsPresent),
    #         "memory": memory
    #     }
    # )
    chain = ConversationChain(
        llm=getLLMModel(llmID),
        chain_type='stuff',
        retriever=getRetriever(vectordb),
        memory=memory,
        #retriever=vectordb.as_retriever(),
        verbose=True,
        chain_type_kwargs={
            "verbose": False,
            "prompt": createPrompt(customerName, customerDistrict, custDetailsPresent),
            "memory": memory
        }
    )
    return chain

def getRetriever(vectordb):
    return vectordb.as_retriever(search_type="mmr", search_kwargs={'k': 2})

def createVectorDB(documents,embeddingModelID):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=150)
    texts = []
    for document in documents:
        tokenized_input = tokenizer.tokenize(document.page_content)
        print("Token Count::::::::::" + str(len(tokenized_input)))
        if (len(tokenized_input) > 1000):
            print("Splitting Content using RTS")
            splitted_doc = text_splitter.split_documents([document])
            texts.extend(splitted_doc)
        #             for text in texts:
        #                 print("splitted content:"+str(len(text.page_content)))
        #                 print(text.page_content)
        elif (len(tokenized_input) < 1000 and len(tokenized_input) > 1):
            texts.append(document)
    # texts = text_splitter.split_documents(documents)
    print("All chunk List START ***********************\n\n")
    pretty_print_docs(texts)
    print("All chunk List END ***********************\n\n")
    embeddings = getEmbeddingModel(embeddingModelID)
    print("Embedding Started >>>>>>>>>>>>>>>>>>", datetime.now().strftime("%H:%M:%S"))
    vectordb = Chroma.from_documents(texts, embeddings, collection_metadata={"hnsw:space": "cosine"})
    print("Vector Store Creation Completed*********************************\n\n")
    return vectordb
    # texts = text_splitter.split_documents(documents)
    # print("All chunk List START ***********************\n\n")
    # pretty_print_docs(texts)
    # print("All chunk List END ***********************\n\n")
    # embeddings = getEmbeddingModel(0)
    # vectordb = Chroma.from_documents(texts, embeddings)
    # return vectordb

def createPrompt(cName, cCity, custDetailsPresent):
    #cProfile = "Customer's Name is " + cName + "\nCustomer's lives in or customer's Resident State or Customer's place is " + cCity + "\n"
    cProfile = "Customer's Name is " + cName + "\nCustomer's lives in the city of " + cCity + "\n"
    print(cProfile)

    # template1 = """You role is of a Professional Customer Support Executive and your name is Jio AIAssist.
    #     You are talking to the below customer whose information is provided in block delimited by <cp></cp>.
    #     Use the following customer related information (delimited by <cp></cp>) and context (delimited by <ctx></ctx>) to answer the question at the end by thinking step by step alongwith reaonsing steps:
    #     If you are unable to answer the question using the context provided, just say that you don't know, don't try to make up an answer.
    #     Use the customer information to replace entities in the question before answering\n
    #     \n"""

    template1 = """You role is of a Professional Customer Support Executive and your name is Jio AIAssist. 
    Use the context (delimited by <ctx></ctx>) to answer the question at the end by thinking step by step alongwith reaonsing steps.
    If you are unable to answer the question using the context provided, just say that you don't know, don't try to make up an answer. 
    Use the customer information to replace entities in the question before answering\n \n"""

    template2 = """
        <ctx>
        {context}
        </ctx>
        <hs>
        {history}
        </hs>
        Question: {question}
        Answer: """

    #prompt_template = template1 + "<cp>\n" + cProfile + "\n</cp>\n" + template2
    prompt_template = template1 + "You are talking to the following customer. <cp>\n" + cProfile + "\n</cp>\n" + template2
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["history", "context", "question"])
    return PROMPT

vectordb = createVectorDB(loadKB(False, False, uploads_dir, None),defaultEmbeddingModelID)

@app.route('/', methods=['GET'])
def test():
    return "Docker hello"

@app.route('/KBUploader')
def KBUpload():
    return render_template("KBTrain.html")

@app.route('/aiassist')
def aiassist():
    return render_template("index.html")

@app.route('/aisearch')
def aisearch():
    return render_template("aisearch.html")

@app.route('/agent/chat/suggestion', methods=['POST'])
def process_json():
    print(f"\n{'*' * 100}\n")
    print("Request Received >>>>>>>>>>>>>>>>>>", datetime.now().strftime("%H:%M:%S"))

    # Retrieve conversation ID from the request (use any suitable ID)
    conversation_id = request.json.get('conversation_id', None)
    
    content_type = request.headers.get('Content-Type')
    if content_type == 'application/json':
        requestQuery = request.get_json()
        print(type(requestQuery))
        custDetailsPresent = False
        customerName = ""
        customerDistrict = ""
        if "custDetails" in requestQuery:
            custDetailsPresent = True
            customerName = requestQuery['custDetails']['cName']
            customerDistrict = requestQuery['custDetails']['cDistrict']

        selectedLLMID=defaultLLMID
        if "llmID" in requestQuery:
            selectedLLMID=(int) (requestQuery['llmID'])

        # Create a conversation ID-specific history list if not exists
        conversation_history_id = f"{conversation_id}_history"
        if conversation_history_id not in globals():
            globals()[conversation_history_id] = []
        conversation_history = globals()[conversation_history_id]

        
        print("chain initiation")
        chainRAG = getRAGChain(customerName, customerDistrict, custDetailsPresent, vectordb,selectedLLMID)
        print("chain created")
        suggestionArray = []
        searchResultArray = []

        for index, query in enumerate(requestQuery['message']):
            # message = answering(query)
            memory.chat_memory.add_user_message(query)
            relevantDoc = vectordb.similarity_search_with_score(query, distance_metric="cos", k=3)
            conversation_history.append(query)
            print("Printing Retriever Docs")
            for doc in getRetriever(vectordb).get_relevant_documents(query):
                searchResult = {}
                print(f"\n{'-' * 100}\n")
                searchResult['documentSource'] = doc.metadata['source']
                searchResult['pageContent'] = doc.page_content
                print(doc)
                print("Document Source>>>>>>  " + searchResult['documentSource'] + "\n\n")
                print("Page Content>>>>>> " + searchResult['pageContent'] + "\n\n")
                print(f"\n{'-' * 100}\n")
                searchResultArray.append(searchResult)
            print("Printing Retriever Docs Ended")

            print("Chain Run Started >>>>>>>>>>>>>>>>>>", datetime.now().strftime("%H:%M:%S"))
            message = chainRAG.run({"query": query})
            print("Chain Run Completed >>>>>>>>>>>>>>>>>>", datetime.now().strftime("%H:%M:%S"))
            print("query:", query)
            print("Response:", message)
            memory.chat_memory.add_ai_message(message)
            if "I don't know" in message:
                message = "Dear Sir/ Ma'am, Could you please ask questions relevant to Jio?"
            responseJSON = {"message": message, "id": index}
            suggestionArray.append(responseJSON)
        print("Response Sent >>>>>>>>>>>>>>>>>>", datetime.now().strftime("%H:%M:%S"))
        return jsonify(suggestions=suggestionArray, searchResult=searchResultArray)
    else:
        return 'Content-Type not supported!'

@app.route('/file_upload', methods=['POST'])
def file_Upload():
    fileprovided = not request.files.getlist('files[]')[0].filename == ''
    urlProvided = not request.form.getlist('weburl')[0] == ''
    embeddingModelProvided = not request.form.getlist('embeddingModelID')[0] == ''
    print("*******")
    print("File Provided:" + str(fileprovided))
    print("URL Provided:" + str(urlProvided))
    print("Embedding Model Provided:" + str(embeddingModelProvided))
    print("*******")

    print(uploads_dir)
    documents = loadKB(fileprovided, urlProvided, uploads_dir, request)
    embeddingModelID = defaultEmbeddingModelID
    if embeddingModelProvided:
        embeddingModelID = int(request.form.getlist('embeddingModelID')[0])
    global vectordb
    vectordb = createVectorDB(documents, embeddingModelID)
    #vectordb=createVectorDB(documents)
    return render_template("aisearch.html")

if __name__ == '__main__':
    app.run(host='0.0.0.0',  port=int(os.environ.get('PORT', 7860)))