fix

indexer.py

@@ -1,44 +1,73 @@
-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
+import json
+import aiohttp
+import asyncio
+import aiofiles
+import urllib.request
+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")
 
-load_dotenv()
+download_progress = {}
 
-def index_repository(token, repo, current_path="", proxies=None):
+async def get_system_proxies():
+    """
+    Retrieves the system's HTTP and HTTPS proxies.
+
+    Returns:
+        dict: A dictionary containing the proxies.
+    """
     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
+        proxies = urllib.request.getproxies()
+        print("System proxies:", proxies)
+        return {
+            "http": proxies.get("http"),
+            "https": proxies.get("http")
+        }
     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()
-    
+        print(f"Error getting system proxies: {e}")
+        return {}
+
+async def get_file_structure(repo, token, path="", proxies=None):
+    """
+    Fetches the file structure of a specified Hugging Face repository.
+
+    Args:
+        repo (str): The name of the repository.
+        token (str): The authorization token for the request.
+        path (str, optional): The specific path in the repository. Defaults to "".
+        proxies (dict, optional): The proxies to use for the request. Defaults to None.
+
+    Returns:
+        list: A list of file structure information.
+    """
+    api_url = f"https://huggingface.co/api/models/{repo}/tree/main/{path}"
+    headers = {'Authorization': f'Bearer {token}'}
+    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 []
+
+async def write_file_structure_to_json(file_structure, file_path):
+    """
+    Writes the file structure to a JSON file.
+
+    Args:
+        file_structure (list): The file structure data.
+        file_path (str): The path where the JSON file will be saved.
+    """
     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}")
+        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}")