utils/base.py

import requests
from typing import List

class BaseAPI:
    def __init__(self, key: str, endpoint: str, model: str, temp: float = 0.3, top_p: float = 0.75, max_tokens: int = 1024):
        self.key = key
        self.endpoint = endpoint
        self.model = model
        self.temp = temp
        self.top_p = top_p
        self.max_tokens = max_tokens

    def get_headers(self):
        return {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.key}'
        }

    def get_payload(self, sys_prompt: str, user_prompt: str, images_base64: List[str]) -> dict:
        _txt = [{"type": "text", "text": user_prompt}]
        _img = [
            {"type": "image_url", "image_url": {"url": f'data:image/jpg;base64,{image}'}}
            for image in images_base64
        ]
        
        sys_msg = {"role": "system", "content": sys_prompt}
        usr_msg = {"role": "user", "content": _txt + _img}

        return {
            'messages': [sys_msg, usr_msg],
            'temperature': self.temp,
            'top_p': self.top_p,
            'max_tokens': self.max_tokens,
            'model': self.model
        }

    def get_caption(self, sys_prompt: str, user_prompt: str, images_base64: List[str], url_suffix: str) -> str:
        headers = self.get_headers()
        payload = self.get_payload(sys_prompt, user_prompt, images_base64)
        url = f'{self.endpoint}{url_suffix}'
        
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        return self.parse_response(response.json())

    def parse_response(self, response: dict) -> str:
        raise NotImplementedError("Subclasses should implement this method")