diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..d8cd1df --- /dev/null +++ b/utils.py @@ -0,0 +1,38 @@ +import json +import hashlib +from pathlib import Path +from functools import wraps + +def with_disk_cache(subdir_name: str): + """ + A decorator to cache function results to a local JSON file. + The cache file is stored in `data//.json`. + """ + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + if not args or not isinstance(args[0], str): + key = hashlib.md5(str(args).encode()).hexdigest() + else: + val = args[0] + # If it's a simple ID, use it for readability. Otherwise hash it to avoid invalid filename characters. + if '/' not in val and '\\' not in val and len(val) < 50: + key = val + else: + key = hashlib.md5(val.encode()).hexdigest() + + cache_p = Path(__file__).parent / 'data' / subdir_name / f"{key}.json" + + if cache_p.is_file(): + try: + return json.loads(cache_p.read_text(encoding="utf-8")) + except Exception: + pass + + result = func(*args, **kwargs) + + cache_p.parent.mkdir(parents=True, exist_ok=True) + cache_p.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") + return result + return wrapper + return decorator diff --git a/utils_ai.py b/utils_ai.py index b51249f..dfe5e81 100644 --- a/utils_ai.py +++ b/utils_ai.py @@ -2,6 +2,8 @@ import os from pathlib import Path from openai import OpenAI +from utils import with_disk_cache + client = OpenAI() def _call_openai_with_prompt(prompt_id: str, prompt_version: str, input_text: str) -> str: @@ -49,6 +51,7 @@ def _call_openai_with_prompt(prompt_id: str, prompt_version: str, input_text: st return str(response) +@with_disk_cache('select_best_torrents') def select_best_torrents(torrents_text: str) -> str: """ Calls the OpenAI API to select the best torrent IDs using a predefined prompt. @@ -63,6 +66,7 @@ def select_best_torrents(torrents_text: str) -> str: ) +@with_disk_cache('generate_rename_mapping') def generate_rename_mapping(directory_text: str) -> dict[str, str]: """ Calls the OpenAI API to generate a renaming mapping for files diff --git a/utils_mteam.py b/utils_mteam.py index fc6e88a..14671a6 100644 --- a/utils_mteam.py +++ b/utils_mteam.py @@ -1,7 +1,10 @@ +import json import requests import tomllib from pathlib import Path +from utils import with_disk_cache + config = tomllib.loads(Path("config.toml").read_text()) def _get_mteam_headers() -> dict: @@ -19,6 +22,7 @@ def _get_mteam_headers() -> dict: } +@with_disk_cache('search_mteam_torrents') def search_mteam_torrents(imdb_url: str, page_number: int = 1, page_size: int = 100) -> dict: """ Search M-Team for torrents using IMDb URL. @@ -47,6 +51,7 @@ def search_mteam_torrents(imdb_url: str, page_number: int = 1, page_size: int = return response.json() +@with_disk_cache('mteam_imdb_info') def mteam_imdb_info(id: str) -> dict: """ Search M-Team for torrents using IMDb URL. diff --git a/workflow.py b/workflow.py index e6a2677..7f5dde4 100644 --- a/workflow.py +++ b/workflow.py @@ -47,12 +47,6 @@ def process_imdb_workflow(imdb_id: str, dl_dir: str = "/data/qb", jellyfin_dir: print(f"\n=== [1] Searching Torrents for {imdb_id} ===") imdb_url = f"https://www.imdb.com/title/{imdb_id}/" search_res = search_mteam_torrents(imdb_url) - - # Save the raw JSON - json_path = f"{imdb_id}.json" - print(f"Saving search results to {json_path}...") - with open(json_path, "w", encoding="utf-8") as f: - json.dump(search_res, f, ensure_ascii=False, indent=2) # Extract the torrent list if "data" in search_res and isinstance(search_res["data"], dict) and "data" in search_res["data"]: @@ -90,6 +84,8 @@ def process_imdb_workflow(imdb_id: str, dl_dir: str = "/data/qb", jellyfin_dir: for tid in selected_ids: print(f"\n=== [3] Downloading .torrent for ID: {tid} ===") torrent_bytes = generate_mteam_download_token(tid) + + # Save straight to local directory torrent_path = f"{tid}.torrent" with open(torrent_path, "wb") as f: f.write(torrent_bytes)