google_ai_studio/job.py

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.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)
            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) -> 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:
        logger.error(f'[+] 上传视频失败: {e}')
        return


def create_model_cache() -> 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:
        logger.error(f'视频创建缓存内容，并返回生成模型异常信息: {e}')
        Common.logger('ai').info(f'视频创建缓存内容，并返回生成模型异常信息: {e}')
        return


def analyze_video(model: genai.GenerativeModel, google_file: File, prompt: str) -> 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:
        logger.error(f'视频处理请求失败: {e}')
        Common.logger('ai').info(f'视频处理请求失败: {e}')
        return


def run():
    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:
        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:
        logger.error('[+] 获取视频时长失败, 跳过任务')
        return
    elif video_duration >= 600:
        logger.error('[+] 视频时长超过10分钟, 跳过任务')
        return

    video_path = download_video(video_link=redis_task['video_path'])
    if not video_path:
        logger.error(f'[+] 视频下载失败, 跳过任务')
        return

    google_file, google_file_state = upload_video(video_path=video_path)
    if not google_file_state:
        return
    elif google_file_state != 'ACTIVE':
        logger.error('[+] 视频上传状态不为 ACTIVE, 跳过任务')
        return

    model = create_model_cache()
    if isinstance(model, str):
        logger.error('[+] 创建模型失败, 跳过任务')
        return

    response = analyze_video(model=model, google_file=google_file, prompt=prompt)
    if isinstance(response, str):
        logger.error('[+] 获取模型响应失败, 跳过任务')
        return

    usage_info, text = str(response.usage_metadata).replace('\n', ', '), response.text.strip()
    logger.info(f'[+] 使用情况: {usage_info}')
    logger.info(f'[+] 模型响应结果: {text}')


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('[+] 任务已停止')