import os import json import server from aiohttp import web from . import downloader as dl # Get the absolute path to the 'data' directory # __file__ is the path to the current script (api.py) # os.path.dirname(__file__) is the directory of the current script (modules) # os.path.join(..., '..', 'data') goes up one level and then into 'data' DATA_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "data")) DANBOORU_PREFIX = "danbooru" E621_PREFIX = "e621" TAGS_SUFFIX = "tags" COOCCURRENCE_SUFFIX = "tags_cooccurrence" def get_csv_file_status(): """ Returns a dictionary of csv file statuses. """ data = { DANBOORU_PREFIX: { "base_tags": False, "extra_tags": [], "base_cooccurrence": False, "extra_cooccurrence": [], }, E621_PREFIX: { "base_tags": False, "extra_tags": [], "base_cooccurrence": False, "extra_cooccurrence": [], }, } for prefix in [DANBOORU_PREFIX, E621_PREFIX]: base_tags_file = f"{prefix}_{TAGS_SUFFIX}.csv" base_cooccurrence_file = f"{prefix}_{COOCCURRENCE_SUFFIX}.csv" tags_base_exists = os.path.exists(os.path.join(DATA_DIR, base_tags_file)) cooccurrence_base_exists = os.path.exists(os.path.join(DATA_DIR, base_cooccurrence_file)) tags_extra_files = [] cooccurrence_extra_files = [] all_csv_files = [f for f in os.listdir(DATA_DIR) if f.startswith(prefix) and f.endswith(".csv")] # Create extra CSV files list for filename in all_csv_files: if filename in [base_tags_file, base_cooccurrence_file]: continue # Skip base files if COOCCURRENCE_SUFFIX in filename.lower(): cooccurrence_extra_files.append(filename) elif TAGS_SUFFIX in filename.lower(): tags_extra_files.append(filename) data[prefix] = { "base_tags": tags_base_exists, "extra_tags": tags_extra_files, "base_cooccurrence": cooccurrence_base_exists, "extra_cooccurrence": cooccurrence_extra_files, } # Return the lists of extra files return data def get_last_check_time_from_metadata(): """ Helper function to get the last remote check timestamp from csv_meta.json. Returns the timestamp string if found, None otherwise. """ try: if not os.path.exists(dl.CSV_META_FILE): return None with open(dl.CSV_META_FILE, "r", encoding="utf-8") as f: metadata = json.load(f) datasets = metadata.get("hf_datasets", []) if datasets and len(datasets) > 0: return datasets[0].get("last_remote_check_timestamp") return None except (IOError, json.JSONDecodeError) as e: print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}") return None # --- API Endpoints --- @server.PromptServer.instance.routes.get("/autocomplete-plus/csv") async def get_csv_list(_request): """ Returns CSV file status. base files: file exists boolean extra files: count of extra files """ csv_file_status = get_csv_file_status() response = { DANBOORU_PREFIX: { "base_tags": csv_file_status[DANBOORU_PREFIX]["base_tags"], "extra_tags": csv_file_status[DANBOORU_PREFIX]["extra_tags"], "base_cooccurrence": csv_file_status[DANBOORU_PREFIX]["base_cooccurrence"], "extra_cooccurrence": csv_file_status[DANBOORU_PREFIX]["extra_cooccurrence"], }, E621_PREFIX: { "base_tags": csv_file_status[E621_PREFIX]["base_tags"], "extra_tags": csv_file_status[E621_PREFIX]["extra_tags"], "base_cooccurrence": csv_file_status[E621_PREFIX]["base_cooccurrence"], "extra_cooccurrence": csv_file_status[E621_PREFIX]["extra_cooccurrence"], }, } # Print csv file status to the console for debugging print(f"""[Autocomplete-Plus] CSV file status: * Danbooru -> base: {response[DANBOORU_PREFIX]["base_tags"]}, extra: [{", ".join(response[DANBOORU_PREFIX]["extra_tags"])}] * E621 -> base: {response[E621_PREFIX]["base_tags"]}, extra: [{", ".join(response[E621_PREFIX]["extra_tags"])}]""") return web.json_response(response) @server.PromptServer.instance.routes.get("/autocomplete-plus/csv/{source}/{suffix}/base") async def get_base_tags_file(request): """ Returns the base tags CSV file. """ source = str(request.match_info["source"]) suffix = str(request.match_info["suffix"]) if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [ TAGS_SUFFIX, COOCCURRENCE_SUFFIX, ]: return web.json_response({"error": "Invalid tag source or suffix"}, status=400) file_path = os.path.join(DATA_DIR, f"{source}_{suffix}.csv") if not os.path.exists(file_path): return web.json_response({"error": "Base tags file not found"}, status=404) return web.FileResponse(file_path) @server.PromptServer.instance.routes.get("/autocomplete-plus/csv/{source}/{suffix}/extra/{index}") async def get_extra_tags_file(request): """ Returns the extra tags CSV file at the specified index. """ try: csv_file_status = get_csv_file_status() source = str(request.match_info["source"]) suffix = str(request.match_info["suffix"]) if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [ TAGS_SUFFIX, COOCCURRENCE_SUFFIX, ]: return web.json_response({"error": "Invalid tag source or suffix"}, status=400) index = int(request.match_info["index"]) if index < 0 or index >= len(csv_file_status[source][f"extra_{suffix}"]): return web.json_response({"error": "Invalid index"}, status=404) file_path = os.path.join(DATA_DIR, csv_file_status[source][f"extra_{suffix}"][index]) if not os.path.exists(file_path): return web.json_response({"error": "Extra tags file not found"}, status=404) return web.FileResponse(file_path) except ValueError: return web.json_response({"error": "Invalid index format"}, status=400) @server.PromptServer.instance.routes.post("/autocomplete-plus/csv/force-check-updates") async def force_check_csv_updates(request): """ Forces a check for CSV file updates from HuggingFace, ignoring cooldown. This allows users to manually trigger an update check at any time. """ try: print("[Autocomplete-Plus] Starting forced check for CSV updates from HuggingFace...") downloader = dl.Downloader() downloader.run_check_and_download(force_check=True) print("[Autocomplete-Plus] Forced check completed successfully.") # Get the updated last check time last_check_time = get_last_check_time_from_metadata() return web.json_response( { "success": True, "message": "Force check completed successfully", "last_check_time": last_check_time, } ) except Exception as e: print(f"[Autocomplete-Plus] Error during forced check: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) @server.PromptServer.instance.routes.get("/autocomplete-plus/csv/last-check-time") async def get_last_check_time(_request): """ Returns the last remote check timestamp from csv_meta.json. Returns null if the file doesn't exist or if there's an error reading it. """ try: if not os.path.exists(dl.CSV_META_FILE): return web.json_response({"last_check_time": None, "message": "csv_meta.json file not found"}) last_check_time = get_last_check_time_from_metadata() if last_check_time is not None: return web.json_response({"last_check_time": last_check_time}) else: return web.json_response({"last_check_time": None, "message": "No datasets found in metadata"}) except (IOError, json.JSONDecodeError) as e: print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}") return web.json_response({"last_check_time": None, "error": str(e)}, status=500)