Add new file

core/utils/text_utils.py

@@ -0,0 +1,226 @@
+# -*- 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