feat:初版
This commit is contained in:
1
plugins/GrokVideo/__init__.py
Normal file
1
plugins/GrokVideo/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# GrokVideo 插件
|
||||
624
plugins/GrokVideo/main.py
Normal file
624
plugins/GrokVideo/main.py
Normal file
@@ -0,0 +1,624 @@
|
||||
"""
|
||||
Grok 视频生成插件
|
||||
|
||||
用户引用图片并发送 /视频 提示词 来生成视频
|
||||
支持队列系统和积分制
|
||||
"""
|
||||
|
||||
import re
|
||||
import tomllib
|
||||
import httpx
|
||||
import xml.etree.ElementTree as ET
|
||||
import asyncio
|
||||
import pymysql
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from typing import Dict, Optional
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from utils.plugin_base import PluginBase
|
||||
from utils.decorators import on_text_message
|
||||
from minio import Minio
|
||||
from minio.error import S3Error
|
||||
|
||||
|
||||
# 定义引用消息装饰器
|
||||
def on_quote_message(priority=50):
|
||||
"""引用消息装饰器"""
|
||||
def decorator(func):
|
||||
setattr(func, '_event_type', 'quote_message') # 修复：应该是 quote_message
|
||||
setattr(func, '_priority', min(max(priority, 0), 99))
|
||||
return func
|
||||
return decorator
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoTask:
|
||||
"""视频生成任务"""
|
||||
user_wxid: str
|
||||
from_wxid: str
|
||||
prompt: str
|
||||
cdnurl: str
|
||||
aeskey: str
|
||||
is_group: bool
|
||||
timestamp: datetime
|
||||
|
||||
|
||||
class GrokVideo(PluginBase):
|
||||
"""Grok 视频生成插件"""
|
||||
|
||||
description = "使用 Grok AI 根据图片和提示词生成视频（支持队列和积分系统）"
|
||||
author = "ShiHao"
|
||||
version = "2.0.0"
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = None
|
||||
self.task_queue: asyncio.Queue = None
|
||||
self.processing_tasks: Dict[str, VideoTask] = {} # 正在处理的任务
|
||||
self.worker_task = None
|
||||
self.minio_client = None
|
||||
|
||||
async def async_init(self):
|
||||
"""插件异步初始化"""
|
||||
config_path = Path(__file__).parent / "config.toml"
|
||||
with open(config_path, "rb") as f:
|
||||
self.config = tomllib.load(f)
|
||||
|
||||
self.api_url = f"{self.config['api']['server_url'].rstrip('/')}/v1/chat/completions"
|
||||
|
||||
# 初始化MinIO客户端
|
||||
self.minio_client = Minio(
|
||||
"101.201.65.129:19000",
|
||||
access_key="admin",
|
||||
secret_key="80012029Lz",
|
||||
secure=False
|
||||
)
|
||||
self.minio_bucket = "wechat"
|
||||
|
||||
# 初始化队列
|
||||
max_queue_size = self.config.get("queue", {}).get("max_queue_size", 10)
|
||||
self.task_queue = asyncio.Queue(maxsize=max_queue_size)
|
||||
|
||||
# 启动工作线程
|
||||
max_concurrent = self.config.get("queue", {}).get("max_concurrent", 1)
|
||||
self.worker_task = asyncio.create_task(self._queue_worker())
|
||||
|
||||
logger.success(f"Grok 视频生成插件已加载")
|
||||
logger.info(f"API: {self.api_url}")
|
||||
logger.info(f"队列配置: 最大并发={max_concurrent}, 最大队列长度={max_queue_size}")
|
||||
logger.info(f"积分系统: {'启用' if self.config.get('points', {}).get('enabled') else '禁用'}")
|
||||
if self.config.get('points', {}).get('enabled'):
|
||||
logger.info(f"每次生成消耗: {self.config['points']['cost']} 积分")
|
||||
|
||||
def get_db_connection(self):
|
||||
"""获取数据库连接"""
|
||||
db_config = self.config["database"]
|
||||
return pymysql.connect(
|
||||
host=db_config["host"],
|
||||
port=db_config["port"],
|
||||
user=db_config["user"],
|
||||
password=db_config["password"],
|
||||
database=db_config["database"],
|
||||
charset=db_config["charset"],
|
||||
autocommit=True
|
||||
)
|
||||
|
||||
def get_user_points(self, wxid: str) -> int:
|
||||
"""获取用户积分"""
|
||||
try:
|
||||
with self.get_db_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
sql = "SELECT points FROM user_signin WHERE wxid = %s"
|
||||
cursor.execute(sql, (wxid,))
|
||||
result = cursor.fetchone()
|
||||
return result[0] if result else 0
|
||||
except Exception as e:
|
||||
logger.error(f"获取用户积分失败: {e}")
|
||||
return 0
|
||||
|
||||
def deduct_points(self, wxid: str, points: int) -> bool:
|
||||
"""扣除用户积分"""
|
||||
try:
|
||||
with self.get_db_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 检查积分是否足够
|
||||
sql_check = "SELECT points FROM user_signin WHERE wxid = %s"
|
||||
cursor.execute(sql_check, (wxid,))
|
||||
result = cursor.fetchone()
|
||||
|
||||
if not result or result[0] < points:
|
||||
return False
|
||||
|
||||
# 扣除积分
|
||||
sql_update = "UPDATE user_signin SET points = points - %s WHERE wxid = %s"
|
||||
cursor.execute(sql_update, (points, wxid))
|
||||
logger.info(f"用户 {wxid} 扣除 {points} 积分")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"扣除用户积分失败: {e}")
|
||||
return False
|
||||
|
||||
def is_admin(self, wxid: str) -> bool:
|
||||
"""检查是否是管理员"""
|
||||
admins = self.config.get("points", {}).get("admins", [])
|
||||
return wxid in admins
|
||||
|
||||
async def upload_video_to_minio(self, local_file: str, original_filename: str = "") -> str:
|
||||
"""上传视频到MinIO"""
|
||||
try:
|
||||
# 生成唯一文件名
|
||||
file_ext = Path(local_file).suffix
|
||||
unique_id = uuid.uuid4().hex
|
||||
|
||||
if original_filename:
|
||||
# 使用原始文件名（去掉扩展名）+ 唯一ID + 扩展名
|
||||
original_name = Path(original_filename).stem
|
||||
# 清理文件名中的特殊字符
|
||||
import re
|
||||
original_name = re.sub(r'[^\w\-_\.]', '_', original_name)
|
||||
filename = f"{original_name}_{unique_id}{file_ext}"
|
||||
else:
|
||||
filename = f"grok_video_{unique_id}{file_ext}"
|
||||
|
||||
object_name = f"videos/{datetime.now().strftime('%Y%m%d')}/{filename}"
|
||||
|
||||
# 上传文件
|
||||
await asyncio.to_thread(
|
||||
self.minio_client.fput_object,
|
||||
self.minio_bucket,
|
||||
object_name,
|
||||
local_file
|
||||
)
|
||||
|
||||
# 返回访问URL
|
||||
url = f"http://101.201.65.129:19000/{self.minio_bucket}/{object_name}"
|
||||
logger.info(f"视频上传成功: {url}")
|
||||
return url
|
||||
|
||||
except S3Error as e:
|
||||
logger.error(f"上传视频到MinIO失败: {e}")
|
||||
return ""
|
||||
|
||||
async def _queue_worker(self):
|
||||
"""队列工作线程"""
|
||||
logger.info("视频生成队列工作线程已启动")
|
||||
while True:
|
||||
try:
|
||||
# 从队列获取任务
|
||||
task_data = await self.task_queue.get()
|
||||
bot, task = task_data
|
||||
|
||||
# 处理任务
|
||||
await self._process_video_task(bot, task)
|
||||
|
||||
# 标记任务完成
|
||||
self.task_queue.task_done()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"队列工作线程错误: {e}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def _process_video_task(self, bot, task: VideoTask):
|
||||
"""处理视频生成任务"""
|
||||
logger.info(f"开始处理视频任务: user={task.user_wxid}, prompt={task.prompt}")
|
||||
|
||||
try:
|
||||
# 下载图片并转换为 base64
|
||||
image_base64 = await self._download_and_encode_image(bot, task.cdnurl, task.aeskey)
|
||||
if not image_base64:
|
||||
await bot.send_text(task.from_wxid, "❌ 无法下载图片")
|
||||
return
|
||||
|
||||
# 调用 Grok API
|
||||
video_url = await self._call_grok_api(task.prompt, image_base64)
|
||||
|
||||
if video_url:
|
||||
# 下载视频
|
||||
video_path = await self._download_video(video_url)
|
||||
|
||||
if video_path:
|
||||
# 上传视频到MinIO
|
||||
logger.info(f"上传视频到MinIO: {video_path}")
|
||||
video_filename = Path(video_path).name
|
||||
minio_url = await self.upload_video_to_minio(video_path, video_filename)
|
||||
|
||||
# 发送视频文件
|
||||
logger.info(f"准备发送视频文件: {video_path}")
|
||||
video_sent = await bot.send_file(task.from_wxid, video_path)
|
||||
|
||||
if video_sent:
|
||||
# 通知MessageLogger记录机器人发送的视频消息
|
||||
try:
|
||||
from plugins.MessageLogger.main import MessageLogger
|
||||
message_logger = MessageLogger.get_instance()
|
||||
if message_logger and minio_url:
|
||||
await message_logger.save_bot_message(
|
||||
task.from_wxid,
|
||||
f"[视频] {task.prompt}",
|
||||
"video",
|
||||
minio_url
|
||||
)
|
||||
logger.info(f"已记录机器人视频消息到数据库: {minio_url}")
|
||||
except Exception as e:
|
||||
logger.warning(f"记录机器人视频消息失败: {e}")
|
||||
|
||||
# 如果启用了积分系统，显示积分信息
|
||||
points_config = self.config.get("points", {})
|
||||
if points_config.get("enabled", False):
|
||||
if self.is_admin(task.user_wxid):
|
||||
# 管理员免费使用
|
||||
success_msg = "✅ 视频生成成功！\n🎖️ 管理员免费使用"
|
||||
else:
|
||||
# 普通用户显示积分消耗
|
||||
cost = points_config.get("cost", 50)
|
||||
remaining_points = self.get_user_points(task.user_wxid)
|
||||
success_msg = f"✅ 视频生成成功！\n💎 本次消耗：{cost} 积分\n💰 剩余积分：{remaining_points}"
|
||||
else:
|
||||
# 积分系统未启用
|
||||
success_msg = "✅ 视频生成成功！"
|
||||
|
||||
await bot.send_text(task.from_wxid, success_msg)
|
||||
logger.success(f"视频文件发送成功: {video_path}")
|
||||
|
||||
# 清理本地文件
|
||||
try:
|
||||
Path(video_path).unlink()
|
||||
logger.info(f"已清理本地视频文件: {video_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"清理本地视频文件失败: {e}")
|
||||
else:
|
||||
await bot.send_text(task.from_wxid, "❌ 视频发送失败")
|
||||
logger.error(f"视频文件发送失败: {video_path}")
|
||||
else:
|
||||
await bot.send_text(task.from_wxid, "❌ 视频下载失败")
|
||||
else:
|
||||
await bot.send_text(task.from_wxid, "❌ 视频生成失败，请稍后再试")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理视频任务失败: {e}")
|
||||
await bot.send_text(task.from_wxid, f"❌ 视频生成失败: {str(e)}")
|
||||
|
||||
@on_quote_message(priority=90)
|
||||
async def handle_video_command(self, bot, message: dict):
|
||||
"""处理 /视频 命令（引用消息）"""
|
||||
content = message.get("Content", "").strip()
|
||||
from_wxid = message.get("FromWxid", "")
|
||||
sender_wxid = message.get("SenderWxid", "")
|
||||
is_group = message.get("IsGroup", False)
|
||||
|
||||
# 获取实际发送者（群聊中使用SenderWxid，私聊使用FromWxid）
|
||||
user_wxid = sender_wxid if is_group else from_wxid
|
||||
|
||||
# 解析 XML 获取标题和引用消息
|
||||
try:
|
||||
root = ET.fromstring(content)
|
||||
title = root.find(".//title")
|
||||
if title is None or not title.text:
|
||||
return
|
||||
|
||||
title_text = title.text.strip()
|
||||
|
||||
# 检查是否是 /视频 命令
|
||||
if not title_text.startswith("/视频"):
|
||||
return
|
||||
|
||||
# 检查是否启用
|
||||
if not self.config["behavior"]["enabled"]:
|
||||
return False
|
||||
|
||||
# 检查群聊过滤
|
||||
if is_group:
|
||||
enabled_groups = self.config["behavior"]["enabled_groups"]
|
||||
disabled_groups = self.config["behavior"]["disabled_groups"]
|
||||
|
||||
if from_wxid in disabled_groups:
|
||||
return False
|
||||
if enabled_groups and from_wxid not in enabled_groups:
|
||||
return False
|
||||
|
||||
# 提取提示词
|
||||
prompt = title_text[3:].strip() # 去掉 "/视频"
|
||||
if not prompt:
|
||||
await bot.send_text(from_wxid, "❌ 请提供提示词，例如：/视频 让太阳升起来")
|
||||
return False
|
||||
|
||||
# 获取引用消息中的图片信息
|
||||
refermsg = root.find(".//refermsg")
|
||||
if refermsg is None:
|
||||
await bot.send_text(from_wxid, "❌ 请引用一张图片后使用此命令")
|
||||
return False
|
||||
|
||||
# 解析引用消息的内容（需要解码 HTML 实体）
|
||||
refer_content = refermsg.find("content")
|
||||
if refer_content is None or not refer_content.text:
|
||||
await bot.send_text(from_wxid, "❌ 引用的消息中没有图片")
|
||||
return False
|
||||
|
||||
# 解码 HTML 实体
|
||||
import html
|
||||
refer_xml = html.unescape(refer_content.text)
|
||||
refer_root = ET.fromstring(refer_xml)
|
||||
|
||||
# 提取图片信息
|
||||
img = refer_root.find(".//img")
|
||||
if img is None:
|
||||
await bot.send_text(from_wxid, "❌ 引用的消息中没有图片")
|
||||
return False
|
||||
|
||||
# 获取图片的 CDN URL 和 AES Key
|
||||
cdnbigimgurl = img.get("cdnbigimgurl", "")
|
||||
aeskey = img.get("aeskey", "")
|
||||
|
||||
if not cdnbigimgurl or not aeskey:
|
||||
await bot.send_text(from_wxid, "❌ 无法获取图片信息")
|
||||
return False
|
||||
|
||||
logger.info(f"收到视频生成请求: user={user_wxid}, prompt={prompt}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"解析引用消息失败: {e}")
|
||||
return
|
||||
|
||||
# 检查积分系统
|
||||
points_config = self.config.get("points", {})
|
||||
if points_config.get("enabled", False):
|
||||
# 检查是否是管理员
|
||||
if not self.is_admin(user_wxid):
|
||||
# 检查积分
|
||||
cost = points_config.get("cost", 50)
|
||||
current_points = self.get_user_points(user_wxid)
|
||||
|
||||
if current_points < cost:
|
||||
await bot.send_text(
|
||||
from_wxid,
|
||||
f"❌ 积分不足！\n💰 当前积分：{current_points}\n💎 需要积分：{cost}\n\n请先签到获取积分~"
|
||||
)
|
||||
return False
|
||||
|
||||
# 扣除积分
|
||||
if not self.deduct_points(user_wxid, cost):
|
||||
await bot.send_text(from_wxid, "❌ 积分扣除失败，请稍后重试")
|
||||
return False
|
||||
|
||||
logger.info(f"用户 {user_wxid} 已扣除 {cost} 积分，剩余 {current_points - cost} 积分")
|
||||
else:
|
||||
logger.info(f"管理员 {user_wxid} 免费使用")
|
||||
|
||||
# 检查队列是否已满
|
||||
if self.task_queue.full():
|
||||
await bot.send_text(from_wxid, f"❌ 队列已满（{self.task_queue.qsize()}/{self.task_queue.maxsize}），请稍后再试")
|
||||
# 如果扣除了积分，需要退还
|
||||
if points_config.get("enabled", False) and not self.is_admin(user_wxid):
|
||||
# 退还积分（这里简化处理，直接加回去）
|
||||
try:
|
||||
with self.get_db_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
sql = "UPDATE user_signin SET points = points + %s WHERE wxid = %s"
|
||||
cursor.execute(sql, (points_config.get("cost", 50), user_wxid))
|
||||
logger.info(f"队列已满，退还积分给用户 {user_wxid}")
|
||||
except Exception as e:
|
||||
logger.error(f"退还积分失败: {e}")
|
||||
return False
|
||||
|
||||
# 创建任务
|
||||
task = VideoTask(
|
||||
user_wxid=user_wxid,
|
||||
from_wxid=from_wxid,
|
||||
prompt=prompt,
|
||||
cdnurl=cdnbigimgurl,
|
||||
aeskey=aeskey,
|
||||
is_group=is_group,
|
||||
timestamp=datetime.now()
|
||||
)
|
||||
|
||||
# 添加到队列
|
||||
try:
|
||||
await self.task_queue.put((bot, task))
|
||||
queue_position = self.task_queue.qsize()
|
||||
|
||||
if queue_position == 1:
|
||||
await bot.send_text(from_wxid, "🎥 正在生成视频，请稍候（预计需要几分钟）...")
|
||||
else:
|
||||
await bot.send_text(
|
||||
from_wxid,
|
||||
f"📋 已加入队列\n🔢 当前排队位置：第 {queue_position} 位\n⏰ 请耐心等待..."
|
||||
)
|
||||
|
||||
logger.success(f"任务已加入队列: user={user_wxid}, position={queue_position}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加任务到队列失败: {e}")
|
||||
await bot.send_text(from_wxid, "❌ 添加任务失败，请稍后重试")
|
||||
# 退还积分
|
||||
if points_config.get("enabled", False) and not self.is_admin(user_wxid):
|
||||
try:
|
||||
with self.get_db_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
sql = "UPDATE user_signin SET points = points + %s WHERE wxid = %s"
|
||||
cursor.execute(sql, (points_config.get("cost", 50), user_wxid))
|
||||
logger.info(f"添加任务失败，退还积分给用户 {user_wxid}")
|
||||
except Exception as e:
|
||||
logger.error(f"退还积分失败: {e}")
|
||||
|
||||
return False # 不阻止后续处理
|
||||
|
||||
async def _download_and_encode_image(self, bot, cdnurl: str, aeskey: str) -> str:
|
||||
"""下载图片并转换为 base64"""
|
||||
try:
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
import base64
|
||||
|
||||
# 创建临时目录
|
||||
temp_dir = Path(__file__).parent / "temp"
|
||||
temp_dir.mkdir(exist_ok=True)
|
||||
|
||||
# 生成临时文件名
|
||||
filename = f"temp_{datetime.now():%Y%m%d_%H%M%S}_{uuid.uuid4().hex[:8]}.jpg"
|
||||
save_path = str((temp_dir / filename).resolve())
|
||||
|
||||
# 使用 CDN 下载 API 下载图片
|
||||
logger.info(f"正在下载图片: {cdnurl[:50]}...")
|
||||
success = await bot.cdn_download(cdnurl, aeskey, save_path, file_type=2)
|
||||
|
||||
if not success:
|
||||
# 如果中图下载失败，尝试原图
|
||||
logger.warning("中图下载失败，尝试下载原图...")
|
||||
success = await bot.cdn_download(cdnurl, aeskey, save_path, file_type=1)
|
||||
|
||||
if not success:
|
||||
logger.error("图片下载失败")
|
||||
return ""
|
||||
|
||||
# 等待文件写入完成并检查文件是否存在
|
||||
import asyncio
|
||||
import os
|
||||
max_wait = 10 # 最多等待10秒
|
||||
wait_time = 0
|
||||
|
||||
while wait_time < max_wait:
|
||||
if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
|
||||
logger.info(f"文件已就绪: {save_path}")
|
||||
break
|
||||
await asyncio.sleep(0.5)
|
||||
wait_time += 0.5
|
||||
|
||||
if not os.path.exists(save_path):
|
||||
logger.error(f"文件下载超时或失败: {save_path}")
|
||||
return ""
|
||||
|
||||
# 读取图片并转换为 base64
|
||||
with open(save_path, "rb") as f:
|
||||
image_data = base64.b64encode(f.read()).decode()
|
||||
|
||||
# 删除临时文件
|
||||
try:
|
||||
Path(save_path).unlink()
|
||||
except:
|
||||
pass
|
||||
|
||||
return f"data:image/jpeg;base64,{image_data}"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"下载图片失败: {e}")
|
||||
return ""
|
||||
|
||||
async def _call_grok_api(self, prompt: str, image_base64: str) -> str:
|
||||
"""调用 Grok API 生成视频"""
|
||||
api_key = self.config["api"]["api_key"]
|
||||
if not api_key:
|
||||
raise Exception("未配置 API Key")
|
||||
|
||||
payload = {
|
||||
"model": self.config["api"]["model_id"],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": image_base64}}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}"
|
||||
}
|
||||
|
||||
timeout = httpx.Timeout(connect=10.0, read=self.config["api"]["timeout"], write=10.0, pool=10.0)
|
||||
|
||||
# 配置代理
|
||||
connector = None
|
||||
proxy_config = self.config.get("proxy", {})
|
||||
if proxy_config.get("enabled", False):
|
||||
from aiohttp_socks import ProxyConnector
|
||||
proxy_type = proxy_config.get("type", "socks5").upper()
|
||||
proxy_host = proxy_config.get("host", "127.0.0.1")
|
||||
proxy_port = proxy_config.get("port", 7890)
|
||||
proxy_username = proxy_config.get("username")
|
||||
proxy_password = proxy_config.get("password")
|
||||
|
||||
# 构建代理 URL
|
||||
if proxy_username and proxy_password:
|
||||
proxy_url = f"{proxy_type}://{proxy_username}:{proxy_password}@{proxy_host}:{proxy_port}"
|
||||
else:
|
||||
proxy_url = f"{proxy_type}://{proxy_host}:{proxy_port}"
|
||||
|
||||
logger.info(f"使用代理: {proxy_type}://{proxy_host}:{proxy_port}")
|
||||
|
||||
# httpx 使用不同的代理配置方式
|
||||
proxy = None
|
||||
if proxy_config.get("enabled", False):
|
||||
proxy_type = proxy_config.get("type", "socks5")
|
||||
proxy_host = proxy_config.get("host", "127.0.0.1")
|
||||
proxy_port = proxy_config.get("port", 7890)
|
||||
proxy = f"{proxy_type}://{proxy_host}:{proxy_port}"
|
||||
logger.info(f"使用代理: {proxy}")
|
||||
|
||||
async with httpx.AsyncClient(timeout=timeout, proxy=proxy) as client:
|
||||
response = await client.post(self.api_url, json=payload, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"API 错误: {response.status_code}")
|
||||
|
||||
result = response.json()
|
||||
|
||||
# 提取视频 URL
|
||||
content = result["choices"][0]["message"]["content"]
|
||||
video_url = self._extract_video_url(content)
|
||||
|
||||
if not video_url:
|
||||
raise Exception("未能从响应中提取视频 URL")
|
||||
|
||||
logger.info(f"获取到视频 URL: {video_url}")
|
||||
return video_url
|
||||
|
||||
def _extract_video_url(self, content: str) -> str:
|
||||
"""从响应内容中提取视频 URL"""
|
||||
# 尝试从 HTML video 标签提取
|
||||
match = re.search(r'<video[^>]*src=["\']([^"\'>]+)["\']', content, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1)
|
||||
|
||||
# 尝试提取直接的 .mp4 URL
|
||||
match = re.search(r'(https?://[^\s<>"\')\]]+\.mp4(?:\?[^\s<>"\')\]]*)?)', content, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1)
|
||||
|
||||
return ""
|
||||
|
||||
async def _download_video(self, video_url: str) -> str:
|
||||
"""下载视频到本地"""
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
|
||||
videos_dir = Path(__file__).parent / "videos"
|
||||
videos_dir.mkdir(exist_ok=True)
|
||||
|
||||
filename = f"grok_{datetime.now():%Y%m%d_%H%M%S}_{uuid.uuid4().hex[:8]}.mp4"
|
||||
file_path = videos_dir / filename
|
||||
|
||||
timeout = httpx.Timeout(connect=10.0, read=300.0, write=10.0, pool=10.0)
|
||||
|
||||
# 配置代理
|
||||
proxy = None
|
||||
proxy_config = self.config.get("proxy", {})
|
||||
if proxy_config.get("enabled", False):
|
||||
proxy_type = proxy_config.get("type", "socks5")
|
||||
proxy_host = proxy_config.get("host", "127.0.0.1")
|
||||
proxy_port = proxy_config.get("port", 7890)
|
||||
proxy = f"{proxy_type}://{proxy_host}:{proxy_port}"
|
||||
logger.debug(f"下载视频使用代理: {proxy}")
|
||||
|
||||
async with httpx.AsyncClient(timeout=timeout, proxy=proxy) as client:
|
||||
response = await client.get(video_url)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
|
||||
return str(file_path.resolve())
|
||||
Reference in New Issue
Block a user