diff --git a/AI/talk_handler.py b/AI/talk_handler.py index 6fdbffd..9073e34 100644 --- a/AI/talk_handler.py +++ b/AI/talk_handler.py @@ -162,6 +162,37 @@ def _set_state_from_conn(conn: sqlite3.Connection, chat_id: int, key: str, value ) +def _get_talk_state_from_conn( + conn: sqlite3.Connection, + chat_id: int, + user_id: int, + key: str, + default: str = "", +) -> str: + row = conn.execute( + "SELECT value FROM talk_state WHERE chat_id = ? AND user_id = ? AND key = ?", + (chat_id, user_id, key), + ).fetchone() + return row["value"] if row else default + + +def _set_talk_state_from_conn( + conn: sqlite3.Connection, + chat_id: int, + user_id: int, + key: str, + value: str, +) -> None: + conn.execute( + """ + INSERT INTO talk_state (chat_id, user_id, key, value) + VALUES (?, ?, ?, ?) + ON CONFLICT(chat_id, user_id, key) DO UPDATE SET value = excluded.value + """, + (chat_id, user_id, key, value), + ) + + def _init_db() -> None: Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True) with _connect_db() as conn: @@ -190,12 +221,42 @@ def _init_db() -> None: ) """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS talk_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + chat_id INTEGER NOT NULL, + user_id INTEGER NOT NULL, + role TEXT NOT NULL, + name TEXT NOT NULL, + text TEXT NOT NULL, + created_at REAL NOT NULL DEFAULT 0 + ) + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS talk_state ( + chat_id INTEGER NOT NULL, + user_id INTEGER NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + PRIMARY KEY (chat_id, user_id, key) + ) + """ + ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_history_chat_id_id ON chat_history(chat_id, id)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_history_chat_id_role_id ON chat_history(chat_id, role, id)" ) + conn.execute( + """ + CREATE INDEX IF NOT EXISTS idx_talk_history_chat_user_id + ON talk_history(chat_id, user_id, id) + """ + ) conn.commit() @@ -235,36 +296,64 @@ def _photo_memory_text(caption: str) -> str: return "[Фото] Без подписи." -def push_message(chat_id: int, role: str, name: str, text: str) -> None: +def push_message( + chat_id: int, + role: str, + name: str, + text: str, + *, + user_id: int | None = None, +) -> None: cleaned_text = _clean_text(text) if not cleaned_text: return with _connect_db() as conn: - conn.execute( - """ - INSERT INTO chat_history (chat_id, role, name, text, created_at) - VALUES (?, ?, ?, ?, ?) - """, - (chat_id, role, name, cleaned_text, time.time()), - ) + if user_id is None: + conn.execute( + """ + INSERT INTO chat_history (chat_id, role, name, text, created_at) + VALUES (?, ?, ?, ?, ?) + """, + (chat_id, role, name, cleaned_text, time.time()), + ) + else: + conn.execute( + """ + INSERT INTO talk_history (chat_id, user_id, role, name, text, created_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + (chat_id, user_id, role, name, cleaned_text, time.time()), + ) conn.commit() -def _get_summary(chat_id: int) -> str: +def _get_summary(chat_id: int, user_id: int | None = None) -> str: with _connect_db() as conn: - return _get_state_from_conn(conn, chat_id, "summary", "") + if user_id is None: + return _get_state_from_conn(conn, chat_id, "summary", "") + return _get_talk_state_from_conn(conn, chat_id, user_id, "summary", "") -def _get_history_rows(chat_id: int) -> list[sqlite3.Row]: +def _get_history_rows(chat_id: int, user_id: int | None = None) -> list[sqlite3.Row]: with _connect_db() as conn: + if user_id is None: + return conn.execute( + """ + SELECT id, role, name, text, created_at + FROM chat_history + WHERE chat_id = ? + ORDER BY id ASC + """, + (chat_id,), + ).fetchall() return conn.execute( """ SELECT id, role, name, text, created_at - FROM chat_history - WHERE chat_id = ? + FROM talk_history + WHERE chat_id = ? AND user_id = ? ORDER BY id ASC """, - (chat_id,), + (chat_id, user_id), ).fetchall() @@ -293,13 +382,25 @@ def _latest_reply_stats(chat_id: int) -> tuple[float, int]: return float(last_assistant["created_at"] or 0.0), int(user_messages_since_reply) -def _store_summary_and_prune(chat_id: int, summary: str, last_row_id: int) -> None: +def _store_summary_and_prune( + chat_id: int, + summary: str, + last_row_id: int, + user_id: int | None = None, +) -> None: with _connect_db() as conn: - _set_state_from_conn(conn, chat_id, "summary", summary) - conn.execute( - "DELETE FROM chat_history WHERE chat_id = ? AND id <= ?", - (chat_id, last_row_id), - ) + if user_id is None: + _set_state_from_conn(conn, chat_id, "summary", summary) + conn.execute( + "DELETE FROM chat_history WHERE chat_id = ? AND id <= ?", + (chat_id, last_row_id), + ) + else: + _set_talk_state_from_conn(conn, chat_id, user_id, "summary", summary) + conn.execute( + "DELETE FROM talk_history WHERE chat_id = ? AND user_id = ? AND id <= ?", + (chat_id, user_id, last_row_id), + ) conn.commit() @@ -317,6 +418,7 @@ def _build_messages( system_prompt: str, current_text: str | None = None, current_name: str | None = None, + user_id: int | None = None, ) -> list[dict[str, Any]]: current_payload = None current_budget = 0 @@ -327,14 +429,14 @@ def _build_messages( } current_budget = len(current_payload["content"]) - summary = _get_summary(chat_id).strip() + summary = _get_summary(chat_id, user_id=user_id).strip() summary_block = "" if summary: summary_block = f"Краткая память чата:\n{summary[:SUMMARY_CHAR_BUDGET]}" used_chars = len(system_prompt) + len(summary_block) + current_budget recent_messages: list[dict[str, str]] = [] - for row in reversed(_get_history_rows(chat_id)[-RECENT_MESSAGES_LIMIT:]): + for row in reversed(_get_history_rows(chat_id, user_id=user_id)[-RECENT_MESSAGES_LIMIT:]): llm_message = _format_row_for_llm(row) if used_chars + len(llm_message["content"]) > PROMPT_CHAR_BUDGET: break @@ -424,11 +526,11 @@ async def _call_llm( return cleaned_content -async def _maybe_refresh_summary(chat_id: int) -> None: +async def _maybe_refresh_summary(chat_id: int, user_id: int | None = None) -> None: async with _summary_lock: # Сворачиваем старую часть истории в summary, чтобы не пихать весь лог в модель. for _ in range(6): - rows = await asyncio.to_thread(_get_history_rows, chat_id) + rows = await asyncio.to_thread(_get_history_rows, chat_id, user_id) if len(rows) <= SUMMARY_TRIGGER_MESSAGES: return @@ -438,7 +540,7 @@ async def _maybe_refresh_summary(chat_id: int) -> None: batch_size = min(available_to_summarize, SUMMARY_BATCH_MESSAGES) rows_to_summarize = rows[:batch_size] - current_summary = await asyncio.to_thread(_get_summary, chat_id) + current_summary = await asyncio.to_thread(_get_summary, chat_id, user_id) transcript = "\n".join( f"{row['name'] if row['role'] == 'user' else BOT_MEMORY_NAME}: " f"{_clip_text(row['text'], SUMMARY_LINE_CHAR_LIMIT)}" @@ -470,10 +572,10 @@ async def _maybe_refresh_summary(chat_id: int) -> None: cleaned_summary = summary.strip() if not cleaned_summary: - logger.warning("Summary refresh produced empty text for chat_id=%s", chat_id) + logger.warning("Summary refresh produced empty text for chat_id=%s user_id=%s", chat_id, user_id) return last_row_id = int(rows_to_summarize[-1]["id"]) - await asyncio.to_thread(_store_summary_and_prune, chat_id, cleaned_summary, last_row_id) + await asyncio.to_thread(_store_summary_and_prune, chat_id, cleaned_summary, last_row_id, user_id) async def _generate_response( @@ -483,11 +585,12 @@ async def _generate_response( current_text: str | None = None, current_name: str | None = None, current_content: Any = None, + user_id: int | None = None, max_tokens: int, temperature: float, top_p: float, ) -> str: - await _maybe_refresh_summary(chat_id) + await _maybe_refresh_summary(chat_id, user_id=user_id) if current_content is None: messages = await asyncio.to_thread( _build_messages, @@ -495,9 +598,10 @@ async def _generate_response( system_prompt, current_text, current_name, + user_id, ) else: - messages = await asyncio.to_thread(_build_messages, chat_id, system_prompt, None, None) + messages = await asyncio.to_thread(_build_messages, chat_id, system_prompt, None, None, user_id) current_payload = {"role": "user", "content": current_content} current_budget = _estimate_content_length(current_content) used_chars = sum(_estimate_content_length(message.get("content", "")) for message in messages) @@ -771,13 +875,15 @@ async def handle_talk(message: Message) -> None: return chat_id = message.chat.id + user_id = message.from_user.id user_name = message.from_user.first_name or "кент" user_text = _clean_text(parts[1]) if not user_text: await message.reply(USER_FALLBACK_TEXT) return - await asyncio.to_thread(push_message, chat_id, "user", user_name, user_text) + # /talk хранит отдельную память по конкретному пользователю, чтобы не мешать чужие истории. + await asyncio.to_thread(push_message, chat_id, "user", user_name, user_text, user_id=user_id) async with _reply_lock: try: @@ -785,6 +891,7 @@ async def handle_talk(message: Message) -> None: response = await _generate_response( chat_id, system_prompt=SYSTEM_PROMPT, + user_id=user_id, max_tokens=220, temperature=0.8, top_p=0.9, @@ -797,13 +904,21 @@ async def handle_talk(message: Message) -> None: normalized_response = _normalize_reply(response) if not normalized_response: logger.warning( - "Talk command produced empty/skip response. chat_id=%s user=%s raw=%s", + "Talk command produced empty/skip response. chat_id=%s user_id=%s user=%s raw=%s", chat_id, + user_id, user_name, _clip_text(response, 200), ) normalized_response = EMPTY_RESPONSE_TEXT - await asyncio.to_thread(push_message, chat_id, "assistant", BOT_MEMORY_NAME, normalized_response) + await asyncio.to_thread( + push_message, + chat_id, + "assistant", + BOT_MEMORY_NAME, + normalized_response, + user_id=user_id, + ) await message.reply(normalized_response, parse_mode=None)