diff --git a/README.md b/README.md index d7bae5e..817962b 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ docker compose up -d --build [ { "text": "text of the replica", + "gender": "male", "speaker_id": 0 }, { @@ -54,7 +55,8 @@ docker compose up -d --build ![Состояния бота](./bot-states.png) -- **INTRO** - начальное состояние сессии сразу после команды `/start`. Выводится сообщение о том, что это за бот и небольшое пользовательское соглашение, уведомляющее о целях сбора данных. Единственная кнопка - `Принять и продолжить`. После нажатия на неё условный переход либо в **NO_MORE_SCENARIOS**, если нету доступных сценариев для озвучивания, либо в **FIRST_REPLICA** . +- **INTRO** - начальное состояние сессии сразу после команды `/start`. Выводится сообщение о том, что это за бот и небольшое пользовательское соглашение, уведомляющее о целях сбора данных. Единственная кнопка - `Принять и продолжить`. После нажатия на неё происходит переход в **SPECIFY_GENDER**. +- **SPECIFY_GENDER** - выводится сообщение с предложением указать пол пользователя. Две кнопки для выбора мужского и женского пола. После нажатия на любую из кнопок происходит условный переход либо в **NO_MORE_SCENARIOS**, если нету доступных сценариев для озвучивания, либо в **FIRST_REPLICA**. - **NO_MORE_SCENARIOS** - выводится сообщение о том, что пока больше нет сценариев для озвучивания. Из **NO_MORE_SCENARIOS** может произойти переход в **FIRST_REPLICA**, когда на сервер загружается новый сценарий, при этом выводится дополнительное уведомление. Самостоятельно пользователь не может покинуть это состояние. - **FIRST_REPLICA** - выводится первая реплика дорожки с дополнительными инструкциями для пользователя (`i = 0`). При отправке аудиосообщения с озвучкой реплики, происходит условный переход в **SHOW_REPLICA** (`i += 1`), если это не последняя реплика в дорожке, либо в **CONFIRM_SAVE**. - **SHOW_REPLICA** - выводится i-ая (0-indexed) реплика дорожки и её номер (1-indexed для пользователя, то есть i + 1). Есть две кнопки. Кнопка "Перезаписать предыдущую реплику", по ней условный переход в **FIRST_REPLICA**, если `i == 1`, либо в **SHOW_REPLICA** (`i -= 1`). Кнопка "Начать заново", по ней переход в **CONFIRM_RESTART**. При отправке аудиосообщения с озвучкой реплики, происходит условный переход в **SHOW_REPLICA** (`i += 1`), если это не последняя реплика в дорожке, либо в **CONFIRM_SAVE**. diff --git a/bot-states.png b/bot-states.png index 21f6239..5d31fe7 100644 Binary files a/bot-states.png and b/bot-states.png differ diff --git a/main.py b/main.py index 334a5ff..f7360aa 100644 --- a/main.py +++ b/main.py @@ -28,6 +28,7 @@ from src.handlers import ( handle_rerecord_previous, handle_restart_track, handle_save_track, + handle_select_gender, handle_select_scenario_delete, handle_unexpected_text, handle_voice_message, @@ -50,6 +51,9 @@ def main() -> None: # Callback query handlers app.add_handler(CallbackQueryHandler(handle_accept_intro, pattern="^accept_intro$")) + app.add_handler( + CallbackQueryHandler(handle_select_gender, pattern="^select_gender:") + ) app.add_handler( CallbackQueryHandler(handle_rerecord_previous, pattern="^rerecord_previous$") ) diff --git a/src/database.py b/src/database.py index 51750d8..48156fa 100644 --- a/src/database.py +++ b/src/database.py @@ -13,6 +13,7 @@ class UserState(Enum): """Состояния пользовательской сессии.""" INTRO = "intro" + SPECIFY_GENDER = "specify_gender" NO_MORE_SCENARIOS = "no_more_scenarios" FIRST_REPLICA = "first_replica" SHOW_REPLICA = "show_replica" @@ -32,6 +33,7 @@ class User: id: int # dataset_speaker_id telegram_id: int created_at: datetime + gender: str | None # "male" или "female" @dataclass @@ -51,6 +53,7 @@ class Replica: speaker_id: int # в рамках сценария replica_index: int # порядок в сценарии (0-indexed) text: str + gender: str # "male" или "female" @dataclass @@ -101,7 +104,8 @@ def init_db() -> None: CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, telegram_id INTEGER UNIQUE NOT NULL, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + gender TEXT ); CREATE TABLE IF NOT EXISTS scenarios ( @@ -115,6 +119,7 @@ def init_db() -> None: speaker_id INTEGER NOT NULL, replica_index INTEGER NOT NULL, text TEXT NOT NULL, + gender TEXT NOT NULL, FOREIGN KEY (scenario_id) REFERENCES scenarios(id), UNIQUE(scenario_id, replica_index) ); @@ -152,14 +157,6 @@ def init_db() -> None: CREATE INDEX IF NOT EXISTS idx_recordings_scenario ON recordings(scenario_id); """) - - # Миграция: добавляем колонку duration если её нет - cursor = conn.execute("PRAGMA table_info(recordings)") - columns = [row[1] for row in cursor.fetchall()] - if "duration" not in columns: - conn.execute("ALTER TABLE recordings ADD COLUMN duration REAL DEFAULT 0.0") - logger.info("Добавлена колонка duration в таблицу recordings") - conn.commit() logger.info("База данных инициализирована") @@ -172,7 +169,8 @@ def get_or_create_user(telegram_id: int) -> User: """Получает или создаёт пользователя по telegram_id.""" with get_connection() as conn: cursor = conn.execute( - "SELECT id, telegram_id, created_at FROM users WHERE telegram_id = ?", + "SELECT id, telegram_id, created_at, gender FROM users " + "WHERE telegram_id = ?", (telegram_id,), ) row = cursor.fetchone() @@ -182,18 +180,22 @@ def get_or_create_user(telegram_id: int) -> User: id=row["id"], telegram_id=row["telegram_id"], created_at=row["created_at"], + gender=row["gender"], ) cursor = conn.execute( "INSERT INTO users (telegram_id) VALUES (?) " - "RETURNING id, telegram_id, created_at", + "RETURNING id, telegram_id, created_at, gender", (telegram_id,), ) row = cursor.fetchone() conn.commit() logger.info(f"Создан новый пользователь: dataset_speaker_id={row['id']}") return User( - id=row["id"], telegram_id=row["telegram_id"], created_at=row["created_at"] + id=row["id"], + telegram_id=row["telegram_id"], + created_at=row["created_at"], + gender=row["gender"], ) @@ -201,7 +203,8 @@ def get_user_by_telegram_id(telegram_id: int) -> User | None: """Получает пользователя по telegram_id.""" with get_connection() as conn: cursor = conn.execute( - "SELECT id, telegram_id, created_at FROM users WHERE telegram_id = ?", + "SELECT id, telegram_id, created_at, gender FROM users " + "WHERE telegram_id = ?", (telegram_id,), ) row = cursor.fetchone() @@ -210,10 +213,22 @@ def get_user_by_telegram_id(telegram_id: int) -> User | None: id=row["id"], telegram_id=row["telegram_id"], created_at=row["created_at"], + gender=row["gender"], ) return None +def update_user_gender(user_id: int, gender: str) -> None: + """Обновляет пол пользователя.""" + with get_connection() as conn: + conn.execute( + "UPDATE users SET gender = ? WHERE id = ?", + (gender, user_id), + ) + conn.commit() + logger.info(f"Обновлён пол пользователя {user_id}: {gender}") + + # === Scenarios CRUD === @@ -258,15 +273,18 @@ def get_all_scenarios() -> list[Scenario]: # === Replicas CRUD === -def create_replicas(scenario_id: str, replicas: list[tuple[int, int, str]]) -> None: - """Создаёт реплики. replicas: [(speaker_id, replica_index, text), ...]""" +def create_replicas( + scenario_id: str, replicas: list[tuple[int, int, str, str]] +) -> None: + """Создаёт реплики. replicas: [(speaker_id, replica_index, text, gender), ...]""" with get_connection() as conn: conn.executemany( - "INSERT INTO replicas (scenario_id, speaker_id, replica_index, text) " - "VALUES (?, ?, ?, ?)", + "INSERT INTO replicas " + "(scenario_id, speaker_id, replica_index, text, gender) " + "VALUES (?, ?, ?, ?, ?)", [ - (scenario_id, speaker_id, idx, text) - for speaker_id, idx, text in replicas + (scenario_id, speaker_id, idx, text, gender) + for speaker_id, idx, text, gender in replicas ], ) conn.commit() @@ -276,8 +294,8 @@ def get_replicas_for_scenario(scenario_id: str) -> list[Replica]: """Получает все реплики сценария.""" with get_connection() as conn: cursor = conn.execute( - "SELECT id, scenario_id, speaker_id, replica_index, text FROM replicas " - "WHERE scenario_id = ? ORDER BY replica_index", + "SELECT id, scenario_id, speaker_id, replica_index, text, gender " + "FROM replicas WHERE scenario_id = ? ORDER BY replica_index", (scenario_id,), ) return [ @@ -287,6 +305,7 @@ def get_replicas_for_scenario(scenario_id: str) -> list[Replica]: speaker_id=row["speaker_id"], replica_index=row["replica_index"], text=row["text"], + gender=row["gender"], ) for row in cursor.fetchall() ] @@ -296,8 +315,9 @@ def get_replicas_for_track(scenario_id: str, speaker_id: int) -> list[Replica]: """Получает реплики для конкретной дорожки (speaker_id в сценарии).""" with get_connection() as conn: cursor = conn.execute( - "SELECT id, scenario_id, speaker_id, replica_index, text FROM replicas " - "WHERE scenario_id = ? AND speaker_id = ? ORDER BY replica_index", + "SELECT id, scenario_id, speaker_id, replica_index, text, gender " + "FROM replicas WHERE scenario_id = ? AND speaker_id = ? " + "ORDER BY replica_index", (scenario_id, speaker_id), ) return [ @@ -307,6 +327,7 @@ def get_replicas_for_track(scenario_id: str, speaker_id: int) -> list[Replica]: speaker_id=row["speaker_id"], replica_index=row["replica_index"], text=row["text"], + gender=row["gender"], ) for row in cursor.fetchall() ] @@ -323,6 +344,16 @@ def get_track_speaker_ids(scenario_id: str) -> list[int]: return [row["speaker_id"] for row in cursor.fetchall()] +def get_scenario_genders(scenario_id: str) -> set[str]: + """Получает список полов, представленных в сценарии.""" + with get_connection() as conn: + cursor = conn.execute( + "SELECT DISTINCT gender FROM replicas WHERE scenario_id = ?", + (scenario_id,), + ) + return {row["gender"] for row in cursor.fetchall()} + + # === Recordings CRUD === diff --git a/src/handlers.py b/src/handlers.py index db97f8a..5128acf 100644 --- a/src/handlers.py +++ b/src/handlers.py @@ -13,6 +13,7 @@ from src.database import ( get_or_create_user, get_replicas_for_track, get_scenario, + get_scenario_genders, get_scenario_stats, get_stats, get_user_audio_duration, @@ -20,6 +21,7 @@ from src.database import ( get_user_stats, get_users_in_state, get_users_with_scenario, + update_user_gender, upsert_user_session, ) from src.decorators import answer_callback, require_state, with_user_and_session @@ -40,6 +42,10 @@ INTRO_TEXT = """👋 Добро пожаловать! Нажмите кнопку ниже, чтобы начать.""" +SPECIFY_GENDER_TEXT = """👤 Укажите ваш пол. + +Это необходимо для подбора подходящих сценариев для озвучивания.""" + NO_MORE_SCENARIOS_TEXT = """📭 Пока нет доступных сценариев для озвучивания. Вы получите уведомление, когда появятся новые сценарии.""" @@ -82,6 +88,15 @@ def get_intro_keyboard() -> InlineKeyboardMarkup: ) +def get_specify_gender_keyboard() -> InlineKeyboardMarkup: + return InlineKeyboardMarkup( + [ + [InlineKeyboardButton("👨 Мужской", callback_data="select_gender:male")], + [InlineKeyboardButton("👩 Женский", callback_data="select_gender:female")], + ] + ) + + def get_show_replica_keyboard() -> InlineKeyboardMarkup: return InlineKeyboardMarkup( [ @@ -343,6 +358,27 @@ async def handle_accept_intro( ) -> None: """Обработчик кнопки 'Принять и продолжить'.""" query = update.callback_query + session.state = UserState.SPECIFY_GENDER + await query.edit_message_text( + SPECIFY_GENDER_TEXT, reply_markup=get_specify_gender_keyboard() + ) + session.last_bot_message_id = query.message.message_id + upsert_user_session(session) + + +@answer_callback +@with_user_and_session +@require_state(UserState.SPECIFY_GENDER) +async def handle_select_gender( + update: Update, context: ContextTypes.DEFAULT_TYPE, user: User, session: UserSession +) -> None: + """Обработчик выбора пола.""" + query = update.callback_query + gender = query.data.replace("select_gender:", "") + + update_user_gender(user.id, gender) + user.gender = gender + track = find_available_track(user.id) if not track: @@ -355,7 +391,10 @@ async def handle_accept_intro( session.speaker_id = speaker_id session.replica_index = 0 await query.edit_message_text(format_first_replica(session)) - logger.info(f"User {user.id} started track {scenario_id}/{speaker_id}") + logger.info( + f"User {user.id} selected gender {gender}, " + f"started track {scenario_id}/{speaker_id}" + ) session.last_bot_message_id = query.message.message_id upsert_user_session(session) @@ -559,6 +598,10 @@ async def handle_exit_admin( # Показываем соответствующее сообщение if session.state == UserState.INTRO: await query.edit_message_text(INTRO_TEXT, reply_markup=get_intro_keyboard()) + elif session.state == UserState.SPECIFY_GENDER: + await query.edit_message_text( + SPECIFY_GENDER_TEXT, reply_markup=get_specify_gender_keyboard() + ) elif session.state == UserState.NO_MORE_SCENARIOS: await query.edit_message_text(NO_MORE_SCENARIOS_TEXT) elif session.state == UserState.FIRST_REPLICA: @@ -605,30 +648,37 @@ async def handle_confirm_upload( del context.user_data["pending_scenario"] + # Получаем полы, представленные в новом сценарии + scenario_genders = get_scenario_genders(pending["id"]) + # Уведомляем и переводим пользователей в NO_MORE_SCENARIOS на новый сценарий for waiting_user_id in get_users_in_state(UserState.NO_MORE_SCENARIOS): try: with get_connection() as conn: row = conn.execute( - "SELECT telegram_id FROM users WHERE id = ?", (waiting_user_id,) + "SELECT telegram_id, gender FROM users WHERE id = ?", + (waiting_user_id,), ).fetchone() if row: - waiting_session = get_user_session(waiting_user_id) - if waiting_session: - track = find_available_track(waiting_user_id) - if track: - scenario_id, speaker_id = track - waiting_session.state = UserState.FIRST_REPLICA - waiting_session.scenario_id = scenario_id - waiting_session.speaker_id = speaker_id - waiting_session.replica_index = 0 - upsert_user_session(waiting_session) + user_gender = row[1] + # Уведомляем только если пол пользователя представлен в сценарии + if user_gender and user_gender in scenario_genders: + waiting_session = get_user_session(waiting_user_id) + if waiting_session: + track = find_available_track(waiting_user_id) + if track: + scenario_id, speaker_id = track + waiting_session.state = UserState.FIRST_REPLICA + waiting_session.scenario_id = scenario_id + waiting_session.speaker_id = speaker_id + waiting_session.replica_index = 0 + upsert_user_session(waiting_session) - msg = ( - f"🎉 Появился новый сценарий!\n\n" - f"{format_first_replica(waiting_session)}" - ) - await context.bot.send_message(row[0], msg) + msg = ( + f"🎉 Появился новый сценарий!\n\n" + f"{format_first_replica(waiting_session)}" + ) + await context.bot.send_message(row[0], msg) except Exception: pass diff --git a/src/scenarios.py b/src/scenarios.py index 683d4b3..6f603a5 100644 --- a/src/scenarios.py +++ b/src/scenarios.py @@ -18,11 +18,11 @@ def load_scenario_from_json(scenario_id: str, json_data: list[dict]) -> int: if get_scenario(scenario_id): raise ValueError(f"Сценарий {scenario_id} уже существует") - replicas_data: list[tuple[int, int, str]] = [] + replicas_data: list[tuple[int, int, str, str]] = [] for idx, item in enumerate(json_data): - if "text" not in item or "speaker_id" not in item: + if "text" not in item or "speaker_id" not in item or "gender" not in item: raise ValueError(f"Некорректный формат реплики #{idx}") - replicas_data.append((item["speaker_id"], idx, item["text"])) + replicas_data.append((item["speaker_id"], idx, item["text"], item["gender"])) create_scenario(scenario_id) create_replicas(scenario_id, replicas_data) @@ -53,6 +53,13 @@ def parse_scenario_file(file_content: bytes) -> tuple[list[dict], str | None]: return [], f"Реплика #{idx}: отсутствует поле 'speaker_id'" if not isinstance(item["speaker_id"], int): return [], f"Реплика #{idx}: 'speaker_id' должен быть числом" + if "gender" not in item: + return [], f"Реплика #{idx}: отсутствует поле 'gender'" + if item["gender"] not in ["male", "female"]: + return ( + [], + f"Реплика #{idx}: 'gender' должен быть 'male' или 'female'", + ) return data, None @@ -69,16 +76,25 @@ def get_scenario_info(json_data: list[dict]) -> dict: def find_available_track(user_id: int) -> tuple[str, int] | None: """ - Находит доступную дорожку для пользователя. + Находит доступную дорожку для пользователя с учётом пола. Приоритет: 1. Дорожки, которые никто не начал озвучивать 2. Дорожки, которые кто-то начал, но не закончил 3. Дорожки с готовой озвучкой (для дополнительных записей) Пользователь не может озвучивать две разные дорожки в одном сценарии. + Учитывается пол пользователя и пол дорожек. Возвращает (scenario_id, speaker_id) или None. """ with get_connection() as conn: + # Получаем пол пользователя + user_gender_row = conn.execute( + "SELECT gender FROM users WHERE id = ?", (user_id,) + ).fetchone() + if not user_gender_row or not user_gender_row[0]: + return None + user_gender = user_gender_row[0] + # Сценарии, в которых пользователь уже записывает дорожку user_scenarios = conn.execute( """ @@ -88,13 +104,16 @@ def find_available_track(user_id: int) -> tuple[str, int] | None: ).fetchall() user_scenario_ids = {row[0] for row in user_scenarios} - # Все дорожки (scenario_id, speaker_id) с количеством реплик + # Все дорожки (scenario_id, speaker_id) с количеством реплик и полом + # Учитываем только дорожки, где пол совпадает с полом пользователя all_tracks = conn.execute( """ - SELECT scenario_id, speaker_id, COUNT(*) as replica_count + SELECT scenario_id, speaker_id, COUNT(*) as replica_count, gender FROM replicas + WHERE gender = ? GROUP BY scenario_id, speaker_id - """ + """, + (user_gender,), ).fetchall() # Статистика записей по дорожкам