Привет Хабр! В прошлый раз мы сделали телеграм-бота с полноценным ИИ. Теперь мы продолжим добавлять новые интересные фичи нашему боту, но в этот раз мы начнем с конца и посмотрим на готовый результат, а потом разберем код и детали реализации.

Дэмо

Первое, что мы сделаем – это добавим небольшое меню с двумя опциями: выбор модели ИИ и отображение уже выбранной модели.

При нажатии кнопки «Выбрать модель» бот отображает список доступных моделей. Поддерживаемые модели можно посмотреть на странице проекта Jlama, но в нашей реализации будет отдельный REST API для управления доступными моделями.

При нажатии «Показать текущую модель» бот выведет название привязанной к чату модели.

Как видим в данном примере наша текущая модель – tjake/Llama-3.1-8B-Instruct-jQ4 и на вопрос «Whats is Java?» будет отвечать именно она. Допустим мы хотим выбрать другую модель, пусть это будет Qwen2.5. Нажимаем кнопку «Выбрать модель» и выбираем Qwen2.5.

Попробуем задать вопрос «What Is C++?».

Теперь на наш вопрос отвечает ИИ Qwen2.5. Убедимся в этом нажав кнопку меню «Показать текущую модель».

Теперь к этому чату будет привязан ИИ Qwen2.5, и на все вопросы будет отвечать он.

Смотрим код

Начнем с метода consume нашего AiChatBot:

    @SneakyThrows
    @Override
    public void consume(Update update) {
        if (update.hasMessage() && update.getMessage().hasText()) {
            long chatId = update.getMessage().getChatId();
            var text = update.getMessage().getText();

            switch (text) {
                case START -> startChat(chatId);
                case CHOSE_MODEL -> showAvailableModels(chatId);
                case SHOW_MODEL -> showCurrentModel(chatId);
                default -> askModel(text, chatId);
            }
        } else if (update.hasCallbackQuery()) {
            String callbackData = update.getCallbackQuery().getData();
            var chatId = update.getCallbackQuery().getMessage().getChatId();
            choseModel(chatId, callbackData);
        }
    }

У нас есть четыре случая обработки входящего текста:

Старт чата с ботом
Выбор модели.
Отображение текущей модели
Вопрос самому ИИ.

При старте чата первое, что нам нужно сделать – это создать объект самого чата Chat и клавиатуру с кнопками выбора модели и показа текущей модели.

    private void startChat(long chatId) throws TelegramApiException {
        chatService.createChat(chatId);
        ReplyKeyboardMarkup keyboardMarkup = new ReplyKeyboardMarkup(
                List.of(
                        new KeyboardRow(new KeyboardButton("Выбрать модель")),
                        new KeyboardRow(new KeyboardButton("Показать текущую модель"))
                )
        );
        keyboardMarkup.setResizeKeyboard(true);

        SendMessage message = SendMessage.builder()
                .chatId(chatId)
                .text("Меню")
                .replyMarkup(keyboardMarkup)
                .build();

        telegramClient.execute(message);
    }

За создание чата отвечает сервис ChatService, метод createChat. Новый чат будет создан только если у текущего пользователя нет уже открытых чатов с ботом. При этом первоначально в качестве модели будет использована модель по умолчанию, указанная в файле application.yaml.

    @Transactional
    public void createChat(long id) {
        var chat = chatRepository.findById(id);
        if(chat.isEmpty()) {
            chatRepository.save(new Chat(id, modelFullName));
            log.info("New chat with id {} has been created", id);
        }
    }

В llm.model-full-name можно указать любую поддерживаемую Jlama модель.

Если мы захотим выбрать другую модель, то в тексте бот получит константу CHOSE_MODEL. Метод showAvailableModels отображает набор кнопок с доступными моделями.

    private void showAvailableModels(long chatId)  {
        List<InlineKeyboardButton> buttons = availableModelService.findAllAvailableModels()
                .stream()
                .map(model -> {
                    var button = new InlineKeyboardButton(model.getName());
                    button.setCallbackData(model.getFullName());

                    return button;
                })
                .toList();
        InlineKeyboardRow row = new InlineKeyboardRow(buttons);
        InlineKeyboardMarkup inlineKeyboardMarkup = new InlineKeyboardMarkup(List.of(row));

        SendMessage message = SendMessage.builder()
                .chatId(chatId)
                .text("Выберите модель")
                .replyMarkup(inlineKeyboardMarkup)
                .build();

        try {
            telegramClient.execute(message);
        } catch (TelegramApiException e) {
            log.error("Error {}", e.getMessage());
        }
    }

Тут мы видим сервис AvailableModelService. Он возвращает список доступных нашему боту моделей. Сам объект AvailableModel довольно простой:

@Data
@Entity
@Table(name = "models")
@AllArgsConstructor
@NoArgsConstructor(access = AccessLevel.PROTECTED)
public class AvailableModel {
    @Id
    @GeneratedValue
    private UUID id;

    @NotEmpty
    private String name;

    @NotEmpty
    private String fullName;

    public AvailableModel(String fullName, String name) {
        this.fullName = fullName;
        this.name = name;
    }
}

AvailableModelService реализует в том числе REST API для управления списком доступных моделей. Мы можем создать нужное нам количество моделей, необязательно все поддерживаемые Jlama. В нашем примере их всего четыре, но ничего не мешает нам создать все возможные для Jlama модели. Разработчики активно добавляют все больше и больше, поэтому в REST API есть определенные смысл – мы сможем добавлять новые модели по мере их появления в Jlama.

@RestController
@RequiredArgsConstructor
@RequestMapping("/models")
public class AvailableModelController {
    private final AvailableModelService availableModelService;
    @PostMapping
    @ResponseStatus(HttpStatus.CREATED)
    public AvailableModel createNewAvailableModel(@RequestBody @Valid AvailableModel model) {
        return availableModelService.createAvailableModel(model);
    }

    @GetMapping("/{id}")
    @ResponseStatus(HttpStatus.CREATED)
    public AvailableModel findModelById(@PathVariable UUID id) {
        return availableModelService.findModelById(id);
    }

    @GetMapping
    @ResponseStatus(HttpStatus.CREATED)
    public AvailableModel findModelByName(@RequestParam String name) {
        return availableModelService.findAvailableModelByName(name);
    }
}

Пока что REST API довольно скромный.

Следует сразу обратить внимание на один момент – нет необходимости заранее выкачивать все доступные модели в рабочую директорию бота. Прежде чем сформировать PromptContext и отправить его в LLM, объект Downloader будет пытаться выкачать саму модель при условии, что ее нет в рабочей директории. Это может немного сказаться на производительности – первый вопрос после переключения ИИ может обрабатываться немного дольше, даст о себе знать время скачивания LLM.

Чтобы отобразить текущую модель бот должен получить константу SHOW_MODEL. Метод showCurrentModel находит нужный чат по его идентификатору и отображает название привязанной к чату модели.

    private void showCurrentModel(long chatId) {
        var chat = chatService.findChatById(chatId);

        SendMessage message = SendMessage.builder()
                .chatId(chatId)
                .text(chat.getModelName())
                .build();

        try {
            telegramClient.execute(message);
        } catch (TelegramApiException e) {
            log.error("Error {}", e.getMessage());
        }
    }

Если бот не получил в тексте каких-либо служебных констант, то текст будет восприниматься, как вопрос ИИ.

    private void askModel(String text, long chatId) {
        var chat = chatService.findChatById(chatId);
        try {
            var answer = model.ask(text, chat.getModelName());

            SendMessage message = SendMessage.builder()
                    .chatId(chatId)
                    .text(answer)
                    .build();

            telegramClient.execute(message);
        } catch (TelegramApiException | IOException e) {
            log.error("Error {}", e.getMessage());
        }
    }

Кнопки с названиями моделей выполнены в виде InlineKeyboardButton. Такая кнопка содержит обратный вызов. В нашей реализации в качестве обратного вызова будет выступать название модели. То есть бот может реагировать на нажатие таких кнопок отдельно от обработки текста. Это реализовано в блоке else if метода consume.

else if (update.hasCallbackQuery()) {
            String callbackData = update.getCallbackQuery().getData();
            var chatId = update.getCallbackQuery().getMessage().getChatId();
            choseModel(chatId, callbackData);
        }

Если в сообщении боту есть обратный вызов, то мы передаем его значение в метод choseModel.

    private void choseModel(long chatId, String model) {
        chatService.changeModel(chatId, model);

        SendMessage message = SendMessage.builder()
                .chatId(chatId)
                .text("Выбрана модель " + model)
                .build();

        try {
            telegramClient.execute(message);
        } catch (TelegramApiException e) {
            log.error("Error {}", e.getMessage());
        }
    }

Метод changeModel меняет старое значение привязанной к чату ИИ на выбранное.

    @Transactional
    public void changeModel(long chatId, String modelName) {
        var chat = chatRepository.findById(chatId)
                .orElseThrow();
        chat.setModelName(modelName);
        log.info("Model {} has been added to chat {}", modelName, chatId);
    }

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

Буду рад любым комментариям и вопросам. Не забудьте подписаться на мой телеграм-канал. В следующей итерации будем учить нашего бота генерировать картинки.

Телеграм-бот с ИИ Jlama: добавляем новые фичи