Update 3 files

- /rag/intent_router.py
- /rag/query_processor.py
- /rag/history_manager.py
This commit is contained in:
Markov Andrey
2026-06-30 16:15:00 +00:00
parent 9d351d037b
commit fbfe178067
3 changed files with 118 additions and 118 deletions

View File

@@ -11,13 +11,13 @@
"""
import logging
from typing import List, Dict, Optional, Any
from typing import List, Dict, Optional, Any, Tuple
from rag.services.postgres_service import PostgresService
from rag.services.giga_client import GigaClient
from rag.utils.text_utils import count_tokens
from rag.functions.hierarchical_summarize import hierarchical_summarize
from rag.config_models import AppConfig
from core.services.postgres_service import PostgresService
from core.services.giga_client import GigaClient
from core.utils.text_utils import count_tokens
from core.functions.hierarchical_summarize import hierarchical_summarize
from core.config_models import AppConfig
logger = logging.getLogger(__name__)
@@ -33,7 +33,7 @@ class HistoryManager:
giga: GigaClient,
config: AppConfig,
default_prompts: Optional[Dict[str, str]] = None
):
) -> None:
"""
Инициализация менеджера истории.
@@ -43,10 +43,14 @@ class HistoryManager:
config: объект конфигурации (AppConfig)
default_prompts: словарь промптов по умолчанию (нужен для hierarchical_summary)
"""
self.db = db
self.giga = giga
self.config = config
self.default_prompts = default_prompts or {}
self.db: PostgresService = db
self.giga: GigaClient = giga
self.config: AppConfig = config
self.default_prompts: Dict[str, str] = default_prompts or {}
# ------------------------------------------------------------------
# Получение и форматирование истории
# ------------------------------------------------------------------
async def get_history(
self,
@@ -65,7 +69,7 @@ class HistoryManager:
Возвращает:
List[Dict[str, str]]: список сообщений с ключами 'role' и 'content'
"""
raw_history = await self.db.get_history(user_jid, room_jid, limit=limit)
raw_history: List[Dict[str, Any]] = await self.db.get_history(user_jid, room_jid, limit=limit)
# Преобразуем в формат, ожидаемый LLM
return [
{"role": rec['role'], "content": rec['content']}
@@ -95,8 +99,8 @@ class HistoryManager:
return []
# Подсчитываем общее количество токенов в истории
history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
total_tokens = count_tokens(history_text)
history_text: str = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
total_tokens: int = count_tokens(history_text)
if total_tokens <= max_tokens:
logger.debug(f"История укладывается в {max_tokens} токенов (фактически {total_tokens})")
@@ -127,14 +131,14 @@ class HistoryManager:
)
# Параметры сжатия из конфига
summarization_config = self.config.summarization
target_tokens = summarization_config.target_tokens_after_summary
chunk_size = summarization_config.chunk_size_tokens
max_depth = summarization_config.max_depth
temperature = self.config.ai.temperature
summarization_config: Dict[str, Any] = getattr(self.config, 'summarization', {})
target_tokens: int = summarization_config.get('target_tokens_after_summary', max_tokens)
chunk_size: int = summarization_config.get('chunk_size_tokens', 500)
max_depth: int = summarization_config.get('max_depth', 2)
temperature: float = getattr(self.config, 'ai_temperature', 0.1)
try:
compressed_text = await hierarchical_summarize(
compressed_text: str = await hierarchical_summarize(
text=history_text,
giga=self.giga,
prompt_template=prompt_template,
@@ -153,10 +157,10 @@ class HistoryManager:
except Exception as e:
logger.error(f"Ошибка при сжатии истории: {e}, выполняем простую обрезку")
# Fallback: простая обрезка по токенам
truncated_history = []
total = 0
truncated_history: List[Dict[str, str]] = []
total: int = 0
for msg in reversed(history):
tokens = count_tokens(msg['content'])
tokens: int = count_tokens(msg['content'])
if total + tokens <= max_tokens:
truncated_history.append(msg)
total += tokens

View File

@@ -12,7 +12,7 @@ import logging
import re
import asyncio
import os
from typing import Optional, Dict, List, Any
from typing import Optional, Dict, List, Any, Union, Tuple
from core.services.giga_client import GigaClient
from core.services.kb_service import KBService
@@ -21,6 +21,7 @@ from core.functions.extract_metrics import extract_metrics
from core.functions.summarize_document import summarize_document
from core.functions.check_consistency import check_consistency
from core.functions.check_spelling import check_spelling
from core.config_models import AppConfig # предположим, что путь такой
logger = logging.getLogger(__name__)
@@ -35,24 +36,17 @@ class IntentRouter:
giga: GigaClient,
kb: KBService,
files: FileService,
config,
config: AppConfig,
default_prompts: Optional[Dict[str, str]] = None
):
) -> None:
"""
Инициализация маршрутизатора.
Аргументы:
giga: клиент GigaChat
kb: сервис базы знаний
files: сервис файлов
config: объект конфигурации
default_prompts: словарь промптов по умолчанию
"""
self.giga = giga
self.kb = kb
self.files = files
self.config = config
self.default_prompts = default_prompts or {}
self.giga: GigaClient = giga
self.kb: KBService = kb
self.files: FileService = files
self.config: AppConfig = config
self.default_prompts: Dict[str, str] = default_prompts or {}
# ------------------------------------------------------------------
# Основной метод маршрутизации
@@ -133,23 +127,23 @@ class IntentRouter:
prompts: Dict[str, str]
) -> Dict[str, Any]:
"""Обработка METRICS: извлечение KPI из базы знаний."""
context = await self.kb.find_relevant_info(
context: str = await self.kb.find_relevant_info(
query, user_jid, room_jid,
top_k=getattr(self.config, 'rag_metrics_top_k', 30)
)
if not context:
return {"answer": "Не найдено данных для извлечения метрик.", "context": "", "sources": []}
metrics_prompt = prompts.get('metrics', '')
metrics = await extract_metrics(
metrics_prompt: str = prompts.get('metrics', '')
metrics: List[Dict] = await extract_metrics(
giga=self.giga,
context=context,
prompt_text=metrics_prompt,
bot_config=self.config
)
if metrics:
lines = [f"- {m.get('metric_name')}: {m.get('value')} {m.get('unit', '')}" for m in metrics[:10]]
answer = "📊 **Извлечённые метрики:**\n" + "\n".join(lines)
lines: List[str] = [f"- {m.get('metric_name')}: {m.get('value')} {m.get('unit', '')}" for m in metrics[:10]]
answer: str = "📊 **Извлечённые метрики:**\n" + "\n".join(lines)
else:
answer = "Не удалось извлечь метрики."
@@ -165,8 +159,8 @@ class IntentRouter:
if not last_file_text:
return {"answer": "Нет документа для суммаризации.", "context": "", "sources": []}
summary_prompt = prompts.get('summary', '')
answer = await summarize_document(
summary_prompt: str = prompts.get('summary', '')
answer: str = await summarize_document(
giga=self.giga,
text=last_file_text,
title="Ваш документ",
@@ -183,19 +177,19 @@ class IntentRouter:
prompts: Dict[str, str]
) -> Dict[str, Any]:
"""Обработка CONTRADICTION: проверка противоречий в базе знаний."""
context = await self.kb.find_relevant_info(
context: str = await self.kb.find_relevant_info(
query, user_jid, room_jid,
top_k=getattr(self.config, 'rag_contradiction_top_k', 10)
)
if not context:
return {"answer": "Недостаточно данных для проверки противоречий.", "context": "", "sources": []}
chunks = [c.strip() for c in context.split("\n\n") if c.strip()]
chunks: List[str] = [c.strip() for c in context.split("\n\n") if c.strip()]
if len(chunks) < 2:
return {"answer": "Недостаточно фрагментов для проверки противоречий.", "context": context, "sources": []}
consistency_prompt = prompts.get('consistency', '')
consistency = await check_consistency(
consistency_prompt: str = prompts.get('consistency', '')
consistency: str = await check_consistency(
giga=self.giga,
chunks=chunks,
query=query,
@@ -203,7 +197,7 @@ class IntentRouter:
bot_config=self.config
)
if "[CONFLICT]" in consistency:
answer = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
answer: str = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
else:
answer = "✅ Противоречий не обнаружено."
@@ -219,10 +213,9 @@ class IntentRouter:
last_file_text: Optional[str]
) -> Dict[str, Any]:
"""Обработка TEMPLATE_FILL: заполнение шаблона документа."""
# Получаем текст шаблона
template_text = last_file_text or ""
template_text: str = last_file_text or ""
if not template_text and last_file_path and os.path.exists(last_file_path):
result = await asyncio.to_thread(self.files.process_any_file, last_file_path)
result: Any = await asyncio.to_thread(self.files.process_any_file, last_file_path)
if isinstance(result, tuple) and len(result) == 2:
template_text = result[0]
else:
@@ -231,16 +224,15 @@ class IntentRouter:
if not template_text:
return {"answer": "❌ Нет шаблона документа для заполнения. Загрузите файл .docx.", "context": "", "sources": []}
# Обрезаем шаблон для экономии токенов
truncated_template = template_text[:5000]
search_query = f"{query}\n{truncated_template}"
context = await self.kb.find_relevant_info(search_query, user_jid, room_jid, top_k=15)
truncated_template: str = template_text[:5000]
search_query: str = f"{query}\n{truncated_template}"
context: str = await self.kb.find_relevant_info(search_query, user_jid, room_jid, top_k=15)
fill_prompt = prompts.get('generate_document', '')
fill_prompt: str = prompts.get('generate_document', '')
if not fill_prompt:
fill_prompt = "Заполни плейсхолдеры, используя базу знаний."
full_prompt = (
full_prompt: str = (
f"Перед тобой шаблон документа. Заполни плейсхолдеры, используя ТОЛЬКО базу знаний.\n\n"
f"[ТЕКСТ ШАБЛОНА]:\n{truncated_template}\n\n"
f"[ДАННЫЕ ИЗ БАЗЫ ЗНАНИЙ]:\n{context}\n\n"
@@ -248,7 +240,7 @@ class IntentRouter:
f"Формат ответа: [SURGICAL_REPLACE]\nинструкция_из_шаблона ||| текст_из_БЗ\n[/SURGICAL_REPLACE]"
)
answer = await self.giga.chat(
answer: str = await self.giga.chat(
history=[],
query=full_prompt,
system_prompt=None,
@@ -266,42 +258,39 @@ class IntentRouter:
if not last_file_path:
return {"answer": "❌ Нет загруженного документа для замены.", "context": "", "sources": []}
# Парсим запрос: "замени слово на слово"
q = re.sub(r'\s+', ' ', query.lower()).strip()
m = re.search(r'замени\s+(.+?)\s+на\s+(.+)', q)
q: str = re.sub(r'\s+', ' ', query.lower()).strip()
m: Optional[re.Match] = re.search(r'замени\s+(.+?)\s+на\s+(.+)', q)
if not m:
m = re.search(r'заменить\s+(.+?)\s+на\s+(.+)', q)
if not m:
return {"answer": "Не удалось распознать, что на что заменять. Используйте формат: замени слово на слово", "context": "", "sources": []}
old_word = m.group(1).strip()
new_word = m.group(2).strip()
old_word: str = m.group(1).strip()
new_word: str = m.group(2).strip()
try:
from mawo_pymorphy3 import create_analyzer
morph = create_analyzer()
# Генерируем все формы старого слова
old_forms = set()
old_forms: set = set()
parsed_old = morph.parse(old_word)[0]
old_forms.add(old_word)
old_forms.add(parsed_old.normal_form)
cases = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
numbers = ['sing', 'plur']
cases: List[str] = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
numbers: List[str] = ['sing', 'plur']
for number in numbers:
for case in cases:
inflected = parsed_old.inflect({case, number})
if inflected:
old_forms.add(inflected.word)
# Кэш для разбора слов
parse_cache = {}
def get_new_form(new_word, old_form_text):
parse_cache: Dict[str, Any] = {}
def get_new_form(new_word: str, old_form_text: str) -> str:
if old_form_text in parse_cache:
parsed_old = parse_cache[old_form_text]
else:
parsed_old = morph.parse(old_form_text)[0]
parse_cache[old_form_text] = parsed_old
tags = set()
tags: set = set()
if parsed_old.tag.case:
tags.add(parsed_old.tag.case)
if parsed_old.tag.number:
@@ -312,14 +301,13 @@ class IntentRouter:
inflected = parsed_new.inflect(tags)
return inflected.word if inflected else new_word
replacements = {}
replacements: Dict[str, str] = {}
for old_form in old_forms:
replacements[old_form] = get_new_form(new_word, old_form)
# Выполняем замену
new_path = await asyncio.to_thread(self.files.surgical_replace, last_file_path, replacements)
new_path: Optional[str] = await asyncio.to_thread(self.files.surgical_replace, last_file_path, replacements)
if new_path:
answer = f"Замена '{old_word}''{new_word}' выполнена. Файл сохранён: {new_path}"
answer: str = f"Замена '{old_word}''{new_word}' выполнена. Файл сохранён: {new_path}"
else:
answer = f"❌ Ошибка при замене '{old_word}''{new_word}'."
return {"answer": answer, "context": "", "sources": []}
@@ -348,7 +336,7 @@ class IntentRouter:
if not last_file_path.lower().endswith('.docx'):
return {"answer": "❌ Проверка орфографии поддерживается только для файлов .docx.", "context": "", "sources": []}
spellcheck_prompt = prompts.get('spellcheck', '')
spellcheck_prompt: str = prompts.get('spellcheck', '')
if not spellcheck_prompt:
return {"answer": "❌ Промпт проверки орфографии не загружен.", "context": "", "sources": []}
@@ -360,7 +348,7 @@ class IntentRouter:
config=self.config
)
if changes:
answer = "📝 **Исправления:**\n" + "\n".join(changes)
answer: str = "📝 **Исправления:**\n" + "\n".join(changes)
else:
answer = "✅ Ошибок не найдено."
return {"answer": answer, "context": "", "sources": []}
@@ -376,14 +364,13 @@ class IntentRouter:
prompts: Dict[str, str]
) -> Dict[str, Any]:
"""Обработка GREETING: простой ответ через GigaChat без поиска."""
# Используем synthesis промпт, если есть, или просто передаём запрос
synthesis_template = prompts.get('synthesis', '')
synthesis_template: str = prompts.get('synthesis', '')
if synthesis_template:
full_query = synthesis_template.format(context="", query=query)
full_query: str = synthesis_template.format(context="", query=query)
else:
full_query = query
answer = await self.giga.chat(
answer: str = await self.giga.chat(
history=history or [],
query=full_query,
system_prompt=system_prompt,

View File

@@ -9,14 +9,14 @@ import logging
import re
from typing import Optional, Dict, List, Any
from .services.giga_client import GigaClient
from .services.kb_service import KBService
from .functions.expand_query import expand_query
from .functions.rerank_context import rerank_context
from .functions.critique_answer import critique_answer
from .utils.text_utils import count_tokens
from .prompt_builder import PromptBuilder
from .config_models import AppConfig
from core.services.giga_client import GigaClient
from core.services.kb_service import KBService
from core.functions.expand_query import expand_query
from core.functions.rerank_context import rerank_context
from core.functions.critique_answer import critique_answer
from core.utils.text_utils import count_tokens
from core.prompt_builder import PromptBuilder
from core.config_models import AppConfig
logger = logging.getLogger(__name__)
@@ -32,12 +32,12 @@ class QueryProcessor:
kb: KBService,
config: AppConfig,
default_prompts: Optional[Dict[str, str]] = None
):
self.giga = giga
self.kb = kb
self.config = config
self.default_prompts = default_prompts or {}
self.prompt_builder = PromptBuilder(config)
) -> None:
self.giga: GigaClient = giga
self.kb: KBService = kb
self.config: AppConfig = config
self.default_prompts: Dict[str, str] = default_prompts or {}
self.prompt_builder: PromptBuilder = PromptBuilder(config)
async def process(
self,
@@ -52,18 +52,31 @@ class QueryProcessor:
) -> Dict[str, Any]:
"""
Выполняет полный RAG-пайплайн для обычного запроса.
Аргументы:
query: текст запроса пользователя
user_jid: JID пользователя
room_jid: JID комнаты (None для личного чата)
prompts: словарь промптов (expand, synthesis, critique, ...)
intent: код намерения (для выбора стратегии)
history: история диалога (уже сжатая, если нужно)
system_prompt: системный промпт
available_tokens_for_context: сколько токенов доступно для контекста
Возвращает:
Словарь с ключами 'answer', 'context', 'sources', 'confidence'
"""
# 1. Расширение запроса
expanded = await expand_query(
expanded: str = await expand_query(
giga=self.giga,
query=query,
prompt_text=prompts.get('expand', ''),
bot_config=self.config
)
search_query = expanded if expanded and expanded != query else query
search_query: str = expanded if expanded and expanded != query else query
# 2. Поиск релевантного контекста в базе знаний
context = await self.kb.find_relevant_info(
context: str = await self.kb.find_relevant_info(
search_query, user_jid, room_jid,
top_k=getattr(self.config, 'rag_default_top_k', 30)
)
@@ -71,20 +84,20 @@ class QueryProcessor:
# 3. Обрезка контекста по токенам
if context:
context_tokens = count_tokens(context)
context_tokens: int = count_tokens(context)
if context_tokens > available_tokens_for_context:
logger.warning(
f"Контекст слишком длинный ({context_tokens} токенов), "
f"обрезаем до {available_tokens_for_context}"
)
max_context_chars = int(available_tokens_for_context * 3.5)
max_context_chars: int = int(available_tokens_for_context * 3.5)
if max_context_chars > 0:
context = context[:max_context_chars]
else:
context = ""
# 4. Переранжирование контекста (если включено и контекст достаточно длинный)
rerank_min_length = getattr(self.config, 'rerank_min_length', 5000)
rerank_min_length: int = getattr(self.config, 'rerank_min_length', 5000)
if intent != "FACT" and len(context) > rerank_min_length:
context = await rerank_context(
bot=None,
@@ -94,8 +107,8 @@ class QueryProcessor:
bot_config=self.config
)
# 5. Формирование промта с помощью PromptBuilder (динамические few-shot, сэндвич)
extra_instructions = ""
# 5. Формирование промта с помощью PromptBuilder
extra_instructions: str = ""
if intent in ("CALCULATION", "PROCEDURE"):
extra_instructions = (
"Пожалуйста, покажи пошаговое решение перед итоговым ответом. "
@@ -103,13 +116,11 @@ class QueryProcessor:
"После всех шагов дай итоговый ответ."
)
# Если synthesis_template задан, добавим его в extra_instructions (но PromptBuilder уже использует стандартные инструкции)
synthesis_template = prompts.get('synthesis', '')
synthesis_template: str = prompts.get('synthesis', '')
if synthesis_template and not extra_instructions:
# Если есть кастомный шаблон, используем его как дополнительную инструкцию
extra_instructions = synthesis_template.format(context=context, query=query) if '{context}' in synthesis_template else synthesis_template
prompt = self.prompt_builder.build_prompt(
prompt: str = self.prompt_builder.build_prompt(
query=query,
intent=intent,
context=context,
@@ -119,21 +130,21 @@ class QueryProcessor:
)
logger.debug(f"Сформированный промт (первые 500 символов): {prompt[:500]}")
# 6. Генерация ответа (без отдельной передачи истории, она уже в промте)
answer = await self.giga.chat(
history=[], # история уже в промте
# 6. Генерация ответа
answer: str = await self.giga.chat(
history=[],
query=prompt,
system_prompt=None, # системный промпт тоже в промте
system_prompt=None,
file_id=None,
temperature=getattr(self.config, 'ai_temperature', 0.1)
)
# 7. Самокритика (если включена)
if getattr(self.config, 'enable_self_critique', False) and context:
critique_prompt = prompts.get('critique', '')
critique_prompt: str = prompts.get('critique', '')
if critique_prompt:
logger.debug("Запуск самокритики")
is_ok = await critique_answer(
is_ok: bool = await critique_answer(
giga=self.giga,
query=query,
context=context,
@@ -143,7 +154,6 @@ class QueryProcessor:
)
if not is_ok:
logger.warning("Ответ не прошёл самокритику, перегенерация")
# Перегенерируем с тем же промтом
answer = await self.giga.chat(
history=[],
query=prompt,
@@ -151,7 +161,6 @@ class QueryProcessor:
file_id=None,
temperature=getattr(self.config, 'ai_temperature', 0.1)
)
# Повторная проверка после перегенерации
if not await critique_answer(
giga=self.giga,
query=query,
@@ -163,7 +172,7 @@ class QueryProcessor:
answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
# 8. Извлечение источников
sources = []
sources: List[str] = []
if context:
for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
sources.append(match.group(1))