Add season number support for TMDB data and GPT prompt integration
This commit is contained in:
@@ -14,8 +14,8 @@ async def main():
|
|||||||
print('WARN: TMDB API key is not set, skip TMDB data integration')
|
print('WARN: TMDB API key is not set, skip TMDB data integration')
|
||||||
else:
|
else:
|
||||||
tmdb = TmdbClient(cfg)
|
tmdb = TmdbClient(cfg)
|
||||||
tmdb_data = await tmdb.get_tmdb_data(cfg.tmdb_id)
|
tmdb_data = await tmdb.get_tmdb_data(cfg.tmdb_id, cfg.season_number)
|
||||||
res = await get_response(cfg, cfg.input, files, cfg.series_name, cfg.year, cfg.tmdb_id, cfg.tvdb_id, tmdb_data)
|
res = await get_response(cfg, cfg.input, files, cfg.series_name, cfg.year, cfg.tmdb_id, cfg.tvdb_id, tmdb_data, cfg.season_number)
|
||||||
for f in res.files:
|
for f in res.files:
|
||||||
print(files[f.index], '->', f.name)
|
print(files[f.index], '->', f.name)
|
||||||
input('Continue?')
|
input('Continue?')
|
||||||
|
|||||||
@@ -40,6 +40,10 @@ class Config:
|
|||||||
def proxy(self) -> Optional[str]:
|
def proxy(self) -> Optional[str]:
|
||||||
return self._args.proxy if self._args.proxy is not None else self._yaml_config.get('proxy')
|
return self._args.proxy if self._args.proxy is not None else self._yaml_config.get('proxy')
|
||||||
|
|
||||||
|
@property
|
||||||
|
def season_number(self) -> Optional[int]:
|
||||||
|
return self._args.season_number
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def series_name(self) -> Optional[str]:
|
def series_name(self) -> Optional[str]:
|
||||||
return self._args.series_name
|
return self._args.series_name
|
||||||
@@ -86,6 +90,7 @@ def get_arg_parser() -> argparse.ArgumentParser:
|
|||||||
parser.add_argument('-T', '--tvdb-id', type=int, help='TVDB ID (optional)')
|
parser.add_argument('-T', '--tvdb-id', type=int, help='TVDB ID (optional)')
|
||||||
parser.add_argument('-H', '--hardlink', action='store_true', help='Use hardlink instead of symlink.')
|
parser.add_argument('-H', '--hardlink', action='store_true', help='Use hardlink instead of symlink.')
|
||||||
parser.add_argument('-n', '--no-tmdb', action='store_true', help='Do not use TMDB API to obtain data.')
|
parser.add_argument('-n', '--no-tmdb', action='store_true', help='Do not use TMDB API to obtain data.')
|
||||||
|
parser.add_argument('-S', '--season-number', type=int, help='Season number (optional)')
|
||||||
parser.add_argument('input', help='Input directory.')
|
parser.add_argument('input', help='Input directory.')
|
||||||
parser.add_argument('output', help='Output directory.')
|
parser.add_argument('output', help='Output directory.')
|
||||||
return parser
|
return parser
|
||||||
|
|||||||
@@ -83,7 +83,7 @@ def parse_result(result: str):
|
|||||||
|
|
||||||
async def get_response(cfg: Config, inp: str, files: List[str],
|
async def get_response(cfg: Config, inp: str, files: List[str],
|
||||||
series_name: str = None, year: int = None, tmdb_id: int = None,
|
series_name: str = None, year: int = None, tmdb_id: int = None,
|
||||||
tvdb_id: int = None, tmdb_data: TmdbData = None) -> Files:
|
tvdb_id: int = None, tmdb_data: TmdbData = None, season_number: int = None) -> Files:
|
||||||
prompt = f'The input directory is `{inp}`.'
|
prompt = f'The input directory is `{inp}`.'
|
||||||
if tmdb_data:
|
if tmdb_data:
|
||||||
prompt += '\n' + gen_tmdb_data(tmdb_data)
|
prompt += '\n' + gen_tmdb_data(tmdb_data)
|
||||||
@@ -96,6 +96,8 @@ async def get_response(cfg: Config, inp: str, files: List[str],
|
|||||||
prompt += f'\nThe TMDB ID is `{tmdb_id}`.'
|
prompt += f'\nThe TMDB ID is `{tmdb_id}`.'
|
||||||
if tvdb_id:
|
if tvdb_id:
|
||||||
prompt += f'\nThe TVDB ID is `{tvdb_id}`.'
|
prompt += f'\nThe TVDB ID is `{tvdb_id}`.'
|
||||||
|
if season_number:
|
||||||
|
prompt += f'\nThe season number is `{season_number}`.'
|
||||||
http_client = httpx.AsyncClient(proxy=cfg.proxy)
|
http_client = httpx.AsyncClient(proxy=cfg.proxy)
|
||||||
client = openai.AsyncClient(api_key=cfg.api_key, base_url=cfg.base_url, http_client=http_client)
|
client = openai.AsyncClient(api_key=cfg.api_key, base_url=cfg.base_url, http_client=http_client)
|
||||||
if is_support_structed_output(cfg.model):
|
if is_support_structed_output(cfg.model):
|
||||||
|
|||||||
Reference in New Issue
Block a user