Split examples, add upload to flask api

Author: logan-markewich

README.md

@@ -18,32 +18,32 @@ pip install -r requirements.txt
 
 
 ## What is included?
-There are two main example files
-- flask_demo.py (localhost:5601)
-  - `python ./flask_demo.py`
+There are two main example folders
+- flask (runs on localhost:5601/2)
+  - `sh launch_app.py `
   - creates a simple api that loads the text from the documents folder
-  - The "/query" endpoint accepts requests that contain a "text" parameter, which is used to query the index
-  - resturns string response containing the query answer
+  - the "/query" endpoint accepts requests that contain a "text" parameter, which is used to query the index
+  - the "/upload" endpoint is a POST endpoint that inserts an attached text file into the index
+  - the index is managed by a seperate server using locks, since inserting a document is a mutable operation and flask is multithreaded
+  - I strongly recommend using a tool like [Postman](https://www.postman.com/downloads/) to test the api - there are example screenshots using postman in the `postman_examples` folder
 
-- streamlit_demo.py (localhost:8501)
+- streamlit (runs on localhost:8501)
   - `streamlit run streamlit_demo.py`
   - creates a simple UI using streamlit
   - loads text from the documents folder (using `st.cache_resource`, so it only loads once)
   - provides an input text-box and a button to run the query
-  - The string response is displayed after it finishes
-  - Want to see this example in action? Check it out [here](https://llama-index.streamlit.app/)
+  - the string response is displayed after it finishes
+  - want to see this example in action? Check it out [here](https://llama-index.streamlit.app/)
 
 
 ## Docker
-Using the local `Dockerfile`, you can run `docker build -t my_tag_name .` to build a python3.11-slim docker image. It ends up being about 980MB.
+Each example contains a `Dockerfile`. You can run `docker build -t my_tag_name .` to build a python3.11-slim docker image inside your desired folder. It ends up being about 600MB.
 
-Inside the `Dockerfile`, you can comment the app you want to run and the port you want to expose, based on if you want streamlit or flask.
+Inside the `Dockerfile`, certain ports are exposed based on which ports the examples need.
 
 When running the image, be sure to include the -p option to access the proper ports (8501, or 5601).
 
 
 ## Contributing
 
-I welcome any suggestions or PRs! If we start adding more examples, it might be good to refactor to have a folder per example type
-
-
+I welcome any suggestions or PRs, or more examples!

documents/.paul_graham_essay.txt.swp

@@ -6,10 +6,6 @@ COPY . .
 
 RUN pip install -r requirements.txt && pip cache purge
 
-# Streamlit
-#CMD ["streamlit", "run", "streamlit_demo.py"]
-#EXPOSE 8501
-
 # Flask
-CMD ["python", "flask_demo.py"]
+CMD ["sh", "launch_app.sh"]
 EXPOSE 5601

Dockerfile renamed to flask/Dockerfile

@@ -0,0 +1,48 @@
+import os
+from multiprocessing.managers import BaseManager
+from flask import Flask, request
+from werkzeug.utils import secure_filename
+
+app = Flask(__name__)
+
+# initialize manager connection
+# NOTE: you might want to handle the password in a less hardcoded way
+manager = BaseManager(('', 5602), b'password')
+manager.register('query_index')
+manager.register('insert_into_index')
+manager.connect()
+
+
+@app.route("/query", methods=["GET"])
+def query_index():
+    global manager
+    query_text = request.args.get("text", None)
+    if query_text is None:
+        return "No text found, please include a ?text=blah parameter in the URL", 400
+    
+    response = manager.query_index(query_text)
+    return str(response), 200
+
+
+@app.route("/upload", methods=["POST"])
+def upload_file():
+    if 'file' not in request.files:
+        return "Please send a POST request with a file", 400
+
+    uploaded_file = request.files["file"]
+    filename = secure_filename(uploaded_file.filename)
+    if request.form.get("filename_as_doc_id", None) is not None:
+        manager.insert_into_index(uploaded_file.read().decode(), doc_id=filename)
+    else:
+        manager.insert_into_index(uploaded_file.read().decode())
+
+    return "File inserted!", 200
+
+
+@app.route("/")
+def home():
+    return "Hello, World! Welcome to the llama_index docker image!"
+
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5601)

documents/paul_graham_essay.txt renamed to flask/documents/paul_graham_essay.txt

@@ -0,0 +1,62 @@
+import os
+
+# NOTE: for local testing only, do NOT deploy with your key hardcoded
+os.environ['OPENAI_API_KEY'] = "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"
+
+from multiprocessing import Lock
+from multiprocessing.managers import BaseManager
+from llama_index import SimpleDirectoryReader, GPTSimpleVectorIndex, Document
+
+index = None
+lock = Lock()
+
+index_name = "./index.json"
+documents_folder = "./documents"
+
+
+def initialize_index():
+    """Create a new global index, or load one from the pre-set path."""
+    global index
+
+    with lock:
+        if os.path.exists(index_name):
+            index = GPTSimpleVectorIndex.load_from_disk(index_name)
+        else:
+            documents = SimpleDirectoryReader(documents_folder).load_data()
+            index = GPTSimpleVectorIndex(documents)
+            index.save_to_disk(index_name)
+
+
+def query_index(query_text):
+    """Query the global index."""
+    global index
+    response = index.query(query_text)
+    return str(response).strip()
+
+
+def insert_into_index(doc_text, doc_id=None):
+    """Insert new document into global index."""
+    global index
+    print(doc_text)
+    document = Document(doc_text)
+    if doc_id is not None:
+        document.doc_id = doc_id
+
+    with lock:
+        index.insert(document)
+
+
+if __name__ == "__main__":
+    # init the global index
+    print("initializing index...")
+    initialize_index()
+
+    # setup server
+    # NOTE: you might want to handle the password in a less hardcoded way
+    manager = BaseManager(('', 5602), b'password')
+    manager.register('query_index', query_index)
+    manager.register('insert_into_index', insert_into_index)
+    server = manager.get_server()
+
+    print("starting server...")
+    server.serve_forever()

flask/flask_demo.py

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# start backend index server
+python ./index_server.py &
+
+# wait for the server to start - if creating a brand new huge index, on startup, increase this further
+sleep 60
+
+# start the flask server
+python ./flask_demo.py