Commit 81f0df5

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-12-07 13:41:21
chore(summary): add customizable model for AI summary
1 parent dae844e
Changed files (3)
src
llm
subtitles
ytdlp
src/llm/summary.py
@@ -98,13 +98,22 @@ DAILY_SUMMARY_PREFIX = "🏪**#爬楼助手**\n"
 CONTEXT_FILENAME = "聊天记录.txt"
 
 
-async def ai_summary(client: Client, message: Message, summary_prefix: str | None = None, **kwargs):
+async def ai_summary(
+    client: Client,
+    message: Message,
+    summary_prefix: str | None = None,
+    summary_model_id: str = GPT.CHAT_SUMMARY_MODEL_ID,
+    summary_model_name: str = GPT.CHAT_SUMMARY_MODEL_NAME,
+    **kwargs,
+):
     """GPT summary of the message history.
 
     Args:
         client (Client): The Pyrogram client.
         message (Message): The trigger message object.
         summary_prefix (str | None): Prefix string of the response message.
+        summary_model_id (str, optional): The model id to use for AI summary.
+        summary_model_name (str, optional): The model name to use for AI summary.
     """
     # send docs if message == "/summary"
     if equal_prefix(message.text, prefix=[PREFIX.AI_SUMMARY, PREFIX.COMBINATION]) and not message.reply_to_message:
@@ -203,8 +212,8 @@ async def ai_summary(client: Client, message: Message, summary_prefix: str | Non
     response = await gpt_response(
         client,
         ai_msg,
-        custom_model_id=GPT.CHAT_SUMMARY_MODEL_ID,
-        custom_model_name=GPT.CHAT_SUMMARY_MODEL_NAME,
+        custom_model_id=summary_model_id,
+        custom_model_name=summary_model_name,
         system_prompt=SYSTEM_PROMPT,
         enable_tools=False,
         include_thoughts=False,
src/subtitles/subtitle.py
@@ -46,6 +46,8 @@ async def get_subtitle(
     *,
     to_telegraph: bool = True,
     ai_summary: bool = True,
+    summary_model_id: str = GPT.SUBTITLE_SUMMARY_MODEL_ID,
+    summary_model_name: str = GPT.SUBTITLE_SUMMARY_MODEL_NAME,
     send_subtitle_as: Literal["file", "str", "none"] = "file",
     **kwargs,
 ):
@@ -147,7 +149,7 @@ async def get_subtitle(
             text=Str(f"/ai {prompt}"),
             reply_to_message=Message(id=rand_number(), chat=subtitle_msg.chat, text=Str(subtitles)),
         )
-        kwargs |= {"include_thoughts": False, "append_grounding": False, "silent": True, "custom_model_id": GPT.SUBTITLE_SUMMARY_MODEL_ID, "custom_model_name": GPT.SUBTITLE_SUMMARY_MODEL_NAME}
+        kwargs |= {"include_thoughts": False, "append_grounding": False, "silent": True, "custom_model_id": summary_model_id, "custom_model_name": summary_model_name}
         res = await gpt_response(client, ai_msg, **kwargs)
         if res.get("texts"):
             await send2tg(client, ai_msg, texts=res["prefix"] + blockquote(res["texts"]), **kwargs)
src/ytdlp/main.py
@@ -50,6 +50,8 @@ async def preview_ytdlp(
     ytdlp_subtitle_target: str | int | None = None,
     ytdlp_send_subtitle: bool = False,
     ytdlp_send_summary: bool = False,
+    summary_model_id: str = GPT.SUBTITLE_SUMMARY_MODEL_ID,
+    summary_model_name: str = GPT.SUBTITLE_SUMMARY_MODEL_NAME,
     to_telegraph: bool = True,
     show_author: bool = True,
     show_title: bool = True,
@@ -78,6 +80,8 @@ async def preview_ytdlp(
         ytdlp_audio_target (str | int, optional): The target chat id to send audio.
         ytdlp_send_subtitle (bool, optional): Send subtitle. Defaults to False.
         ytdlp_send_summary (bool, optional): Send AI summary. Defaults to False.
+        summary_model_id (str, optional): The model id to use for AI summary.
+        summary_model_name (str, optional): The model name to use for AI summary.
         to_telegraph (bool, optional): Whether to publish the subtitle or transcription to telegraph.
     """
     logger.trace(f"{url=} {kwargs=}")
@@ -147,7 +151,7 @@ async def preview_ytdlp(
             text=Str(f"{strings_list(PREFIX.GPT)[0]} {prompt}"),
             reply_to_message=Message(id=rand_number(), chat=message.chat, text=Str(subtitles)),
         )
-        params = {"include_thoughts": False, "append_grounding": False, "silent": True, "custom_model_id": GPT.SUBTITLE_SUMMARY_MODEL_ID, "custom_model_name": GPT.SUBTITLE_SUMMARY_MODEL_NAME}
+        params = {"include_thoughts": False, "append_grounding": False, "silent": True, "custom_model_id": summary_model_id, "custom_model_name": summary_model_name}
         aires = await gpt_response(client, ai_msg, **params)
         if aires.get("texts"):
             summary = f"🤖<b>{aires['model_name']}总结:</b>\n{markdown.markdown(aires['texts'])}\n"