g0ldyy/comet
1
0
Fork 0
mirror of https://github.com/g0ldyy/comet.git synced 2026-01-12 01:16:12 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement smart file selection algorithm and refine torrent record management for multi-episode content

Browse Source
This commit is contained in:
g0ldyy
2026-01-07 20:34:43 +01:00
parent f392b73e42
commit 479b5ae8ff
7 changed files with 313 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comet/core/log_levels.py
+6
View File
@@ -68,6 +68,12 @@ CUSTOM_LOG_LEVELS = {
"loguru_color": "<fg #FFD700>",
"no": 35,
},
"PLAYBACK": {
"color": "#00FA9A",
"icon": "▶️",
"loguru_color": "<fg #00FA9A>",
"no": 35,
},
}
ALL_LOG_LEVELS = {**STANDARD_LOG_LEVELS, **CUSTOM_LOG_LEVELS}
comet/debrid/stremthru.py
+154 -53
View File
@@ -2,13 +2,14 @@ import asyncio
from urllib.parse import quote, unquote
import aiohttp
from RTN import parse, title_match
from RTN import normalize_title, parse, title_match
from comet.core.execution import get_executor
from comet.core.logger import logger
from comet.core.models import settings
from comet.debrid.exceptions import DebridAuthError
from comet.services.debrid_cache import cache_availability
from comet.services.filtering import quick_alias_match
from comet.services.torrent_manager import torrent_update_queue
from comet.utils.parsing import is_video
@@ -217,6 +218,18 @@ class StremThru:
sources: list = None,
aliases: dict = None,
):
"""
Smart file selection algorithm with scoring system.
Priority order (highest to lowest):
1. Exact season + episode match with single episode file (+1000)
2. Exact season + episode match with multi-episode file (+500)
3. Episode match without season info (+200)
4. Exact filename match with requested torrent_name (+100)
5. Title alias match (+50)
6. Index match from original selection (+25)
7. Fallback: largest video file (+file_size as tiebreaker)
"""
try:
magnet_uri = f"magnet:?xt=urn:btih:{hash}&dn={quote(torrent_name)}"
@@ -236,82 +249,170 @@ class StremThru:
name = unquote(name)
torrent_name = unquote(torrent_name)
name_parsed = parse(name)
target_file = None
ez_aliases = aliases.get("ez", [])
if ez_aliases:
ez_aliases_normalized = [normalize_title(a) for a in ez_aliases]
debrid_files = magnet["data"]["files"]
debrid_files_parsed = []
files = []
# Filter to video files only, excluding samples
video_files = []
filenames_to_parse = []
for file in debrid_files:
filename = file["name"]
filename_lower = filename.lower()
if "sample" in filename.lower():
if "sample" in filename_lower:
continue
if not is_video(filename):
continue
filename_parsed = parse(filename)
video_files.append(file)
filenames_to_parse.append(filename)
file_season = (
filename_parsed.seasons[0] if filename_parsed.seasons else None
if not video_files:
logger.warning(f"No video files found in torrent {hash}")
return
loop = asyncio.get_running_loop()
parsed_results = await loop.run_in_executor(
get_executor(), batch_parse, filenames_to_parse
)
file_episode = (
filename_parsed.episodes[0] if filename_parsed.episodes else None
)
file_index = file["index"] if file["index"] != -1 else None
file_size = file["size"] if file["size"] != -1 else None
file = {
"index": file_index,
"title": filename,
"size": file_size,
"season": file_season,
"episode": file_episode,
"link": file.get("link", None),
}
scored_files = []
debrid_files_parsed.append(file)
if not filename_parsed.parsed_title or not title_match(
name_parsed.parsed_title,
filename_parsed.parsed_title,
aliases=aliases,
for file, filename, parsed in zip(
video_files, filenames_to_parse, parsed_results
):
file_index = file["index"] if file.get("index", -1) != -1 else None
file_size = file["size"] if file.get("size", -1) != -1 else 0
file_link = file.get("link")
if not file_link:
continue
file["info_hash"] = hash
file["parsed"] = filename_parsed
files.append(file)
file_season = parsed.seasons[0] if parsed.seasons else None
file_episode = parsed.episodes[0] if parsed.episodes else None
for file in debrid_files_parsed:
if file["title"] == torrent_name:
target_file = file
break
# Calculate score
score = 0
match_reason = []
if season == file["season"] and episode == file["episode"]:
target_file = file
break
# Season + Episode matching (highest priority)
if season is not None and episode is not None:
season_matches = (not parsed.seasons) or (season in parsed.seasons)
episode_matches = parsed.episodes and episode in parsed.episodes
if not target_file:
for file in files:
if str(file["index"]) == index:
target_file = file
break
if season_matches and episode_matches:
if len(parsed.episodes) == 1:
score += 1000 # Perfect single episode match
match_reason.append("exact_episode")
else:
score += 500 # Multi-episode file containing our episode
match_reason.append("multi_episode")
elif episode_matches:
score += 200 # Episode matches but season doesn't
match_reason.append("episode_only")
if len(files) > 0:
asyncio.create_task(cache_availability(self.real_debrid_name, files))
# Exact filename match
if filename == torrent_name:
score += 100
match_reason.append("exact_name")
if not target_file and len(debrid_files) > 0:
files_with_link = [
file for file in debrid_files if file.get("link", None)
]
if len(files_with_link) > 0:
target_file = max(files_with_link, key=lambda x: x["size"])
# Title/alias matching
if parsed.parsed_title:
# Quick alias match first
if ez_aliases and quick_alias_match(
normalize_title(filename), ez_aliases_normalized
):
score += 50
match_reason.append("alias")
elif title_match(name, parsed.parsed_title, aliases=aliases):
score += 50
match_reason.append("title")
if not target_file:
# Index match from original selection
if file_index is not None and str(file_index) == str(index):
score += 25
match_reason.append("index")
# Use file size as tiebreaker (larger files preferred)
# Normalize to 0-10 range to not overwhelm other scores
size_score = min(
file_size / (10 * 1024 * 1024 * 1024), 10
) # Cap at 10GB
score += size_score
enriched_file = {
"index": file_index,
"title": filename,
"size": file_size if file_size > 0 else None,
"season": file_season,
"episode": file_episode,
"link": file_link,
"parsed": parsed,
"score": score,
"match_reason": match_reason,
}
scored_files.append(enriched_file)
if not scored_files:
logger.log(
"PLAYBACK",
f"No valid video files with links found in torrent {hash}",
)
return
# Sort by score descending
scored_files.sort(key=lambda x: x["score"], reverse=True)
# Select best file
target_file = scored_files[0]
logger.log(
"PLAYBACK",
f"File selection for {hash}: selected '{target_file['title']}' "
f"(score={target_file['score']:.1f}, reasons={target_file['match_reason']}) "
f"from {len(scored_files)} candidates",
)
all_files_for_cache = []
for f in scored_files:
if f["season"] is not None or f["episode"] is not None:
all_files_for_cache.append(
{
"info_hash": hash,
"index": f["index"],
"title": f["title"],
"size": f["size"],
"season": f["season"] or season,
"episode": f["episode"],
"parsed": f["parsed"],
}
)
# Also ensure the selected file is cached with the REQUESTED season/episode
# This handles cases where filename doesn't contain S/E info but user requested it
if season is not None or episode is not None:
all_files_for_cache.append(
{
"info_hash": hash,
"index": target_file["index"],
"title": target_file["title"],
"size": target_file["size"],
"season": season,
"episode": episode,
"parsed": target_file["parsed"],
}
)
if all_files_for_cache:
asyncio.create_task(
cache_availability(self.real_debrid_name, all_files_for_cache)
)
link = await self.session.post(
f"{self.base_url}/link/generate?client_ip={self.client_ip}",
json={"link": target_file["link"]},
comet/services/debrid.py
+19 -8
View File
@@ -80,7 +80,7 @@ class DebridService:
for hash in info_hashes:
torrents[hash]["cached"] = False
if self.debrid_service == "torrent" or len(torrents) == 0:
if len(torrents) == 0:
return
rows = await get_cached_availability(
@@ -91,13 +91,24 @@ class DebridService:
info_hash = row["info_hash"]
torrents[info_hash]["cached"] = True
if row["parsed"] is not None:
torrents[info_hash]["parsed"] = ParsedData(
**orjson.loads(row["parsed"])
)
if row["file_index"] is not None:
torrents[info_hash]["fileIndex"] = row["file_index"]
if row["title"] is not None:
torrents[info_hash]["title"] = row["title"]
try:
torrents[info_hash]["fileIndex"] = int(row["file_index"])
except ValueError:
pass
if row["size"] is not None:
torrents[info_hash]["size"] = row["size"]
# Only update title/parsed if the cached file has resolution info
# Otherwise keep the original torrent info which may have better quality data
# E.g. torrent "[Group] Show S01 1080p" vs file "Show - 02.mkv"
if row["parsed"] is not None:
cached_parsed = ParsedData(**orjson.loads(row["parsed"]))
if (
cached_parsed.resolution != "unknown"
or torrents[info_hash]["parsed"].resolution == "unknown"
):
torrents[info_hash]["parsed"] = cached_parsed
if row["title"] is not None:
torrents[info_hash]["title"] = row["title"]
comet/services/debrid_cache.py
+28 -10
View File
@@ -116,31 +116,43 @@ async def cache_availability(debrid_service: str, availability: list):
async def get_cached_availability(
debrid_service: str, info_hashes: list, season: int = None, episode: int = None
):
base_query = f"""
SELECT info_hash, file_index, title, size, parsed
select_clause = "SELECT info_hash, file_index, title, size, parsed"
if debrid_service == "torrent" and settings.DATABASE_TYPE == "postgresql":
select_clause = (
"SELECT DISTINCT ON (info_hash) info_hash, file_index, title, size, parsed"
)
base_from_where = f"""
FROM debrid_availability
WHERE info_hash IN (SELECT CAST(value as TEXT) FROM {"json_array_elements_text" if settings.DATABASE_TYPE == "postgresql" else "json_each"}(:info_hashes))
AND debrid_service = :debrid_service
AND timestamp + :cache_ttl >= :current_time
"""
params = {
"info_hashes": orjson.dumps(info_hashes).decode("utf-8"),
"debrid_service": debrid_service,
"cache_ttl": settings.DEBRID_CACHE_TTL,
"current_time": time.time(),
"season": season,
"episode": episode,
}
if debrid_service != "torrent":
base_from_where += " AND debrid_service = :debrid_service"
params["debrid_service"] = debrid_service
group_by_clause = ""
if debrid_service == "torrent" and settings.DATABASE_TYPE == "sqlite":
group_by_clause = " GROUP BY info_hash"
if debrid_service == "offcloud":
query = (
base_query
+ """
season_episode_filter = """
AND ((CAST(:season as INTEGER) IS NULL AND season IS NULL) OR season = CAST(:season as INTEGER))
AND ((CAST(:episode as INTEGER) IS NULL AND episode IS NULL) OR episode = CAST(:episode as INTEGER))
"""
query = (
select_clause + base_from_where + season_episode_filter + group_by_clause
)
results = await database.fetch_all(query, params)
found_hashes = {r["info_hash"] for r in results}
@@ -149,20 +161,26 @@ async def get_cached_availability(
if remaining_hashes:
null_title_params = {
"info_hashes": orjson.dumps(remaining_hashes).decode("utf-8"),
"debrid_service": debrid_service,
"cache_ttl": settings.DEBRID_CACHE_TTL,
"current_time": time.time(),
}
null_title_query = base_query + " AND title IS NULL"
if debrid_service != "torrent":
null_title_params["debrid_service"] = debrid_service
null_title_query = (
select_clause + base_from_where + " AND title IS NULL" + group_by_clause
)
null_results = await database.fetch_all(null_title_query, null_title_params)
results.extend(null_results)
else:
query = (
base_query
select_clause
+ base_from_where
+ """
AND ((CAST(:season as INTEGER) IS NULL AND season IS NULL) OR season = CAST(:season as INTEGER))
AND ((CAST(:episode as INTEGER) IS NULL AND episode IS NULL) OR episode = CAST(:episode as INTEGER))
"""
+ group_by_clause
)
results = await database.fetch_all(query, params)
comet/services/filtering.py
+21 -5
View File
@@ -1,4 +1,4 @@
from RTN import parse, title_match
from RTN import normalize_title, parse, title_match
from comet.core.logger import logger
from comet.core.models import settings
@@ -13,12 +13,22 @@ else:
pass
def quick_alias_match(text_normalized: str, ez_aliases_normalized: list[str]):
return any(alias in text_normalized for alias in ez_aliases_normalized)
def filter_worker(torrents, title, year, year_end, aliases, remove_adult_content):
results = []
ez_aliases = aliases.get("ez", [])
if ez_aliases:
ez_aliases_normalized = [normalize_title(a) for a in ez_aliases]
for torrent in torrents:
torrent_title = torrent["title"]
if "sample" in torrent_title.lower() or torrent_title == "":
torrent_title_lower = torrent_title.lower()
if "sample" in torrent_title_lower or torrent_title == "":
_log_exclusion(f"🚫 Rejected (Sample/Empty) | {torrent_title}")
continue
@@ -28,9 +38,15 @@ def filter_worker(torrents, title, year, year_end, aliases, remove_adult_content
_log_exclusion(f"🔞 Rejected (Adult) | {torrent_title}")
continue
if not parsed.parsed_title or not title_match(
title, parsed.parsed_title, aliases=aliases
):
if not parsed.parsed_title:
_log_exclusion(f"❌ Rejected (No Parsed Title) | {torrent_title}")
continue
alias_matched = ez_aliases and quick_alias_match(
normalize_title(torrent_title), ez_aliases_normalized
)
if not alias_matched:
if not title_match(title, parsed.parsed_title, aliases=aliases):
_log_exclusion(
f"❌ Rejected (Title Mismatch) | {torrent_title} | Parsed: {parsed.parsed_title} | Expected: {title}"
)
comet/services/orchestration.py
+30 -12
View File
@@ -104,7 +104,7 @@ class TorrentManager:
if self.is_kitsu:
rows = await database.fetch_all(
"""
SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed
SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed, episode
FROM torrents
WHERE media_id = :media_id
AND (episode IS NULL OR episode = CAST(:episode as INTEGER))
@@ -117,7 +117,7 @@ class TorrentManager:
else:
rows = await database.fetch_all(
"""
SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed
SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed, episode
FROM torrents
WHERE media_id = :media_id
AND ((season IS NOT NULL AND season = CAST(:season as INTEGER)) OR (season IS NULL AND CAST(:season as INTEGER) IS NULL))
@@ -130,7 +130,15 @@ class TorrentManager:
},
)
rows = sorted(rows, key=lambda r: (r["episode"] is not None, r["episode"]))
for row in rows:
parsed_data = ParsedData(**orjson.loads(row["parsed"]))
if row["episode"] is None and parsed_data.episodes:
if self.episode not in parsed_data.episodes:
continue
info_hash = row["info_hash"]
self.torrents[info_hash] = {
"fileIndex": row["file_index"],
@@ -139,35 +147,45 @@ class TorrentManager:
"size": row["size"],
"tracker": row["tracker"],
"sources": orjson.loads(row["sources"]),
"parsed": ParsedData(**orjson.loads(row["parsed"])),
"parsed": parsed_data,
}
async def cache_torrents(self):
for torrent in self.ready_to_cache:
if self.is_kitsu:
cache_season = self.season
cache_seasons = [self.season]
else:
cache_season = (
torrent["parsed"].seasons[0]
cache_seasons = (
torrent["parsed"].seasons
if torrent["parsed"].seasons
else self.season
else [self.season]
)
parsed_episodes = (
torrent["parsed"].episodes if torrent["parsed"].episodes else [None]
)
if len(parsed_episodes) > 1:
episode = None
else:
episode = parsed_episodes[0]
for season in cache_seasons:
file_info = {
"info_hash": torrent["infoHash"],
"index": torrent["fileIndex"],
"title": torrent["title"],
"size": torrent["size"],
"season": cache_season,
"episode": torrent["parsed"].episodes[0]
if torrent["parsed"].episodes
else None,
"season": season,
"episode": episode,
"parsed": torrent["parsed"],
"seeders": torrent["seeders"],
"tracker": torrent["tracker"],
"sources": torrent["sources"],
}
await torrent_update_queue.add_torrent_info(file_info, self.media_only_id)
await torrent_update_queue.add_torrent_info(
file_info, self.media_only_id
)
async def filter_manager(self, scraper_name: str, torrents: list):
if len(torrents) == 0:
comet/services/torrent_manager.py
+15 -10
View File
@@ -125,10 +125,14 @@ async def add_torrent(
parsed: ParsedData,
):
try:
parsed_season = parsed.seasons[0] if parsed.seasons else search_season
parsed_episode = parsed.episodes[0] if parsed.episodes else None
seasons_to_process = parsed.seasons if parsed.seasons else [search_season]
parsed_episodes = parsed.episodes if parsed.episodes else [None]
if parsed_episode is not None:
episode_to_insert = parsed_episodes[0] if len(parsed_episodes) == 1 else None
for season in seasons_to_process:
# Only delete season-only entry if we are inserting a SINGLE specific episode.
if episode_to_insert is not None:
await database.execute(
"""
DELETE FROM torrents
@@ -138,12 +142,12 @@ async def add_torrent(
""",
{
"info_hash": info_hash,
"season": parsed_season,
"season": season,
},
)
logger.log(
"SCRAPER",
f"Deleted season-only entry for S{parsed_season:02d} of {info_hash}",
f"Deleted season-only entry for S{season:02d} of {info_hash} in favor of specific episode",
)
await _upsert_torrent_record(
@@ -151,8 +155,8 @@ async def add_torrent(
"media_id": media_id,
"info_hash": info_hash,
"file_index": file_index,
"season": parsed_season,
"episode": parsed_episode,
"season": season,
"episode": episode_to_insert,
"title": title,
"seeders": seeders,
"size": size,
@@ -164,9 +168,10 @@ async def add_torrent(
)
additional = ""
if parsed_season:
additional += f" - S{parsed_season:02d}"
additional += f"E{parsed_episode:02d}" if parsed_episode else ""
if seasons_to_process:
additional += f" - S{seasons_to_process[0]:02d}"
if parsed_episodes and parsed_episodes[0] is not None:
additional += f"E{parsed_episodes[0]:02d}"
logger.log("SCRAPER", f"Added torrent for {media_id} - {title}{additional}")
except Exception as e: