[+] Caching
This commit is contained in:
@@ -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/<subdir_name>/<key>.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
|
||||||
@@ -2,6 +2,8 @@ import os
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
|
|
||||||
|
from utils import with_disk_cache
|
||||||
|
|
||||||
client = OpenAI()
|
client = OpenAI()
|
||||||
|
|
||||||
def _call_openai_with_prompt(prompt_id: str, prompt_version: str, input_text: str) -> str:
|
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)
|
return str(response)
|
||||||
|
|
||||||
|
|
||||||
|
@with_disk_cache('select_best_torrents')
|
||||||
def select_best_torrents(torrents_text: str) -> str:
|
def select_best_torrents(torrents_text: str) -> str:
|
||||||
"""
|
"""
|
||||||
Calls the OpenAI API to select the best torrent IDs using a predefined prompt.
|
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]:
|
def generate_rename_mapping(directory_text: str) -> dict[str, str]:
|
||||||
"""
|
"""
|
||||||
Calls the OpenAI API to generate a renaming mapping for files
|
Calls the OpenAI API to generate a renaming mapping for files
|
||||||
|
|||||||
@@ -1,7 +1,10 @@
|
|||||||
|
import json
|
||||||
import requests
|
import requests
|
||||||
import tomllib
|
import tomllib
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
from utils import with_disk_cache
|
||||||
|
|
||||||
config = tomllib.loads(Path("config.toml").read_text())
|
config = tomllib.loads(Path("config.toml").read_text())
|
||||||
|
|
||||||
def _get_mteam_headers() -> dict:
|
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:
|
def search_mteam_torrents(imdb_url: str, page_number: int = 1, page_size: int = 100) -> dict:
|
||||||
"""
|
"""
|
||||||
Search M-Team for torrents using IMDb URL.
|
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()
|
return response.json()
|
||||||
|
|
||||||
|
|
||||||
|
@with_disk_cache('mteam_imdb_info')
|
||||||
def mteam_imdb_info(id: str) -> dict:
|
def mteam_imdb_info(id: str) -> dict:
|
||||||
"""
|
"""
|
||||||
Search M-Team for torrents using IMDb URL.
|
Search M-Team for torrents using IMDb URL.
|
||||||
|
|||||||
+2
-6
@@ -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} ===")
|
print(f"\n=== [1] Searching Torrents for {imdb_id} ===")
|
||||||
imdb_url = f"https://www.imdb.com/title/{imdb_id}/"
|
imdb_url = f"https://www.imdb.com/title/{imdb_id}/"
|
||||||
search_res = search_mteam_torrents(imdb_url)
|
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
|
# Extract the torrent list
|
||||||
if "data" in search_res and isinstance(search_res["data"], dict) and "data" in search_res["data"]:
|
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:
|
for tid in selected_ids:
|
||||||
print(f"\n=== [3] Downloading .torrent for ID: {tid} ===")
|
print(f"\n=== [3] Downloading .torrent for ID: {tid} ===")
|
||||||
torrent_bytes = generate_mteam_download_token(tid)
|
torrent_bytes = generate_mteam_download_token(tid)
|
||||||
|
|
||||||
|
# Save straight to local directory
|
||||||
torrent_path = f"{tid}.torrent"
|
torrent_path = f"{tid}.torrent"
|
||||||
with open(torrent_path, "wb") as f:
|
with open(torrent_path, "wb") as f:
|
||||||
f.write(torrent_bytes)
|
f.write(torrent_bytes)
|
||||||
|
|||||||
Reference in New Issue
Block a user