rubyrails117/Upload-Docs-react-test

1from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext2from llama_index.vector_stores import PineconeVectorStore3from llama_index.storage.storage_context import StorageContext4from pinecone import Pinecone5import logging6import sys7import os8import openai9from pymongo.mongo_client import MongoClient10from datetime import datetime11 12logging.basicConfig(stream=sys.stdout, level=logging.INFO)13logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))14 15 16pinecone_key = os.environ['PINECONE_KEY']17globalkey=os.environ["global_key"]18mongo_key=os.environ["MONGO_KEY"]19 20def isBrainFound(brainName):21    if(brainName=="global-tekmed"):22        pinecone.init(globalkey,23                      environment="us-west4-gcp")24    else:    25        Pinecone(api_key=pinecone_key)26    active_indexes = pinecone.list_indexes()27    print(active_indexes)28    brainName = brainName.lower()29    if brainName in active_indexes:30        return True31    else:32        return False33 34def add_file(collection,username, filename):35    document = {36        "brain": username,37        "filename": filename,38        "namespace":None,39        "timestamp": datetime.utcnow()40    }41    collection.insert_one(document)42 43 44def delete_file(collection,username, filename):45    query = {46        "brain": username,47        "filename": filename48    }49    collection.delete_one(query)50 51def updateBrain(files,brainName):52    if(brainName=="unchain3d-demo" or brainName=="Unchain3d-demo"):53        return "updated"54    55    # if(isBrainFound(brainName)==False):56    #     return "brain not found"57    58    print(len(files))59    print("Updating brain")60    print(brainName)61    62    # if(brainName=="global-tekmed"):63    #     pinecone.init(globalkey,64    #                   environment="us-west4-gcp")65    # else:    66    #     Pinecone(api_key=pinecone_key)67    pc=Pinecone(api_key=pinecone_key)68   69    documents = []70    corrupt = []71    newfiles = []72 73 74    i = 075    for file in files:76        loader = SimpleDirectoryReader(input_files=[file.name])77        try:78            document1 = loader.load_data()79            j=180            for doc in document1:81                doc.doc_id=os.path.basename(file.name)82                doc.extra_info={"filename":os.path.basename(file.name),"pageno":j}83                documents.append(doc)84                j=j+185            newfiles.append(os.path.basename(file.name))86        except Exception as e:87            print(e)88            if (str(e) == "Cannot read an empty file"):89                return "Please Wait! Files are uploading, Try again Later!"90            corrupt.append(os.path.basename(file.name))91        i = i+192        print(i)93 94    pindex = pc.Index(brainName)95 96    # try:97    #     print(pindex.describe_index_stats())98    # except Exception as e:99    #     print(e)100 101    vector_store = PineconeVectorStore(pinecone_index=pindex)102    storage_context = StorageContext.from_defaults(vector_store=vector_store)103    service_context = ServiceContext.from_defaults(chunk_size=512,chunk_overlap=20)104    index = VectorStoreIndex.from_documents([], storage_context=storage_context,service_context=service_context)105 106    client = MongoClient(mongo_key)107    db = client['nbrain']108    collection = db['files']109 110    # for prevfile in newfiles:111    #     index.delete_ref_doc(prevfile, delete_from_docstore=True)112    #     delete_file(collection,brainName,prevfile)113 114    uploadedFiles=[]115    j = 1116    117    fileProcessing=documents[0].doc_id118    for doc in documents:119        try:120            index.insert(doc)121            if doc.doc_id not in uploadedFiles:122                print(doc.doc_id)123                add_file(collection,brainName,doc.doc_id)124                uploadedFiles.append(doc.doc_id)125                print(j)126                j = j+1127            if(fileProcessing!=doc.doc_id):128                # yield fileProcessing129                fileProcessing=doc.doc_id130        except Exception as e:131            if doc.doc_id not in corrupt:132                corrupt.append(doc.doc_id)133            print("ERROR : "+str(e))134 135 136    print("Brain Updated: "+brainName)137    # try:138    #     print(pindex.describe_index_stats())139    # except Exception as e:140    #     print(e)141 142    print(corrupt)143 144    if (len(corrupt) > 0):145        print("Brain Updated! Below files are corrupt/unformatted, and not added to the brain. "+ str(corrupt))146        return """Brain Updated! 147        Below files are corrupt/unformatted, and not added to the brain.148         """ + str(corrupt)149 150    return "updated"151