import aiohttp

from .data import APIEndpoints, Payload
from .exception import LanguageModelException
from .var import LanguageModelError, LlamaAPI


class LanguageModel:
    def __init__(self, api_base: str, exception_message: str = LanguageModelError.GENERIC_EXCEPTION):
        self.endpoints = APIEndpoints(
            api_base,
            LlamaAPI.CHAT_COMPLETIONS
        )
        self.exception_message = exception_message
        
    async def make_request(self, payload: Payload) -> str:
        data = payload.render()
        
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(self.endpoints.chat_completion, json = data) as response:
                    if response.status == 200:
                        result = await response.json()
                        return result["choices"][0]["message"]["content"]
                    raise LanguageModelException(f"Request failed with status code {response.status}!")
            except (aiohttp.ClientError, LanguageModelException) as e:
                print(f"Exception occurred: {e}")
                return self.exception_message