Add td_api::fixRichMessageWithAi.

6 files changed, 26 insertions, 0 deletions

diff --git a/td/generate/scheme/td_api.tl b/td/generate/scheme/td_api.tl
index e5f89efff..0dc9ce3e6 100644
--- a/td/generate/scheme/td_api.tl
+++ b/td/generate/scheme/td_api.tl
@@ -11952,6 +11952,10 @@ composeRichMessageWithAi message:inputRichMessage translate_to_language_code:str
 //@text The original text
 fixTextWithAi text:formattedText = FixedText;
 
+//@description Fixes rich message using an AI model. May return an error with a message "AICOMPOSE_FLOOD_PREMIUM" if Telegram Premium is required to send further requests
+//@message The original message
+fixRichMessageWithAi message:inputRichMessage = RichMessage;
+
 //@description Recognizes speech in a video note or a voice note message
 //@chat_id Identifier of the chat to which the message belongs
 //@message_id Identifier of the message. Use messageProperties.can_recognize_speech to check whether the message is suitable
diff --git a/td/telegram/Requests.cpp b/td/telegram/Requests.cpp
index 3d9547254..b6b368c10 100644
--- a/td/telegram/Requests.cpp
+++ b/td/telegram/Requests.cpp
@@ -2906,6 +2906,12 @@ void Requests::on_request(uint64 id, td_api::fixTextWithAi &request) {
   td_->translation_manager_->proofread_message_with_ai(std::move(request.text_), std::move(promise));
 }
 
+void Requests::on_request(uint64 id, td_api::fixRichMessageWithAi &request) {
+  CHECK_IS_USER();
+  CREATE_REQUEST_PROMISE();
+  td_->translation_manager_->proofread_rich_message_with_ai(std::move(request.message_), std::move(promise));
+}
+
 void Requests::on_request(uint64 id, const td_api::recognizeSpeech &request) {
   CHECK_IS_USER();
   CREATE_OK_REQUEST_PROMISE();
diff --git a/td/telegram/Requests.h b/td/telegram/Requests.h
index bb191e2b0..49afc1d8a 100644
--- a/td/telegram/Requests.h
+++ b/td/telegram/Requests.h
@@ -352,6 +352,8 @@ class Requests {
 
   void on_request(uint64 id, td_api::fixTextWithAi &request);
 
+  void on_request(uint64 id, td_api::fixRichMessageWithAi &request);
+
   void on_request(uint64 id, const td_api::recognizeSpeech &request);
 
   void on_request(uint64 id, const td_api::rateSpeechRecognition &request);
diff --git a/td/telegram/TranslationManager.cpp b/td/telegram/TranslationManager.cpp
index c21c40717..32f540a26 100644
--- a/td/telegram/TranslationManager.cpp
+++ b/td/telegram/TranslationManager.cpp
@@ -695,6 +695,13 @@ void TranslationManager::compose_rich_message_with_ai(td_api::object_ptr<td_api:
       ->send(true, input_rich_message, translate_to_language_code, std::move(input_tone), emojify, false);
 }
 
+void TranslationManager::proofread_rich_message_with_ai(td_api::object_ptr<td_api::inputRichMessage> &&message,
+                                                        Promise<td_api::object_ptr<td_api::richMessage>> &&promise) {
+  TRY_RESULT_PROMISE(promise, input_rich_message, get_input_rich_message(std::move(message)));
+  td_->create_handler<ComposeRichMessageWithAiQuery>(std::move(promise))
+      ->send(true, input_rich_message, string(), nullptr, false, true);
+}
+
 string TranslationManager::get_ai_compose_tones_key() {
   return "ai_compose_styles";
 }
diff --git a/td/telegram/TranslationManager.h b/td/telegram/TranslationManager.h
index 7d3c6919a..38e94e6bb 100644
--- a/td/telegram/TranslationManager.h
+++ b/td/telegram/TranslationManager.h
@@ -65,6 +65,9 @@ class TranslationManager final : public Actor {
   void proofread_message_with_ai(td_api::object_ptr<td_api::formattedText> &&text,
                                  Promise<td_api::object_ptr<td_api::fixedText>> &&promise);
 
+  void proofread_rich_message_with_ai(td_api::object_ptr<td_api::inputRichMessage> &&message,
+                                      Promise<td_api::object_ptr<td_api::richMessage>> &&promise);
+
   void reload_ai_compose_tones(Promise<Unit> &&promise);
 
   void on_get_ai_compose_tones(telegram_api::object_ptr<telegram_api::aicompose_Tones> &&tones_ptr);
diff --git a/td/telegram/cli.cpp b/td/telegram/cli.cpp
index 7d5be32bc..f04464048 100644
--- a/td/telegram/cli.cpp
+++ b/td/telegram/cli.cpp
@@ -5242,6 +5242,10 @@ class CliClient final : public Actor {
       string text;
       get_args(args, text);
       send_request(td_api::make_object<td_api::fixTextWithAi>(as_formatted_text(text)));
+    } else if (op == "frmwa") {
+      string message;
+      get_args(args, message);
+      send_request(td_api::make_object<td_api::fixRichMessageWithAi>(as_input_rich_message(message)));
     } else if (op == "rs") {
       ChatId chat_id;
       MessageId message_id;