From 901cde448060420ce1b79a03e860444ace4e8b86 Mon Sep 17 00:00:00 2001
From: Mykhailo Moskalets <69348201+13MrBlackCat13@users.noreply.github.com>
Date: Thu, 9 Nov 2023 16:02:10 +0100
Subject: [PATCH] Update main.py

---
 main.py | 38 ++++++++++++++++++--------------------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/main.py b/main.py
index c4d27cc..a854311 100644
--- a/main.py
+++ b/main.py
@@ -4,17 +4,13 @@
 import shutil
 from dotenv import load_dotenv
 
-# Загрузка переменных окружения
-load_dotenv()
-API_URL = os.getenv("API_URL", "http://localhost:5000/api/emailcategory/eml")
-SUPER_TOKEN = os.getenv("SUPER_TOKEN", "devtest")
-TOKEN_GROUP = os.getenv("TOKEN_GROUP", "group1")
-UNSORTED_DIR = os.getenv("UNSORTED_DIR", "Unsorted")
-PROBABILITY_THRESHOLD = float(os.getenv("PROBABILITY_THRESHOLD", "0.6"))
+# Configuration
+from config import Config
 
-# Настройка логгера
+# Logger setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
+# Utility functions
 def read_eml_file(path):
     try:
         with open(path, "rb") as f:
@@ -23,6 +19,16 @@ def read_eml_file(path):
         logging.error(f"Error reading file {path}: {e}")
         return None
 
+def move_file_to_category(file_path, file_name, category):
+    sorted_folder = os.path.join("Sorted", category)
+    if not os.path.exists(sorted_folder):
+        os.makedirs(sorted_folder)
+
+    new_path = os.path.join(sorted_folder, file_name)
+    shutil.move(file_path, new_path)
+    logging.info(f"Moved {file_name} to {sorted_folder}")
+
+# API interaction
 def process_file(file_path, file_name):
     eml_content = read_eml_file(file_path)
     if eml_content is None:
@@ -31,7 +37,7 @@ def process_file(file_path, file_name):
     data = {"model_name": "all", "text_only": "true"}
 
     try:
-        response = requests.post(API_URL, files={"file": (file_name, eml_content, "message/rfc822")}, params=data)
+        response = requests.post(Config.API_URL, files={"file": (file_name, eml_content, "message/rfc822")}, params=data)
         if response.status_code == 200:
             handle_response(response.json(), file_path, file_name)
         else:
@@ -44,21 +50,13 @@ def handle_response(response_data, file_path, file_name):
     category_probabilities = {k: v for k, v in avg_probabilities.items() if k != 'Model'}
 
     for category, probability in category_probabilities.items():
-        if float(probability) >= PROBABILITY_THRESHOLD:
+        if float(probability) >= Config.PROBABILITY_THRESHOLD:
             move_file_to_category(file_path, file_name, category)
             break
 
-def move_file_to_category(file_path, file_name, category):
-    sorted_folder = os.path.join("Sorted", category)
-    if not os.path.exists(sorted_folder):
-        os.makedirs(sorted_folder)
-
-    new_path = os.path.join(sorted_folder, file_name)
-    shutil.move(file_path, new_path)
-    logging.info(f"Moved {file_name} to {sorted_folder}")
-
+# Main execution
 def main():
-    for root, dirs, files in os.walk(UNSORTED_DIR):
+    for root, dirs, files in os.walk(Config.UNSORTED_DIR):
         for file in files:
             if file.endswith(".eml"):
                 process_file(os.path.join(root, file), file)