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, ad_in_video_data 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 = { "一、基础信息": { "视频主题": "", "视频关键词": "", "产品卖点":[] }, "二、主体和场景": { "视频主体": "", "人物数量": "", "人物年龄段": [], "人物性别": [], "场景描述": [] }, "三、情感与风格": { "情感倾向": "", "视频风格": "" }, "四、视频传播性与时效": { "片尾引导": "", "传播强度": "", "时效性": "", "适用时间": "" }, "五、用户画像": { "价值类型": "", "用户价值点": "", "性别": "", "年龄": "", "观众收入": "" }, "六、音画信息": { "背景音类型": "", "背景音风格": "", "语音类型": "", "字幕": "", "颜色": "", "logo": "" }, "七、封面信息": { "封面主体": "", "人物个数": "", "文字数量": "", "文字关键字": [], "封面主题": "" }, "八、剪辑信息": { "视频开场风格": "", "展现形式": "" }, "九、类目": { "视频一级分类": "", "视频二级分类": ["品类- 、分数-","品类- 、分数-","品类- 、分数-"] }, "十、视频时长": { "时长": "", }, "八、视频宽高": { "宽高": "", "宽高比": "" } } 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:ad_video_recommend' 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=60) 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.ad_logging(str(redis_task['ad_id']), redis_task['creative_code'], redis_task['creative_title'], redis_task['material_address'], redis_task['click_button_text'], redis_task['creative_logo_address'], redis_task['update_time'], f"[+] 上传视频失败: {e}") logger.error(f'[+] 上传视频失败: {e}') ad_in_video_data(redis_task) 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.ad_logging(str(redis_task['ad_id']), redis_task['creative_code'], redis_task['creative_title'], redis_task['material_address'], redis_task['click_button_text'], redis_task['creative_logo_address'], redis_task['update_time'], f"[+] 视频创建缓存内容，并返回生成模型异常信息: {e}") logger.error(f'视频创建缓存内容，并返回生成模型异常信息: {e}') ad_in_video_data(redis_task) Common.logger('ai').info(f'视频创建缓存内容，并返回生成模型异常信息: {e}') return def analyze_video(model: genai.GenerativeModel, google_file: File, prompt: str, redis_task) -> Optional[GenerateContentResponse]: try: session = model.start_chat(history=[]) content = { 'parts': [ google_file, f'{prompt}\n输出返回格式样例:\n{SAMPLE_DATA}', ], } return session.send_message(content=content) except Exception as e: AliyunLogger.ad_logging(str(redis_task['ad_id']), redis_task['creative_code'], redis_task['creative_title'], redis_task['material_address'], redis_task['click_button_text'], redis_task['creative_logo_address'], redis_task['update_time'], f"[+] 视频处理请求失败: {e}") ad_in_video_data(redis_task) 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.ad_feishu_list() # video_duration = get_video_duration(video_link=redis_task[3]) # if not video_duration: # AliyunLogger.logging( str( redis_task[0] ), redis_task[1], 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['material_address']) if not video_path: AliyunLogger.ad_logging( str(redis_task['ad_id']), redis_task['creative_code'], redis_task['creative_title'], redis_task['material_address'], redis_task['click_button_text'], redis_task['creative_logo_address'], redis_task['update_time'], "[+] 视频下载失败, 跳过任务" ) 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) 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.ad_logging(str(redis_task['ad_id']), redis_task['creative_code'], redis_task['creative_title'], redis_task['material_address'], redis_task['click_button_text'], redis_task['creative_logo_address'], redis_task['update_time'], 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('[+] 任务已停止')