[+] 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 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
|
||||
|
||||
@@ -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.
|
||||
|
||||
+2
-6
@@ -48,12 +48,6 @@ def process_imdb_workflow(imdb_id: str, dl_dir: str = "/data/qb", jellyfin_dir:
|
||||
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"]
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user