28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
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 |