diff --git a/chunker.py b/chunker.py index 61c2792..ec031c3 100644 --- a/chunker.py +++ b/chunker.py @@ -18,6 +18,8 @@ def generate_chunks_semantic(content): nodes = splitter.build_semantic_nodes_from_documents(document) + os.remove("./temp.txt") + output = [] for i, node in enumerate(nodes): diff --git a/main.py b/main.py index 4b2dbde..c38008f 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ import ollama import pymupdf import os +import shutil from minio import Minio from minio.error import S3Error import chunker @@ -15,7 +16,7 @@ import redis_conn from tqdm import tqdm def retrieve_file_contents(path, pPullFromS3): - pulled_bucket_data = "C:\\MINIO_TEMP_DATA" + pulled_bucket_data = os.path.join(".", "MINIO_TEMP_DATA") #Minio Connection minioClient = Minio( @@ -34,13 +35,7 @@ def retrieve_file_contents(path, pPullFromS3): } 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 + path = pull_files_from_s3(pulled_bucket_data, minioClient, ) lib.read_files(path, files := []) @@ -50,19 +45,7 @@ def retrieve_file_contents(path, pPullFromS3): for file in tqdm(files): content = file_extraction_functions[file[0]](file[1]) 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", - ) - - os.remove(f"./{file[2]}.txt") - + push_files_into_s3(file, content, minioClient) contents.append({ "type": file[0], @@ -71,8 +54,7 @@ def retrieve_file_contents(path, pPullFromS3): "content": content }) - #if os.path.exists(pulled_bucket_data): - # os.remove(pulled_bucket_data) + shutil.rmtree(pulled_bucket_data) return contents @@ -86,6 +68,7 @@ def create_embeddings(pContent): host="localhost", port="5555" ) + cur = conn.cursor() create_table_query = ''' @@ -130,6 +113,29 @@ def insert_data_mongo(id, filepath,pChunk): if result.acknowledged: collection.create_index("doc_id") + +def push_files_into_s3(file, content, client): + with open(f"./{file[2]}.txt", "w", encoding="UTF-8") as future_s3_file: + future_s3_file.writelines(content) + future_s3_file.flush() + + client.fput_object( + bucket_name="datafiles", + object_name=f"{file[2]}.txt", + file_path=f"./{file[2]}.txt", + ) + os.remove(f"./{file[2]}.txt") + +def pull_files_from_s3(pPath, client): + if not os.path.exists(pPath): + os.makedirs(pPath) + file_objects = client.list_objects("datafiles") + for file in file_objects: + local_path = os.path.join(pPath, file.object_name) + client.fget_object("datafiles", file.object_name, local_path) + return pPath + + def reload_files(pPullFromS3): pgvec_conn.flush_pg() redis_conn.flush_redis()