fckbot/rag/utils/text_utils.py

# -*- coding: utf-8 -*-
"""
Утилиты для работы с текстом:
- Разбиение на чанки (токенизация) с помощью tiktoken для индексации документов.
- Подсчёт количества токенов в тексте (для обрезки истории диалога).
- Fallback-разбиение по символам, если tiktoken не установлен.

Используется в:
- KBService при индексации документов (split_into_chunks).
- RAG-пайплайне для обрезки истории диалога по токенам (count_tokens).

Особенности:
- Разбиение на предложения использует библиотеку razdel для корректной обработки
  аббревиатур (т.е., и.т.д.) и других сложных случаев русского языка.
- Если razdel недоступен, используется fallback-регулярное выражение.
- Если tiktoken недоступен, используется грубая оценка токенов (3 символа на токен).
"""

import re
import logging
from typing import List

# ---- Инициализация токенизатора tiktoken (для точного подсчёта токенов) ----
try:
    import tiktoken
    # Используем кодировку cl100k_base, которая используется в моделях GPT-4, GPT-3.5 и GigaChat
    TOKENIZER = tiktoken.get_encoding("cl100k_base")
except ImportError:
    TOKENIZER = None
    logging.warning("tiktoken не установлен. Для точного подсчёта токенов установите: pip install tiktoken")
# ----------------------------------------------------------------------------

# ---- Импорт razdel для умного разбиения по предложениям (русский язык) ----
try:
    from razdel import sentenize
    HAS_RAZDEL = True
except ImportError:
    HAS_RAZDEL = False
    logging.warning("razdel не установлен. Разбиение по предложениям будет неточным. Установите: pip install razdel")
# ----------------------------------------------------------------------------

logger = logging.getLogger(__name__)


def count_tokens(text: str) -> int:
    """
    Подсчитывает количество токенов в тексте, используя tiktoken (если доступен).
    Если tiktoken недоступен, использует грубую оценку: 1 токен ≈ 3 символа.

    Эта функция используется в RAG-пайплайне для обрезки истории диалога
    до заданного лимита токенов (max_context_tokens).

    Аргументы:
        text (str): текст для подсчёта токенов.

    Возвращает:
        int: количество токенов (или оценка).
    """
    if not text:
        return 0

    if TOKENIZER is not None:
        # Точный подсчёт через tiktoken
        return len(TOKENIZER.encode(text))
    else:
        # Грубая оценка: для русского языка в среднем 1 токен ≈ 2.5–3 символа.
        # Используем 3 как консервативную оценку (чтобы не превысить лимит).
        return len(text) // 3 + 1


def split_into_chunks(
    text: str,
    chunk_size_tokens: int = 200,
    overlap_tokens: int = 50,
    approx_chunk_chars: int = 500,
    approx_overlap_chars: int = 100,
    strategy: str = "recursive_split_by_sentences"
) -> List[str]:
    """
    Разбивает текст на чанки (фрагменты) для последующей индексации в векторной БД.
    Сначала пытается использовать tiktoken для точного разбиения по токенам.
    Если tiktoken недоступен, использует приближённое разбиение по символам.

    Аргументы:
        text (str): исходный текст.
        chunk_size_tokens (int): целевой размер чанка в токенах (по умолчанию 200).
        overlap_tokens (int): перекрытие между соседними чанками в токенах (по умолчанию 50).
        approx_chunk_chars (int): размер чанка в символах для fallback-режима (по умолчанию 500).
        approx_overlap_chars (int): перекрытие в символах для fallback-режима (по умолчанию 100).
        strategy (str): стратегия разбиения (пока только "recursive_split_by_sentences").

    Возвращает:
        List[str]: список чанков (текстовых фрагментов).
    """
    if not text or not text.strip():
        return []

    # Нормализуем пробелы: заменяем множественные пробелы, табуляции и переносы на один пробел.
    # Это упрощает дальнейшую обработку и улучшает качество разбиения.
    text = re.sub(r'\s+', ' ', text).strip()

    # Если токенизатор недоступен, используем fallback-разбиение по символам.
    if TOKENIZER is None:
        logger.debug("tiktoken недоступен, используем разбиение по символам")
        return _split_by_chars(text, approx_chunk_chars, approx_overlap_chars)

    # Основной метод: разбиение по предложениям с учётом токенов.
    return _split_by_tokens(text, chunk_size_tokens, overlap_tokens)


def _split_by_tokens(text: str, chunk_size: int, overlap: int) -> List[str]:
    """
    Разбивает текст на чанки, используя tiktoken для точного подсчёта токенов.
    Алгоритм:
      1. Разбить текст на предложения (с помощью razdel для русского языка).
      2. Накопить предложения, пока сумма токенов не превысит chunk_size.
      3. Когда превышает – сохранить текущий чанк и создать перекрытие из последних предложений (overlap токенов).
      4. Продолжить с оставшимся хвостом.
      5. В конце объединить очень короткие чанки (< chunk_size/3) с предыдущим,
         чтобы избежать множества мелких фрагментов.

    Аргументы:
        text (str): исходный текст.
        chunk_size (int): максимальный размер чанка в токенах.
        overlap (int): количество токенов перекрытия между соседними чанками.

    Возвращает:
        List[str]: список чанков.
    """
    # Вспомогательная функция для подсчёта токенов в строке.
    def _count_tokens(s: str) -> int:
        return len(TOKENIZER.encode(s))

    # ---- 1. Разбиение на предложения ----
    if HAS_RAZDEL:
        # razdel возвращает итератор объектов Sentence, у которых есть атрибут .text
        sentences = [sent.text for sent in sentenize(text)]
    else:
        # Fallback: разбиение по точкам, вопросительным и восклицательным знакам.
        # Это может быть неточным для аббревиатур (например, "т.е.", "и.т.д.").
        sentences = re.split(r'(?<=[.!?])\s+', text)
    # -------------------------------------

    chunks = []
    current_chunk = []      # Список предложений текущего чанка
    current_len = 0         # Суммарное количество токенов в текущем чанке

    for sent in sentences:
        sent_tokens = _count_tokens(sent)   # Токены в текущем предложении

        # Проверяем, помещается ли предложение в текущий чанк
        if current_len + sent_tokens <= chunk_size:
            # Помещается – добавляем
            current_chunk.append(sent)
            current_len += sent_tokens
        else:
            # Не помещается – сохраняем текущий чанк
            if current_chunk:
                chunks.append(' '.join(current_chunk))

            # --- Формируем перекрытие для следующего чанка ---
            # Берём последние предложения текущего чанка, сумма токенов которых
            # не превышает overlap (чтобы сохранить контекст между чанками).
            overlap_sentences = []
            overlap_len = 0
            # Идём с конца текущего чанка (от последнего предложения к первому)
            for prev_sent in reversed(current_chunk):
                prev_len = _count_tokens(prev_sent)
                if overlap_len + prev_len <= overlap:
                    overlap_sentences.insert(0, prev_sent)   # Вставляем в начало, чтобы сохранить порядок
                    overlap_len += prev_len
                else:
                    break
            # -------------------------------------------------

            # Новый чанк начинается с перекрытия + текущее предложение
            current_chunk = overlap_sentences + [sent]
            current_len = overlap_len + sent_tokens

    # Добавляем последний чанк, если он не пуст
    if current_chunk:
        chunks.append(' '.join(current_chunk))

    # ---- Склеиваем очень короткие чанки с предыдущим ----
    # Это предотвращает создание большого количества мелких фрагментов,
    # которые могут ухудшить качество поиска.
    merged = []
    for chunk in chunks:
        # Если текущий чанк слишком короткий (< 1/3 от chunk_size) и есть предыдущий чанк,
        # присоединяем его к предыдущему.
        if merged and _count_tokens(chunk) < chunk_size // 3:
            merged[-1] = merged[-1] + " " + chunk
        else:
            merged.append(chunk)
    # -------------------------------------------------

    return merged


def _split_by_chars(text: str, chunk_size: int, overlap: int) -> List[str]:
    """
    Fallback-метод: разбиение по символам (при отсутствии tiktoken).
    Использует простое скользящее окно с фиксированной длиной chunk_size и перекрытием overlap.

    Аргументы:
        text (str): исходный текст.
        chunk_size (int): размер чанка в символах.
        overlap (int): перекрытие в символах.

    Возвращает:
        List[str]: список чанков.
    """
    if chunk_size <= 0:
        return [text]

    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        start = end - overlap
        if start < 0:
            start = 0
        if start >= len(text):
            break
    return chunks