Merge pull request #2 from mairror/feat/add_prediction

feat: add prediction

Author: blopezpi

.isort.cfg

@@ -0,0 +1,2 @@
+[settings]
+known_third_party = commands,config,dotenv,errors,requests,ruamel,telegram,utils

.pre-commit-config.yaml

@@ -36,6 +36,7 @@ repos:
     rev: v5.4.2
     hooks:
     -   id: isort
+        args: ["--profile", "black", "--filter-files"]
 -   repo: https://github.com/asottile/pyupgrade
     rev: v2.7.2
     hooks:

README.md

@@ -1,5 +1,5 @@
 # telegram-bot
-Bot telegram 
+Bot telegram
 
 ## Converting Anaconda environment to pip requirements
 
@@ -23,3 +23,4 @@ There isn't any framework to do unit test with our code, you can check it in thi
 - [Telegram Bot Testing Suite](https://github.com/python-telegram-bot/ptbtest)
 - [Telegram Bot Testing Suite RTD](https://ptbtestsuite.readthedocs.io/en/master/?badge=master)
 - https://core.telegram.org/bots/api
+- https://api.slack.com/reference/surfaces/formatting

src/commands.py

@@ -1,15 +1,11 @@
-from telegram import Update
-from utils.logging import telegram_bot
-
-import telegram
-from telegram.ext import (
-    CallbackContext
-)
-import requests
+import json
 from typing import Dict
-from config.settings import API_URL, API_KEY, API_UPLOAD_PATH
-
 
+import requests
+from config.settings import API_KEY, API_PREDICT_PATH, API_UPLOAD_PATH, API_URL
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
+from telegram.ext import CallbackContext
+from utils.logging import telegram_bot
 
 
 def send_photo_to_api(files: Dict):
@@ -24,93 +20,187 @@ def send_photo_to_api(files: Dict):
     """
     data = {"source": "telegram"}
     headers = {"X-Api-Key": API_KEY}
-    telegram_bot.debug(f"Sending the file object to the API.")
+    telegram_bot.debug("Sending the file object to the API.")
     try:
-        r = requests.post(API_URL + API_UPLOAD_PATH, files=files, data=data, headers=headers)
+        r = requests.post(
+            API_URL + API_UPLOAD_PATH, files=files, data=data, headers=headers
+        )
         if r.status_code == 201:
             telegram_bot.info(f"File sucessfully uploaded: {r.text}")
     except Exception as e:
         telegram_bot.error(f"There is an error when send the file: {e}.")
 
 
-
 def start(update: Update, context: CallbackContext) -> None:
     """
     Name: start
     Description:
-        Start callable function when /start command is used on the telegram bot. 
+        Start callable function when /start command is used on the telegram bot.
         Send a reply when using that command.
-    Inputs: 
+    Inputs:
         :update: type(Update): update handler object.
         :context: type(CallbackContext).
     Outputs:
-        None 
+        None
     """
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Welcome to mairror! Do you need to guess your age?")
+    context.bot.send_message(
+        chat_id=update.effective_chat.id,
+        text="Welcome to mairror! Do you need to guess your age?",
+    )
+
 
 def text(update: Update, context: CallbackContext) -> None:
     """
     Name: text
     Description:
-        When a text message is sent in the telegram bot, the context send that message. 
-    Inputs: 
+        When a text message is sent in the telegram bot, the context send that message.
+    Inputs:
         :update: type(Update): update handler object.
         :context: type(CallbackContext).
     Outputs:
-        None 
+        None
     """
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Sorry, I didn't understand you.")
+    context.bot.send_message(
+        chat_id=update.effective_chat.id, text="Sorry, I didn't understand you."
+    )
+
 
 def unknown(update: Update, context: CallbackContext) -> None:
     """
     Name: unknown
     Description:
-        When an unknown command is used the bot is replied with this message. 
-    Inputs: 
+        When an unknown command is used the bot is replied with this message.
+    Inputs:
         :update: type(Update): update handler object.
         :context: type(CallbackContext).
     Outputs:
-        None 
+        None
     """
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Sorry, I didn't understand that command.")
+    context.bot.send_message(
+        chat_id=update.effective_chat.id,
+        text="Sorry, I didn't understand that command.",
+    )
 
 
 def help_command(update: Update, context: CallbackContext) -> None:
     """
     Name: unknown
     Description:
-        When used the /help command, you can get the command helper for the telegram bot. 
-    Inputs: 
+        When used the /help command, you can get the command helper for the telegram bot.
+    Inputs:
         :update: type(Update): update handler object.
         :context: type(CallbackContext).
     Outputs:
-        None 
+        None
     """
-    update.message.reply_text('''
+    update.message.reply_text(
+        """
 <b>Mairror Bot:</b>
 
 You can find <a href="https://github.com/mairror/telegram-bot/blob/main/README.md">here</a> this bot's documentation.
 Authors of this bot:
 
     - <a href="https://www.linkedin.com/in/%E2%9C%85-borja-l-422666a9">Borja</a>
     - <a href="https://www.linkedin.com/in/aacecan">Alex</a>
-        ''', parse_mode="HTML", disable_web_page_preview=True)
+        """,
+        parse_mode="HTML",
+        disable_web_page_preview=True,
+    )
+
+
+def build_text_prediction(prediction: Dict):
+    newline = "\n"
+    text = f"""
+<b>PREDICTION:</b>\n
+{newline.join([f"Face {count+1}: (Age: {pred['age']}), (Gender: {pred['gender']})"
+                for count, pred in enumerate(prediction["predictions"]) ])}
+"""
+    return text
+
+
+def predict(image):
+    data = {"image_id": image}
+    print(data)
+    headers = {"X-Api-Key": API_KEY}
+    telegram_bot.debug(f"Predict image {image}.")
+    try:
+        r = requests.post(
+            API_URL + API_PREDICT_PATH, data=json.dumps(data), headers=headers
+        )
+        if r.status_code == 200:
+            telegram_bot.info(f"Sucessfully predicted: {r.text}")
+            return build_text_prediction(json.loads(r.text))
+        else:
+            telegram_bot.error(f"Error querying the API: {r.text}.")
+            return "There was an error predicting your image."
+    except Exception as e:
+        telegram_bot.error(f"There is an error when get tyhe prediction: {e}.")
+
+
+def keyboard():
+    keyboard = [
+        [
+            InlineKeyboardButton("Yes!!!", callback_data="yes"),
+            InlineKeyboardButton("No", callback_data="no"),
+        ],
+    ]
+
+    reply_markup = InlineKeyboardMarkup(keyboard)
+    return reply_markup
+
+
+def button(update: Update, context: CallbackContext) -> None:
+    """
+    Name: photo
+    Description:
+        Used to send a photo into the API when you uploaded in the telegram bot.
+    Inputs:
+        :update: type(Update): update handler object.
+        :context: type(CallbackContext).
+    Outputs:
+        None
+    """
+    query = update.callback_query
+    query.answer()
+    if query.data == "yes":
+        query.edit_message_text(text="I'm a great guesser!")
+    else:
+        query.edit_message_text(
+            text="Sorry, I'll do the best of me. I need more images to improve the prediction."
+        )
 
 
 def photo(update: Update, context: CallbackContext) -> None:
     """
     Name: photo
     Description:
-        Used to send a photo into the API when you uploaded in the telegram bot. 
-    Inputs: 
+        Used to send a photo into the API when you uploaded in the telegram bot.
+    Inputs:
         :update: type(Update): update handler object.
         :context: type(CallbackContext).
     Outputs:
-        None 
+        None
     """
+    image = f"{update.message.photo[-1].file_unique_id}_{update.message.chat.id}.jpg"
     files = {
         "file": (
             # AQADFrkxG-3bCVFy_3235057.jpg
-            f"{update.message.photo[-1].file_unique_id}_{update.message.chat.id}.jpg",
-                        bytes(update.message.photo[-1].get_file().download_as_bytearray()))}
+            image,
+            bytes(update.message.photo[-1].get_file().download_as_bytearray()),
+        )
+    }
+
     send_photo_to_api(files)
+
+    context.bot.send_chat_action(chat_id=update.effective_chat.id, action="typing")
+
+    prediction_result = predict("raw/" + image)
+    if prediction_result:
+        context.bot.send_message(
+            chat_id=update.effective_chat.id,
+            text=prediction_result,
+            parse_mode="HTML",
+            disable_web_page_preview=True,
+        )
+    reply_markup = keyboard()
+    update.message.reply_text("Have I guessed your age?", reply_markup=reply_markup)

src/config/settings.py

@@ -1,6 +1,6 @@
-from dotenv import load_dotenv
 import os
 
+from dotenv import load_dotenv
 
 if os.path.exists(".env"):
     load_dotenv()
@@ -10,3 +10,5 @@
 API_URL = os.getenv("API_URL", "http://localhost:8000")
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
 API_UPLOAD_PATH = os.getenv("API_UPLOAD_PATH", "/images/upload")
+API_PREDICT_PATH = os.getenv("API_PREDICT_PATH", "/images/predict")
+SLACK_WEBHOOK_URI = os.getenv("SLACK_WEBHOOK_URI", "xxxxx")

src/errors.py

@@ -0,0 +1,51 @@
+import json
+import traceback
+
+import requests
+from config.settings import SLACK_WEBHOOK_URI
+from telegram import Update
+from telegram.ext import CallbackContext
+from utils.logging import telegram_bot
+
+
+def error_handler(update: object, context: CallbackContext) -> None:
+    """
+    Name: error_handler
+    Description:
+        Manage any error from the dispatcher logging it and send it to the developer chat id.
+    Inputs:
+        :update: type(Update): update handler object.
+        :context: type(CallbackContext).
+    Outputs:
+        None
+    """
+    telegram_bot.error(
+        msg="Exception while handling an update:", exc_info=context.error
+    )
+
+    tb_list = traceback.format_exception(
+        None, context.error, context.error.__traceback__
+    )
+    tb_string = "".join(tb_list)
+
+    update_str = update.to_dict() if isinstance(update, Update) else str(update)
+    message = f"""
+<!here>
+An exception was raised while handling an update:\n
+*UPDATE JSON DATA*
+```
+{json.dumps(update_str, indent=2, ensure_ascii=False)}
+```
+*TRACEBACK*
+```
+{tb_string}
+```
+"""
+
+    data = {
+        "blocks": [{"type": "section", "text": {"type": "mrkdwn", "text": message}}]
+    }
+
+    headers = {"Content-type": "application/json"}
+
+    requests.post(SLACK_WEBHOOK_URI, data=json.dumps(data), headers=headers)

src/main.py

@@ -1,13 +1,14 @@
+from commands import button, help_command, photo, start, text, unknown
+from config.settings import BOT_TELEGRAM_TOKEN
+from errors import error_handler
 from telegram.ext import (
-    Updater, 
-    CommandHandler, 
-    MessageHandler, 
-    Filters
+    CallbackQueryHandler,
+    CommandHandler,
+    Filters,
+    MessageHandler,
+    Updater,
 )
 from utils.logging import telegram_bot
-from config.settings import BOT_TELEGRAM_TOKEN
-from commands import start, photo, text, help_command, unknown
-
 
 
 def main():
@@ -22,30 +23,35 @@ def main():
         - Photos -> These are required.
         - Gifs, Videos, Text, etc -> Not allowed
     """
-    updater = Updater(
-        BOT_TELEGRAM_TOKEN,
-        use_context=True,
-        workers=32)
+    telegram_bot.info("Starting telegram bot.")
+    updater = Updater(BOT_TELEGRAM_TOKEN, use_context=True, workers=32)
 
     dispatcher = updater.dispatcher
-    
+    dispatcher.add_error_handler(error_handler)
+
     dispatcher.add_handler(CommandHandler("start", start))
-    
-    dispatcher.add_handler(MessageHandler(Filters.photo & ~Filters.command, photo, run_async=True))
-    
-    text_handler = MessageHandler(Filters.text & (~Filters.command), text, run_async=True)
+
+    dispatcher.add_handler(
+        MessageHandler(Filters.photo & ~Filters.command, photo, run_async=True)
+    )
+
+    text_handler = MessageHandler(
+        Filters.text & (~Filters.command), text, run_async=True
+    )
     dispatcher.add_handler(text_handler)
-    
+
     dispatcher.add_handler(CommandHandler("help", help_command))
-    
+
+    updater.dispatcher.add_handler(CallbackQueryHandler(button))
+
     unknown_handler = MessageHandler(Filters.command, unknown, run_async=True)
     dispatcher.add_handler(unknown_handler)
 
+    telegram_bot.info("Start polling.")
 
     updater.start_polling()
     updater.idle()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
-