puuuush

structure change

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()