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:21:07 +00:00
parent be9af917ab
commit 2008bb1283
1 changed files with 153 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
rag/rag_orchestrator.py
View File

@@ -8,8 +8,8 @@
Этот модуль не зависит от XMPP и может использоваться как в ботах,
так и в отдельном RAG-ядре (API-сервисе).
Добавлено управление токенами: резервирование места для ответа и промтов,
обрезка истории и контекста с учётом доступного пространства.
Улучшено управление токенами: теперь история и контекст обрезаются с учётом
лимитов модели, резервирования для ответа и промптов.
"""
import asyncio
@@ -78,6 +78,68 @@ class RAGOrchestrator:
self.default_prompts = default_prompts or {}
logger.info("RAGOrchestrator инициализирован")
# ------------------------------------------------------------------
# Вспомогательный метод для расчёта токенов с резервированием
# ------------------------------------------------------------------
def _prepare_prompt_parts(
self,
synthesis_template: str,
system_prompt: Optional[str],
query: str,
context: str,
max_total_tokens: int = 8192, # для GigaChat-Max
reserved_for_answer: int = 1000,
reserved_for_overhead: int = 200 # буфер для форматирования
) -> Dict[str, Any]:
"""
Подготавливает части промта и обрезает историю/контекст по токенам.
Возвращает словарь с ключами:
- history: отфильтрованная история (список сообщений)
- context: отфильтрованный контекст (строка)
- prompt_tokens: количество токенов в промте (без истории и контекста)
- total_used: общее использованное количество токенов
- available_for_history_and_context: сколько токенов выделено для истории и контекста
"""
# 1. Подсчёт токенов в статичных частях промта
system_tokens = count_tokens(system_prompt) if system_prompt else 0
synthesis_tokens = count_tokens(synthesis_template)
query_tokens = count_tokens(query)
# 2. Резервирование
# Вычитаем системный промпт, шаблон, запрос, буфер и место для ответа
prompt_tokens = system_tokens + synthesis_tokens + query_tokens
available_for_history_and_context = (
max_total_tokens
- prompt_tokens
- reserved_for_answer
- reserved_for_overhead
)
if available_for_history_and_context <= 0:
logger.warning(
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,
"system_tokens": system_tokens,
"synthesis_tokens": synthesis_tokens,
"query_tokens": query_tokens,
}
# ------------------------------------------------------------------
# Основной метод обработки запроса
# ------------------------------------------------------------------
async def process_query(
self,
query: str,
@@ -115,46 +177,64 @@ class RAGOrchestrator:
# ----- 1. Подготовка промптов -----
if prompts is None:
prompts = self.default_prompts.copy()
synthesis_template = prompts.get('synthesis', '{context}\n\n{query}\n\nОтвет:')
synthesis_template = prompts.get('synthesis', '')
if not synthesis_template:
synthesis_template = "{context}\n\n{query}\n\nОтвет:"
system_prompt = prompts.get('system', None)
# ----- 2. Получение истории диалога из БД -----
# ----- 2. Получение истории диалога из БД (необрезанной) -----
history = await self.db.get_history(user_jid, room_jid, limit=100)
# ----- 3. Управление токенами -----
# 3.1. Резервирование места для ответа и промтов
max_context_tokens = getattr(self.config, 'max_context_tokens', 3000)
# ----- 3. Расчёт лимитов токенов (НОВОЕ) -----
# Параметры из конфига
max_model_tokens = getattr(self.config, 'max_model_tokens', 8192) # GigaChat-Max default
reserved_for_answer = getattr(self.config, 'reserved_for_answer_tokens', 1000)
reserved_for_prompt = 0
if system_prompt:
reserved_for_prompt += count_tokens(system_prompt)
reserved_for_prompt += count_tokens(synthesis_template)
# Дополнительный резерв на служебные токены
overhead = 100
available_for_history_and_context = max_context_tokens - reserved_for_answer - reserved_for_prompt - overhead
reserved_for_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200)
if available_for_history_and_context < 100:
logger.warning(f"Доступное пространство для истории и контекста слишком мало: {available_for_history_and_context} токенов. Увеличьте max_context_tokens или уменьшите резервы.")
available_for_history_and_context = max(100, available_for_history_and_context)
# Используем вспомогательный метод для расчёта
token_info = self._prepare_prompt_parts(
synthesis_template=synthesis_template,
system_prompt=system_prompt,
query=query,
context="", # пока пусто, позже добавим
max_total_tokens=max_model_tokens,
reserved_for_answer=reserved_for_answer,
reserved_for_overhead=reserved_for_overhead
)
available_for_history_and_context = token_info["available_for_history_and_context"]
logger.debug(f"Управление токенами: max={max_context_tokens}, резерв ответа={reserved_for_answer}, промт={reserved_for_prompt}, доступно={available_for_history_and_context}")
logger.debug(
f"Доступно для истории и контекста: {available_for_history_and_context} токенов, "
f"промпт: {token_info['prompt_tokens']} токенов"
)
# 3.2. Обрезка истории по токенам (сохраняем последние сообщения)
total_tokens = 0
# ----- 4. Обрезка истории по токенам (с учётом лимита) -----
# Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние
truncated_history = []
total_history_tokens = 0
history_limit = min(available_for_history_and_context // 2, 2000) # выделяем половину для истории
for record in reversed(history):
tokens = count_tokens(record['content'])
if total_tokens + tokens <= available_for_history_and_context:
if total_history_tokens + tokens <= history_limit:
truncated_history.append(record)
total_tokens += tokens
total_history_tokens += tokens
else:
break
truncated_history.reverse()
# Сразу обрезаем историю для использования в дальнейшем
formatted_history = truncated_history
logger.debug(f"История обрезана до {len(truncated_history)} сообщений, {total_tokens} токенов")
truncated_history.reverse() # восстанавливаем хронологический порядок
# ----- 4. Классификация намерений -----
# Преобразуем в формат, ожидаемый GigaChat
formatted_history = [
{"role": rec['role'], "content": rec['content']}
for rec in truncated_history
]
logger.debug(
f"История обрезана: {len(truncated_history)} сообщений, "
f"{total_history_tokens} токенов (лимит {history_limit})"
)
# ----- 5. Классификация намерений (если не переопределена) -----
intent = intent_override
if intent is None and getattr(self.config, 'enable_intent_classification', True):
intent_prompt = prompts.get('intent', '')
@@ -170,19 +250,19 @@ class RAGOrchestrator:
else:
intent = intent or "GENERAL"
# Принудительная установка SURGICAL
# ----- 6. Принудительная установка SURGICAL по ключевым словам -----
keywords = getattr(self.config, 'surgical_keywords', [])
if any(kw in query.lower() for kw in keywords) and last_file_path:
intent = "SURGICAL"
logger.info(f"Принудительный Intent: SURGICAL (есть файл и ключевые слова)")
# ----- 5. Обработка специализированных намерений -----
# ----- 7. Обработка специализированных намерений -----
answer = None
context = None
ctx_for_critique = None
synthesis_template_for_critique = None
# --- 5.1. METRICS ---
# --- 7.1. METRICS ---
if intent == "METRICS":
context = await self.kb.find_relevant_info(
query, user_jid, room_jid,
@@ -204,7 +284,7 @@ class RAGOrchestrator:
else:
answer = "Не удалось извлечь метрики."
# --- 5.2. SUMMARY ---
# --- 7.2. SUMMARY ---
elif intent == "SUMMARY":
if not last_file_text:
answer = "Нет документа для суммаризации."
@@ -218,7 +298,7 @@ class RAGOrchestrator:
bot_config=self.config
)
# --- 5.3. CONTRADICTION ---
# --- 7.3. CONTRADICTION ---
elif intent == "CONTRADICTION":
context = await self.kb.find_relevant_info(
query, user_jid, room_jid,
@@ -244,7 +324,7 @@ class RAGOrchestrator:
else:
answer = "✅ Противоречий не обнаружено."
# --- 5.4. TEMPLATE_FILL ---
# --- 7.4. TEMPLATE_FILL ---
elif intent == "TEMPLATE_FILL":
template_text = last_file_text or ""
if not template_text and last_file_path and os.path.exists(last_file_path):
@@ -277,7 +357,7 @@ class RAGOrchestrator:
temperature=getattr(self.config, 'ai_temperature', 0.1)
)
# --- 5.5. SPELLCHECK ---
# --- 7.5. SPELLCHECK ---
elif intent == "SPELLCHECK":
if room_jid is not None:
answer = "⚠️ Проверка орфографии доступна только в личном чате."
@@ -290,7 +370,7 @@ class RAGOrchestrator:
# TODO: адаптировать check_spelling для работы через HTTP
answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
# --- 5.6. GREETING ---
# --- 7.6. GREETING ---
elif intent == "GREETING":
answer = await self.giga.chat(
history=formatted_history,
@@ -300,7 +380,7 @@ class RAGOrchestrator:
temperature=getattr(self.config, 'ai_temperature', 0.1)
)
# --- 5.7. SURGICAL ---
# --- 7.7. SURGICAL ---
elif intent == "SURGICAL":
if not last_file_path:
answer = "❌ Нет загруженного документа для замены."
@@ -353,9 +433,9 @@ class RAGOrchestrator:
except ImportError:
answer = "❌ Библиотека mawo-pymorphy3 не установлена. Установите её для морфологической замены."
# ----- 6. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
# ----- 8. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
if answer is None:
# 6.1. Расширение запроса
# 8.1. Расширение запроса
expanded = await expand_query(
giga=self.giga,
query=query,
@@ -364,14 +444,37 @@ class RAGOrchestrator:
)
search_query = expanded if expanded and expanded != query else query
# 6.2. Поиск релевантного контекста в базе знаний
# 8.2. Поиск релевантного контекста в базе знаний
context = await self.kb.find_relevant_info(
search_query, user_jid, room_jid,
top_k=getattr(self.config, 'rag_default_top_k', 30)
)
logger.info(f"Найден контекст длиной {len(context)} символов (room={room_jid})")
# 6.3. Переранжирование контекста (если включено)
# 8.3. Обрезка контекста по токенам (НОВОЕ)
# Оставшиеся токены после истории и промптов – используем для контекста
max_context_tokens = available_for_history_and_context - total_history_tokens
# Защита от отрицательных значений
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)
if max_context_chars > 0:
context = context[:max_context_chars]
else:
context = ""
# 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(
@@ -382,53 +485,20 @@ class RAGOrchestrator:
bot_config=self.config
)
# 6.4. Управление токенами для контекста и ответа
# Вычисляем, сколько токенов уже занято историей
used_tokens = total_tokens + count_tokens(synthesis_template) + count_tokens(system_prompt or '')
# Вычисляем, сколько токенов осталось для контекста
remaining_for_context = available_for_history_and_context - used_tokens
if remaining_for_context < 0:
# Если даже без контекста превышаем – обрезаем историю жёстко
logger.warning(f"Перерасход токенов: история уже {used_tokens} токенов, доступно {available_for_history_and_context}")
# Урезаем историю до минимального размера (например, 5 сообщений)
if len(truncated_history) > 5:
truncated_history = truncated_history[-5:]
remaining_for_context = available_for_history_and_context - count_tokens(''.join([h['content'] for h in truncated_history])) - count_tokens(synthesis_template) - count_tokens(system_prompt or '')
if remaining_for_context < 0:
remaining_for_context = 100 # минимум
# Обрезаем контекст до remaining_for_context токенов
if remaining_for_context > 0 and context:
context_tokens = count_tokens(context)
if context_tokens > remaining_for_context:
# Обрезаем контекст по границам предложений (упрощённо – по символам)
# В идеале нужно использовать иерархическое резюмирование, но пока оставим простую обрезку
chars_per_token = 3 # приблизительно
max_chars = int(remaining_for_context * chars_per_token)
if len(context) > max_chars:
context = context[:max_chars]
# Обрезаем до последнего полного предложения
last_period = context.rfind('.')
if last_period > 0:
context = context[:last_period+1]
logger.warning(f"Контекст обрезан до {len(context)} символов (приблизительно {remaining_for_context} токенов)")
elif remaining_for_context <= 0:
context = ""
logger.warning("Нет места для контекста в промте, ответ будет без контекста")
# 6.5. Синтез ответа с добавлением цепочки рассуждений (CoT)
synthesis_template_local = synthesis_template
# 8.5. Синтез ответа с добавлением цепочки рассуждений (CoT)
synthesis_template = synthesis_template
if intent in ("CALCULATION", "PROCEDURE"):
cot_instruction = (
"\n\nПожалуйста, покажи пошаговое решение перед итоговым ответом. "
"Опиши каждый шаг вычислений или действий в логической последовательности. "
"После всех шагов дай итоговый ответ."
)
synthesis_template_local += cot_instruction
synthesis_template += cot_instruction
full_query = synthesis_template_local.format(context=context, query=query)
full_query = synthesis_template.format(context=context, query=query)
logger.debug(f"Полный запрос к GigaChat (первые 500 символов): {full_query[:500]}")
# 6.6. Генерация ответа через GigaChat
# 8.6. Генерация ответа
answer = await self.giga.chat(
history=formatted_history,
query=full_query,
@@ -437,11 +507,10 @@ class RAGOrchestrator:
temperature=getattr(self.config, 'ai_temperature', 0.1)
)
# Сохраняем контекст для самокритики
ctx_for_critique = context
synthesis_template_for_critique = synthesis_template_local
synthesis_template_for_critique = synthesis_template
# ----- 7. Самокритика (если включена) -----
# ----- 9. Самокритика (если включена) -----
if (getattr(self.config, 'enable_self_critique', False) and
intent not in ("METRICS", "SUMMARY", "CONTRADICTION") and
ctx_for_critique is not None and
@@ -480,18 +549,18 @@ class RAGOrchestrator:
):
answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
# ----- 8. Сохранение истории диалога в БД -----
# ----- 10. Сохранение истории диалога в БД -----
await self.db.add_history(user_jid, "user", query, room_jid)
if answer:
await self.db.add_history(user_jid, "assistant", answer, room_jid)
# ----- 9. Извлечение источников из контекста -----
# ----- 11. Извлечение источников из контекста -----
sources = []
if context:
for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
sources.append(match.group(1))
# ----- 10. Формирование результата -----
# ----- 12. Формирование результата -----
return {
"answer": answer,
"intent": intent,