diff --git a/main.py b/main.py index 9c2c0ef..4b2dbde 100644 --- a/main.py +++ b/main.py @@ -14,13 +14,14 @@ import pgvec_conn import redis_conn from tqdm import tqdm -count = 0 - -def retrieve_file_contents(path): +def retrieve_file_contents(path, pPullFromS3): + pulled_bucket_data = "C:\\MINIO_TEMP_DATA" #Minio Connection minioClient = Minio( "localhost:9000", + access_key= "PVFOeJbx87rQyi0WXF1X", + secret_key = "Am8Cd9auYGbEGuEXfJtnWEPsMwJCx9N58NCNHCgs", secure=False, ) @@ -32,6 +33,15 @@ def retrieve_file_contents(path): "mp3": lambda path: lib.extract_mp3_content(path), } + if pPullFromS3: + if not os.path.exists(pulled_bucket_data): + os.makedirs(pulled_bucket_data) + file_objects = minioClient.list_objects("datafiles") + for file in file_objects: + local_path = os.path.join(pulled_bucket_data, file.object_name) + minioClient.fget_object("datafiles", file.object_name, local_path) + path = pulled_bucket_data + lib.read_files(path, files := []) contents = [] @@ -39,27 +49,31 @@ def retrieve_file_contents(path): print("parsing files...") for file in tqdm(files): content = file_extraction_functions[file[0]](file[1]) - with open(f"./{file[2]}.txt", "w", encoding="UTF-8") as future_s3_file: - future_s3_file.writelines(file[1] + "\n" + content) - future_s3_file.flush() - #future_s3_file.close() + if not pPullFromS3: + with open(f"./{file[2]}.txt", "w", encoding="UTF-8") as future_s3_file: + future_s3_file.writelines(content) + future_s3_file.flush() + #future_s3_file.close() - minioClient.fput_object( - bucket_name="datafiles", - object_name=f"{file[2]}.txt", - file_path=f"./{file[2]}.txt", - ) + minioClient.fput_object( + bucket_name="datafiles", + object_name=f"{file[2]}.txt", + file_path=f"./{file[2]}.txt", + ) - os.remove(f"./{file[2]}.txt") + os.remove(f"./{file[2]}.txt") contents.append({ "type": file[0], - "path": file[1], + "path": file[1] if not pPullFromS3 else "Diese Datei kommt ursprünglich nicht von dir, sondern aus einem S3-Bucket. ", "filename": file[2], "content": content }) + #if os.path.exists(pulled_bucket_data): + # os.remove(pulled_bucket_data) + return contents @@ -116,15 +130,17 @@ def insert_data_mongo(id, filepath,pChunk): if result.acknowledged: collection.create_index("doc_id") -def reload_files(): - # path = input("Please provider path to folder: ") +def reload_files(pPullFromS3): pgvec_conn.flush_pg() redis_conn.flush_redis() mongo_conn.flush_mongo() + if not (pPullFromS3): + provided_path = input("Please Provide Full Qualified Path: ") + contents = retrieve_file_contents(provided_path, pPullFromS3) + else: + provided_path = "" + contents = retrieve_file_contents(provided_path, pPullFromS3) - provided_path = input("Please Provide Full Qualified Path: ") - - contents = retrieve_file_contents(provided_path) create_embeddings(contents) @@ -143,10 +159,11 @@ def prompt_cycle(): def get_user_action(): - user_action = input("Reload Files (r), Add Files (a), Prompt (p): ") + user_action = input("Reload Files (r), Load Files From S3 (l), Add Files (a), Prompt (p): ") { - "r": lambda: reload_files(), + "r": lambda: reload_files(False), + "l": lambda: reload_files(True), "a": lambda: add_files(), "p": lambda: prompt_cycle() }[user_action]()