import json import os import time import uuid from typing import Literal, Optional, Tuple import cv2 import google.generativeai as genai import requests import schedule from google.generativeai.types import (File, GenerateContentResponse, HarmBlockThreshold, HarmCategory) from loguru import logger from common.aliyun_log import AliyunLogger from common.common_log import Common from common.feishu_data import Material from common.redis import SyncRedisHelper ENV = os.getenv('ENV', 'dev') API_KEY = os.getenv('API_KEY') TASK_TYPE = os.getenv('TASK_TYPE') PROXY_ADDR = 'http://localhost:1081' CACHE_DIR = '/app/cache/' if ENV == 'prod' else os.path.expanduser('~/Downloads/') SAMPLE_DATA = { "一、基础信息": { "关键维度": "", "内容选题": "", "视频主题": "", "视频关键词":"" }, "二、主体和场景": { "视频主体": "", "视频场景": [] }, "三、情感与风格": { "情感倾向":"", "视频风格":"" }, "四、视频传播性与画像": { "片尾引导": {}, "传播性判断": "", "视频用户画像": {} }, "五、音画细节": { "音频细节": {}, "视频水印": {}, "视频字幕": {}, "视频口播": "" }, "六、封面信息": { "封面主体": "", "人物个数": "", "文字数量": "", "文字关键字":[], "封面主题":"" }, "七、人物与场景": { "知名人物":"", "人物年龄段":"", "场景描述": "" }, "八、时效性与分类": { "时效性":"", "视频一级分类": "", "二级分类": ["品类- 、分数-", "品类- 、分数-", "品类- 、分数-"] }, "九、标题理解": { } } if ENV == 'dev': os.environ['http_proxy'] = PROXY_ADDR os.environ['https_proxy'] = PROXY_ADDR def get_redis_task(task_type: Literal['recommend', 'top']) -> Optional[bytes]: redis_key = f'task:video_ai_{task_type}' redis_task: bytes = SyncRedisHelper().get_client().rpop(redis_key) if redis_task: logger.success(f'[+] 获取到 {task_type} 类型任务: {redis_task.decode()}') else: logger.error(f'[+] 未获取到 {task_type} 类型任务') return redis_task def get_video_duration(video_link: str) -> int: cap = cv2.VideoCapture(video_link) if cap.isOpened(): rate = cap.get(5) frame_num = cap.get(7) duration = int(frame_num / rate) return duration return 0 def download_video(video_link: str) -> Optional[str]: file_path = os.path.join(CACHE_DIR, f'{str(uuid.uuid4())}.mp4') for _ in range(3): try: response = requests.get(url=video_link, timeout=360) if response.status_code == 200: with open(file_path, 'wb') as f: f.write(response.content) logger.info(f'[+] 视频链接: {video_link}, 存储地址: {file_path}') return file_path except Exception: time.sleep(1) continue return def upload_video(video_path: str, redis_task) -> Optional[Tuple[File, str]]: try: file = genai.upload_file(path=video_path) while True: if file.state.name == 'PROCESSING': time.sleep(1) file = genai.get_file(name=file.name) else: return file, file.state.name except Exception as e: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], '', redis_task['type'], redis_task['partition'], f"[+] 上传视频失败: {e}" ) logger.error(f'[+] 上传视频失败: {e}') return def create_model_cache(redis_task) -> Optional[genai.GenerativeModel]: try: model = genai.GenerativeModel( model_name='gemini-1.5-flash', generation_config={'response_mime_type': 'application/json'}, safety_settings={HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE}, ) logger.info('[+] 创建缓存模型成功') return model except Exception as e: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], '', redis_task['type'], redis_task['partition'], f"[+] 视频创建缓存内容，并返回生成模型异常信息: {e}" ) logger.error(f'视频创建缓存内容，并返回生成模型异常信息: {e}') Common.logger('ai').info(f'视频创建缓存内容，并返回生成模型异常信息: {e}') return def analyze_video(model: genai.GenerativeModel, google_file: File, prompt: str, redis_task, title) -> Optional[GenerateContentResponse]: try: session = model.start_chat(history=[]) if title: content = { 'parts': [ google_file, f'{prompt}\n输出返回格式样例:\n{SAMPLE_DATA}', ], } return session.send_message(content=content) else: content = { 'parts': [ google_file, f'{prompt}\n九、标题理解：请对"{title}"标题进行分析概括3～5个词[]":\n输出返回格式样例:\n{SAMPLE_DATA}', ], } return session.send_message(content=content) except Exception as e: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], '', redis_task['type'], redis_task['partition'], f"[+] 视频处理请求失败: {e}" ) logger.error(f'视频处理请求失败: {e}') Common.logger('ai').info(f'视频处理请求失败: {e}') return def run(): video_path = None try: if not API_KEY: logger.error('[+] 请在环境变量中新增 API_KEY') return if not TASK_TYPE: logger.error('[+] 请在环境变量中新增 TASK_TYPE, 可选值: recommend | top') return genai.configure(api_key=API_KEY) redis_task = get_redis_task(task_type=TASK_TYPE) if not redis_task: time.sleep(10) return redis_task = json.loads(redis_task) mark, prompt = Material.feishu_list() video_duration = get_video_duration(video_link=redis_task['video_path']) if not video_duration: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], "", redis_task['type'], redis_task['partition'], "[+] 获取视频时长失败, 跳过任务" ) logger.error('[+] 获取视频时长失败, 跳过任务') return elif video_duration >= 600: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], "", redis_task['type'], redis_task['partition'], "[+] 视频时长超过10分钟, 跳过任务" ) logger.error('[+] 视频时长超过10分钟, 跳过任务') return video_path = download_video(video_link=redis_task['video_path']) if not video_path: AliyunLogger.logging( str( redis_task['video_id'] ), redis_task['title'], redis_task['video_path'], "", redis_task['type'], redis_task['partition'], "[+] 视频下载失败, 跳过任务" ) logger.error(f'[+] 视频下载失败, 跳过任务') if os.path.exists(video_path): os.remove(video_path) logger.info(f"文件已删除: {video_path}") return google_file, google_file_state = upload_video(video_path=video_path, redis_task=redis_task) if not google_file_state: return elif google_file_state != 'ACTIVE': logger.error('[+] 视频上传状态不为 ACTIVE, 跳过任务') genai.delete_file(google_file) if os.path.exists(video_path): os.remove(video_path) logger.info(f"文件已删除: {video_path}") return model = create_model_cache(redis_task=redis_task) if isinstance(model, str): logger.error('[+] 创建模型失败, 跳过任务') genai.delete_file(google_file) if os.path.exists(video_path): os.remove(video_path) logger.info(f"文件已删除: {video_path}") return response = analyze_video(model=model, google_file=google_file, prompt=prompt, redis_task=redis_task, title=redis_task['title']) if isinstance(response, str): logger.error('[+] 获取模型响应失败, 跳过任务') genai.delete_file(google_file) if os.path.exists(video_path): os.remove(video_path) logger.info(f"文件已删除: {video_path}") return text = response.text.strip() cleaned_text = text.replace("```json", '').replace("```", '').strip() AliyunLogger.logging( str(redis_task['video_id']), redis_task['title'], redis_task['video_path'], str(mark), redis_task['type'], redis_task['partition'], str(cleaned_text)) logger.info(f'[+] 模型响应结果: {text}') if os.path.exists(video_path): os.remove(video_path) logger.info(f"文件已删除: {video_path}") genai.delete_file(google_file) except KeyboardInterrupt: if os.path.exists(video_path): os.remove(video_path) if __name__ == '__main__': logger.info(f'[+] 任务已启动 -> API_KEY: {API_KEY}, TASK_TYPE: {TASK_TYPE}') schedule.every(interval=1).seconds.do(run) while True: try: schedule.run_pending() time.sleep(1) except KeyboardInterrupt: break logger.info('[+] 任务已停止')