Sending requests to yandexGPT. Yousing service tokens to generate IAM one and connect to API through IAM tokens

YandexGPT.py

import credentials

from time import time
import jwt

import requests


class YandexGPT:
    def __init__(self, service_account_id, key_id, private_key, folder_id, default_system_msg=None,
                 iam_token_live_time=60):
        self.service_account_id = service_account_id
        self.key_id = key_id
        self.private_key = private_key
        self.folder_id = folder_id

        self.default_system_msg = default_system_msg

        self.iam_token_live_time = iam_token_live_time

        self.time_since_iam_generation = 0

        self._jwt = None
        self.update_jws_token()

        self._iam_token = None
        self.update_iam_token()

    def update_jws_token(self, live_time=60):
        now = int(time())
        payload = {
            'aud': 'https://iam.api.cloud.yandex.net/iam/v1/tokens',
            'iss': self.service_account_id,
            'iat': now,
            'exp': now + live_time
        }

        self._jwt = jwt.encode(
            payload,
            self.private_key,
            algorithm='PS256',
            headers={'kid': self.key_id})

    def update_iam_token(self):
        self.time_since_iam_generation = time()
        response = requests.post(url='https://iam.api.cloud.yandex.net/iam/v1/tokens',
                                 json={'jwt': self._jwt})
        json = response.json()
        if response.status_code == 200:
            self._iam_token = json['iamToken']
            return self._iam_token

        code = json['code']
        if code == 16:
            # если jws токен протух
            self.update_jws_token()
            return self.update_iam_token()
        else:
            # если произошло невероятное
            raise ConnectionError(json['error'])

    @property
    def iam_token(self):
        if self.time_since_iam_generation - time() > self.iam_token_live_time:
            self.update_iam_token()
        return self._iam_token

    def completion(self, user_msg, system_msg=None, ai_model='yandexgpt/latest', temperature=0.6, maxTokens=20):
        if system_msg is None:
            system_msg = self.default_system_msg

        messages = [{'role': 'system', 'text': system_msg}] if system_msg else []
        messages.append({'role': 'user', 'text': user_msg})

        payload = {
            "modelUri": f"gpt://{self.folder_id}/{ai_model}",
            "completionOptions": {
                "stream": False,
                "temperature": temperature,
                "maxTokens": maxTokens
            },
            "messages": messages
        }
        kwargs = {'url': 'https://llm.api.cloud.yandex.net/foundationModels/v1/completion',
                'json': payload,
                'headers': {'Authorization': f"Bearer {self.iam_token}",
                            'x-folder-id': self.folder_id}
                }
        status_code = 429
        while status_code == 429:  # 429 вызывается, если одновременно отправлено 2 запроса
            response = requests.post(**kwargs)
            status_code = response.status_code
        json = response.json()
        if status_code != 200:
            raise ConnectionAbortedError(json['error'])
        return json['result']['alternatives'][0]['message']['text']


if __name__ == '__main__':
    a = YandexGPT(service_account_id=credentials.service_account_id,
                  key_id=credentials.key_id,
                  private_key=credentials.private_key,
                  folder_id=credentials.folder_id)
    print(a.completion('Что такое ламинат?', maxTokens=20))