nestedconversationbot.py

  1#!/usr/bin/env python
  2# pylint: disable=unused-argument, wrong-import-position
  3# This program is dedicated to the public domain under the CC0 license.
  4
  5"""
  6First, a few callback functions are defined. Then, those functions are passed to
  7the Application and registered at their respective places.
  8Then, the bot is started and runs until we press Ctrl-C on the command line.
  9
 10Usage:
 11Example of a bot-user conversation using nested ConversationHandlers.
 12Send /start to initiate the conversation.
 13Press Ctrl-C on the command line or send a signal to the process to stop the
 14bot.
 15"""
 16
 17import logging
 18from typing import Any, Dict, Tuple
 19
 20from telegram import __version__ as TG_VER
 21
 22try:
 23    from telegram import __version_info__
 24except ImportError:
 25    __version_info__ = (0, 0, 0, 0, 0)  # type: ignore[assignment]
 26
 27if __version_info__ < (20, 0, 0, "alpha", 1):
 28    raise RuntimeError(
 29        f"This example is not compatible with your current PTB version {TG_VER}. To view the "
 30        f"{TG_VER} version of this example, "
 31        f"visit https://docs.python-telegram-bot.org/en/v{TG_VER}/examples.html"
 32    )
 33from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
 34from telegram.ext import (
 35    Application,
 36    CallbackQueryHandler,
 37    CommandHandler,
 38    ContextTypes,
 39    ConversationHandler,
 40    MessageHandler,
 41    filters,
 42)
 43
 44# Enable logging
 45logging.basicConfig(
 46    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
 47)
 48logger = logging.getLogger(__name__)
 49
 50# State definitions for top level conversation
 51SELECTING_ACTION, ADDING_MEMBER, ADDING_SELF, DESCRIBING_SELF = map(chr, range(4))
 52# State definitions for second level conversation
 53SELECTING_LEVEL, SELECTING_GENDER = map(chr, range(4, 6))
 54# State definitions for descriptions conversation
 55SELECTING_FEATURE, TYPING = map(chr, range(6, 8))
 56# Meta states
 57STOPPING, SHOWING = map(chr, range(8, 10))
 58# Shortcut for ConversationHandler.END
 59END = ConversationHandler.END
 60
 61# Different constants for this example
 62(
 63    PARENTS,
 64    CHILDREN,
 65    SELF,
 66    GENDER,
 67    MALE,
 68    FEMALE,
 69    AGE,
 70    NAME,
 71    START_OVER,
 72    FEATURES,
 73    CURRENT_FEATURE,
 74    CURRENT_LEVEL,
 75) = map(chr, range(10, 22))
 76
 77
 78# Helper
 79def _name_switcher(level: str) -> Tuple[str, str]:
 80    if level == PARENTS:
 81        return "Father", "Mother"
 82    return "Brother", "Sister"
 83
 84
 85# Top level conversation callbacks
 86async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
 87    """Select an action: Adding parent/child or show data."""
 88    text = (
 89        "You may choose to add a family member, yourself, show the gathered data, or end the "
 90        "conversation. To abort, simply type /stop."
 91    )
 92
 93    buttons = [
 94        [
 95            InlineKeyboardButton(text="Add family member", callback_data=str(ADDING_MEMBER)),
 96            InlineKeyboardButton(text="Add yourself", callback_data=str(ADDING_SELF)),
 97        ],
 98        [
 99            InlineKeyboardButton(text="Show data", callback_data=str(SHOWING)),
100            InlineKeyboardButton(text="Done", callback_data=str(END)),
101        ],
102    ]
103    keyboard = InlineKeyboardMarkup(buttons)
104
105    # If we're starting over we don't need to send a new message
106    if context.user_data.get(START_OVER):
107        await update.callback_query.answer()
108        await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
109    else:
110        await update.message.reply_text(
111            "Hi, I'm Family Bot and I'm here to help you gather information about your family."
112        )
113        await update.message.reply_text(text=text, reply_markup=keyboard)
114
115    context.user_data[START_OVER] = False
116    return SELECTING_ACTION
117
118
119async def adding_self(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
120    """Add information about yourself."""
121    context.user_data[CURRENT_LEVEL] = SELF
122    text = "Okay, please tell me about yourself."
123    button = InlineKeyboardButton(text="Add info", callback_data=str(MALE))
124    keyboard = InlineKeyboardMarkup.from_button(button)
125
126    await update.callback_query.answer()
127    await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
128
129    return DESCRIBING_SELF
130
131
132async def show_data(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
133    """Pretty print gathered data."""
134
135    def pretty_print(data: Dict[str, Any], level: str) -> str:
136        people = data.get(level)
137        if not people:
138            return "\nNo information yet."
139
140        return_str = ""
141        if level == SELF:
142            for person in data[level]:
143                return_str += f"\nName: {person.get(NAME, '-')}, Age: {person.get(AGE, '-')}"
144        else:
145            male, female = _name_switcher(level)
146
147            for person in data[level]:
148                gender = female if person[GENDER] == FEMALE else male
149                return_str += (
150                    f"\n{gender}: Name: {person.get(NAME, '-')}, Age: {person.get(AGE, '-')}"
151                )
152        return return_str
153
154    user_data = context.user_data
155    text = f"Yourself:{pretty_print(user_data, SELF)}"
156    text += f"\n\nParents:{pretty_print(user_data, PARENTS)}"
157    text += f"\n\nChildren:{pretty_print(user_data, CHILDREN)}"
158
159    buttons = [[InlineKeyboardButton(text="Back", callback_data=str(END))]]
160    keyboard = InlineKeyboardMarkup(buttons)
161
162    await update.callback_query.answer()
163    await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
164    user_data[START_OVER] = True
165
166    return SHOWING
167
168
169async def stop(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
170    """End Conversation by command."""
171    await update.message.reply_text("Okay, bye.")
172
173    return END
174
175
176async def end(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
177    """End conversation from InlineKeyboardButton."""
178    await update.callback_query.answer()
179
180    text = "See you around!"
181    await update.callback_query.edit_message_text(text=text)
182
183    return END
184
185
186# Second level conversation callbacks
187async def select_level(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
188    """Choose to add a parent or a child."""
189    text = "You may add a parent or a child. Also you can show the gathered data or go back."
190    buttons = [
191        [
192            InlineKeyboardButton(text="Add parent", callback_data=str(PARENTS)),
193            InlineKeyboardButton(text="Add child", callback_data=str(CHILDREN)),
194        ],
195        [
196            InlineKeyboardButton(text="Show data", callback_data=str(SHOWING)),
197            InlineKeyboardButton(text="Back", callback_data=str(END)),
198        ],
199    ]
200    keyboard = InlineKeyboardMarkup(buttons)
201
202    await update.callback_query.answer()
203    await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
204
205    return SELECTING_LEVEL
206
207
208async def select_gender(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
209    """Choose to add mother or father."""
210    level = update.callback_query.data
211    context.user_data[CURRENT_LEVEL] = level
212
213    text = "Please choose, whom to add."
214
215    male, female = _name_switcher(level)
216
217    buttons = [
218        [
219            InlineKeyboardButton(text=f"Add {male}", callback_data=str(MALE)),
220            InlineKeyboardButton(text=f"Add {female}", callback_data=str(FEMALE)),
221        ],
222        [
223            InlineKeyboardButton(text="Show data", callback_data=str(SHOWING)),
224            InlineKeyboardButton(text="Back", callback_data=str(END)),
225        ],
226    ]
227    keyboard = InlineKeyboardMarkup(buttons)
228
229    await update.callback_query.answer()
230    await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
231
232    return SELECTING_GENDER
233
234
235async def end_second_level(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
236    """Return to top level conversation."""
237    context.user_data[START_OVER] = True
238    await start(update, context)
239
240    return END
241
242
243# Third level callbacks
244async def select_feature(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
245    """Select a feature to update for the person."""
246    buttons = [
247        [
248            InlineKeyboardButton(text="Name", callback_data=str(NAME)),
249            InlineKeyboardButton(text="Age", callback_data=str(AGE)),
250            InlineKeyboardButton(text="Done", callback_data=str(END)),
251        ]
252    ]
253    keyboard = InlineKeyboardMarkup(buttons)
254
255    # If we collect features for a new person, clear the cache and save the gender
256    if not context.user_data.get(START_OVER):
257        context.user_data[FEATURES] = {GENDER: update.callback_query.data}
258        text = "Please select a feature to update."
259
260        await update.callback_query.answer()
261        await update.callback_query.edit_message_text(text=text, reply_markup=keyboard)
262    # But after we do that, we need to send a new message
263    else:
264        text = "Got it! Please select a feature to update."
265        await update.message.reply_text(text=text, reply_markup=keyboard)
266
267    context.user_data[START_OVER] = False
268    return SELECTING_FEATURE
269
270
271async def ask_for_input(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
272    """Prompt user to input data for selected feature."""
273    context.user_data[CURRENT_FEATURE] = update.callback_query.data
274    text = "Okay, tell me."
275
276    await update.callback_query.answer()
277    await update.callback_query.edit_message_text(text=text)
278
279    return TYPING
280
281
282async def save_input(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
283    """Save input for feature and return to feature selection."""
284    user_data = context.user_data
285    user_data[FEATURES][user_data[CURRENT_FEATURE]] = update.message.text
286
287    user_data[START_OVER] = True
288
289    return await select_feature(update, context)
290
291
292async def end_describing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
293    """End gathering of features and return to parent conversation."""
294    user_data = context.user_data
295    level = user_data[CURRENT_LEVEL]
296    if not user_data.get(level):
297        user_data[level] = []
298    user_data[level].append(user_data[FEATURES])
299
300    # Print upper level menu
301    if level == SELF:
302        user_data[START_OVER] = True
303        await start(update, context)
304    else:
305        await select_level(update, context)
306
307    return END
308
309
310async def stop_nested(update: Update, context: ContextTypes.DEFAULT_TYPE) -> str:
311    """Completely end conversation from within nested conversation."""
312    await update.message.reply_text("Okay, bye.")
313
314    return STOPPING
315
316
317def main() -> None:
318    """Run the bot."""
319    # Create the Application and pass it your bot's token.
320    application = Application.builder().token("TOKEN").build()
321
322    # Set up third level ConversationHandler (collecting features)
323    description_conv = ConversationHandler(
324        entry_points=[
325            CallbackQueryHandler(
326                select_feature, pattern="^" + str(MALE) + "$|^" + str(FEMALE) + "$"
327            )
328        ],
329        states={
330            SELECTING_FEATURE: [
331                CallbackQueryHandler(ask_for_input, pattern="^(?!" + str(END) + ").*$")
332            ],
333            TYPING: [MessageHandler(filters.TEXT & ~filters.COMMAND, save_input)],
334        },
335        fallbacks=[
336            CallbackQueryHandler(end_describing, pattern="^" + str(END) + "$"),
337            CommandHandler("stop", stop_nested),
338        ],
339        map_to_parent={
340            # Return to second level menu
341            END: SELECTING_LEVEL,
342            # End conversation altogether
343            STOPPING: STOPPING,
344        },
345    )
346
347    # Set up second level ConversationHandler (adding a person)
348    add_member_conv = ConversationHandler(
349        entry_points=[CallbackQueryHandler(select_level, pattern="^" + str(ADDING_MEMBER) + "$")],
350        states={
351            SELECTING_LEVEL: [
352                CallbackQueryHandler(select_gender, pattern=f"^{PARENTS}$|^{CHILDREN}$")
353            ],
354            SELECTING_GENDER: [description_conv],
355        },
356        fallbacks=[
357            CallbackQueryHandler(show_data, pattern="^" + str(SHOWING) + "$"),
358            CallbackQueryHandler(end_second_level, pattern="^" + str(END) + "$"),
359            CommandHandler("stop", stop_nested),
360        ],
361        map_to_parent={
362            # After showing data return to top level menu
363            SHOWING: SHOWING,
364            # Return to top level menu
365            END: SELECTING_ACTION,
366            # End conversation altogether
367            STOPPING: END,
368        },
369    )
370
371    # Set up top level ConversationHandler (selecting action)
372    # Because the states of the third level conversation map to the ones of the second level
373    # conversation, we need to make sure the top level conversation can also handle them
374    selection_handlers = [
375        add_member_conv,
376        CallbackQueryHandler(show_data, pattern="^" + str(SHOWING) + "$"),
377        CallbackQueryHandler(adding_self, pattern="^" + str(ADDING_SELF) + "$"),
378        CallbackQueryHandler(end, pattern="^" + str(END) + "$"),
379    ]
380    conv_handler = ConversationHandler(
381        entry_points=[CommandHandler("start", start)],
382        states={
383            SHOWING: [CallbackQueryHandler(start, pattern="^" + str(END) + "$")],
384            SELECTING_ACTION: selection_handlers,
385            SELECTING_LEVEL: selection_handlers,
386            DESCRIBING_SELF: [description_conv],
387            STOPPING: [CommandHandler("start", start)],
388        },
389        fallbacks=[CommandHandler("stop", stop)],
390    )
391
392    application.add_handler(conv_handler)
393
394    # Run the bot until the user presses Ctrl-C
395    application.run_polling()
396
397
398if __name__ == "__main__":
399    main()

State Diagram

flowchart TB %% Documentation: https://mermaid-js.github.io/mermaid/#/flowchart A(("/start")):::entryPoint -->|Hi! I'm FamilyBot...| B((SELECTING_ACTION)):::state B --> C("Show Data"):::userInput C --> |"(List of gathered data)"| D((SHOWING)):::state D --> E("Back"):::userInput E --> B B --> F("Add Yourself"):::userInput F --> G(("DESCRIBING_SELF")):::state G --> H("Add info"):::userInput H --> I((SELECT_FEATURE)):::state I --> |"Please select a feature to update. <br /> - Name <br /> - Age <br /> - Done"|J("(choice)"):::userInput J --> |"Okay, tell me."| K((TYPING)):::state K --> L("(text)"):::userInput L --> |"[saving]"|I I --> M("Done"):::userInput M --> B B --> N("Add family member"):::userInput R --> I W --> |"See you around!"|End(("END")):::termination Y(("ANY STATE")):::state --> Z("/stop"):::userInput Z -->|"Okay, bye."| End B --> W("Done"):::userInput subgraph nestedConversation[Nested Conversation: Add Family Member] direction BT N --> O(("SELECT_LEVEL")):::state O --> |"Add... <br /> - Add Parent <br /> - Add Child <br />"|P("(choice)"):::userInput P --> Q(("SELECT_GENDER")):::state Q --> |"- Mother <br /> - Father <br /> / <br /> - Sister <br /> - Brother"| R("(choice)"):::userInput Q --> V("Show Data"):::userInput Q --> T(("SELECTING_ACTION")):::state Q --> U("Back"):::userInput U --> T O --> U O --> V V --> S(("SHOWING")):::state V --> T end classDef userInput fill:#2a5279, color:#ffffff, stroke:#ffffff classDef state fill:#222222, color:#ffffff, stroke:#ffffff classDef entryPoint fill:#009c11, stroke:#42FF57, color:#ffffff classDef termination fill:#bb0007, stroke:#E60109, color:#ffffff style nestedConversation fill:#999999, stroke-width:2px, stroke:#333333