Add td_api::translateRichMessage.

5 files changed, 148 insertions, 0 deletions

diff --git a/td/generate/scheme/td_api.tl b/td/generate/scheme/td_api.tl
index c712b229a..6bb416c9e 100644
--- a/td/generate/scheme/td_api.tl
+++ b/td/generate/scheme/td_api.tl
@@ -11908,6 +11908,12 @@ removeTextCompositionStyle name:string = Ok;
 //@tone Tone of the translation; must be one of "", "formal", "neutral", "casual"; defaults to "neutral"
 translateText text:formattedText to_language_code:string tone:string = FormattedText;
 
+//@description Translates a rich message to the given language
+//@message Rich message to translate
+//@to_language_code Language code of the language to which the message is translated. See translateText.to_language_code for the list of supported values
+//@tone Tone of the translation; must be one of "", "formal", "neutral", "casual"; defaults to "neutral"
+translateRichMessage message:inputRichMessage to_language_code:string tone:string = RichMessage;
+
 //@description Extracts text or caption of the given message and translates it to the given language; must not be used in secret chats. If the current user is a Telegram Premium user, then text formatting is preserved
 //@chat_id Identifier of the chat to which the message belongs
 //@message_id Identifier of the message
diff --git a/td/telegram/Requests.cpp b/td/telegram/Requests.cpp
index a4977e767..fe80abffd 100644
--- a/td/telegram/Requests.cpp
+++ b/td/telegram/Requests.cpp
@@ -2845,6 +2845,15 @@ void Requests::on_request(uint64 id, td_api::translateText &request) {
                                             std::move(promise));
 }
 
+void Requests::on_request(uint64 id, td_api::translateRichMessage &request) {
+  CHECK_IS_USER();
+  CLEAN_INPUT_STRING(request.to_language_code_);
+  CLEAN_INPUT_STRING(request.tone_);
+  CREATE_REQUEST_PROMISE();
+  td_->translation_manager_->translate_rich_message(std::move(request.message_), request.to_language_code_,
+                                                    request.tone_, std::move(promise));
+}
+
 void Requests::on_request(uint64 id, td_api::translateMessageText &request) {
   CHECK_IS_USER();
   CLEAN_INPUT_STRING(request.to_language_code_);
diff --git a/td/telegram/Requests.h b/td/telegram/Requests.h
index 98f98208f..093623a0b 100644
--- a/td/telegram/Requests.h
+++ b/td/telegram/Requests.h
@@ -338,6 +338,8 @@ class Requests {
 
   void on_request(uint64 id, td_api::translateText &request);
 
+  void on_request(uint64 id, td_api::translateRichMessage &request);
+
   void on_request(uint64 id, td_api::translateMessageText &request);
 
   void on_request(uint64 id, td_api::summarizeMessage &request);
diff --git a/td/telegram/TranslationManager.cpp b/td/telegram/TranslationManager.cpp
index ad4244a4a..ce0148621 100644
--- a/td/telegram/TranslationManager.cpp
+++ b/td/telegram/TranslationManager.cpp
@@ -82,6 +82,59 @@ class TranslateTextQuery final : public Td::ResultHandler {
   }
 };
 
+class TranslateRichMessageQuery final : public Td::ResultHandler {
+  Promise<vector<telegram_api::object_ptr<telegram_api::richMessage>>> promise_;
+
+ public:
+  explicit TranslateRichMessageQuery(Promise<vector<telegram_api::object_ptr<telegram_api::richMessage>>> &&promise)
+      : promise_(std::move(promise)) {
+  }
+
+  void send(vector<telegram_api::object_ptr<telegram_api::InputRichMessage>> &&rich_messages,
+            MessageFullId message_full_id, const string &to_language_code, string tone) {
+    int32 flags = 0;
+    if (tone == "neutral") {
+      tone.clear();
+    }
+    if (!tone.empty()) {
+      flags |= telegram_api::messages_translateRichMessage::TONE_MASK;
+    }
+    if (message_full_id.get_message_id().is_valid()) {
+      CHECK(rich_messages.size() == 1u);
+      flags |= telegram_api::messages_translateRichMessage::PEER_MASK;
+      auto input_peer = td_->dialog_manager_->get_input_peer(message_full_id.get_dialog_id(), AccessRights::Read);
+      CHECK(input_peer != nullptr);
+      auto message_ids = {message_full_id.get_message_id().get_server_message_id().get()};
+      send_query(G()->net_query_creator().create(telegram_api::messages_translateRichMessage(
+          flags, std::move(input_peer), std::move(message_ids), Auto(), to_language_code, tone)));
+    } else {
+      flags |= telegram_api::messages_translateRichMessage::TEXT_MASK;
+      send_query(G()->net_query_creator().create(telegram_api::messages_translateRichMessage(
+          flags, nullptr, vector<int32>{}, std::move(rich_messages), to_language_code, tone)));
+    }
+  }
+
+  void on_result(BufferSlice packet) final {
+    auto result_ptr = fetch_result<telegram_api::messages_translateRichMessage>(packet);
+    if (result_ptr.is_error()) {
+      return on_error(result_ptr.move_as_error());
+    }
+
+    auto ptr = result_ptr.move_as_ok();
+    LOG(INFO) << "Receive result for TranslateRichMessageQuery: " << to_string(ptr);
+    promise_.set_value(std::move(ptr->result_));
+  }
+
+  void on_error(Status status) final {
+    if (status.message() == "INPUT_TEXT_EMPTY") {
+      vector<telegram_api::object_ptr<telegram_api::richMessage>> result;
+      result.push_back(telegram_api::make_object<telegram_api::richMessage>());
+      return promise_.set_value(std::move(result));
+    }
+    promise_.set_error(std::move(status));
+  }
+};
+
 class GetAiComposeTonesQuery final : public Td::ResultHandler {
   Promise<Unit> promise_;
 
@@ -507,6 +560,65 @@ void TranslationManager::on_get_translated_texts(vector<telegram_api::object_ptr
       get_formatted_text_object(td_->user_manager_.get(), formatted_text, skip_bot_commands, max_media_timestamp));
 }
 
+Result<TranslationManager::InputRichMessage> TranslationManager::get_input_rich_message(
+    td_api::object_ptr<td_api::inputRichMessage> &&message) const {
+  TRY_RESULT(rich_message, RichMessage::get_rich_message(td_, DialogId(), std::move(message), false));
+
+  InputRichMessage input_rich_message;
+  input_rich_message.skip_bot_commands_ = !rich_message.has_bot_commands();
+  input_rich_message.message_ = std::move(rich_message);
+  return std::move(input_rich_message);
+}
+
+void TranslationManager::translate_rich_message(td_api::object_ptr<td_api::inputRichMessage> &&message,
+                                                const string &to_language_code, const string &tone,
+                                                Promise<td_api::object_ptr<td_api::richMessage>> &&promise) {
+  TRY_RESULT_PROMISE(promise, input_rich_message, get_input_rich_message(std::move(message)));
+  translate_rich_message(std::move(input_rich_message), MessageFullId(), to_language_code, tone, std::move(promise));
+}
+
+void TranslationManager::translate_rich_message(InputRichMessage &&rich_message, MessageFullId message_full_id,
+                                                const string &to_language_code, const string &tone,
+                                                Promise<td_api::object_ptr<td_api::richMessage>> &&promise) {
+  auto input_rich_message = rich_message.message_.get_input_rich_message(td_);
+  if (input_rich_message == nullptr) {
+    return promise.set_error(400, "Invalid rich message specified");
+  }
+  vector<telegram_api::object_ptr<telegram_api::InputRichMessage>> input_rich_messages;
+  input_rich_messages.push_back(std::move(input_rich_message));
+
+  if (tone != string() && tone != "formal" && tone != "neutral" && tone != "casual") {
+    return promise.set_error(400, "Invalid tone specified");
+  }
+
+  auto query_promise = PromiseCreator::lambda(
+      [actor_id = actor_id(this), skip_bot_commands = rich_message.skip_bot_commands_, promise = std::move(promise)](
+          Result<vector<telegram_api::object_ptr<telegram_api::richMessage>>> result) mutable {
+        if (result.is_error()) {
+          return promise.set_error(result.move_as_error());
+        }
+        send_closure(actor_id, &TranslationManager::on_get_translated_rich_messages, result.move_as_ok(),
+                     skip_bot_commands, std::move(promise));
+      });
+
+  td_->create_handler<TranslateRichMessageQuery>(std::move(query_promise))
+      ->send(std::move(input_rich_messages), message_full_id, to_language_code, tone);
+}
+
+void TranslationManager::on_get_translated_rich_messages(
+    vector<telegram_api::object_ptr<telegram_api::richMessage>> rich_messages, bool skip_bot_commands,
+    Promise<td_api::object_ptr<td_api::richMessage>> &&promise) {
+  TRY_STATUS_PROMISE(promise, G()->close_status());
+  if (rich_messages.size() != 1u) {
+    if (rich_messages.empty()) {
+      return promise.set_error(500, "Translation failed");
+    }
+    return promise.set_error(500, "Receive invalid number of results");
+  }
+  auto rich_message = RichMessage(td_, std::move(rich_messages[0]), DialogId());
+  promise.set_value(rich_message.get_rich_message_object(td_, skip_bot_commands));
+}
+
 void TranslationManager::compose_message_with_ai(td_api::object_ptr<td_api::formattedText> &&text,
                                                  const string &translate_to_language_code, const string &tone,
                                                  bool emojify,
diff --git a/td/telegram/TranslationManager.h b/td/telegram/TranslationManager.h
index 85ff2f31a..34c2d963a 100644
--- a/td/telegram/TranslationManager.h
+++ b/td/telegram/TranslationManager.h
@@ -10,6 +10,7 @@
 #include "td/telegram/CustomEmojiId.h"
 #include "td/telegram/MessageEntity.h"
 #include "td/telegram/MessageFullId.h"
+#include "td/telegram/RichMessage.h"
 #include "td/telegram/td_api.h"
 #include "td/telegram/telegram_api.h"
 
@@ -41,6 +42,18 @@ class TranslationManager final : public Actor {
   void translate_text(InputText &&text, MessageFullId message_full_id, const string &to_language_code,
                       const string &tone, Promise<td_api::object_ptr<td_api::formattedText>> &&promise);
 
+  struct InputRichMessage {
+    RichMessage message_;
+    bool skip_bot_commands_ = true;
+  };
+
+  void translate_rich_message(td_api::object_ptr<td_api::inputRichMessage> &&message, const string &to_language_code,
+                              const string &tone, Promise<td_api::object_ptr<td_api::richMessage>> &&promise);
+
+  void translate_rich_message(InputRichMessage &&richmessage, MessageFullId message_full_id,
+                              const string &to_language_code, const string &tone,
+                              Promise<td_api::object_ptr<td_api::richMessage>> &&promise);
+
   void compose_message_with_ai(td_api::object_ptr<td_api::formattedText> &&text,
                                const string &translate_to_language_code, const string &tone, bool emojify,
                                Promise<td_api::object_ptr<td_api::formattedText>> &&promise);
@@ -78,10 +91,16 @@ class TranslationManager final : public Actor {
 
   Result<InputText> get_input_text(td_api::object_ptr<td_api::formattedText> &&text) const;
 
+  Result<InputRichMessage> get_input_rich_message(td_api::object_ptr<td_api::inputRichMessage> &&text) const;
+
   void on_get_translated_texts(vector<telegram_api::object_ptr<telegram_api::textWithEntities>> texts,
                                bool skip_bot_commands, int32 max_media_timestamp,
                                Promise<td_api::object_ptr<td_api::formattedText>> &&promise);
 
+  void on_get_translated_rich_messages(vector<telegram_api::object_ptr<telegram_api::richMessage>> rich_messages,
+                                       bool skip_bot_commands,
+                                       Promise<td_api::object_ptr<td_api::richMessage>> &&promise);
+
   void do_create_tone(const string &title, CustomEmojiId custom_emoji_id, const string &prompt, bool show_creator,
                       Promise<td_api::object_ptr<td_api::textCompositionStyle>> &&promise);