diff --git a/bot/bot_templates/main.txt b/bot/bot_templates/main.txt index 3051763..5f839ee 100644 --- a/bot/bot_templates/main.txt +++ b/bot/bot_templates/main.txt @@ -2,7 +2,7 @@ import asyncio import logging from aiogram import Bot, Dispatcher, F import re -from aiogram.types import Message +from aiogram.types import Message, CallbackQuery from aiogram.client.session.aiohttp import AiohttpSession from aiogram.fsm.storage.memory import MemoryStorage from aiogram.client.telegram import TelegramAPIServer diff --git a/bot/services.py b/bot/services.py index cfd4296..fb243cd 100644 --- a/bot/services.py +++ b/bot/services.py @@ -11,7 +11,8 @@ def generate_code(bot: Bot) -> str: bot_id=bot.id, component_type=Component.ComponentType.TRIGGER, ) - bot_component_codes = "" + bot_component_codes = [] + raw_state_check = "" for component in components: if component.component_content_type.model != "onmessage": raise ValidationError( @@ -19,21 +20,31 @@ def generate_code(bot: Bot) -> str: ) for next_component in component.get_all_next_components(): - # if next_component.id == component.id: - # continue object = next_component.component_content_type.model_class().objects.get( pk=next_component.pk, ) - bot_component_codes += object.generate_code() - bot_component_codes += "\n" * 2 + code_result = object.generate_code() + if isinstance(code_result, tuple): + keyboard, callback_code = code_result + if keyboard: + bot_component_codes.append(keyboard) + if callback_code: + bot_component_codes.append(callback_code) + else: + if "raw_state" in code_result: + raw_state_check += code_result + "\n" + else: + bot_component_codes.append(code_result) + + if raw_state_check: + bot_component_codes.append(raw_state_check) with open("bot/bot_templates/main.txt") as f: base = f.read() + code = base.format( - FUNCTION_CODES=bot_component_codes, + FUNCTION_CODES="\n\n".join(bot_component_codes), TOKEN=bot.token, BASE_URL=settings.BALE_API_URL, ) - print("FOOOOO", code) - return black.format_str(code, mode=black.Mode()) diff --git a/bot/test/tests.py b/bot/test/tests.py index 71133ef..b9a1370 100644 --- a/bot/test/tests.py +++ b/bot/test/tests.py @@ -1,14 +1,17 @@ -import unittest -from io import BytesIO - -from django.core.files.base import ContentFile from django.test import TestCase from django.urls import reverse from PIL import Image from rest_framework import status from bot.models import Bot -from component.models import Component, Markup, OnMessage, SetState +from component.models import ( + CodeComponent, + Component, + Markup, + OnMessage, + SetState, + SwitchComponent, +) from component.telegram.models import SendMessage, SendPhoto from iam.models import IamUser from iam.utils import create_token_for_iamuser @@ -35,64 +38,180 @@ def setUp(self): component_type=Component.ComponentType.TRIGGER, text="/start", ) - send_message_comp = SendMessage.objects.create( + send_message_start = SendMessage.objects.create( bot=self.bot, - chat_id=693259126, - text="Hello, World!", + chat_id=".from_user.id", + text="Hello welcome to the bot", position_x=1, position_y=1, previous_component=on_message_component, ) - SendMessage.objects.create( + CodeComponent.objects.create( bot=self.bot, - chat_id=693259126, - text="Bye World", + code="print('This is log message to validate code')", position_x=1, position_y=1, previous_component=on_message_component, ) - image = Image.new("RGB", (100, 100), color="blue") - image_io = BytesIO() - image.save(image_io, "PNG") - image_content = ContentFile(image_io.getvalue(), "test.png") - photo_component = SendPhoto.objects.create( + send_support = SendMessage.objects.create( bot=self.bot, - chat_id=".chat.id", - photo=image_content, + chat_id=".from_user.id", + text="Send your message then i want to send to admin", + position_x=1, + position_y=1, + ) + SetState.objects.create( + bot=self.bot, + state="support", + position_x=1, + position_y=1, + previous_component=send_support, + ) + + send_message_help = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="This is test help message", + position_x=1, + position_y=1, + ) + + send_message_accept_terms = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="Some text to accept terms and conditions... Do you accept?, send Yes or No", position_x=1, position_y=1, - previous_component=on_message_component, + ) + SetState.objects.create( + bot=self.bot, + state="accept_terms", + position_x=1, + position_y=1, + previous_component=send_message_accept_terms, ) Markup.objects.create( - parent_component=photo_component, + parent_component=send_message_start, markup_type=Markup.MarkupType.ReplyKeyboard, - buttons=[["Button 1", "Button 2"], ["Button 3", "Button 4"]], + buttons=[ + [{"value": "Help", "next_component": send_message_help.id}], + [ + { + "value": "Accept terms and conditions", + "next_component": send_message_accept_terms.id, + }, + ], + [{"value": "Support", "next_component": send_support.id}], + ], ) - SetState.objects.create( + handle_support = OnMessage.objects.create( bot=self.bot, - state="state", + state="support", position_x=1, position_y=1, - previous_component=on_message_component, + ) + send_message_to_admin = SendMessage.objects.create( + bot=self.bot, + chat_id="693259126", + text="text send from user", + position_x=1, + position_y=1, + previous_component=handle_support, + ) + send_message_to_user = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="your message send to admin", + position_x=1, + position_y=1, + previous_component=send_message_to_admin, ) - on_state_component = OnMessage.objects.create( + handle_accept_terms = OnMessage.objects.create( bot=self.bot, - state="state", + state="accept_terms", + position_x=1, + position_y=1, + ) + send_message_accept_terms_yes = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="You accepted terms and conditions, send your phone number start with +98", position_x=1, position_y=1, ) - SendMessage.objects.create( + send_message_accept_terms_no = SendMessage.objects.create( bot=self.bot, - chat_id=693259126, - text="State set", + chat_id=".from_user.id", + text="You did not accept terms and conditions", position_x=1, position_y=1, - previous_component=on_state_component, + ) + + SwitchComponent.objects.create( + bot=self.bot, + position_x=1, + position_y=1, + previous_component=handle_accept_terms, + expression=".text", + values=["Yes", "No"], + next_components=[ + send_message_accept_terms_yes.id, + send_message_accept_terms_no.id, + ], + ) + + phone_number_hanndle = OnMessage.objects.create( + bot=self.bot, + regex=True, + text="^\\+98\\d+$", + position_x=1, + position_y=1, + ) + send_message_phone_number = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="Is your phone number correct?", + position_x=1, + position_y=1, + previous_component=phone_number_hanndle, + ) + phone_number_handle_yes_no = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text=".text", + position_x=1, + position_y=1, + previous_component=phone_number_hanndle, + ) + + handle_yes = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="OK, thanks", + position_x=1, + position_y=1, + ) + + handle_no = SendMessage.objects.create( + bot=self.bot, + chat_id=".from_user.id", + text="OK, you lose..., send /start to start again", + position_x=1, + position_y=1, + ) + + Markup.objects.create( + parent_component=phone_number_handle_yes_no, + markup_type=Markup.MarkupType.InlineKeyboard, + buttons=[ + [{"value": "Yes", "next_component": handle_yes.id}], + [{"value": "No", "next_component": handle_no.id}], + ], ) def test_create_bot(self): @@ -106,5 +225,5 @@ def test_create_bot(self): ) self.assertEqual(response.status_code, status.HTTP_200_OK, response.content) # print(response.content.decode()) - with open("code.py", "w") as f: + with open("code_1.py", "w") as f: f.write(response.content.decode()) diff --git a/component/models.py b/component/models.py index adc56f2..8247571 100644 --- a/component/models.py +++ b/component/models.py @@ -1,5 +1,7 @@ from typing import Any +from django.core.exceptions import ValidationError + from component.telegram.models import * @@ -30,20 +32,45 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # ] def generate_code(self) -> str: - dict_key = "" + if len(self.values) != len(self.next_components): + raise ValidationError( + "Values and next_components must have the same length", + ) + + base_code = "" code = [ f"async def {self.code_function_name}(message: Message, **kwargs):", f" value = message{self.expression}", f" match value:", ] - for i in range(len(self.values)): - value = self.values[i] - next_component = self.next_components[i] - code += [ - f" case {value}:", - f" await {next_component.code_function_name}(message, **kwargs)", - ] - return code + + for value, next_component in zip(self.values, self.next_components): + next_component = Component.objects.get( + pk=next_component, + ) + next_component = ( + next_component.component_content_type.model_class().objects.get( + pk=next_component.pk, + ) + ) + base_code += next_component.generate_code() + "\n" + code.extend( + [ + f" case '{value}':", + f" await {next_component.code_function_name}(message, **kwargs)", + ], + ) + + # Add default case + code.extend( + [ + " case _:", + " pass # No matching case found", + ], + ) + + base_code += "\n".join(code) + return base_code @property def required_fields(self) -> list: @@ -64,6 +91,29 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: def required_fields(self) -> list: return ["code"] + def _format_code_component(self, underlying_object) -> list[str]: + if not underlying_object.code: + return [" pass # No code provided"] + + try: + import black + + formatted_code = black.format_str(underlying_object.code, mode=black.Mode()) + return [f" {formatted_code}"] + except Exception as e: + return [ + f" # Original code failed black formatting: {str(e)}", + f" # {underlying_object.code}", + " pass", + ] + + def generate_code(self) -> str: + code = [ + f"async def {self.code_function_name}(message: Message, **kwargs):", + *self._format_code_component(self), + ] + return "\n".join(code) + class SetState(Component): """Use this method to set a state. Returns True on success.""" @@ -85,11 +135,8 @@ def required_fields(self) -> list: return ["state"] def generate_code(self) -> str: - - underlying_object: ( - SetState - ) = self.component_content_type.model_class().objects.get( - pk=self.pk, + underlying_object: SetState = ( + self.component_content_type.model_class().objects.get(pk=self.pk) ) code = [ @@ -139,35 +186,62 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: help_text="Optional comma separated list of states to match. If not specified, matches any state.", ) - def generate_code(self) -> str: + def _build_text_filter(self, text: str, regex: bool, case_sensitive: bool) -> str: + """Build text matching filter based on configuration.""" + if regex: + # For regex, use regexp function + return f"F.text.regexp(r'{text}')" + else: + # For exact text matching, handle case sensitivity + if case_sensitive: + return f"F.text == '{text}'" + else: + return f"F.text.lower() == '{text.lower()}'" - underlying_object: ( - OnMessage - ) = self.component_content_type.model_class().objects.get( - pk=self.pk, + def _build_state_filter(self, state_string: str) -> str: + """Build state matching filter from comma-separated state list.""" + if not state_string: + return "" + + # Parse and clean state list + states = [s.strip() for s in state_string.split(",") if s.strip()] + if not states: + return "" + + # Build lambda filter for state matching + state_list = [f"'{state}'" for state in states] + return f"lambda _, raw_state: raw_state in [{', '.join(state_list)}]" + + def generate_code(self) -> str: + underlying_object: OnMessage = ( + self.component_content_type.model_class().objects.get(pk=self.pk) ) if underlying_object.next_component.count() == 0: return "" filters = [] if underlying_object.text: - if underlying_object.regex: - filters.append(f"regexp=re.compile(r'{underlying_object.text})'") - else: - filters.append( - f"F.text{'.lower()' if underlying_object.case_sensitive else ''} == '{underlying_object.text}'", - ) - if underlying_object.state: - state_list = [f"'{s.strip()}'" for s in underlying_object.state.split(",")] - filters.append( - f"lambda _, raw_state: raw_state in [{','.join(state_list)}]", + text_filter = self._build_text_filter( + underlying_object.text, + underlying_object.regex, + underlying_object.case_sensitive, ) + filters.append(text_filter) + + if underlying_object.state: + state_filter = self._build_state_filter(underlying_object.state) + if state_filter: + filters.append(state_filter) + filter_str = ", ".join(filters) if filters else "" code = [ - f"@dp.message({','.join(filters)})", + f"@dp.message({filter_str})", f"async def {self.code_function_name}(message: Message, **kwargs):", ] + if underlying_object.state: + code.append(" await kwargs['state'].clear()") + for next_component in underlying_object.next_component.all(): next_component = ( next_component.component_content_type.model_class().objects.get( @@ -200,19 +274,154 @@ class MarkupType(models.TextChoices): def validate(self) -> None: buttons = self.buttons - assert isinstance(buttons, list) - for row in buttons: - assert isinstance(row, list) - for button in row: - assert isinstance(button, dict) - assert "value" in button - assert "next_component" in button - assert isinstance(button["value"], str) - assert isinstance(button["next_component"], int) - - def get_callback_data(self, cell: str) -> str: - return f"{self.prent_component.id}-{cell}" + if not isinstance(buttons, list): + raise ValidationError("Buttons must be a list") + + for row_idx, row in enumerate(buttons): + if not isinstance(row, list): + raise ValidationError(f"Row {row_idx} must be a list") + + for button_idx, button in enumerate(row): + if not isinstance(button, dict): + raise ValidationError( + f"Button at row {row_idx}, column {button_idx} must be a dict", + ) + if "value" not in button: + raise ValidationError( + f"Button at row {row_idx}, column {button_idx} must have 'value' key", + ) + if not isinstance(button["value"], str): + raise ValidationError( + f"Button value at row {row_idx}, column {button_idx} must be a string", + ) + + def get_callback_data(self, cell: dict) -> str: + if not isinstance(cell, dict) or "value" not in cell: + raise ValidationError("Cell must be a dict with 'value' key") + + value = cell["value"].replace(" ", "_") + return f"{self.parent_component.id}-{value}" def save(self, *args: Any, **kwargs: Any) -> None: self.validate() super().save(*args, **kwargs) + + def _get_markup_config(self) -> tuple[str, str, str]: + """Returns the configuration for the markup type.""" + config_map = { + self.MarkupType.ReplyKeyboard: ( + "ReplyKeyboardMarkup", + "KeyboardButton", + "resize_keyboard=True, one_time_keyboard=False, keyboard", + ), + self.MarkupType.InlineKeyboard: ( + "InlineKeyboardMarkup", + "InlineKeyboardButton", + "inline_keyboard", + ), + } + + if self.markup_type not in config_map: + raise NotImplementedError(f"Unknown markup type: {self.markup_type}") + + return config_map[self.markup_type] + + def _generate_button_args(self, cell: dict) -> dict: + """Generates the arguments for a button.""" + args = {"text": cell["value"]} + if self.markup_type == self.MarkupType.InlineKeyboard: + args["callback_data"] = self.get_callback_data(cell) + return args + + def _generate_button_code(self, button_class: str, args: dict) -> str: + """Generates the code for a single button.""" + button_lines = [f"{button_class}("] + for k, v in args.items(): + button_lines.append(f' {k}="{v}",') + button_lines.append("),") + return "\n".join(button_lines) + + def _generate_callback_handlers(self, cell: dict) -> tuple[list[str], list[str]]: + """Generates callback handlers for a cell if needed.""" + base_code = [] + callback_code = [] + + first_next_component = cell.get("next_component") + if not first_next_component: + return base_code, callback_code + + try: + component_obj = Component.objects.get(id=first_next_component) + object = component_obj.component_content_type.model_class().objects.get( + pk=first_next_component, + ) + except Component.DoesNotExist: + raise ValidationError( + f"Component with id {first_next_component} does not exist", + ) + + if self.markup_type == self.MarkupType.InlineKeyboard: + callback_data = self.get_callback_data(cell) + callback_code.extend( + [ + f"@dp.callback_query(lambda callback_query: callback_query.data == '{callback_data}')", + f"async def {object.code_function_name}_callback(callback_query: CallbackQuery, **kwargs):", + f" await {object.code_function_name}(callback_query, **kwargs)", + ], + ) + else: + callback_code.extend( + [ + f"@dp.message(F.text == '{cell['value']}')", + f"async def {object.code_function_name}_handler(message: Message, **kwargs):", + f" await {object.code_function_name}(message, **kwargs)", + ], + ) + + # Generate code for all next components + try: + for next_component in Component.objects.get( + id=first_next_component, + ).get_all_next_components(): + next_obj = ( + next_component.component_content_type.model_class().objects.get( + pk=next_component.pk, + ) + ) + base_code.append(next_obj.generate_code()) + except AttributeError: + # get_all_next_components method might not exist + pass + + return base_code, callback_code + + def generate_code(self) -> tuple[str, list[str]]: + """Generates the keyboard markup code and callback handlers.""" + base_code = [] + callback_code = [] + + keyword_class, button_class, variable_name = self._get_markup_config() + + keyboard_rows = [] + for row in self.buttons: + row_buttons = [] + for cell in row: + cell_base_code, cell_callback_code = self._generate_callback_handlers( + cell, + ) + base_code.extend(cell_base_code) + callback_code.extend(cell_callback_code) + + button_args = self._generate_button_args(cell) + row_buttons.append( + self._generate_button_code(button_class, button_args), + ) + + keyboard_rows.append("[\n" + "".join(row_buttons) + "\n]") + + keyboard_buttons = "[\n" + ",\n".join(keyboard_rows) + "\n]" + keyboard = f"{keyword_class}({variable_name} = {keyboard_buttons})" + + callback_code.append("\n\n") + base_code.extend(callback_code) + return keyboard, "\n".join(base_code) diff --git a/component/telegram/models.py b/component/telegram/models.py index 7cc753f..78fb6b6 100644 --- a/component/telegram/models.py +++ b/component/telegram/models.py @@ -84,31 +84,6 @@ def _get_method_name(self, class_name: str) -> str: method += c.lower() return method.lstrip("_") - def _generate_keyboard_code(self, keyboard) -> list[str]: - if not isinstance(keyboard, InlineKeyboardMarkup): - return [] - - code = [" builder = InlineKeyboardBuilder()"] - for k in keyboard.inline_keyboard.all(): - code.append( - f" builder.button(text='{k.text}', callback_data='{k.callback_data}')", - ) - code.append(" keyboard = builder.as_markup()") - return code - - def _format_code_component(self, underlying_object) -> list[str]: - try: - import black - - formatted_code = black.format_str(underlying_object.code, mode=black.Mode()) - return [f" {formatted_code}"] - except Exception as e: - return [ - f" # Original code failed black formatting: {str(e)}", - f" # {underlying_object.code}", - " pass", - ] - def _get_component_params( self, underlying_object, @@ -164,48 +139,20 @@ def generate_code(self) -> str: f"async def {underlying_object.code_function_name}(input_data: Message, **kwargs):", ] - # Handle code component - if underlying_object.__class__.__name__ == "CodeComponent": - code.extend(self._format_code_component(underlying_object)) - return "\n".join(code) - keyboard = None + callback_code = "" # Check if markup exists before accessing it if hasattr(underlying_object, "markup") and underlying_object.markup: markup = underlying_object.markup - match markup.markup_type: - case markup.MarkupType.ReplyKeyboard: - keyword_class = "ReplyKeyboardMarkup" - button_class = "KeyboardButton" - case markup.MarkupType.InlineKeyboard: - keyword_class = "InlineKeyboardMarkup" - button_class = "InlineKeyboardButton" - case _: - raise NotImplementedError(f"Unknown markup {markup.markup_type}") - keyboard_buttons = "[" - for row in markup.buttons: - keyboard_buttons += "[\n" - - for cell in row: - args = {"text": cell} - if markup.markup_type == markup.MarkupType.InlineKeyboard: - args["callback_data"] = markup.get_callback_data(cell) - - keyboard_buttons += f"{button_class}(\n" - for k, v in args.items(): - keyboard_buttons += f'{k} = "{v}"\n' - keyboard_buttons += f")" - - keyboard_buttons += "]" - keyboard_buttons += "]" - keyboard = f"{keyword_class}(resize_keyboard=True, one_time_keyboard=False, keyboard = {keyboard_buttons})" - + keyboard, callback_code = markup.generate_code() + code.append(f" keyboard = {keyboard}") # Generate parameters and method call params_str = self._get_component_params( underlying_object, keyboard, file_params, ) + code.append(f" await bot.{method}({params_str})") # Handle next components @@ -219,7 +166,10 @@ def generate_code(self) -> str: f" await {next_component.code_function_name}(input_data, **kwargs)", ) - return "\n".join(code) + new_code = "\n".join(code) + new_code += "\n\n" + callback_code + + return new_code def get_all_next_components(self) -> List["Component"]: ans = {}