main
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3import asyncio
4
5from loguru import logger
6from pyrogram.client import Client
7from pyrogram.types import Message
8
9from ai.chat_summary import ai_chat_summary
10from ai.main import ai_image_generation, ai_text_generation, ai_video_generation
11from asr.voice_recognition import voice_to_text
12from bridge.ocr import send_to_ocr_bridge
13from config import FAVORITE, PREFIX, PROXY
14from danmu.entrypoint import query_danmu
15from database.r2 import del_cf_r2
16from history.query import query_chat_history
17from messages.details import show_msg_info
18from messages.help import social_media_help
19from messages.modify import parse_kwargs
20from messages.parser import parse_msg
21from messages.sender import send2tg
22from messages.utils import delete_message, equal_prefix, startswith_prefix
23from networking import match_social_media_link
24from others.convert_chinese import chinese_conversion
25from others.convert_img_file import convert_image_file
26from others.download_external import download_url_in_message
27from others.extract_audio import extract_audio_file
28from others.favorite import save_favorite, send_favorite
29from others.ffmpeg import ffmpeg_cut, ffmpeg_h264, ffprobe
30from others.search_google import search_google
31from others.search_ytb import search_youtube
32from others.tmdb import search_tmdb
33from others.version import get_bot_version, update_bot
34from others.watermark import add_watermark
35from preview.arxiv import preview_arxiv
36from preview.bilibili import preview_bilibili
37from preview.douyin import preview_douyin
38from preview.github import preview_github
39from preview.instagram import preview_instagram
40from preview.netease import preview_music163
41from preview.reddit import preview_reddit
42from preview.spotify import preview_spotify
43from preview.twitter import preview_twitter
44from preview.v2ex import preview_v2ex
45from preview.wechat import preview_wechat
46from preview.weibo import preview_weibo
47from preview.xiaohongshu import preview_xhs
48from price.entrypoint import get_asset_price
49from quotly.quotly import quote_message
50from subtitles.subtitle import get_subtitle
51from summarize.main import ai_summary
52from tts.tts import text_to_speech
53from utils import to_int, true
54from ytdlp.main import preview_ytdlp
55from ytdlp.utils import ProxyError
56
57
58async def process_message(
59 client: Client,
60 message: Message,
61 target_chat: int | str | None = None,
62 reply_msg_id: int = 0,
63 *,
64 ai_text: bool = True, # Text Generation
65 ai_img: bool = True, # AI Image Generation
66 ai_video: bool = True, # AI Video Generation
67 ai_summarize: bool = True, # GPT Summary
68 chat_summary: bool = True, # AI Chat Summary.
69 asr: bool = True, # Voice to Text
70 audio_extract: bool = True, # Extract Audio from message. /audio -> m4a file. /voice -> wav file
71 danmu: bool = True, # Query Danmu database.
72 favorite: bool = True, # Send & Save message to favorite.
73 google_search: bool = True, # Search on Google.
74 ocr: bool = True, # OCR image file.
75 history: bool = True, # Query Chat History.
76 price: bool = True, # Query Stock / Fiat / Crypto price.
77 subtitle: bool = True, # Download YouTube subtitle.
78 tts: bool = True, # Text to Speech.
79 quotly: bool = True, # Quote message.
80 convert_chinese: bool = True, # Traditional Chinese <-> Simplified Chinese.
81 wget: bool = True, # Download external URL.
82 ytb: bool = True, # Search on YouTube.
83 tmdb: bool = True, # Search on TMDB.
84 ffmpeg: bool = True, # FFMpeg processing
85 watermark: bool = True, # Add watermark to video.
86 version: bool = True, # Show bot version with dependencies.
87 convert_img: bool = True, # Convert image file.
88 show_progress: bool = True, # Show progress bar.
89 detail_progress: bool = False, # Show detailed progress bar.
90 **kwargs,
91):
92 """Call utility functions to handle the message.
93
94 Args:
95 client (Client): The Pyrogram client.
96 message (Message): The trigger message object.
97 target_chat (int | str, optional): Send result to this telegram target chat. If not set, send to the trigger message's chat.
98 reply_msg_id (int, optional): If set to integer > 0, the result is sent as a reply message to this message_id.
99 If set to 0, reply to the trigger message itself.
100 If set to -1, do not send as a reply message.
101 """
102 kwargs |= {"target_chat": target_chat, "reply_msg_id": reply_msg_id, "show_progress": show_progress, "detail_progress": detail_progress}
103 if true(kwargs.get("target_chat")):
104 kwargs["target_chat"] = to_int(kwargs["target_chat"])
105 message, msg_kwargs = parse_kwargs(message)
106 kwargs |= msg_kwargs # merge the kwargs from the message text
107 if ai_text:
108 await ai_text_generation(client, message, **kwargs) # /ai
109 if ai_img:
110 await ai_image_generation(client, message, **kwargs) # /gen
111 if ai_video:
112 await ai_video_generation(client, message, **kwargs) # /gvid
113 if ai_summarize:
114 await ai_summary(client, message, **kwargs) # /summary
115 if chat_summary:
116 await ai_chat_summary(client, message, **kwargs) # /chatsum
117 if asr:
118 await voice_to_text(client, message, **kwargs) # /asr
119 if audio_extract:
120 await extract_audio_file(client, message, **kwargs) # /audio
121 if subtitle:
122 await get_subtitle(client, message, **kwargs) # /subtitle
123 if wget:
124 await download_url_in_message(client, message, **kwargs) # /wget
125 if google_search:
126 await search_google(client, message, **kwargs) # /google
127 if ytb:
128 await search_youtube(client, message, **kwargs) # /ytb
129 if ocr:
130 await send_to_ocr_bridge(client, message, **kwargs) # /ocr
131 if price:
132 await get_asset_price(client, message, **kwargs) # /price
133 if history:
134 await query_chat_history(client, message, **kwargs) # /history
135 if danmu:
136 await query_danmu(client, message, **kwargs) # /danmu
137 if favorite:
138 await save_favorite(client, message, **kwargs) # /save
139 await send_favorite(client, message, **kwargs) # /fav
140 if tts:
141 await text_to_speech(client, message, **kwargs) # /tts
142 if convert_chinese:
143 await chinese_conversion(client, message, **kwargs) # /sc
144 if quotly:
145 await quote_message(client, message, **kwargs) # /quote
146 if tmdb:
147 await search_tmdb(client, message, **kwargs) # /tmdb
148 if convert_img:
149 await convert_image_file(client, message, **kwargs)
150 if ffmpeg:
151 await ffmpeg_cut(client, message, **kwargs)
152 await ffmpeg_h264(client, message, **kwargs)
153 await ffprobe(client, message, **kwargs)
154 if watermark:
155 await add_watermark(client, message, **kwargs)
156 if version:
157 await get_bot_version(client, message, **kwargs)
158 await update_bot(message)
159
160 await show_msg_info(client, message, **kwargs)
161 await preview_social_media(client, message, **kwargs)
162
163
164async def preview_social_media(
165 client: Client,
166 message: Message,
167 *,
168 need_prefix: bool = True, # Need prefix to parse social media. (/dl)
169 prepend_sender_user: bool = False,
170 social: bool = True, # Parse social media link
171 douyin: bool = True, # Parse Douyin
172 tiktok: bool = True, # Parse TikTok
173 instagram: bool = True, # Parse Instagram
174 twitter: bool = True, # Parse Twitter
175 weibo: bool = True, # Parse Weibo
176 reddit: bool = True, # Parse Reddit
177 github: bool = True, # Parse Github
178 xhs: bool = True, # Parse XiaoHongShu
179 v2ex: bool = True, # Parse V2EX
180 music163: bool = True, # Parse Music163
181 spotify: bool = True, # Parse Spotify
182 ytdlp: bool = True, # Parse YT-DLP link
183 ytdlp_bilibili: bool = True, # Parse YT-DLP Bilibili link
184 ytdlp_youtube: bool = True, # Parse YT-DLP YouTube link
185 arxiv: bool = True, # Parse arXiv
186 **kwargs,
187):
188 """Preview social media link in the message.
189
190 Args:
191 client (Client): The Pyrogram client.
192 message (Message): The trigger message object.
193 target_chat (int | str, optional): Send result to this telegram target chat. If not set, send to the trigger message's chat.
194 reply_msg_id (int, optional): If set to integer > 0, the result is sent as a reply message to this message_id.
195 If set to 0, reply to the trigger message itself.
196 If set to -1, do not send as a reply message.
197 need_prefix (bool, optional): Need to start with PREFIX to call this funciton. Defaults to True.
198 cmd_prefix (str, optional): prefix to call this function.
199 prepend_sender_user (bool, optional): Prepend the sender's username to the message. Defaults to False.
200 show_progress (bool, optional): Show a progress message on Telegram. Defaults to True.
201 detail_progress (bool, optional): Show detailed progress (Only if show_proress is set to True). Defaults to False.
202 """
203 if not social:
204 return None
205 # these commands are handled in `process_message`
206 ignore_prefix = [
207 PREFIX.AI_IMG_GENERATION,
208 PREFIX.AI_SUMMARY,
209 PREFIX.AI_TEXT_GENERATION,
210 PREFIX.ASR,
211 PREFIX.AUDIO,
212 PREFIX.CHAT_SUMMARY,
213 PREFIX.COMBINATION,
214 PREFIX.CONVERT_TO_SC,
215 PREFIX.CONVERT_TO_TC,
216 PREFIX.CONVERT,
217 PREFIX.CRYPTO,
218 PREFIX.DANMU,
219 PREFIX.FAYAN,
220 PREFIX.OCR,
221 PREFIX.PRICE,
222 PREFIX.SEARCH_GOOGLE,
223 PREFIX.SEARCH_YOUTUBE,
224 PREFIX.STOCK,
225 PREFIX.SUBTITLE,
226 PREFIX.TMDB,
227 PREFIX.TTS,
228 PREFIX.VOICE,
229 PREFIX.WGET,
230 FAVORITE.SAVE_PREFIX,
231 FAVORITE.SEND_PREFIX,
232 ]
233
234 info = parse_msg(message)
235 this_msg = message
236 this_texts = info["text"] # texts of the trigger message
237 if startswith_prefix(this_texts, prefix=ignore_prefix):
238 return None
239 if need_prefix and not startswith_prefix(this_texts, prefix=[PREFIX.SOCIAL_MEDIA, "/help", "/retry"]):
240 return None
241
242 # message only contains prefix command
243 if equal_prefix(this_texts, prefix=[PREFIX.SOCIAL_MEDIA, "/help", "/retry"]):
244 # without reply, send docs if message only contains prefix command
245 if not message.reply_to_message:
246 docs = social_media_help(message)
247 await delete_message(message)
248 helps = await send2tg(client, message, texts=docs, **kwargs)
249 await asyncio.sleep(30)
250 return await delete_message(helps)
251
252 # with reply, treat the reply_msg as the trigger to preview social media link
253 message = message.reply_to_message
254 info = parse_msg(message, silent=True, use_cache=False) # parse again
255
256 # add send_from_user.
257 if prepend_sender_user:
258 # Caution: this format should be consistent with `save_messages` function in `message.database.py`
259 kwargs["send_from_user"] = f"👤[@{info['full_name']}](tg://user?id={info['uid']})//"
260
261 warn_msg = None
262 try:
263 matched = await match_social_media_link(info["text"]) # match "platform" and "url" (and other info)
264 if matched["platform"]:
265 logger.success(f"Matched: {matched}")
266 kwargs |= matched
267 if startswith_prefix(this_texts, prefix="/retry"):
268 await del_cf_r2(matched["db_key"])
269
270 if douyin and matched["platform"] == "douyin":
271 return await preview_douyin(client, message, **kwargs)
272 if tiktok and matched["platform"] == "tiktok":
273 return await preview_douyin(client, message, **kwargs)
274 if instagram and matched["platform"] == "instagram":
275 return await preview_instagram(client, message, **kwargs)
276 if twitter and matched["platform"] in ["x", "twitter", "fxtwitter", "fixupx"]:
277 return await preview_twitter(client, message, **kwargs)
278 if weibo and matched["platform"] == "weibo":
279 return await preview_weibo(client, message, **kwargs)
280 if xhs and matched["platform"] == "xiaohongshu":
281 return await preview_xhs(client, message, **kwargs)
282 if xhs and matched["platform"] == "wechat":
283 return await preview_wechat(client, message, **kwargs)
284 if github and matched["platform"] == "github":
285 return await preview_github(client, message, **kwargs)
286 if reddit and matched["platform"] == "reddit":
287 return await preview_reddit(client, message, **kwargs)
288 if music163 and matched["platform"] == "music163":
289 return await preview_music163(client, message, **kwargs)
290 if spotify and matched["platform"] == "spotify":
291 return await preview_spotify(client, message, **kwargs)
292 if v2ex and matched["platform"] == "v2ex":
293 return await preview_v2ex(client, message, **kwargs)
294 if arxiv and matched["platform"] == "arxiv":
295 return await preview_arxiv(client, message, **kwargs)
296 if matched["platform"].startswith("bilibili-"): # this is not bilibili video, for videos, use yt-dlp
297 return await preview_bilibili(client, message, **kwargs)
298
299 sent_messages = []
300 try:
301 if ytdlp_bilibili and matched["platform"] == "bilibili":
302 sent_messages = await preview_ytdlp(client, message, **kwargs)
303 if ytdlp_youtube and matched["platform"] == "youtube":
304 sent_messages = await preview_ytdlp(client, message, **kwargs)
305 if ytdlp and matched["platform"] == "ytdlp":
306 sent_messages = await preview_ytdlp(client, message, **kwargs)
307 except ProxyError:
308 logger.error(f"🚫{matched['platform']}代理错误")
309 if PROXY.YTDLP_FALLBACK:
310 logger.warning(f"🔄使用备用代理{PROXY.YTDLP_FALLBACK}")
311 sent_messages = await preview_ytdlp(client, message, proxy=PROXY.YTDLP_FALLBACK, **kwargs)
312
313 if not sent_messages and any(keyword in info["text"] for keyword in ["facebook.com", "threads.com", "youtube.com", "bilibili.com"]):
314 if kwargs.get("show_progress"):
315 warn_msg = await message.reply_text("⚠️该链接不被支持", quote=True)
316 await asyncio.sleep(10)
317 await delete_message(warn_msg)
318 # if ytdlp failed, try to download directly
319 elif (
320 not sent_messages
321 and startswith_prefix(this_texts, prefix=PREFIX.SOCIAL_MEDIA)
322 and matched["platform"]
323 not in [
324 "bilibili",
325 "douyin",
326 "github",
327 "instagram",
328 "music163",
329 "reddit",
330 "spotify",
331 "tiktok",
332 "v2ex",
333 "weibo",
334 "x",
335 "xiaohongshu",
336 "youtube",
337 ]
338 ):
339 if kwargs.get("show_progress"):
340 warn_msg = await client.send_message(info["cid"], text="⚠️暂时不支持解析链接, 尝试直接下载该网页")
341 await download_url_in_message(client, this_msg, extra_prefix=PREFIX.SOCIAL_MEDIA, **kwargs)
342
343 except Exception as e:
344 logger.exception(e)
345 finally:
346 await delete_message(warn_msg)
347 return sent_messages