diploma/dataset-tg-bot
Code Actions Releases

Пол пользователя

Browse Source
This commit is contained in:
Артём Тищенко
2026-02-04 13:00:20 +03:00
parent 05c4773a50
commit d51bbe04ba
6 changed files with 154 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@@ -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**.

BIN
bot-states.png
View File

Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 209 KiB

After

Width:  |  Height:  |  Size: 219 KiB

4
main.py
View File

@@ -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$")
)

77
src/database.py
View File

@@ -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 ===

84
src/handlers.py
View File

@@ -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

33
src/scenarios.py
View File

@@ -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()
# Статистика записей по дорожкам