Error handle and thread safety

server/privateGPT.py

@@ -1,4 +1,4 @@
-from flask import Flask,jsonify, render_template, flash, redirect, url_for, Markup, request
+from flask import Flask, jsonify, render_template, flash, redirect, url_for, Markup, request
 from flask_cors import CORS
 from dotenv import load_dotenv
 from langchain.chains import RetrievalQA
@@ -44,7 +44,6 @@
 model_n_ctx = os.environ.get('MODEL_N_CTX')
 llm = None
 
-from constants import CHROMA_SETTINGS
 
 class MyElmLoader(UnstructuredEmailLoader):
     """Wrapper to fallback to text/plain when default does not work"""
@@ -57,7 +56,7 @@ def load(self) -> List[Document]:
             except ValueError as e:
                 if 'text/html content not found in email' in str(e):
                     # Try plain text
-                    self.unstructured_kwargs["content_source"]="text/plain"
+                    self.unstructured_kwargs["content_source"] = "text/plain"
                     doc = UnstructuredEmailLoader.load(self)
                 else:
                     raise
@@ -107,60 +106,81 @@ def load_documents(source_dir: str) -> List[Document]:
         )
     return [load_single_document(file_path) for file_path in all_files]
 
+
 @app.route('/ingest', methods=['GET'])
 def ingest_data():
-    # Load environment variables
+    #  Load environment variables
     persist_directory = os.environ.get('PERSIST_DIRECTORY')
     source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
     embeddings_model_name = os.environ.get('EMBEDDINGS_MODEL_NAME')
 
-    # Load documents and split in chunks
-    print(f"Loading documents from {source_directory}")
-    chunk_size = 500
-    chunk_overlap = 50
-    documents = load_documents(source_directory)
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    texts = text_splitter.split_documents(documents)
-    print(f"Loaded {len(documents)} documents from {source_directory}")
-    print(f"Split into {len(texts)} chunks of text (max. {chunk_size} characters each)")
+    #  Load documents and split in chunks
+    try:
+        print(f"Loading documents from {source_directory}")
+        chunk_size = 500
+        chunk_overlap = 50
+        documents = load_documents(source_directory)
+    except Exception as e:
+        return jsonify(error=f"Failed to load documents: {str(e)}"), 500
+
+    try:
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+        texts = text_splitter.split_documents(documents)
+        print(f"Loaded {len(documents)} documents from {source_directory}")
+        print(
+            f"Split into {len(texts)} chunks of text (max. {chunk_size} characters each)")
+    except Exception as e:
+        return jsonify(error=f"Failed to split texts: {str(e)}"), 500
 
     # Create embeddings
-    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
-
+    try:
+        embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
+    except Exception as e:
+        return jsonify(error=f"Failed to create embeddings: {str(e)}"), 500
+
     # Create and store locally vectorstore
-    db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
-    db.persist()
-    db = None
+    try:
+        db = Chroma.from_documents(
+            texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
+        db.persist()
+        db = None
+    except Exception as e:
+        return jsonify(error=f"Failed to create and store vectorstore: {str(e)}"), 500
+
     return jsonify(response="Success")
-
+
+
 @app.route('/get_answer', methods=['POST'])
 def get_answer():
     query = request.json
     embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
-    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
+    db = Chroma(persist_directory=persist_directory,
+                embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
     retriever = db.as_retriever()
-    if llm==None:
-        return "Model not downloaded", 400    
-    qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
-    if query!=None and query!="":
+    if llm == None:
+        return "Model not downloaded", 400
+    qa = RetrievalQA.from_chain_type(
+        llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
+    if query != None and query != "":
         res = qa(query)
         answer, docs = res['result'], res['source_documents']
-        
-        source_data =[]
+
+        source_data = []
         for document in docs:
-             source_data.append({"name":document.metadata["source"]})
+            source_data.append({"name": document.metadata["source"]})
 
-        return jsonify(query=query,answer=answer,source=source_data)
+        return jsonify(query=query, answer=answer, source=source_data)
 
-    return "Empty Query",400
+    return "Empty Query", 400
 
 
 @app.route('/upload_doc', methods=['POST'])
 def upload_doc():
-    
+
     if 'document' not in request.files:
         return jsonify(response="No document file found"), 400
-    
+
     document = request.files['document']
     if document.filename == '':
         return jsonify(response="No selected file"), 400
@@ -171,42 +191,60 @@ def upload_doc():
 
     return jsonify(response="Document upload successful")
 
+
 @app.route('/download_model', methods=['GET'])
 def download_and_save():
-    url = 'https://gpt4all.io/models/ggml-gpt4all-j-v1.3-groovy.bin'  # Specify the URL of the resource to download
-    filename = 'ggml-gpt4all-j-v1.3-groovy.bin'  # Specify the name for the downloaded file
-    models_folder = 'models'  # Specify the name of the folder inside the Flask app root
+    try:
+        url = 'https://gpt4all.io/models/ggml-gpt4all-j-v1.3-groovy.bin'
+        filename = 'ggml-gpt4all-j-v1.3-groovy.bin'
+        models_folder = 'models'
+        n_ctx = 1024  # Replace this with correct value
+
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raises an exception for 4xx and 5xx status codes
+    except requests.RequestException as e:
+        return jsonify(error=f"Failed to download the model: {str(e)}"), 500
 
-    if not os.path.exists(models_folder):
-        os.makedirs(models_folder)
-    response = requests.get(url,stream=True)
     total_size = int(response.headers.get('content-length', 0))
     bytes_downloaded = 0
     file_path = f'{models_folder}/{filename}'
+
     if os.path.exists(file_path):
-        return jsonify(response="Download completed")
-
-    with open(file_path, 'wb') as file:
-        for chunk in response.iter_content(chunk_size=4096):
-            file.write(chunk)
-            bytes_downloaded += len(chunk)
-            progress = round((bytes_downloaded / total_size) * 100, 2)
-            print(f'Download Progress: {progress}%')
-    global llm
-    callbacks = [StreamingStdOutCallbackHandler()]
-    llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
+        return jsonify(response="Download already completed")
+
+    try:
+        with open(file_path, 'wb') as file:
+            for chunk in response.iter_content(chunk_size=4096):
+                file.write(chunk)
+                bytes_downloaded += len(chunk)
+                progress = round((bytes_downloaded / total_size) * 100, 2)
+                print(f'Download Progress: {progress}%')
+    except IOError as e:
+        return jsonify(error=f"Failed to save the model: {str(e)}"), 500
+
+    try:
+        callbacks = [StreamingStdOutCallbackHandler()]
+        llm = GPT4All(model=file_path, n_ctx=n_ctx, backend='gptj',
+                      callbacks=callbacks, verbose=False)
+    except Exception as e:
+        return jsonify(error=f"Failed to load the model: {str(e)}"), 500
+
     return jsonify(response="Download completed")
 
+
 def load_model():
-    filename = 'ggml-gpt4all-j-v1.3-groovy.bin'  # Specify the name for the downloaded file
+    # Specify the name for the downloaded file
+    filename = 'ggml-gpt4all-j-v1.3-groovy.bin'
     models_folder = 'models'  # Specify the name of the folder inside the Flask app root
     file_path = f'{models_folder}/{filename}'
     if os.path.exists(file_path):
         global llm
         callbacks = [StreamingStdOutCallbackHandler()]
-        llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
+        llm = GPT4All(model=model_path, n_ctx=model_n_ctx,
+                      backend='gptj', callbacks=callbacks, verbose=False)
+
 
 if __name__ == "__main__":
-  load_model()
-  print("LLM0", llm)
-  app.run(host="0.0.0.0", debug = False)
+    load_model()
+    print("LLM0", llm)
+    app.run(host="0.0.0.0", debug=False)