mirror of https://github.com/arc53/DocsGPT
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
import os
|
|
import datetime
|
|
from flask import Blueprint, request, send_from_directory
|
|
from pymongo import MongoClient
|
|
from werkzeug.utils import secure_filename
|
|
|
|
|
|
from application.core.settings import settings
|
|
mongo = MongoClient(settings.MONGO_URI)
|
|
db = mongo["docsgpt"]
|
|
conversations_collection = db["conversations"]
|
|
vectors_collection = db["vectors"]
|
|
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
|
|
internal = Blueprint('internal', __name__)
|
|
@internal.route("/api/download", methods=["get"])
|
|
def download_file():
|
|
user = secure_filename(request.args.get("user"))
|
|
job_name = secure_filename(request.args.get("name"))
|
|
filename = secure_filename(request.args.get("file"))
|
|
save_dir = os.path.join(current_dir, settings.UPLOAD_FOLDER, user, job_name)
|
|
return send_from_directory(save_dir, filename, as_attachment=True)
|
|
|
|
|
|
|
|
@internal.route("/api/upload_index", methods=["POST"])
|
|
def upload_index_files():
|
|
"""Upload two files(index.faiss, index.pkl) to the user's folder."""
|
|
if "user" not in request.form:
|
|
return {"status": "no user"}
|
|
user = secure_filename(request.form["user"])
|
|
if "name" not in request.form:
|
|
return {"status": "no name"}
|
|
job_name = secure_filename(request.form["name"])
|
|
save_dir = os.path.join(current_dir, "indexes", user, job_name)
|
|
if settings.VECTOR_STORE == "faiss":
|
|
if "file_faiss" not in request.files:
|
|
print("No file part")
|
|
return {"status": "no file"}
|
|
file_faiss = request.files["file_faiss"]
|
|
if file_faiss.filename == "":
|
|
return {"status": "no file name"}
|
|
if "file_pkl" not in request.files:
|
|
print("No file part")
|
|
return {"status": "no file"}
|
|
file_pkl = request.files["file_pkl"]
|
|
if file_pkl.filename == "":
|
|
return {"status": "no file name"}
|
|
# saves index files
|
|
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir)
|
|
file_faiss.save(os.path.join(save_dir, "index.faiss"))
|
|
file_pkl.save(os.path.join(save_dir, "index.pkl"))
|
|
# create entry in vectors_collection
|
|
vectors_collection.insert_one(
|
|
{
|
|
"user": user,
|
|
"name": job_name,
|
|
"language": job_name,
|
|
"location": save_dir,
|
|
"date": datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
|
|
"model": settings.EMBEDDINGS_NAME,
|
|
"type": "local",
|
|
}
|
|
)
|
|
return {"status": "ok"} |