From d7a48327914aff218385e1b6c505af4084e4138a Mon Sep 17 00:00:00 2001 From: Azalea <22280294+hykilpikonna@users.noreply.github.com> Date: Mon, 9 Mar 2026 00:23:16 -0400 Subject: [PATCH] [+] Workflow --- workflow.py | 154 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 workflow.py diff --git a/workflow.py b/workflow.py new file mode 100644 index 0000000..e6a2677 --- /dev/null +++ b/workflow.py @@ -0,0 +1,154 @@ +import json +import time +from pathlib import Path + +from utils_mteam import ( + mteam_imdb_info, + search_mteam_torrents, + format_mteam_torrent, + generate_mteam_download_token, +) +from utils_qb import ( + get_qb_client, + download_torrent, + get_torrent_file_tree +) + +from utils_ai import ( + select_best_torrents, + generate_rename_mapping, + apply_rename_mapping +) + + +def format_file_tree(file_tree: list) -> str: + """Helper to convert the qB file tree output into simple relative paths for the LLM prompt""" + lines = [] + for f in file_tree: + lines.append(f.get("name", "")) + return "\n".join(lines) + + +def process_imdb_workflow(imdb_id: str, dl_dir: str = "/data/qb", jellyfin_dir: str = "/data/jellyfin"): + """ + Workflow to automatically find, download, and map torrents for an IMDb ID into a Jellyfin library. + """ + print(f"=== [0] Fetching IMDB info for {imdb_id} ===") + imdb_info = mteam_imdb_info(imdb_id) + if 'data' not in imdb_info: + print("Failed to get IMDB info") + return + + title = imdb_info['data'].get('title', 'Unknown_Title') + year = imdb_info['data'].get('year', '') + title_dir = f"{title} ({year})" + print(f"Found Title: {title_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"]: + torrents = search_res["data"]["data"] + elif "data" in search_res and isinstance(search_res["data"], list): + torrents = search_res["data"] + elif isinstance(search_res, list): + torrents = search_res + else: + torrents = [] + + if not torrents: + print("No torrents found.") + return + + # Format the torrents text + formatted_torrents = [] + for t in torrents: + if isinstance(t, dict): + formatted_torrents.append(format_mteam_torrent(t)) + torrents_text = "\n\n".join(formatted_torrents) + + print(f"\n=== [2] Selecting best torrents using LLM ===") + selected_ids_str = select_best_torrents(torrents_text) + selected_ids = [tid.strip() for tid in selected_ids_str.split() if tid.strip()] + print(f"Selected torrent IDs: {selected_ids}") + + if not selected_ids: + print("No torrents selected.") + return + + qb = get_qb_client() + jellyfin_base = Path(jellyfin_dir) / f"{title_dir} [{imdb_id}]" + + for tid in selected_ids: + print(f"\n=== [3] Downloading .torrent for ID: {tid} ===") + torrent_bytes = generate_mteam_download_token(tid) + torrent_path = f"{tid}.torrent" + with open(torrent_path, "wb") as f: + f.write(torrent_bytes) + print(f"Saved .torrent to {torrent_path}") + + print(f"\n=== [4] Adding torrent to qBittorrent ===") + download_torrent(qb, torrent_path, dl_dir) + + print(f"\n=== [5] Waiting for download to finish ===") + # Wait slightly for qB to process the adding request + time.sleep(3) + + # Best effort to grab the hash of the torrent we just added + # (Assuming it's the most recently added sorted descending) + recent_torrents = qb.torrents_info(sort="added_on", reverse=True) + if not recent_torrents: + print("Could not find any torrents in qB!") + continue + + target_torrent = recent_torrents[0] + t_hash = target_torrent.hash + t_name = target_torrent.name + print(f"Tracking torrent: {t_name} (Hash: {t_hash})") + + while True: + info = qb.torrents_info(hashes=t_hash) + if not info: + print("Torrent disappeared from qB!") + break + + t_info = info[0] + progress = t_info.progress + state = t_info.state + print(f"Progress: {progress * 100:.1f}% (State: {state})") + + # Progress of 1.0 means 100%. Alternatively, check the state. + if progress >= 1.0 or state in ('uploading', 'pausedUP', 'stalledUP', 'forcedUP'): + break + time.sleep(5) + print("Download complete!") + + print(f"\n=== [6] Generating rename mapping using LLM ===") + file_tree = get_torrent_file_tree(qb, t_hash) + file_tree_str = format_file_tree(file_tree) + + prompt_text = f"Base directory: `{title_dir}`\n\n{file_tree_str}" + print(f"Sending paths to LLM...") + mapping = generate_rename_mapping(prompt_text) + print("Generated Mapping:") + for src, dst in mapping.items(): + print(f" {src} -->> {dst}") + + print(f"\n=== [7] Creating symbolic links ===") + apply_rename_mapping(mapping, base_src_dir=dl_dir, base_dst_dir=jellyfin_base) + print(f"Finished processing torrent: {tid}") + +if __name__ == "__main__": + import sys + + # Allow executing directly with `uv run workflow.py tt7742120` + test_id = sys.argv[1] if len(sys.argv) > 1 else "tt38872297" + process_imdb_workflow(test_id)