diff --git a/rag/rag_orchestrator.py b/rag/rag_orchestrator.py index cdf8fc0..9e49576 100644 --- a/rag/rag_orchestrator.py +++ b/rag/rag_orchestrator.py @@ -10,6 +10,10 @@ Улучшено управление токенами: теперь история и контекст обрезаются с учётом лимитов модели, резервирования для ответа и промптов. + +ДОБАВЛЕНО: +- Иерархическое резюмирование истории диалога. +- Иерархическое резюмирование больших документов при индексации. """ import asyncio @@ -89,9 +93,9 @@ class RAGOrchestrator: system_prompt: Optional[str], query: str, context: str, - max_total_tokens: int = 8192, # для GigaChat-Max + max_total_tokens: int = 8192, reserved_for_answer: int = 1000, - reserved_for_overhead: int = 200 # буфер для форматирования + reserved_for_overhead: int = 200 ) -> Dict[str, Any]: """ Подготавливает части промта и обрезает историю/контекст по токенам. @@ -109,7 +113,6 @@ class RAGOrchestrator: query_tokens = count_tokens(query) # 2. Резервирование - # Вычитаем системный промпт, шаблон, запрос, буфер и место для ответа prompt_tokens = system_tokens + synthesis_tokens + query_tokens available_for_history_and_context = ( max_total_tokens @@ -123,12 +126,8 @@ class RAGOrchestrator: f"Недостаточно токенов для истории и контекста: {available_for_history_and_context}. " f"Увеличьте max_total_tokens или уменьшите размер промптов." ) - # Если отрицательное – устанавливаем минимум 100 токенов для истории и контекста available_for_history_and_context = max(available_for_history_and_context, 100) - # 3. Обрезка истории (уже есть в process_query, но мы её переместим сюда) - # Для простоты пока оставим обрезку в process_query, но используем этот метод для расчёта лимита - return { "available_for_history_and_context": available_for_history_and_context, "prompt_tokens": prompt_tokens, @@ -137,6 +136,67 @@ class RAGOrchestrator: "query_tokens": query_tokens, } + # ------------------------------------------------------------------ + # Метод для сжатия больших документов при индексации + # ------------------------------------------------------------------ + + async def _compress_document_if_needed(self, text: str, file_name: str) -> str: + """ + Если текст превышает порог (max_tokens_for_document), сжимает его + с помощью иерархического резюмирования. + Возвращает либо исходный текст, либо сжатый. + """ + # Читаем настройки из конфига + summarization_config = getattr(self.config, 'summarization', {}) + enabled = summarization_config.get('enable_hierarchical_summarization', True) + if not enabled: + return text + + max_tokens = summarization_config.get('max_tokens_for_document', 8000) + target_tokens = summarization_config.get('target_tokens_after_summary', 3000) + chunk_size = summarization_config.get('chunk_size_tokens', 500) + max_depth = summarization_config.get('max_depth', 2) + + # Подсчитываем токены + text_tokens = count_tokens(text) + if text_tokens <= max_tokens: + logger.debug(f"Документ '{file_name}' ({text_tokens} токенов) не требует сжатия") + return text + + logger.info(f"Документ '{file_name}' ({text_tokens} токенов) превышает лимит {max_tokens}, применяем иерархическое резюмирование") + + # Загружаем промпт для резюмирования + summary_prompt = self.default_prompts.get('hierarchical_summary', '') + if not summary_prompt: + try: + prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt' + with open(prompt_path, 'r', encoding='utf-8') as f: + summary_prompt = f.read() + except Exception as e: + logger.warning(f"Не удалось загрузить промпт hierarchical_summary: {e}, используем стандартный") + summary_prompt = ( + "Ты — профессиональный реферант. Кратко изложи суть текста, сохранив ключевые факты, цифры и логические связи.\n" + "Объём: не более 30% от исходного, но не менее 2 предложений.\n" + "Стиль: деловой, нейтральный, без оценок.\n\n" + "ТЕКСТ:\n{text}\n\nРЕЗЮМЕ:" + ) + + try: + compressed = await hierarchical_summarize( + text=text, + giga=self.giga, + prompt_template=summary_prompt, + target_tokens=target_tokens, + chunk_size_tokens=chunk_size, + max_depth=max_depth, + temperature=getattr(self.config, 'ai_temperature', 0.1) + ) + logger.info(f"Документ '{file_name}' сжат с {text_tokens} до {count_tokens(compressed)} токенов") + return compressed + except Exception as e: + logger.error(f"Ошибка при сжатии документа '{file_name}': {e}, используем исходный текст") + return text + # ------------------------------------------------------------------ # Основной метод обработки запроса # ------------------------------------------------------------------ @@ -186,18 +246,16 @@ class RAGOrchestrator: # ----- 2. Получение истории диалога из БД (необрезанной) ----- history = await self.db.get_history(user_jid, room_jid, limit=100) - # ----- 3. Расчёт лимитов токенов (НОВОЕ) ----- - # Параметры из конфига - max_model_tokens = getattr(self.config, 'max_model_tokens', 8192) # GigaChat-Max default + # ----- 3. Расчёт лимитов токенов ----- + max_model_tokens = getattr(self.config, 'max_model_tokens', 8192) reserved_for_answer = getattr(self.config, 'reserved_for_answer_tokens', 1000) reserved_for_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200) - # Используем вспомогательный метод для расчёта token_info = self._prepare_prompt_parts( synthesis_template=synthesis_template, system_prompt=system_prompt, query=query, - context="", # пока пусто, позже добавим + context="", max_total_tokens=max_model_tokens, reserved_for_answer=reserved_for_answer, reserved_for_overhead=reserved_for_overhead @@ -205,25 +263,30 @@ class RAGOrchestrator: available_for_history_and_context = token_info["available_for_history_and_context"] logger.debug( - f"Доступно для истории и контекста: {available_for_history_and_context} токенов, " - f"промпт: {token_info['prompt_tokens']} токенов" + f"Доступно для истории и контекста: {available_for_history_and_context} токенов" ) - # ----- 4. Обрезка истории (с использованием иерархического резюмирования) ----- + # ----- 4. Сжатие истории (иерархическое резюмирование) ----- max_history_tokens = min(available_for_history_and_context // 2, 2000) # Если история слишком длинная, сжимаем её иерархически history_text = "\n".join([f"{rec['role']}: {rec['content']}" for rec in history]) if count_tokens(history_text) > max_history_tokens: - logger.info(f"История слишком длинная ({count_tokens(history_text)} токенов), применяем иерархическое резюмирование") + logger.info( + f"История слишком длинная ({count_tokens(history_text)} токенов), " + f"применяем иерархическое резюмирование" + ) summary_prompt = prompts.get('hierarchical_summary', '') if not summary_prompt: - # Загружаем из файла, если не передано try: - with open(self.config.prompts_dir / 'hierarchical_summary.txt', 'r', encoding='utf-8') as f: + prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt' + with open(prompt_path, 'r', encoding='utf-8') as f: summary_prompt = f.read() except Exception: - summary_prompt = "Кратко изложи суть диалога, сохранив ключевые факты и вопросы.\n\nТЕКСТ:\n{text}" + summary_prompt = ( + "Кратко изложи суть диалога, сохранив ключевые факты и вопросы.\n\n" + "ТЕКСТ:\n{text}" + ) try: compressed_history = await hierarchical_summarize( text=history_text, @@ -234,9 +297,12 @@ class RAGOrchestrator: max_depth=2, temperature=getattr(self.config, 'ai_temperature', 0.1) ) - # Заменяем историю одним сжатым сообщением - formatted_history = [{"role": "system", "content": "Сжатая история диалога:\n" + compressed_history}] - logger.info(f"История сжата с {count_tokens(history_text)} до {count_tokens(compressed_history)} токенов") + formatted_history = [ + {"role": "system", "content": "Сжатая история диалога:\n" + compressed_history} + ] + logger.info( + f"История сжата с {count_tokens(history_text)} до {count_tokens(compressed_history)} токенов" + ) except Exception as e: logger.error(f"Ошибка при иерархическом резюмировании истории: {e}") # Fallback: простая обрезка @@ -255,14 +321,15 @@ class RAGOrchestrator: for rec in truncated_history ] else: - # Если история укладывается, просто форматируем - truncated_history = history[-10:] # ограничим последними 10 сообщениями для простоты, но можно взять больше + # Если история укладывается, форматируем последние сообщения + # Берём максимум 20 сообщений, чтобы не перегружать + truncated_history = history[-20:] formatted_history = [ {"role": rec['role'], "content": rec['content']} for rec in truncated_history ] - # ----- 5. Классификация намерений (если не переопределена) ----- + # ----- 5. Классификация намерений ----- intent = intent_override if intent is None and getattr(self.config, 'enable_intent_classification', True): intent_prompt = prompts.get('intent', '') @@ -395,7 +462,6 @@ class RAGOrchestrator: elif not last_file_path.lower().endswith('.docx'): answer = "❌ Проверка орфографии поддерживается только для файлов .docx." else: - # TODO: адаптировать check_spelling для работы через HTTP answer = "⚠️ Проверка орфографии требует доработки (передача промптов)." # --- 7.6. GREETING --- @@ -479,30 +545,24 @@ class RAGOrchestrator: ) logger.info(f"Найден контекст длиной {len(context)} символов (room={room_jid})") - # 8.3. Обрезка контекста по токенам (НОВОЕ) + # 8.3. Обрезка контекста по токенам # Оставшиеся токены после истории и промптов – используем для контекста - max_context_tokens = available_for_history_and_context - total_history_tokens - # Защита от отрицательных значений + max_context_tokens = available_for_history_and_context - sum(count_tokens(msg['content']) for msg in formatted_history) max_context_tokens = max(max_context_tokens, 0) if max_context_tokens > 0 and context: - # Если контекст слишком длинный, обрезаем по токенам - # Используем функцию count_tokens для оценки context_tokens = count_tokens(context) if context_tokens > max_context_tokens: - # Простая обрезка: берём первые max_context_tokens токенов - # В будущем можно использовать иерархическое резюмирование logger.warning( f"Контекст слишком длинный ({context_tokens} токенов), обрезаем до {max_context_tokens}" ) - # Обрезаем по символам (грубая оценка: 1 токен ≈ 3 символа) - max_context_chars = int(max_context_tokens * 3) + max_context_chars = int(max_context_tokens * 3.5) if max_context_chars > 0: context = context[:max_context_chars] else: context = "" - # 8.4. Переранжирование контекста (если включено) + # 8.4. Переранжирование контекста rerank_min_length = getattr(self.config, 'rerank_min_length', 5000) if intent != "FACT" and len(context) > rerank_min_length: context = await rerank_context( @@ -597,6 +657,10 @@ class RAGOrchestrator: "confidence": None } + # ------------------------------------------------------------------ + # Индексация документа (с возможным сжатием) + # ------------------------------------------------------------------ + async def index_document( self, file_name: str, @@ -611,11 +675,16 @@ class RAGOrchestrator: ) -> Dict[str, Any]: """ Индексация документа в базу знаний. - Вызывается из HTTP-эндпоинта /rag/index. + Перед индексацией применяет иерархическое резюмирование, + если документ слишком большой. """ + # Сжимаем документ, если он слишком большой + compressed_text = await self._compress_document_if_needed(file_text, file_name) + + # Передаём сжатый текст в сервис базы знаний doc_id, chunk_count = await self.kb.add_document( file_name=file_name, - file_text=file_text, + file_text=compressed_text, user_jid=user_jid, is_global=is_global, title=title,