LongArticlesJob/applications/api/google_ai_api.py

"""
@author: luojunhui
"""

from google import genai
from google.genai.errors import ClientError
from tqdm import tqdm


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

    def upload_file(self, file_path: str):
        """
        file_path: 文件路径
        """
        tqdm.write("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
            tqdm.write("success uploaded file: {}".format(file_path))
            return file_name, file_state, expire_time
        except Exception as e:
            tqdm.write("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 fetch_info_from_google_ai(self, prompt, video_file):
        """
        获取视频文本
        prompt: 提示词
        video_file: <class 'google.genai.types.File'>
        """
        try:
            response = self.client.models.generate_content(
                model='gemini-1.5-flash',
                contents=[
                    video_file,
                    prompt
                ]
            )
            return {
                'code': 200,
                'text': response.text
            }
        except ClientError as e:
            return {
                'code': e.response.status_code,
                'text': 'resource exhausted'
            }
        except Exception as e:
            return {
                'code': 500,
                'text': str(e)
            }


    def delete_video(self, file_name: str):
        """
        删除视频
        """
        self.client.files.delete(name=file_name)

    def get_file_list(self):
        """
        获取文件列表
        """
        file_list = self.client.files.list()
        return file_list