# azure_api.py
import requests
from typing import List

class AzureAPI:
    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
        self.version = "2024-02-15-preview"

    def get_caption(self, prompt: str, user_prompt: str, images_base64: List[str]) -> str:
        headers = {
            'content-type': 'application/json',
            "api-key": self.key,
        }
        system_msg = {"role": "system", "content": prompt}
        user_msg = [{"type": "text", "text": user_prompt}]
        img_msg = [
            {"type": "image_url", "image_url": {"url": f'data:image/jpg;base64,{image}'}}
            for image in images_base64
        ]
        payload = {
            'messages': [system_msg, {"role": "user", "content": user_msg + img_msg}],
            'temperature': self.temp,
            'top_p': self.top_p,
            'max_tokens': self.max_tokens,
            'model': self.model
        }
        url = f'{self.endpoint}/openai/deployments/{self.model}/chat/completions?api-version={self.version}'
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']