[+] Workflow

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)