AutoLove: апдейты девушке с YandexGPT

Чтение YAML и JSON

def read_yaml(path: str) -> dict:
    with open(path, 'r') as stream:
        return yaml.safe_load(stream)


def read_json(path: str) -> dict:
    with open(path, 'r') as stream:
        return json.load(stream)

Преобразование списков строк формата «00:00» в списки datetime.time

def yaml_to_datetime_time(times_str: list) -> list:
    return [
        datetime.time(
            *map(
                int,
                time_str.split(':')
            )
        )
        for time_str
        in times_str
    ]

class YandexGPT:
    available_models = [
        'yandexgpt',
        'yandexgpt-lite',
        'summarization'
    ]

    def __init__(
            self,
            model_type: str = 'yandexgpt',
            iam_token: str = None,
            catalog_id: str = None,
            yandex_cloud_config_file_path: str = None,
            yandex_gpt_key_file_path: str = None,
            iam_url: str = "https://iam.api.cloud.yandex.net/iam/v1/tokens"
    ) -> None:
        # setting config files
        self._yandex_cloud_config_file_path = yandex_cloud_config_file_path
        self._yandex_gpt_key_file_path = yandex_gpt_key_file_path
        self._iam_url = iam_url
        # setting model type
        if model_type not in self.available_models:
            raise ValueError(f"Model type must be one of {self.available_models}")
        else:
            self.model_type = model_type
        # setting IAM token
        if not iam_token:
            self._set_iam_token()
        else:
            self._iam_token = iam_token
        # setting catalog id
        if not catalog_id:
            self._set_catalog_id()
        else:
            self._catalog_id = catalog_id

    def _set_iam_token(self) -> None:
        # reading yaml config
        config = read_yaml(self._yandex_cloud_config_file_path)
        # reading json key
        key = read_json(self._yandex_gpt_key_file_path)
        # generating jwt token
        jwt_token = self._generate_jwt_token(
            service_account_id=config['ServiceAccountID'],
            private_key=key['private_key'],
            key_id=config['ServiceAccountKeyID'],
            url=self._iam_url
        )
        # sending request to get IAM token and setting IAM token
        self._iam_token = self._swap_jwt_to_iam(
            jwt_token=jwt_token,
            url=self._iam_url
        )

    @staticmethod
    def _swap_jwt_to_iam(
            jwt_token: str,
            url: str = "https://iam.api.cloud.yandex.net/iam/v1/tokens"
    ) -> str:
        # sending request to get IAM token
        headers = {
            "Content-Type": "application/json"
        }
        data = {
            "jwt": jwt_token
        }
        response = requests.post(url, headers=headers, json=data)
        # checking response
        if response.status_code == 200:
            return response.json()['iamToken']
        else:
            raise Exception(
                f"Failed to get IAM token. "
                f"Status code: {response.status_code}"
                f"\n"
                f"{response.text}"
            )

    @staticmethod
    def _generate_jwt_token(
            service_account_id: str,
            private_key: str,
            key_id: str,
            url: str = 'https://iam.api.cloud.yandex.net/iam/v1/tokens'
    ) -> str:
        # generating jwt token
        now = int(time.time())
        payload = {
            'aud': url,
            'iss': service_account_id,
            'iat': now,
            'exp': now + 360
        }
        encoded_token = jwt.encode(
            payload,
            private_key,
            algorithm='PS256',
            headers={'kid': key_id}
        )
        return encoded_token

    def _set_catalog_id(self) -> None:
        # reading yaml config file
        config = read_yaml(self._yandex_cloud_config_file_path)
        # setting catalog id from config file
        self._catalog_id = config['CatalogID']

    def send_completion_request(
            self,
            messages: List[dict],
            temperature: float = 0.6,
            max_tokens: int = 1000,
            stream: bool = False,
            completion_url: str = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
    ) -> dict:
        # checking if IAM token and catalog id is set
        if not self._iam_token or not self._catalog_id:
            raise Exception("IAM token and catalog id must be set to send completion requests")
        # sending request
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self._iam_token}",
            "x-folder-id": self._catalog_id
        }
        data = {
            "modelUri": f"gpt://{self._catalog_id}/{self.model_type}/latest",
            "completionOptions": {
                "stream": stream,
                "temperature": temperature,
                "maxTokens": str(max_tokens)
            },
            "messages": messages
        }
        response = requests.post(completion_url, headers=headers, json=data)
        # checking response
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Failed to send completion request. Status code: {response.status_code}\n{response.text}")

Я предусмотрел возможность передачи параметров класса напрямую, а также возможность изменить логику инициализации из конфиг файлов  для того чтобы каждый желающий мог адаптировать данный функционал под конкретные сценарии использования.

class AutoChat(Telegram, YandexGPT):
    def __init__(
            self,
            target_username: str = None,
            sending_time: List[datetime.time] = None,
            chat_system_prompt: List[str] = None,
            chat_mode: bool = None,
            timezone: pytz.timezone = None,
            telegram_config_file_path: str = None,
            auto_chat_config_file_path: str = None,
            *args,
            **kwargs,
    ) -> None:
        # initializing telegram
        # checking if api_id, api_hash, phone, database_encryption_key, files_directory and login in **kwargs
        if all(
                key
                in kwargs
                for key
                in [
                    'api_id',
                    'api_hash',
                    'phone',
                    'database_encryption_key',
                    'files_directory',
                    'login'
                ]
        ):
            Telegram.__init__(
                *args,
                **kwargs,
            )
        # if there is no then calling custom telegram init method
        else:
            self._init_telegram(telegram_config_file_path)
        # initializing yandex gpt
        # checking if iam_token and catalog_id in **kwargs
        if all(
                key
                in kwargs
                for key
                in [
                    'iam_token',
                    'catalog_id'
                ]
        ):
            YandexGPT.__init__(
                *args,
                **kwargs
            )
        # if there is no then calling custom yandex gpt init method
        else:
            self._init_yandex_gpt(
                *args,
                **kwargs
            )
        # initializing auto chat
        # checking if target_username, sending_time, chat_system_prompt, chat_mode, timezone is set
        if all([
            target_username,
            sending_time,
            chat_system_prompt,
            chat_mode,
            timezone
        ]):
            self._target_user_id = self._get_target_user_id(target_username)
            self._sending_time = yaml_to_datetime_time(sending_time)
            self._chat_system_prompt = chat_system_prompt
            self._chat_mode = chat_mode
            self._timezone = timezone
        # if there is no then calling custom auto chat init method
        else:
            self._init_autochat(auto_chat_config_file_path)

    def _init_yandex_gpt(
            self,
            *args,
            **kwargs
    ) -> None:
        # checking if yandex_cloud_config_file_path and yandex_gpt_key_file_path in **kwargs
        if not all(
                key
                in kwargs
                for key
                in [
                    'yandex_cloud_config_file_path',
                    'yandex_gpt_key_file_path'
                ]
        ):
            raise ValueError(
                "AutoChat args must contain either"
                " 'yandex_cloud_config_file_path' and "
                "'yandex_gpt_key_file_path' or "
                "'iam_token' and 'catalog_id'"
            )
        # if they are than calling custom yandex gpt init method
        else:
            YandexGPT.__init__(
                self,
                *args,
                **kwargs
            )

    def _init_telegram(
            self,
            telegram_config_file_path: str = None,
    ) -> None:
        # checking if telegram_config_file_path is set
        if not telegram_config_file_path:
            raise ValueError(
                "AutoChat args must contain either"
                " 'telegram_config_file_path' or "
                "'api_id', 'api_hash', 'phone', "
                "'database_encryption_key', "
                "'files_directory' and 'login'"
            )
        # if it is than calling custom telegram init method
        else:
            # reading config file
            telegram_config = read_yaml(telegram_config_file_path)
            # initializing Telegram class
            Telegram.__init__(
                self,
                api_id=int(telegram_config['ApiId']),
                api_hash=telegram_config['ApiHash'],
                phone=telegram_config['PhoneNumber'],
                database_encryption_key=telegram_config['DatabaseEncryptionKey'],
                files_directory=os.path.join(
                    os.getcwd(),
                    telegram_config['FilesDirectory']
                ),
                login=telegram_config['Login'],
            )

    def _init_autochat(
            self,
            auto_chat_config_file_path: str = None
    ) -> None:
        # checking if auto_chat_config_file_path is set
        if not auto_chat_config_file_path:
            raise ValueError(
                "AutoChat args must contain ether 'auto_chat_config_file_path'"
                " or 'target_username', 'sending_time', 'chat_system_prompt',"
                " 'chat_mode' and 'timezone'"
            )
        # if it is than calling custom auto chat init method
        # reading config file
        auto_chat_config = read_yaml(auto_chat_config_file_path)
        # setting autochat
        self._target_user_id = self._get_target_user_id(auto_chat_config['TargetUsername'])
        self._sending_time = yaml_to_datetime_time(auto_chat_config['SendingTime'])
        self._chat_system_prompt = auto_chat_config['ChatSystemPrompt']
        self._chat_mode = auto_chat_config['ChatMode']
        self._timezone = pytz.timezone(auto_chat_config['TimeZone'])

    def _get_target_user_id(
            self,
            target_username: str
    ) -> int:
        # getting target user id
        return int(self._get_contact_by_username(target_username)['id'])

    def _get_client_contacts(self) -> List[dict]:
        # getting client chats
        chats = self.get_chats()
        chats.wait()
        # getting contacts from private chats
        contacts = []
        for chat_id in chats.update['chat_ids']:
            chat = self.get_chat(chat_id)
            chat.wait()
            if chat.update['type']['@type'] == 'chatTypePrivate':
                contact = self.get_user(chat_id)
                contact.wait()
                contacts.append(contact.update)
        # returning contacts
        return contacts

    def _get_contact_by_username(
            self,
            target_username: str
    ) -> dict:
        # getting list of all contacts
        contacts = self._get_client_contacts()
        # finding target contact by username
        for contact in contacts:
            if contact['username'] == target_username:
                return contact
        # raise ValueError if there is no such contact
        raise ValueError(f"Contact with username {target_username} not found")

    def send_message_to_target_username(
            self,
            text: Union[str, Element],
            entities: Union[List[dict], None] = None,
    ) -> Union[AsyncResult, None]:
        # sending message to target contact
        return self.send_message(
            chat_id=self._target_user_id,
            text=text,
            entities=entities
        )

    def which_sending_time(self) -> int:
        # getting current time without seconds and microseconds
        current_time = (
            datetime.datetime
            .now(self._timezone)
            .replace(second=0, microsecond=0)
            .time()
        )
        # checking if current time is in sending time list and returning its index or -1 if it is not
        for index, sending_time_unit in enumerate(self._sending_time):
            if current_time == sending_time_unit:
                return index
        return -1

    def check_time_and_send(self):
        # checking if current time is in sending time list
        which_sending_time = self.which_sending_time()
        # if it is than sending message
        if which_sending_time > -1:
            prompt = self._chat_system_prompt[which_sending_time]
            text = self.send_completion_request(
                messages=[
                    {
                        "role": "system",
                        "text": prompt
                    }
                ]
            )['result']['alternatives'][0]['message']['text']
            self.send_message_to_target_username(text=text)

Как видно, я также вынес весь код связанный с инициализацией целевых параметров в отдельные методы и оставил возможность загрузить готовые параметры вместо конфигов при их наличии.