LongArticlesJob/applications/api/google_ai_api.py

"""
@author: luojunhui
"""

from google import genai


class GoogleAIAPI(object):
    """
    Google 视频内容理解API
    """
    def __init__(self):
        self.api_key = 'AIzaSyBAFFL6yHa4kUK-YcuAw8tbiSHQB6oJG34'
        self.client = genai.Client(api_key=self.api_key)

    def upload_file(self, file_path: str):
        """
        file_path: 文件路径
        """
        print("start uploading file: {}".format(file_path))
        try:
            video_file = self.client.files.upload(file=file_path)
            file_name = video_file.name
            file_state = video_file.state.name
            expire_time = video_file.expiration_time
            print("success uploaded file: {}".format(file_path))
            return file_name, file_state, expire_time
        except Exception as e:
            print("fail to upload file: {} because {}".format(file_path, e))
            return None

    def get_file_status(self, file_name: str,):
        """
        获取文件的状态
        file_name: 文件名称
        """
        try:
            video_file = self.client.files.get(name=file_name)
            state = video_file.state.name
            return state
        except Exception as e:
            print(e)
            return None

    def get_google_file(self, file_name: str):
        """
        获取文件
        file_name: 文件名称
        """
        try:
            video_file = self.client.files.get(name=file_name)
            return video_file
        except Exception as e:
            print(e)
            return None

    def get_video_text(self, prompt, video_file):
        """
        获取视频文本
        prompt: 提示词
        video_file: <class 'google.genai.types.File'>
        """
        response = self.client.models.generate_content(
            model='gemini-1.5-pro',
            contents=[
                video_file,
                prompt
            ]
        )
        return response.text