0.0.2.9 V Beta

LoadBalancer.py

@@ -14,7 +14,7 @@ download_progress = {}
 
 class LoadBalancer:
     def __init__(self, cache_dir, index_file, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
-        self.version = "0.0.2.7 V Beta"
+        self.version = "0.0.2.9 V Beta"
         self.instances = []
         self.instances_health = {}
         self.polling_interval = polling_interval

hf_scrapper.py

@@ -1,17 +1,19 @@
 import os
-import requests
 import json
+import aiohttp
+import aiofiles
+import asyncio
 import urllib.request
-import time
-from requests.exceptions import RequestException
-from tqdm import tqdm
+from aiohttp import ClientSession, ClientTimeout
+from aiohttp.client_exceptions import ClientError
+from tqdm.asyncio import tqdm
 
 CACHE_DIR = os.getenv("CACHE_DIR")
 CACHE_JSON_PATH = os.path.join(CACHE_DIR, "cached_films.json")
 
 download_progress = {}
 
-def get_system_proxies():
+async def get_system_proxies():
     """
     Retrieves the system's HTTP and HTTPS proxies.
 
@@ -29,7 +31,7 @@ def get_system_proxies():
         print(f"Error getting system proxies: {e}")
         return {}
 
-def get_file_structure(repo, token, path="", proxies=None):
+async def get_file_structure(repo, token, path="", proxies=None):
     """
     Fetches the file structure of a specified Hugging Face repository.
 
@@ -44,16 +46,18 @@ def get_file_structure(repo, token, path="", proxies=None):
     """
     api_url = f"https://huggingface.co/api/models/{repo}/tree/main/{path}"
     headers = {'Authorization': f'Bearer {token}'}
-    print(f"Fetching file structure from URL: {api_url} with proxies: {proxies}")
-    try:
-        response = requests.get(api_url, headers=headers, proxies=proxies)
-        response.raise_for_status()
-        return response.json()
-    except RequestException as e:
-        print(f"Error fetching file structure: {e}")
-        return []
+    timeout = ClientTimeout(total=10)
+    async with ClientSession(timeout=timeout) as session:
+        print(f"Fetching file structure from URL: {api_url} with proxies: {proxies}")
+        try:
+            async with session.get(api_url, headers=headers, proxy=proxies.get("http")) as response:
+                response.raise_for_status()
+                return await response.json()
+        except ClientError as e:
+            print(f"Error fetching file structure: {e}")
+            return []
 
-def write_file_structure_to_json(file_structure, file_path):
+async def write_file_structure_to_json(file_structure, file_path):
     """
     Writes the file structure to a JSON file.
 
@@ -62,12 +66,8 @@ def write_file_structure_to_json(file_structure, file_path):
         file_path (str): The path where the JSON file will be saved.
     """
     try:
-        with open(file_path, 'w') as json_file:
-            json.dump(file_structure, json_file, indent=2)
+        async with aiofiles.open(file_path, 'w') as json_file:
+            await json_file.write(json.dumps(file_structure, indent=2))
         print(f'File structure written to {file_path}')
     except IOError as e:
-        print(f"Error writing file structure to JSON: {e}")
-
-if __name__ == "__main__":
-    file_url = "https://huggingface.co/Unicone-Studio/jellyfin_media/resolve/main/films/Funky%20Monkey%202004/Funky%20Monkey%20(2004)%20Web-dl%201080p.mp4"
-    token = os.getenv("TOKEN")
+        print(f"Error writing file structure to JSON: {e}")

indexer.py

@@ -1,4 +1,5 @@
 import json
+import logging
 from hf_scrapper import get_system_proxies, get_file_structure, write_file_structure_to_json
 from dotenv import load_dotenv
 import os
@@ -6,26 +7,38 @@ import os
 load_dotenv()
 
 def index_repository(token, repo, current_path="", proxies=None):
-    file_structure = get_file_structure(repo, token, current_path, proxies)
-    full_structure = []
-    for item in file_structure:
-        if item['type'] == 'directory':
-            sub_directory_structure = index_repository(token, repo, item['path'], proxies)
-            full_structure.append({
-                "type": "directory",
-                "path": item['path'],
-                "contents": sub_directory_structure
-            })
-        else:
-            full_structure.append(item)
-    return full_structure
+    try:
+        file_structure = get_file_structure(repo, token, current_path, proxies)
+        full_structure = []
+        for item in file_structure:
+            if item['type'] == 'directory':
+                sub_directory_structure = index_repository(token, repo, item['path'], proxies)
+                full_structure.append({
+                    "type": "directory",
+                    "path": item['path'],
+                    "contents": sub_directory_structure
+                })
+            else:
+                full_structure.append(item)
+        return full_structure
+    except Exception as e:
+        logging.error(f"Error indexing repository: {e}")
+        raise
 
 def indexer():
     token = os.getenv("TOKEN")
     repo = os.getenv("REPO")
     output_path = os.getenv("INDEX_FILE")
+    
+    if not token or not repo or not output_path:
+        logging.error("Environment variables TOKEN, REPO, or INDEX_FILE are not set.")
+        return
+
     proxies = get_system_proxies()
-    full_structure = index_repository(token, repo, "", proxies)
-    write_file_structure_to_json(full_structure, output_path)
-    print(f"Full file structure for repository '{repo}' has been indexed and saved to {output_path}")
-    
+    
+    try:
+        full_structure = index_repository(token, repo, "", proxies)
+        write_file_structure_to_json(full_structure, output_path)
+        logging.info(f"Full file structure for repository '{repo}' has been indexed and saved to {output_path}")
+    except Exception as e:
+        logging.error(f"Error during indexing: {e}")

requirements.txt

@@ -1,5 +1,7 @@
 fastapi
 aiofiles
+aiohttp
+tqdm
 uvicorn
 requests
 python-dotenv