utils/google.py

from .base import BaseAPI
from typing import List

class GoogleAPI(BaseAPI):
    def get_headers(self):
        return {
            'Authorization': f'Bearer {self.key}'
        }

    def get_caption(self, prompt: str, user_prompt: str, images_base64: List[str]) -> str:
        url_suffix = f'/v1/models/{self.model}:generateContent?key={self.key}'
        return super().get_caption(prompt, user_prompt, images_base64, url_suffix)

    def parse_response(self, response: dict) -> str:
        return response['predictions'][0]['content']