markov/fckbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update file rag_orchestrator.py

Browse Source
This commit is contained in:
Markov Andrey
2026-06-30 12:54:10 +00:00
parent 405dea3fa3
commit cf5cc568b4
1 changed files with 107 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

145
rag/rag_orchestrator.py
View File

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