liuwei_x_x/abot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化IO问题，使用异步方案进行视频下载等操作。

Browse Source
This commit is contained in:
liuwei
2025-06-16 10:33:26 +08:00
parent 02a387628c
commit ed324eaa24
3 changed files with 170 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
plugins/message_summary/main.py
View File

@@ -4,7 +4,8 @@ import time
from pathlib import Path
from typing import Dict, Any, Tuple, Optional, List
import requests
import aiohttp
from loguru import logger
from base.plugin_common.message_plugin_interface import MessagePluginInterface
from base.plugin_common.plugin_interface import PluginStatus
@@ -205,13 +206,12 @@ class MessageSummaryPlugin(MessagePluginInterface):
}
try:
# 发送POST请求
response = requests.post(self._api_url, headers=headers, json=data)
async with aiohttp.ClientSession() as session:
async with session.post(self._api_url, headers=headers, json=data) as response:
response.raise_for_status() # 检查请求是否成功
response_data = await response.json()
# 解析响应
response_data = response.json()
self.LOG.info(f"Dify API响应状态码: {response.status_code}")
self.LOG.info(f"Dify API响应状态码: {response.status}")
self.LOG.debug(f"响应数据: {json.dumps(response_data, ensure_ascii=False, indent=2)}")
# 提取回答内容
@@ -244,7 +244,7 @@ class MessageSummaryPlugin(MessagePluginInterface):
# 返回文本内容和图片路径
return answer, spath
except requests.exceptions.RequestException as e:
except aiohttp.ClientError as e:
self.LOG.error(f"请求Dify API时出错: {e}")
return f"生成总结时出错: {str(e)}", None

45
plugins/video/main.py
View File

@@ -1,7 +1,9 @@
import time
from loguru import logger
import os
import aiohttp
import aiofiles
import asyncio
from loguru import logger
import requests
from typing import Dict, Any, List, Optional, Tuple
from pathlib import Path
@@ -133,7 +135,7 @@ class VideoPlugin(MessagePluginInterface):
save_path = os.path.join(self.download_dir, video_filename)
self.LOG.info(f"开始下载视频到: {save_path}")
file_abspath, first_frame = self._download_stream(" http://api.yujn.cn/api/heisis.php?type=video", save_path)
file_abspath, first_frame = await self._download_stream(" http://api.yujn.cn/api/heisis.php?type=video", save_path)
if not file_abspath or not os.path.exists(file_abspath) or not file_abspath.endswith("mp4"):
self.LOG.error(f"视频下载失败，文件路径: {file_abspath}")
@@ -154,9 +156,9 @@ class VideoPlugin(MessagePluginInterface):
sender)
return False, f"处理出错: {e}"
def _download_stream(self, url, save_path):
async def _download_stream(self, url, save_path):
"""
从指定URL读取视频流并保存到本地
从指定URL读取视频流并保存到本地（异步版本）
:param url: 视频流的URL
:param save_path: 本地保存路径（包含文件名，例如 "video.mp4"）
"""
@@ -170,10 +172,11 @@ class VideoPlugin(MessagePluginInterface):
"Connection": "keep-alive",
"Referer": "https://api.guiguiya.com/"
}
response = requests.get(url, stream=True, timeout=30, headers=headers, allow_redirects=True)
async with aiohttp.ClientSession() as session:
async with session.get(url, headers=headers, allow_redirects=True) as response:
# 检查请求是否成功
response.raise_for_status() # 如果状态码不是200，将抛出异常
response.raise_for_status()
# 确保保存路径的目录存在
save_dir = os.path.dirname(save_path)
@@ -187,7 +190,8 @@ class VideoPlugin(MessagePluginInterface):
if "video" not in content_type and "application/octet-stream" not in content_type:
self.LOG.warning(f"警告: 返回的可能不是视频流，Content-Type: {content_type}")
self.LOG.warning(f"响应内容预览: {response.text[:100]}") # 打印前100字符查看
text = await response.text()
self.LOG.warning(f"响应内容预览: {text[:100]}") # 打印前100字符查看
return None, None
# 以二进制写入模式保存流数据
@@ -196,10 +200,10 @@ class VideoPlugin(MessagePluginInterface):
self.LOG.info(f"预期下载大小: {total_size} 字节")
downloaded_size = 0
with open(save_path, "wb") as file:
for chunk in response.iter_content(chunk_size=1024): # 分块读取，每块1KB
async with aiofiles.open(save_path, "wb") as file:
async for chunk in response.content.iter_chunked(1024): # 分块读取，每块1KB
if chunk: # 过滤空块
file.write(chunk)
await file.write(chunk)
downloaded_size += len(chunk)
self.LOG.info(f"视频已下载到: {save_path}, 大小: {downloaded_size} 字节")
@@ -222,8 +226,8 @@ class VideoPlugin(MessagePluginInterface):
self.LOG.error(f"下载的文件太小，可能不是有效视频: {file_size} 字节")
# 尝试读取文件内容以诊断问题
try:
with open(save_path, 'rb') as f:
content_preview = f.read(200)
async with aiofiles.open(save_path, 'rb') as f:
content_preview = await f.read(200)
self.LOG.warning(f"文件内容预览(十六进制): {content_preview.hex()[:100]}")
except Exception as e:
self.LOG.error(f"读取文件内容失败: {e}")
@@ -231,7 +235,7 @@ class VideoPlugin(MessagePluginInterface):
# 加入首帧下载
first_frame_path = os.path.join(self.download_dir, f"frame_{int(time.time())}.jpg")
first_frame = self._get_first_frame(save_path, first_frame_path)
first_frame = await self._get_first_frame(save_path, first_frame_path)
if not first_frame or not os.path.exists(first_frame):
self.LOG.warning(f"无法提取首帧，使用默认图片")
@@ -239,7 +243,7 @@ class VideoPlugin(MessagePluginInterface):
return os.path.abspath(save_path), first_frame
except requests.RequestException as e:
except aiohttp.ClientError as e:
self.LOG.error(f"请求失败: {e}")
except IOError as e:
self.LOG.error(f"文件写入失败: {e}")
@@ -247,15 +251,18 @@ class VideoPlugin(MessagePluginInterface):
self.LOG.error(f"下载视频时发生未知错误: {e}")
return None, None
def _get_first_frame(self, video_path, output_path):
async def _get_first_frame(self, video_path, output_path):
"""
提取视频的第一帧并保存为图片
提取视频的第一帧并保存为图片（异步版本）
:param video_path: 视频文件路径
:param output_path: 输出图片路径
:return: 输出图片的绝对路径，如果失败则返回None
"""
try:
self.LOG.info(f"开始提取视频首帧: {video_path}")
# 使用 asyncio.to_thread 包装 OpenCV 操作
def extract_frame():
# 打开视频文件
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
@@ -282,6 +289,10 @@ class VideoPlugin(MessagePluginInterface):
cap.release()
return os.path.abspath(output_path)
# 在线程池中执行 OpenCV 操作
result = await asyncio.to_thread(extract_frame)
return result
except Exception as e:
self.LOG.error(f"提取视频首帧时出错: {e}")
return None

21
plugins/video_man/main.py
View File

@@ -1,5 +1,7 @@
import os
import time
import asyncio
import aiofiles
from typing import Dict, Any, List, Optional, Tuple
import aiohttp
@@ -170,17 +172,17 @@ class VideoManPlugin(MessagePluginInterface):
self.LOG.error(f"无法下载视频，HTTP状态码: {video_response.status}")
return None
# 保存视频
with open(save_path, "wb") as file:
# 使用 aiofiles 异步保存视频
async with aiofiles.open(save_path, "wb") as file:
async for chunk in video_response.content.iter_chunked(1024):
if chunk: # 过滤空块
file.write(chunk)
await file.write(chunk)
abs_path = os.path.abspath(save_path)
self.LOG.info(f"视频已下载至: {abs_path}")
first_frame_path = os.path.join(self.download_dir, f"frame_{int(time.time())}.jpg")
first_frame = self._get_first_frame(save_path, first_frame_path)
first_frame = await self._get_first_frame(save_path, first_frame_path)
return abs_path, first_frame
@@ -193,15 +195,18 @@ class VideoManPlugin(MessagePluginInterface):
return None
def _get_first_frame(self, video_path, output_path):
async def _get_first_frame(self, video_path, output_path):
"""
提取视频的第一帧并保存为图片
提取视频的第一帧并保存为图片（异步版本）
:param video_path: 视频文件路径
:param output_path: 输出图片路径
:return: 输出图片的绝对路径，如果失败则返回None
"""
try:
self.LOG.info(f"开始提取视频首帧: {video_path}")
# 使用 asyncio.to_thread 包装 OpenCV 操作
def extract_frame():
# 打开视频文件
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
@@ -228,6 +233,10 @@ class VideoManPlugin(MessagePluginInterface):
cap.release()
return os.path.abspath(output_path)
# 在线程池中执行 OpenCV 操作
result = await asyncio.to_thread(extract_frame)
return result
except Exception as e:
self.LOG.error(f"提取视频首帧时出错: {e}")
return None