Commit 29b5f38
Changed files (6)
src/llm/contexts.py
@@ -14,7 +14,7 @@ from pyrogram.client import Client
from pyrogram.types import Message
from asr.utils import GEMINI_AUDIO_EXT, downsampe_audio
-from config import GPT
+from config import DOWNLOAD_DIR, GPT
from llm.utils import BOT_TIPS, clean_context, convert_md
from messages.parser import parse_msg
from utils import read_text
@@ -81,13 +81,14 @@ async def single_gpt_context(client: Client, message: Message) -> dict:
for msg in messages:
info = parse_msg(msg, silent=True, use_cache=False)
sender = info["fwd_full_name"] or info["full_name"]
+ media_path = DOWNLOAD_DIR + "/" + info["file_name"]
try:
if info["mtype"] == "photo":
res = await base64_media(client, msg)
contexts.append({"type": "image_url", "image_url": {"url": f"data:image/{res['ext']};base64,{res['base64']}"}})
elif info["mtype"] == "document":
if info["mime_type"].startswith("text/") or Path(info["file_name"]).suffix in extra_txt_extensions:
- fpath: str = await client.download_media(msg) # type: ignore
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
contexts.append(
{
"type": "text",
@@ -95,7 +96,7 @@ async def single_gpt_context(client: Client, message: Message) -> dict:
}
)
if Path(info["file_name"]).suffix in extra_markdown_extensions:
- fpath: str = await client.download_media(msg) # type: ignore
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
text = convert_md(fpath)
Path(fpath).unlink(missing_ok=True)
contexts.append(
@@ -144,9 +145,10 @@ async def single_gemini_context(client: Client, message: Message, app: genai.Cli
for msg in messages:
info = parse_msg(msg, silent=True, use_cache=False)
sender = info["fwd_full_name"] or info["full_name"]
+ media_path = DOWNLOAD_DIR + "/" + info["file_name"]
try:
if info["mtype"] in ["video", "photo", "audio", "voice"] or info["mime_type"] in gemini_mime_types or any(info["file_name"].endswith(ext) for ext in gemini_extensions):
- fpath: str = await client.download_media(msg, in_memory=False) # type: ignore # type: ignore
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
if info["mtype"] in ["audio", "voice"] and Path(fpath).suffix not in GEMINI_AUDIO_EXT:
audio_path = await downsampe_audio(fpath)
fpath = audio_path.as_posix()
@@ -160,10 +162,10 @@ async def single_gemini_context(client: Client, message: Message, app: genai.Cli
Path(fpath).unlink(missing_ok=True)
elif info["mtype"] == "document":
if info["mime_type"].startswith("text/") or Path(info["file_name"]).suffix in txt_extensions:
- fpath: str = await client.download_media(msg, in_memory=False) # type: ignore # type: ignore
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
parts.append(Part.from_text(text=f"[filename]: {info['file_name']}\n[file content]:\n{read_text(fpath).strip()}"))
if Path(info["file_name"]).suffix in extra_markdown_extensions:
- fpath: str = await client.download_media(msg) # type: ignore
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
text = convert_md(fpath)
Path(fpath).unlink(missing_ok=True)
parts.append(Part.from_text(text=f"[filename]: {info['file_name']}\n[file content]:\n{text.strip()}"))
src/messages/parser.py
@@ -7,6 +7,7 @@ from zoneinfo import ZoneInfo
from glom import Coalesce, glom
from loguru import logger
+from pathvalidate import sanitize_filename
from pyrogram.enums import MessageEntityType
from pyrogram.types import Chat, Message
@@ -111,7 +112,7 @@ def parse_msg(message: Message, *, silent: bool = False, verbose: bool = False,
"handle": str(handle),
"datetime": dt,
"time": str(time),
- "file_name": str(file_name),
+ "file_name": sanitize_filename(file_name, replacement_text="_"),
"file_id": str(file_id),
"mime_type": str(mime_type),
"file_size": int(file_size),
src/subtitles/subtitle.py
@@ -11,7 +11,7 @@ from pyrogram.types import Message
from pyrogram.types.messages_and_media.message import Str
from asr.voice_recognition import asr_file
-from config import ASR, PREFIX, PROVIDER, READING_SPEED, TEXT_LENGTH, cache
+from config import ASR, DOWNLOAD_DIR, PREFIX, PROVIDER, READING_SPEED, TEXT_LENGTH, cache
from llm.gpt import gpt_response
from messages.parser import parse_msg
from messages.progress import modify_progress
@@ -89,7 +89,8 @@ async def get_subtitle(
if this_info["mtype"] in ["audio", "video"] or reply_info.get("mtype", "") in ["audio", "video"]:
await modify_progress(text=error + "\n正在通过ASR识别字幕", force_update=True, **kwargs)
msg = message if this_info["mtype"] in ["audio", "video"] else message.reply_to_message
- fpath: str = await client.download_media(msg) # type: ignore
+ media_path = f"{DOWNLOAD_DIR}/{this_info['file_name'] or reply_info.get('file_name', '')}"
+ fpath: str = await client.download_media(msg, media_path) # type: ignore
prompt = f"请转录{matched['platform'].title()}视频作者【{vinfo['author']}】的一期节目的音频。\n该期节目标题: {vinfo['title']}\n节目简介: {description}"
res = await asr_file(fpath, engine=asr_engine, prompt=prompt, client=client, message=message, silent=True, **kwargs)
if res.get("error"):
src/tts/tts.py
@@ -5,7 +5,7 @@ from pathlib import Path
from pyrogram.client import Client
from pyrogram.types import Message
-from config import PREFIX, TTS
+from config import DOWNLOAD_DIR, PREFIX, TTS
from messages.parser import parse_msg
from messages.sender import send2tg
from messages.utils import blockquote, equal_prefix, set_reaction, startswith_prefix
@@ -49,7 +49,7 @@ async def text_to_speech(client: Client, message: Message, **kwargs):
# file
if info["mtype"] == "document":
if info["mime_type"].startswith("text/") or Path(info["file_name"]).suffix.lower() in [".txt", ".md"]:
- fpath: str = await client.download_media(message, in_memory=False) # type: ignore
+ fpath: str = await client.download_media(message, f"{DOWNLOAD_DIR}/{info['file_name']}") # type: ignore
texts = read_text(fpath).strip()
else:
await reaction_msg.reply(text="不支持该文件格式, 请以 `.txt` 格式发送", quote=True)
pyproject.toml
@@ -19,6 +19,7 @@ dependencies = [
"markitdown[docx,pdf,pptx,xls,xlsx]==0.1.2",
"openai==1.97.0",
"orjson==3.11.0",
+ "pathvalidate==3.3.1",
"pillow-heif==1.0.0",
"pillow>=11.2.1",
"puremagic==1.30",
uv.lock
@@ -241,6 +241,7 @@ dependencies = [
{ name = "markitdown", extra = ["docx", "pdf", "pptx", "xls", "xlsx"] },
{ name = "openai" },
{ name = "orjson" },
+ { name = "pathvalidate" },
{ name = "pillow" },
{ name = "pillow-heif" },
{ name = "puremagic" },
@@ -288,6 +289,7 @@ requires-dist = [
{ name = "markitdown", extras = ["docx", "pdf", "pptx", "xls", "xlsx"], specifier = "==0.1.2" },
{ name = "openai", specifier = "==1.97.0" },
{ name = "orjson", specifier = "==3.11.0" },
+ { name = "pathvalidate", specifier = "==3.3.1" },
{ name = "pillow", specifier = ">=11.2.1" },
{ name = "pillow-heif", specifier = "==1.0.0" },
{ name = "puremagic", specifier = "==1.30" },
@@ -1619,6 +1621,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c6/ac/dac4a63f978e4dcb3c6d3a78c4d8e0192a113d288502a1216950c41b1027/parso-0.8.4-py2.py3-none-any.whl", hash = "sha256:a418670a20291dacd2dddc80c377c5c3791378ee1e8d12bffc35420643d43f18", size = 103650, upload-time = "2024-04-05T09:43:53.299Z" },
]
+[[package]]
+name = "pathvalidate"
+version = "3.3.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fa/2a/52a8da6fe965dea6192eb716b357558e103aea0a1e9a8352ad575a8406ca/pathvalidate-3.3.1.tar.gz", hash = "sha256:b18c07212bfead624345bb8e1d6141cdcf15a39736994ea0b94035ad2b1ba177", size = 63262, upload-time = "2025-06-15T09:07:20.736Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305, upload-time = "2025-06-15T09:07:19.117Z" },
+]
+
[[package]]
name = "pdfminer-six"
version = "20250506"
@@ -2389,14 +2400,14 @@ wheels = [
[[package]]
name = "s3transfer"
-version = "0.13.0"
+version = "0.13.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "botocore" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/ed/5d/9dcc100abc6711e8247af5aa561fc07c4a046f72f659c3adea9a449e191a/s3transfer-0.13.0.tar.gz", hash = "sha256:f5e6db74eb7776a37208001113ea7aa97695368242b364d73e91c981ac522177", size = 150232, upload-time = "2025-05-22T19:24:50.245Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/6d/05/d52bf1e65044b4e5e27d4e63e8d1579dbdec54fce685908ae09bc3720030/s3transfer-0.13.1.tar.gz", hash = "sha256:c3fdba22ba1bd367922f27ec8032d6a1cf5f10c934fb5d68cf60fd5a23d936cf", size = 150589, upload-time = "2025-07-18T19:22:42.31Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/18/17/22bf8155aa0ea2305eefa3a6402e040df7ebe512d1310165eda1e233c3f8/s3transfer-0.13.0-py3-none-any.whl", hash = "sha256:0148ef34d6dd964d0d8cf4311b2b21c474693e57c2e069ec708ce043d2b527be", size = 85152, upload-time = "2025-05-22T19:24:48.703Z" },
+ { url = "https://files.pythonhosted.org/packages/6d/4f/d073e09df851cfa251ef7840007d04db3293a0482ce607d2b993926089be/s3transfer-0.13.1-py3-none-any.whl", hash = "sha256:a981aa7429be23fe6dfc13e80e4020057cbab622b08c0315288758d67cabc724", size = 85308, upload-time = "2025-07-18T19:22:40.947Z" },
]
[[package]]