Update 86 files
- /rag/commands/expert.py - /core/services/__init__.py - /core/services/embedding_service.py - /core/services/file_service.py - /core/services/giga_client.py - /core/services/kb_service.py - /core/services/postgres_service.py - /core/services/qdrant_service.py - /core/services/reranker_service.py - /core/functions/__init__.py - /core/functions/check_consistency.py - /core/functions/check_spelling.py - /core/functions/critique_answer.py - /core/functions/expand_query.py - /core/functions/extract_metrics.py - /core/functions/file_processor.py - /core/functions/generate_document.py - /core/functions/intent_classify.py - /core/functions/rerank_context.py - /core/functions/summarize_document.py - /core/utils/__init__.py - /core/utils/arg_parser.py - /core/utils/config_loader.py - /core/utils/layout_converter.py - /core/utils/logger.py - /core/utils/text_utils.py - /core/utils/web_utils.py - /bots/commands/__init__.py - /bots/commands/base.py - /bots/commands/create.py - /core/commands/global_remove.py - /core/commands/help.py - /core/commands/info.py - /core/commands/kb.py - /core/commands/learn.py - /core/commands/other.py - /core/commands/registry.py - /core/commands/stats.py - /core/commands/template.py - /core/handlers/__init__.py - /core/handlers/message_handler.py - /core/workers/__init__.py - /core/workers/indexing_worker.py - /core/xmpp/__init__.py - /core/xmpp/client.py - /bots/commands/global_remove.py - /bots/commands/help.py - /bots/commands/info.py - /bots/commands/kb.py - /bots/commands/registry.py - /bots/commands/other.py - /bots/commands/template.py - /bots/commands/learn.py - /bots/commands/stats.py - /bots/handlers/message_handler.py - /bots/handlers/__init__.py - /bots/workers/__init__.py - /bots/workers/indexing_worker.py - /bots/xmpp/__init__.py - /bots/xmpp/client.py - /rag/services/qdrant_service.py - /rag/services/giga_client.py - /rag/services/embedding_service.py - /rag/services/reranker_service.py - /rag/services/__init__.py - /rag/services/file_service.py - /rag/services/postgres_service.py - /rag/services/kb_service.py - /rag/functions/check_spelling.py - /rag/functions/__init__.py - /rag/functions/extract_metrics.py - /rag/functions/generate_document.py - /rag/functions/expand_query.py - /rag/functions/critique_answer.py - /rag/functions/file_processor.py - /rag/functions/check_consistency.py - /rag/functions/summarize_document.py - /rag/functions/rerank_context.py - /rag/functions/intent_classify.py - /rag/utils/__init__.py - /rag/utils/layout_converter.py - /rag/utils/text_utils.py - /rag/utils/arg_parser.py - /rag/utils/config_loader.py - /rag/utils/logger.py - /rag/utils/web_utils.py
This commit is contained in:
Markov Andrey
12
rag/services/__init__.py
Normal file
12
rag/services/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Пакет services – долгоживущие компоненты для работы с внешними системами.
|
||||
Содержит:
|
||||
- PostgresService – пул соединений с PostgreSQL,
|
||||
- QdrantService – векторная БД,
|
||||
- EmbeddingService – эмбеддинги GigaChat,
|
||||
- KBService – база знаний (индексация, поиск),
|
||||
- FileService – извлечение текста из файлов,
|
||||
- GigaClient – клиент GigaChat,
|
||||
- RerankerService – локальный кросс-энкодер для переранжирования.
|
||||
"""
|
||||
195
rag/services/embedding_service.py
Normal file
195
rag/services/embedding_service.py
Normal file
@@ -0,0 +1,195 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис генерации эмбеддингов через GigaChat API.
|
||||
Поддерживает одиночные запросы с кэшированием (lru_cache) и пакетную обработку.
|
||||
Все параметры (модель, таймаут, размер кэша, verify_ssl) вынесены в конфиг.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import List, Optional, Any
|
||||
|
||||
from gigachat import GigaChat
|
||||
from gigachat.exceptions import AuthenticationError, RateLimitError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from fastembed import SparseTextEmbedding
|
||||
HAS_SPARSE = True
|
||||
except ImportError:
|
||||
HAS_SPARSE = False
|
||||
logging.warning("fastembed не установлен, sparse-эмбеддинги отключены")
|
||||
|
||||
class EmbeddingService:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "GigaChat-Embeddings",
|
||||
timeout: int = 30,
|
||||
cache_size: int = 4096,
|
||||
verify_ssl: bool = False
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.cache_size = cache_size
|
||||
self.verify_ssl = verify_ssl
|
||||
self._giga = None
|
||||
# Кэширование одиночных эмбеддингов (синхронная обёртка с lru_cache)
|
||||
self.embed_sync = lru_cache(maxsize=cache_size)(self._embed_sync_impl)
|
||||
self._sparse_model = None
|
||||
|
||||
def _get_client(self) -> GigaChat:
|
||||
"""Ленивое создание клиента GigaChat (с нужными параметрами)."""
|
||||
if self._giga is None:
|
||||
self._giga = GigaChat(
|
||||
credentials=self.api_key,
|
||||
model=self.model,
|
||||
verify_ssl_certs=self.verify_ssl,
|
||||
timeout=self.timeout
|
||||
)
|
||||
logger.debug(f"Клиент GigaChat создан (модель {self.model}, timeout {self.timeout}s)")
|
||||
return self._giga
|
||||
|
||||
def _get_sparse_model(self):
|
||||
if not HAS_SPARSE:
|
||||
return None
|
||||
if self._sparse_model is None:
|
||||
self._sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25")
|
||||
return self._sparse_model
|
||||
|
||||
# ---------- Синхронное получение эмбеддинга (с обрезкой) ----------
|
||||
def _embed_sync_impl(self, text: str) -> Optional[List[float]]:
|
||||
"""Синхронный вызов GigaChat API для одного текста. Результат кэшируется."""
|
||||
if not text or not text.strip():
|
||||
logger.warning("Попытка получить эмбеддинг пустого текста")
|
||||
return None
|
||||
|
||||
# Обрезка до 800 символов (безопасно для лимита 514 токенов GigaChat)
|
||||
MAX_CHARS = 800
|
||||
if len(text) > MAX_CHARS:
|
||||
text = text[:MAX_CHARS]
|
||||
logger.debug(f"Текст обрезан до {MAX_CHARS} символов для эмбеддинга")
|
||||
|
||||
client = self._get_client()
|
||||
try:
|
||||
# Совместимость с разными версиями SDK: пробуем input=[text] или просто [text]
|
||||
try:
|
||||
response = client.embeddings(input=[text])
|
||||
except TypeError:
|
||||
response = client.embeddings([text])
|
||||
if response.data and len(response.data) > 0:
|
||||
logger.debug(f"Эмбеддинг получен, размерность {len(response.data[0].embedding)}")
|
||||
return response.data[0].embedding
|
||||
else:
|
||||
logger.error(f"Пустой ответ эмбеддингов для текста: {text[:100]}")
|
||||
return None
|
||||
except (AuthenticationError, RateLimitError) as e:
|
||||
logger.error(f"Ошибка авторизации/лимитов GigaChat: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка GigaChat при эмбеддинге: {e}")
|
||||
raise
|
||||
|
||||
async def embed(self, text: str, retries: int = 3, delay: float = 1.0) -> Optional[List[float]]:
|
||||
"""Асинхронная обёртка с повторными попытками."""
|
||||
if not isinstance(text, str):
|
||||
logger.error(f"embed получил не строку, а {type(text)}: {text!r}")
|
||||
return None
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
result = await loop.run_in_executor(None, self.embed_sync, text)
|
||||
if result is not None:
|
||||
return result
|
||||
else:
|
||||
raise ValueError("Эмбеддинг вернул None")
|
||||
except Exception as e:
|
||||
wait = delay * (2 ** attempt)
|
||||
logger.warning(
|
||||
f"Попытка {attempt+1}/{retries} получить эмбеддинг не удалась: {e}. "
|
||||
f"Повтор через {wait:.1f} сек."
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
logger.error(f"Не удалось получить эмбеддинг после {retries} попыток")
|
||||
return None
|
||||
|
||||
# ---------- Пакетная обработка (для индексации) ----------
|
||||
def _embed_batch_sync(self, texts: List[str]) -> Optional[List[List[float]]]:
|
||||
if not texts:
|
||||
return []
|
||||
MAX_CHARS = 800
|
||||
truncated_texts = [t[:MAX_CHARS] if t and len(t) > MAX_CHARS else (t or "") for t in texts]
|
||||
client = self._get_client()
|
||||
for attempt in range(5):
|
||||
try:
|
||||
try:
|
||||
response = client.embeddings(input=truncated_texts)
|
||||
except TypeError:
|
||||
response = client.embeddings(truncated_texts)
|
||||
if response.data and len(response.data) == len(truncated_texts):
|
||||
sorted_embeddings = [None] * len(truncated_texts)
|
||||
for item in response.data:
|
||||
idx = item.index
|
||||
if idx < len(sorted_embeddings):
|
||||
sorted_embeddings[idx] = item.embedding
|
||||
return sorted_embeddings
|
||||
else:
|
||||
logger.error(f"Неполный ответ эмбеддингов: ожидалось {len(truncated_texts)}, получено {len(response.data) if response.data else 0}")
|
||||
return None
|
||||
except RateLimitError as e:
|
||||
if attempt < 4:
|
||||
wait = 2 ** attempt # 1,2,4,8 секунд
|
||||
logger.warning(f"Rate limit, повтор через {wait} секунд")
|
||||
time.sleep(wait) # синхронный sleep, так как мы в синхронной функции
|
||||
else:
|
||||
logger.error("Превышено количество попыток при 429")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка пакетного эмбеддинга: {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
async def embed_batch(self, texts: List[str], retries: int = 3, delay: float = 1.0) -> Optional[List[List[float]]]:
|
||||
if not texts:
|
||||
return []
|
||||
BATCH_SIZE = 50
|
||||
results = []
|
||||
for i in range(0, len(texts), BATCH_SIZE):
|
||||
batch = texts[i:i+BATCH_SIZE]
|
||||
loop = asyncio.get_running_loop()
|
||||
vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch)
|
||||
if vectors is None:
|
||||
for attempt in range(retries):
|
||||
await asyncio.sleep(delay * (2 ** attempt))
|
||||
vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch)
|
||||
if vectors is not None:
|
||||
break
|
||||
else:
|
||||
return None
|
||||
results.extend(vectors)
|
||||
# ДОБАВИТЬ: пауза между батчами
|
||||
await asyncio.sleep(0.5) # 500 мс
|
||||
return results
|
||||
|
||||
async def embed_sparse_batch(self, texts: List[str]) -> List[Optional[Any]]:
|
||||
"""Возвращает список объектов SparseEmbedding для каждого текста"""
|
||||
model = self._get_sparse_model()
|
||||
if model is None:
|
||||
# Возвращаем список из None такой же длины, чтобы не нарушать индексацию
|
||||
return [None] * len(texts)
|
||||
loop = asyncio.get_running_loop()
|
||||
embeddings = await loop.run_in_executor(None, list, model.embed(texts))
|
||||
return embeddings
|
||||
|
||||
async def embed_sparse(self, text: str) -> Optional[Any]:
|
||||
model = self._get_sparse_model()
|
||||
if model is None:
|
||||
return None
|
||||
loop = asyncio.get_running_loop()
|
||||
embeddings = await loop.run_in_executor(None, lambda: list(model.embed([text])))
|
||||
if embeddings and len(embeddings) > 0:
|
||||
return embeddings[0]
|
||||
return None
|
||||
658
rag/services/file_service.py
Normal file
658
rag/services/file_service.py
Normal file
@@ -0,0 +1,658 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис обработки файлов: извлечение текста из различных форматов,
|
||||
хирургическая замена в DOCX, создание DOCX из текста,
|
||||
транскрипция аудио через SaluteSpeech.
|
||||
Все методы синхронные (используются в asyncio.to_thread).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import zipfile
|
||||
from datetime import datetime
|
||||
from typing import List, Tuple, Optional, Dict, Any
|
||||
|
||||
import requests
|
||||
from docx import Document as DocxReader
|
||||
from docx import Document as DocxWriter
|
||||
from docx.shared import Pt
|
||||
|
||||
# Опциональные библиотеки
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
HAS_PDF = True
|
||||
except ImportError:
|
||||
HAS_PDF = False
|
||||
|
||||
try:
|
||||
import openpyxl
|
||||
HAS_EXCEL = True
|
||||
except ImportError:
|
||||
HAS_EXCEL = False
|
||||
|
||||
try:
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
HAS_OCR = True
|
||||
except ImportError:
|
||||
HAS_OCR = False
|
||||
|
||||
try:
|
||||
import py7zr
|
||||
HAS_7Z = True
|
||||
except ImportError:
|
||||
HAS_7Z = False
|
||||
|
||||
try:
|
||||
from pptx import Presentation
|
||||
HAS_PPTX = True
|
||||
except ImportError:
|
||||
HAS_PPTX = False
|
||||
|
||||
# Для чанкинга используем tiktoken, но здесь оставим только извлечение текста,
|
||||
# а чанкинг вынесен в core/utils/text_utils.py
|
||||
from core.utils.config_loader import BotConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileService:
|
||||
def __init__(self, config: BotConfig):
|
||||
self.config = config
|
||||
self.temp_dir = config.temp_dir
|
||||
self.temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def cleanup_old_files(self, days: int = 3):
|
||||
"""Удаляет старые файлы из temp_dir."""
|
||||
now = datetime.now().timestamp()
|
||||
for item in self.temp_dir.iterdir():
|
||||
if item.is_file() and (now - item.stat().st_mtime) > days * 86400:
|
||||
try:
|
||||
item.unlink()
|
||||
logger.info(f"Удалён старый файл: {item.name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка удаления {item.name}: {e}")
|
||||
|
||||
# ---------- Извлечение текста из разных форматов ----------
|
||||
def extract_text_from_docx(self, file_path: str) -> str:
|
||||
"""
|
||||
Извлекает текст из DOCX, включая колонтитулы, таблицы и сноски.
|
||||
Возвращает текст, разделённый переносами строк.
|
||||
"""
|
||||
try:
|
||||
doc = DocxReader(file_path)
|
||||
full_text = []
|
||||
|
||||
# Вспомогательная функция для извлечения текста из XML-элемента
|
||||
def get_xml_text(element):
|
||||
text_nodes = element.xpath('.//*[local-name()="t"]')
|
||||
return "".join([node.text for node in text_nodes if node.text])
|
||||
|
||||
# Колонтитулы
|
||||
for section in doc.sections:
|
||||
for header in [section.header, section.first_page_header, section.even_page_header]:
|
||||
if header:
|
||||
h_text = get_xml_text(header._element)
|
||||
if h_text.strip():
|
||||
full_text.append(f"[КОЛОНТИТУЛ ВЕРХНИЙ]: {h_text}")
|
||||
for footer in [section.footer, section.first_page_footer, section.even_page_footer]:
|
||||
if footer:
|
||||
f_text = get_xml_text(footer._element)
|
||||
if f_text.strip():
|
||||
full_text.append(f"[КОЛОНТИТУЛ НИЖНИЙ]: {f_text}")
|
||||
|
||||
# Тело документа
|
||||
for block in doc.element.body:
|
||||
if block.tag.endswith('p'):
|
||||
text = get_xml_text(block)
|
||||
if text.strip():
|
||||
full_text.append(text)
|
||||
elif block.tag.endswith('tbl'):
|
||||
from docx.table import Table
|
||||
t = Table(block, doc)
|
||||
full_text.append("\n[ТАБЛИЦА]")
|
||||
for row in t.rows:
|
||||
row_data = [get_xml_text(cell._element).strip().replace('\n', ' ') for cell in row.cells]
|
||||
# Удаляем дублирование, если ячейки пустые
|
||||
clean_row = []
|
||||
for item in row_data:
|
||||
if not clean_row or item != clean_row[-1]:
|
||||
clean_row.append(item)
|
||||
full_text.append(" | ".join(clean_row))
|
||||
full_text.append("[КОНЕЦ ТАБЛИЦЫ]\n")
|
||||
|
||||
# Сноски
|
||||
try:
|
||||
for footnote in doc.footnotes:
|
||||
fn_text = get_xml_text(footnote._element)
|
||||
if fn_text.strip():
|
||||
full_text.append(f"[СНОСКА]: {fn_text}")
|
||||
except:
|
||||
pass
|
||||
|
||||
return "\n".join(full_text)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения DOCX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_pdf(self, file_path: str) -> str:
|
||||
"""Извлекает текст из PDF через pypdf."""
|
||||
if not HAS_PDF:
|
||||
return "[Ошибка: библиотека pypdf не установлена]"
|
||||
try:
|
||||
reader = PdfReader(file_path)
|
||||
pages = []
|
||||
for i, page in enumerate(reader.pages):
|
||||
text = page.extract_text()
|
||||
if text.strip():
|
||||
pages.append(f"--- Страница {i+1} ---\n{text}")
|
||||
return "\n".join(pages)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения PDF {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_xlsx(self, file_path: str) -> str:
|
||||
"""Извлекает текст из Excel (XLSX) через openpyxl."""
|
||||
if not HAS_EXCEL:
|
||||
return "[Ошибка: библиотека openpyxl не установлена]"
|
||||
try:
|
||||
wb = openpyxl.load_workbook(file_path, data_only=True)
|
||||
output = []
|
||||
for sheet in wb.sheetnames:
|
||||
ws = wb[sheet]
|
||||
output.append(f"\n--- ЛИСТ: {sheet} ---")
|
||||
for row in ws.iter_rows(values_only=True):
|
||||
if any(c is not None for c in row):
|
||||
output.append(" | ".join(str(c) if c is not None else "" for c in row))
|
||||
return "\n".join(output)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения XLSX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_csv(self, file_path: str) -> str:
|
||||
"""Извлекает текст из CSV."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
reader = csv.reader(f)
|
||||
return "\n".join("\t".join(row) for row in reader)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения CSV {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_txt(self, file_path: str) -> str:
|
||||
"""Извлекает текст из TXT."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения TXT {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_json(self, file_path: str) -> str:
|
||||
"""Извлекает текст из JSON (дамп в виде строки)."""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения JSON {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_image(self, file_path: str) -> str:
|
||||
"""Извлекает текст из изображения через OCR (Tesseract)."""
|
||||
if not HAS_OCR:
|
||||
return ""
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
text = pytesseract.image_to_string(img, lang='rus+eng')
|
||||
return f"--- ТЕКСТ С ИЗОБРАЖЕНИЯ ---\n{text}" if text.strip() else ""
|
||||
except Exception as e:
|
||||
logger.debug(f"OCR не удался: {e}")
|
||||
return ""
|
||||
|
||||
def extract_text_from_pptx(self, file_path: str) -> str:
|
||||
"""
|
||||
Извлекает текст из презентации .pptx:
|
||||
- текст всех слайдов (из текстовых полей, таблиц)
|
||||
- заметки к слайдам
|
||||
Возвращает строку с разделением по слайдам.
|
||||
"""
|
||||
if not HAS_PPTX:
|
||||
return "[Ошибка: библиотека python-pptx не установлена]"
|
||||
try:
|
||||
prs = Presentation(file_path)
|
||||
full_text = []
|
||||
for slide_num, slide in enumerate(prs.slides, 1):
|
||||
slide_content = []
|
||||
# Текст из фигур (надписи, автофигуры)
|
||||
for shape in slide.shapes:
|
||||
if hasattr(shape, "text") and shape.text:
|
||||
slide_content.append(shape.text.strip())
|
||||
# Таблицы
|
||||
if shape.has_table:
|
||||
for row in shape.table.rows:
|
||||
row_text = [cell.text.strip() for cell in row.cells]
|
||||
slide_content.append(" | ".join(row_text))
|
||||
# Заметки к слайду
|
||||
if slide.has_notes_slide and slide.notes_slide.notes_text_frame:
|
||||
notes = slide.notes_slide.notes_text_frame.text.strip()
|
||||
if notes:
|
||||
slide_content.append(f"[ЗАМЕТКИ]: {notes}")
|
||||
if slide_content:
|
||||
full_text.append(f"--- Слайд {slide_num} ---\n" + "\n".join(slide_content))
|
||||
return "\n\n".join(full_text)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка чтения PPTX {file_path}: {e}")
|
||||
return ""
|
||||
|
||||
async def transcribe_audio(self, file_path: str) -> str:
|
||||
"""
|
||||
Асинхронно транскрибирует аудиофайл через SaluteSpeech API.
|
||||
Все блокирующие вызовы (subprocess.run, requests.post) обёрнуты в asyncio.to_thread.
|
||||
Требует наличия ffmpeg и переменной окружения SALUTE_SPEECH_AUTH.
|
||||
"""
|
||||
import uuid
|
||||
|
||||
# 1. Получаем токен асинхронно
|
||||
token = await self._get_salute_token()
|
||||
if not token:
|
||||
return "[Ошибка авторизации в SaluteSpeech]"
|
||||
|
||||
# 2. Проверяем наличие ffmpeg
|
||||
ffmpeg_bin = shutil.which('ffmpeg') or '/usr/local/bin/ffmpeg'
|
||||
curl_bin = shutil.which('curl') or '/usr/local/bin/curl'
|
||||
if not shutil.which('ffmpeg') and not os.path.exists(ffmpeg_bin):
|
||||
return "[Ошибка: на сервере не установлен ffmpeg]"
|
||||
|
||||
pcm_output = file_path + ".raw"
|
||||
|
||||
try:
|
||||
# ---- ОБЁРТКА: конвертация ffmpeg (subprocess.run) ----
|
||||
def _run_ffmpeg():
|
||||
cmd_ffmpeg = [
|
||||
ffmpeg_bin, '-y', '-i', file_path,
|
||||
'-ac', '1', '-ar', '16000', '-f', 's16le', '-filter:a', 'volume=2.0', pcm_output
|
||||
]
|
||||
return subprocess.run(cmd_ffmpeg, capture_output=True)
|
||||
|
||||
result = await asyncio.to_thread(_run_ffmpeg)
|
||||
if result.returncode != 0:
|
||||
err_msg = result.stderr.decode()
|
||||
logger.error(f"FFmpeg ошибка: {err_msg}")
|
||||
return f"[Ошибка конвертации аудио: {err_msg[:100]}]"
|
||||
|
||||
if not os.path.exists(pcm_output) or os.path.getsize(pcm_output) == 0:
|
||||
return "[Ошибка: пустой аудиофайл после конвертации]"
|
||||
|
||||
# 3. Отправка запроса к SaluteSpeech через curl (subprocess.run)
|
||||
url = "https://smartspeech.sber.ru/rest/v1/speech:recognize"
|
||||
|
||||
def _run_curl():
|
||||
curl_cmd = [
|
||||
curl_bin, '-s', '-k', '--http1.1', '-X', 'POST', url,
|
||||
'-H', f'Authorization: Bearer {token}',
|
||||
'-H', 'Content-Type: audio/x-pcm;bit=16;rate=16000',
|
||||
'-H', f'RqUID: {str(uuid.uuid4())}',
|
||||
'--data-binary', f'@{pcm_output}'
|
||||
]
|
||||
return subprocess.run(curl_cmd, capture_output=True, text=True)
|
||||
|
||||
result_curl = await asyncio.to_thread(_run_curl)
|
||||
|
||||
# Удаляем временный PCM-файл
|
||||
if os.path.exists(pcm_output):
|
||||
os.remove(pcm_output)
|
||||
|
||||
if result_curl.returncode != 0:
|
||||
return f"[Ошибка сетевого соединения (curl): {result_curl.returncode}]"
|
||||
|
||||
# 4. Парсинг ответа (синхронная операция – быстрая, можно оставить без обёртки)
|
||||
try:
|
||||
response_data = json.loads(result_curl.stdout)
|
||||
if 'result' in response_data:
|
||||
if response_data['result']:
|
||||
text = " ".join(response_data['result'])
|
||||
return f"--- РАСПОЗНАННЫЙ ГОЛОС ---\n{text}"
|
||||
else:
|
||||
return "😶 Тишина на записи."
|
||||
error_msg = response_data.get('message', 'Неизвестная ошибка API')
|
||||
return f"--- ОШИБКА ГОЛОСА ---\n[Голос не распознан: {error_msg}]"
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка парсинга JSON: {result_curl.stdout}")
|
||||
return f"[Ошибка парсинга ответа: {e}]"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка в transcribe_audio: {e}")
|
||||
return f"[Ошибка распознавания: {e}]"
|
||||
|
||||
async def _get_salute_token(self) -> Optional[str]:
|
||||
"""
|
||||
Асинхронно получает токен доступа к SaluteSpeech по Basic авторизации.
|
||||
Блокирующий вызов requests.post обёрнут в asyncio.to_thread.
|
||||
"""
|
||||
import uuid
|
||||
salute_auth = self.config.salute_speech_auth
|
||||
if not salute_auth:
|
||||
logger.error("Переменная окружения SALUTE_SPEECH_AUTH не задана")
|
||||
return None
|
||||
url = "https://ngw.devices.sberbank.ru.:9443/api/v2/oauth"
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'RqUID': str(uuid.uuid4()),
|
||||
'Authorization': f'Basic {salute_auth}'
|
||||
}
|
||||
payload = {'scope': 'SALUTE_SPEECH_PERS'}
|
||||
try:
|
||||
# ---- ОБЁРТКА: синхронный requests.post выполняется в отдельном потоке ----
|
||||
def _sync_request():
|
||||
return requests.post(url, headers=headers, data=payload, verify=False)
|
||||
response = await asyncio.to_thread(_sync_request)
|
||||
token = response.json().get('access_token')
|
||||
logger.info("SaluteSpeech token получен")
|
||||
return token
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка получения токена SaluteSpeech: {e}")
|
||||
return None
|
||||
|
||||
def process_archive(self, file_path: str) -> Tuple[List[Tuple[str, str]], int]:
|
||||
"""
|
||||
Распаковывает ZIP/7z и рекурсивно обрабатывает каждый файл.
|
||||
Возвращает список (имя_файла, извлечённый_текст) и общее количество обработанных файлов.
|
||||
"""
|
||||
results = []
|
||||
total_count = 0
|
||||
extract_path = self.temp_dir / f"ext_{int(datetime.now().timestamp()*1000000)}"
|
||||
extract_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
max_total_bytes = self.config.max_file_size_mb * 2 * 1024 * 1024 # для архивов лимит удвоен
|
||||
MAX_FILES = self.config.max_archive_files
|
||||
total_files = 0
|
||||
total_size = 0
|
||||
|
||||
try:
|
||||
# ---------- ZIP ----------
|
||||
if file_path.endswith('.zip'):
|
||||
with zipfile.ZipFile(file_path, 'r') as z:
|
||||
# 1. Предварительный подсчёт суммарного несжатого размера всех файлов в архиве
|
||||
total_uncompressed = 0
|
||||
for info in z.infolist():
|
||||
if not info.is_dir():
|
||||
total_uncompressed += info.file_size
|
||||
if total_uncompressed > max_total_bytes:
|
||||
logger.warning(
|
||||
f"Архив {file_path} превышает лимит {max_total_bytes // (1024*1024)} МБ "
|
||||
f"(несжатый размер {total_uncompressed // (1024*1024)} МБ)"
|
||||
)
|
||||
return [], 0
|
||||
|
||||
# 2. Распаковка и обработка
|
||||
for info in z.infolist():
|
||||
if info.is_dir():
|
||||
continue
|
||||
total_files += 1
|
||||
if total_files > MAX_FILES:
|
||||
logger.warning(f"Архив {file_path} содержит более {MAX_FILES} файлов, обработка прервана")
|
||||
return [], 0
|
||||
|
||||
# Имя файла в правильной кодировке
|
||||
try:
|
||||
fname = info.filename.encode('cp437').decode('utf-8')
|
||||
except:
|
||||
try:
|
||||
fname = info.filename.encode('cp437').decode('cp866')
|
||||
except:
|
||||
fname = info.filename
|
||||
|
||||
target = extract_path / os.path.basename(fname)
|
||||
with open(target, 'wb') as f:
|
||||
f.write(z.read(info))
|
||||
|
||||
text, count = self.process_any_file(str(target))
|
||||
if text:
|
||||
size = len(text.encode('utf-8'))
|
||||
if total_size + size > max_total_bytes:
|
||||
logger.warning(f"Архив {file_path} превышает лимит {self.config.max_file_size_mb*2} МБ")
|
||||
return [], 0
|
||||
results.append((fname, text))
|
||||
total_count += count
|
||||
total_size += size
|
||||
|
||||
# ---------- 7Z ----------
|
||||
elif file_path.endswith('.7z') and HAS_7Z:
|
||||
with py7zr.SevenZipFile(file_path, 'r') as z:
|
||||
# Для 7z нет простого способа получить суммарный размер до распаковки,
|
||||
# поэтому распаковываем с контролем в процессе
|
||||
z.extractall(extract_path)
|
||||
for root, dirs, files in os.walk(extract_path):
|
||||
for f in files:
|
||||
total_files += 1
|
||||
if total_files > MAX_FILES:
|
||||
logger.warning(f"Архив {file_path} содержит более {MAX_FILES} файлов, обработка прервана")
|
||||
return [], 0
|
||||
fpath = os.path.join(root, f)
|
||||
text, count = self.process_any_file(fpath)
|
||||
if text:
|
||||
size = len(text.encode('utf-8'))
|
||||
if total_size + size > max_total_bytes:
|
||||
logger.warning(f"Архив {file_path} превышает лимит {self.config.max_file_size_mb*2} МБ")
|
||||
return [], 0
|
||||
results.append((f, text))
|
||||
total_count += count
|
||||
total_size += size
|
||||
else:
|
||||
# Неподдерживаемый формат архива (но по логике сюда не должны попасть)
|
||||
logger.warning(f"Неподдерживаемый тип архива: {file_path}")
|
||||
return [], 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка распаковки архива {file_path}: {e}")
|
||||
return [], 0
|
||||
finally:
|
||||
shutil.rmtree(extract_path, ignore_errors=True)
|
||||
|
||||
return results, total_count
|
||||
|
||||
def process_any_file(self, file_path: str) -> Tuple[str, int]:
|
||||
"""
|
||||
Главный метод: определяет тип файла по расширению и вызывает соответствующий
|
||||
метод извлечения текста. Возвращает (текст, количество_файлов_внутри).
|
||||
Для архивов возвращает список файлов и общее количество.
|
||||
"""
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
text = ""
|
||||
count = 0
|
||||
try:
|
||||
size = os.path.getsize(file_path)
|
||||
max_size = self.config.max_file_size_mb * 1024 * 1024
|
||||
if size > max_size:
|
||||
logger.warning(f"Файл {file_path} превышает лимит {self.config.max_file_size_mb} МБ")
|
||||
return f"[Ошибка: файл превышает максимальный размер {self.config.max_file_size_mb} МБ]", 0
|
||||
if ext == '.docx':
|
||||
text = self.extract_text_from_docx(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.pdf':
|
||||
text = self.extract_text_from_pdf(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.xlsx':
|
||||
text = self.extract_text_from_xlsx(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.csv':
|
||||
text = self.extract_text_from_csv(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.txt':
|
||||
text = self.extract_text_from_txt(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext == '.json':
|
||||
text = self.extract_text_from_json(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'):
|
||||
text = self.extract_text_from_image(file_path)
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'):
|
||||
# Асинхронный метод transcribe_audio запускаем синхронно через get_event_loop
|
||||
# (так как process_any_file выполняется в отдельном потоке)
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
text = loop.run_until_complete(self.transcribe_audio(file_path))
|
||||
count = 1 if text else 0
|
||||
elif ext in ('.zip', '.7z'):
|
||||
return self.process_archive(file_path)
|
||||
elif ext == '.pptx':
|
||||
text = self.extract_text_from_pptx(file_path)
|
||||
count = 1 if text else 0
|
||||
else:
|
||||
logger.warning(f"Неподдерживаемый тип файла: {ext}")
|
||||
return "", 0
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка обработки файла {file_path}: {e}")
|
||||
return "", 0
|
||||
return text, count
|
||||
|
||||
def surgical_replace(self, file_path: str, replacements: Dict[str, str]) -> Optional[str]:
|
||||
"""
|
||||
Умная замена слов в DOCX с учётом грамматики (падеж, число, род).
|
||||
Сохраняет форматирование. При ошибке выполняет простую замену.
|
||||
"""
|
||||
import re
|
||||
from datetime import datetime
|
||||
from docx import Document
|
||||
|
||||
# ---- Попытка умной замены через mawo-pymorphy3 ----
|
||||
try:
|
||||
from mawo_pymorphy3 import create_analyzer
|
||||
morph = create_analyzer()
|
||||
logger.info("Морфологический анализатор загружен")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось загрузить mawo-pymorphy3: {e}. Выполняю простую замену.")
|
||||
# Простая замена
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
for old, new in replacements.items():
|
||||
if old in run.text:
|
||||
run.text = run.text.replace(old, new)
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Простая замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
except Exception as e2:
|
||||
logger.error(f"Ошибка простой замены: {e2}")
|
||||
return None
|
||||
|
||||
# ---- Умная замена ----
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
logger.info(f"Документ загружен, параграфов: {len(doc.paragraphs)}")
|
||||
|
||||
# Кэш для разобранных слов
|
||||
parse_cache = {}
|
||||
|
||||
# Предвычисление форм старых слов
|
||||
old_forms = {}
|
||||
for old_word, new_word in replacements.items():
|
||||
old_low = old_word.lower()
|
||||
parsed = morph.parse(old_low)[0]
|
||||
forms = {parsed.normal_form}
|
||||
cases = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
|
||||
numbers = ['sing', 'plur']
|
||||
for number in numbers:
|
||||
for case in cases:
|
||||
inflected = parsed.inflect({case, number})
|
||||
if inflected:
|
||||
forms.add(inflected.word)
|
||||
old_forms[old_low] = (new_word, forms)
|
||||
|
||||
def get_new_form(new_word, old_form_text):
|
||||
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()
|
||||
if parsed_old.tag.case:
|
||||
tags.add(parsed_old.tag.case)
|
||||
if parsed_old.tag.number:
|
||||
tags.add(parsed_old.tag.number)
|
||||
if parsed_old.tag.gender:
|
||||
tags.add(parsed_old.tag.gender)
|
||||
parsed_new = morph.parse(new_word)[0]
|
||||
inflected = parsed_new.inflect(tags)
|
||||
return inflected.word if inflected else new_word
|
||||
|
||||
total_runs = 0
|
||||
modified_runs = 0
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
total_runs += 1
|
||||
original = run.text
|
||||
if not original:
|
||||
continue
|
||||
tokens = re.split(r'(\W+)', original)
|
||||
modified = False
|
||||
for i, token in enumerate(tokens):
|
||||
if token.isalpha():
|
||||
token_low = token.lower()
|
||||
for old_low, (new_word, forms) in old_forms.items():
|
||||
if token_low in forms:
|
||||
new_form = get_new_form(new_word, token)
|
||||
tokens[i] = new_form
|
||||
modified = True
|
||||
break
|
||||
if modified:
|
||||
run.text = ''.join(tokens)
|
||||
modified_runs += 1
|
||||
|
||||
logger.info(f"Умная замена: обработано {total_runs} runs, изменено {modified_runs}")
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Умная замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при умной замене: {e}. Выполняю простую замену.")
|
||||
# Простая замена (fallback)
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
new_path = self.temp_dir / f"surgical_{int(datetime.now().timestamp())}.docx"
|
||||
for paragraph in doc.paragraphs:
|
||||
for run in paragraph.runs:
|
||||
for old, new in replacements.items():
|
||||
if old in run.text:
|
||||
run.text = run.text.replace(old, new)
|
||||
doc.save(str(new_path))
|
||||
logger.info(f"Простая замена выполнена: {new_path}")
|
||||
return str(new_path)
|
||||
except Exception as e2:
|
||||
logger.error(f"Ошибка простой замены: {e2}")
|
||||
return None
|
||||
|
||||
def create_docx(self, content: str, title: str, user_name: str, mode: str = "CREATE") -> Optional[str]:
|
||||
"""
|
||||
Создаёт новый DOCX из текста и сохраняет во временную папку.
|
||||
Возвращает путь к созданному файлу или None.
|
||||
"""
|
||||
try:
|
||||
fname = f"output_{user_name}_{int(datetime.now().timestamp())}.docx"
|
||||
target = self.temp_dir / fname
|
||||
doc = DocxWriter()
|
||||
doc.add_heading(title, 0)
|
||||
for paragraph in content.split('\n'):
|
||||
if paragraph.strip():
|
||||
p = doc.add_paragraph(paragraph.strip())
|
||||
p.runs[0].font.size = Pt(11)
|
||||
doc.save(str(target))
|
||||
return str(target)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка создания DOCX: {e}")
|
||||
return None
|
||||
109
rag/services/giga_client.py
Normal file
109
rag/services/giga_client.py
Normal file
@@ -0,0 +1,109 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Простой клиент для GigaChat API.
|
||||
Только вызовы chat и upload_file, без бизнес-логики.
|
||||
Используется AIService и другими модулями.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
from gigachat import GigaChat
|
||||
from gigachat.models import Chat, Messages, MessagesRole
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GigaClient:
|
||||
def __init__(self, api_key: str, model: str, temperature: float, timeout: int, verify_ssl: bool = False):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.temperature = temperature
|
||||
self.timeout = timeout
|
||||
self.verify_ssl = verify_ssl
|
||||
self._client = None
|
||||
|
||||
def _get_client(self) -> GigaChat:
|
||||
"""Ленивое создание клиента (создаётся при первом запросе)."""
|
||||
if self._client is None:
|
||||
self._client = GigaChat(
|
||||
credentials=self.api_key,
|
||||
model=self.model,
|
||||
verify_ssl_certs=self.verify_ssl,
|
||||
timeout=self.timeout
|
||||
)
|
||||
logger.debug("GigaChat client initialized")
|
||||
return self._client
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
history: List[Dict[str, str]],
|
||||
query: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
file_id: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Отправляет запрос к GigaChat и возвращает ответ.
|
||||
|
||||
Аргументы:
|
||||
history: список предыдущих сообщений (каждое с ключами role и content).
|
||||
query: текущий запрос пользователя.
|
||||
system_prompt: системный промпт (если задан, добавляется первым сообщением).
|
||||
file_id: идентификатор загруженного файла (для Vision).
|
||||
|
||||
Возвращает:
|
||||
str: ответ модели или сообщение об ошибке.
|
||||
"""
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append(Messages(role=MessagesRole.SYSTEM, content=system_prompt))
|
||||
for msg in history[-10:]: # ограничим историю 10 сообщениями для экономии токенов
|
||||
role = MessagesRole.USER if msg['role'] == 'user' else MessagesRole.ASSISTANT
|
||||
messages.append(Messages(role=role, content=msg['content']))
|
||||
messages.append(Messages(role=MessagesRole.USER, content=query))
|
||||
|
||||
attachments = [file_id] if file_id else None
|
||||
|
||||
actual_temp = temperature if temperature is not None else self.temperature
|
||||
|
||||
try:
|
||||
# Запускаем синхронный вызов giga.chat в отдельном потоке
|
||||
response = await asyncio.to_thread(
|
||||
self._get_client().chat,
|
||||
Chat(
|
||||
messages=messages,
|
||||
model=self.model,
|
||||
temperature=actual_temp,
|
||||
n=1,
|
||||
attachments=attachments
|
||||
)
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при вызове GigaChat: {e}", exc_info=True)
|
||||
return "Извините, произошла ошибка при обращении к сервису ИИ."
|
||||
|
||||
async def upload_file(self, file_path: str) -> Optional[str]:
|
||||
"""
|
||||
Загружает файл в GigaChat для Vision-анализа (изображения).
|
||||
Возвращает file_id для использования в chat().
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
# Загружаем файл синхронно в потоке
|
||||
uploaded = await asyncio.to_thread(self._get_client().upload_file, f, "general")
|
||||
# В разных версиях SDK поле может называться id_, id или file_id
|
||||
if hasattr(uploaded, "id_"):
|
||||
return uploaded.id_
|
||||
if hasattr(uploaded, "id"):
|
||||
return uploaded.id
|
||||
if hasattr(uploaded, "file_id"):
|
||||
return uploaded.file_id
|
||||
# fallback: пробуем преобразовать в словарь
|
||||
data = uploaded.dict() if hasattr(uploaded, "dict") else {}
|
||||
return data.get("id_") or data.get("id") or data.get("file_id")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка загрузки файла в GigaChat: {e}", exc_info=True)
|
||||
return None
|
||||
426
rag/services/kb_service.py
Normal file
426
rag/services/kb_service.py
Normal file
@@ -0,0 +1,426 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис базы знаний (Knowledge Base Service).
|
||||
Реализует RAG-пайплайн:
|
||||
- Индексация документов (разбиение на чанки, получение эмбеддингов, сохранение в Qdrant и PostgreSQL).
|
||||
- Поиск релевантной информации по запросу (семантический поиск в Qdrant с фильтром по room_jid).
|
||||
- Управление доступом к документам (личные / глобальные / комнатные).
|
||||
- Поддержка обновления версий документов в комнатах (замена по хешу).
|
||||
|
||||
ВНИМАНИЕ: Поиск работает в двух режимах:
|
||||
1. Семантический (dense) — для общих вопросов, использует эмбеддинги GigaChat
|
||||
2. Полнотекстовый (keyword) — для точных запросов (пункт 1.1, фразы в кавычках)
|
||||
|
||||
Использует PostgreSQL (asyncpg) через PostgresService и Qdrant через QdrantService.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
|
||||
# Импорт сервисов
|
||||
from core.services.postgres_service import PostgresService
|
||||
from core.services.qdrant_service import QdrantService
|
||||
from core.services.embedding_service import EmbeddingService
|
||||
from core.utils.text_utils import split_into_chunks
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KBService:
|
||||
"""
|
||||
База знаний: индексация документов, поиск, управление.
|
||||
Абстрагирует работу с PostgreSQL и Qdrant для RAG.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: PostgresService,
|
||||
qdrant: QdrantService,
|
||||
embedding: EmbeddingService,
|
||||
collection_name: str,
|
||||
chunk_size_tokens: int,
|
||||
overlap_tokens: int,
|
||||
approx_chunk_chars: int,
|
||||
approx_overlap_chars: int,
|
||||
):
|
||||
"""
|
||||
Инициализация KBService.
|
||||
|
||||
Аргументы:
|
||||
db: сервис PostgreSQL (asyncpg)
|
||||
qdrant: сервис Qdrant
|
||||
embedding: сервис эмбеддингов GigaChat
|
||||
collection_name: имя коллекции в Qdrant
|
||||
chunk_size_tokens: размер чанка в токенах (для tiktoken)
|
||||
overlap_tokens: перекрытие между чанками в токенах
|
||||
approx_chunk_chars: запасной размер чанка в символах (если tiktoken недоступен)
|
||||
approx_overlap_chars: запасное перекрытие в символах
|
||||
"""
|
||||
self.db = db
|
||||
self.qdrant = qdrant
|
||||
self.embedding = embedding
|
||||
self.collection_name = collection_name
|
||||
self.chunk_size_tokens = chunk_size_tokens
|
||||
self.overlap_tokens = overlap_tokens
|
||||
self.approx_chunk_chars = approx_chunk_chars
|
||||
self.approx_overlap_chars = approx_overlap_chars
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Поиск релевантной информации
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str:
|
||||
"""
|
||||
Находит релевантные фрагменты в БЗ.
|
||||
|
||||
Режимы поиска:
|
||||
- Если запрос содержит номер пункта (например, "1.1") или кавычки —
|
||||
используется полнотекстовый поиск (search_keywords)
|
||||
- Иначе — семантический поиск (dense через GigaChat эмбеддинги)
|
||||
|
||||
При ошибке эмбеддинга или пустом результате использует fallback-поиск по source_name.
|
||||
|
||||
Аргументы:
|
||||
query: поисковый запрос пользователя
|
||||
user_jid: JID пользователя (для фильтрации личных документов)
|
||||
room_jid: JID комнаты (None для личного чата) – для фильтрации комнатных документов
|
||||
top_k: количество возвращаемых чанков
|
||||
|
||||
Возвращает:
|
||||
str: строка с объединёнными фрагментами вида "[источник: имя]\nтекст"
|
||||
"""
|
||||
try:
|
||||
# Определяем, является ли запрос "точным"
|
||||
# Точный запрос: содержит номер пункта (1.1, 2.3.4) или фразу в кавычках
|
||||
is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query
|
||||
|
||||
if is_exact_query:
|
||||
# РЕЖИМ 1: Полнотекстовый поиск для точных запросов
|
||||
logger.debug(f"Используем полнотекстовый поиск для запроса: {query}")
|
||||
search_results = self.qdrant.search_keywords(query, user_jid, room_jid, top_k=top_k)
|
||||
else:
|
||||
# Получаем dense-вектор запроса
|
||||
query_vector = await self.embedding.embed(query)
|
||||
if query_vector is None:
|
||||
logger.error("Не удалось получить dense-эмбеддинг запроса, переключаемся на fallback")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
# Получаем sparse-вектор запроса (опционально, может быть None)
|
||||
sparse_vector = await self.embedding.embed_sparse(query)
|
||||
if sparse_vector is None:
|
||||
logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск")
|
||||
|
||||
# Гибридный поиск
|
||||
search_results = self.qdrant.search_hybrid(
|
||||
dense_vector=query_vector,
|
||||
sparse_vector=sparse_vector,
|
||||
user_jid=user_jid,
|
||||
room_jid=room_jid,
|
||||
top_k=top_k
|
||||
)
|
||||
|
||||
if not search_results:
|
||||
logger.info(f"По запросу '{query[:50]}' ничего не найдено в Qdrant")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
# Формируем контекст из найденных чанков
|
||||
context_parts = []
|
||||
for res in search_results:
|
||||
payload = res.get('payload', {})
|
||||
source_name = payload.get('source_name', 'Неизвестный источник')
|
||||
chunk_text = payload.get('text', '')
|
||||
if chunk_text:
|
||||
context_parts.append(f"[источник: {source_name}]\n{chunk_text}")
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
logger.debug(f"Сформирован контекст длиной {len(context)} символов")
|
||||
return context
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Ошибка при поиске в RAG: {e}")
|
||||
return await self._fallback_search(query, user_jid, room_jid)
|
||||
|
||||
async def _fallback_search(self, query: str, user_jid: str, room_jid: str = None) -> str:
|
||||
"""
|
||||
Полнотекстовый поиск по source_name в PostgreSQL.
|
||||
Используется, если Qdrant недоступен или не вернул результатов.
|
||||
|
||||
Аргументы:
|
||||
query: исходный запрос
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
str: строка с информацией о найденных документах
|
||||
"""
|
||||
if not query or len(query.strip()) < 3:
|
||||
return ""
|
||||
|
||||
# Очищаем запрос от спецсимволов для полнотекстового поиска
|
||||
clean_query = re.sub(r'[^\w\s]', ' ', query)
|
||||
|
||||
try:
|
||||
rows = await self.db.fulltext_search(clean_query, user_jid, room_jid)
|
||||
if not rows:
|
||||
return ""
|
||||
|
||||
parts = [
|
||||
f"[источник: {row['source_name']}]\n(Документ доступен, но текст не проиндексирован в Qdrant. Повторите запрос позже.)"
|
||||
for row in rows
|
||||
]
|
||||
return "\n\n".join(parts)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при fallback поиске: {e}")
|
||||
return ""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Индексация документов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_document(
|
||||
self,
|
||||
file_name: str,
|
||||
file_text: str,
|
||||
user_jid: str,
|
||||
is_global: bool,
|
||||
title: Optional[str] = None,
|
||||
metadata: Optional[Dict] = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None,
|
||||
update_if_exists: bool = True
|
||||
) -> Tuple[int, int]:
|
||||
"""
|
||||
Индексирует документ в базу знаний.
|
||||
|
||||
Этапы индексации:
|
||||
1. Разбивает текст на чанки
|
||||
2. Получает эмбеддинги для каждого чанка через GigaChat
|
||||
3. Сохраняет метаданные в PostgreSQL
|
||||
4. Сохраняет векторы и payload в Qdrant
|
||||
5. Если update_if_exists=True и room_jid задан, проверяет наличие документа
|
||||
с таким же source_name в комнате и при изменении хеша заменяет старую версию.
|
||||
|
||||
Аргументы:
|
||||
file_name: исходное имя файла (или URL)
|
||||
file_text: извлечённый текст документа
|
||||
user_jid: JID владельца документа
|
||||
is_global: глобальный ли документ
|
||||
title: отображаемое название (если не указано, используется file_name)
|
||||
metadata: дополнительные метаданные (сохраняются в JSONB)
|
||||
room_jid: JID комнаты (если документ привязан к комнате)
|
||||
file_hash: SHA-256 хеш содержимого (для проверки изменений)
|
||||
update_if_exists: заменять ли существующий документ в комнате при изменении хеша
|
||||
|
||||
Возвращает:
|
||||
Tuple[int, int]: (doc_id, количество_проиндексированных_чанков)
|
||||
"""
|
||||
source_name = title if title else file_name
|
||||
doc_meta = metadata or {}
|
||||
if file_hash:
|
||||
doc_meta['file_hash'] = file_hash
|
||||
|
||||
# --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) ---
|
||||
if update_if_exists and room_jid and not is_global:
|
||||
existing = await self.db.find_document_by_name_and_room(source_name, room_jid)
|
||||
if existing:
|
||||
old_doc_id = existing['id']
|
||||
old_hash = existing.get('file_hash')
|
||||
if file_hash and old_hash == file_hash:
|
||||
logger.info(f"Документ {source_name} не изменился (хеш совпадает), пропускаем индексацию")
|
||||
return old_doc_id, 0
|
||||
# Удаляем старый документ из Qdrant и PostgreSQL
|
||||
logger.info(f"Обновление документа {source_name} в комнате {room_jid}: удаляем старую версию doc_id={old_doc_id}")
|
||||
self.qdrant.delete_by_doc_id(old_doc_id)
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM document_access WHERE doc_id = $1", old_doc_id)
|
||||
await conn.execute("DELETE FROM documents WHERE id = $1", old_doc_id)
|
||||
|
||||
# --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL ---
|
||||
doc_id = await self.db.add_document(
|
||||
source_name=source_name,
|
||||
owner_jid=user_jid,
|
||||
is_global=is_global,
|
||||
collection_name=self.collection_name,
|
||||
metadata=doc_meta,
|
||||
room_jid=room_jid,
|
||||
file_hash=file_hash
|
||||
)
|
||||
|
||||
# --- РАЗБИЕНИЕ НА ЧАНКИ ---
|
||||
chunks = split_into_chunks(
|
||||
file_text,
|
||||
chunk_size_tokens=self.chunk_size_tokens,
|
||||
overlap_tokens=self.overlap_tokens,
|
||||
approx_chunk_chars=self.approx_chunk_chars,
|
||||
approx_overlap_chars=self.approx_overlap_chars,
|
||||
strategy="recursive_split_by_sentences"
|
||||
)
|
||||
if not chunks:
|
||||
return doc_id, 0
|
||||
|
||||
# --- ПАКЕТНОЕ ПОЛУЧЕНИЕ ЭМБЕДДИНГОВ (dense + sparse) ---
|
||||
dense_vectors = await self.embedding.embed_batch(chunks)
|
||||
if dense_vectors is None or len(dense_vectors) != len(chunks):
|
||||
logger.error(f"Не удалось получить dense-эмбеддинги для документа {doc_id}")
|
||||
return doc_id, 0
|
||||
|
||||
sparse_vectors = await self.embedding.embed_sparse_batch(chunks)
|
||||
# sparse_vectors может быть None или короче, но мы продолжим (sparse опционально)
|
||||
|
||||
# --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT ---
|
||||
chunk_count = 0
|
||||
for idx, chunk in enumerate(chunks):
|
||||
dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None
|
||||
if dense_vec is None:
|
||||
logger.error(f"Нет dense-вектора для чанка {idx}, пропускаем")
|
||||
continue
|
||||
|
||||
sparse_vec = None
|
||||
if sparse_vectors and idx < len(sparse_vectors):
|
||||
sparse_vec = sparse_vectors[idx]
|
||||
|
||||
payload = {
|
||||
"doc_id": doc_id,
|
||||
"user_jid": user_jid,
|
||||
"is_global": is_global,
|
||||
"chunk_index": idx,
|
||||
"text": chunk,
|
||||
"source_name": source_name,
|
||||
"room_jid": room_jid if room_jid else "",
|
||||
}
|
||||
await asyncio.to_thread(
|
||||
self.qdrant.add_chunk,
|
||||
dense=dense_vec,
|
||||
sparse=sparse_vec,
|
||||
payload=payload
|
||||
)
|
||||
chunk_count += 1
|
||||
|
||||
# --- ДЛЯ ЛИЧНЫХ ДОКУМЕНТОВ ПРЕДОСТАВЛЯЕМ ДОСТУП ВЛАДЕЛЬЦУ ---
|
||||
if not is_global and not room_jid:
|
||||
await self.db.add_access(doc_id, user_jid)
|
||||
|
||||
logger.info(f"Документ {source_name} (doc_id={doc_id}) проиндексирован, чанков: {chunk_count}")
|
||||
return doc_id, chunk_count
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Получение списков документов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
|
||||
"""
|
||||
Возвращает список названий документов (source_name), доступных пользователю в данном контексте.
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
List[str]: список названий
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name FROM documents WHERE room_jid = $1 ORDER BY source_name",
|
||||
room_jid
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
|
||||
user_jid
|
||||
)
|
||||
return [row['source_name'] for row in rows]
|
||||
|
||||
async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Возвращает список документов с указанием, является ли документ глобальным.
|
||||
Используется командой !kb для раздельного отображения глобальных и локальных документов.
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
|
||||
Возвращает:
|
||||
List[Dict[str, Any]]: список словарей [{"source_name": ..., "is_global": bool}]
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name, is_global FROM documents WHERE room_jid = $1 ORDER BY source_name",
|
||||
room_jid
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
|
||||
user_jid
|
||||
)
|
||||
return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Очистка баз знаний
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def clear_user_kb(self, user_jid: str) -> None:
|
||||
"""
|
||||
Удаляет все личные документы пользователя (не глобальные и не комнатные).
|
||||
|
||||
Аргументы:
|
||||
user_jid: JID пользователя
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
|
||||
user_jid
|
||||
)
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
# Удаляем из Qdrant
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
# Удаляем из PostgreSQL (используем ANY для массового удаления)
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info(f"Личная БЗ пользователя {user_jid} очищена (документов: {len(doc_ids)})")
|
||||
|
||||
async def clear_room_kb(self, room_jid: str) -> None:
|
||||
"""
|
||||
Удаляет все документы указанной комнаты.
|
||||
|
||||
Аргументы:
|
||||
room_jid: JID комнаты
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})")
|
||||
|
||||
async def clear_global_kb(self) -> None:
|
||||
"""
|
||||
Удаляет все глобальные документы.
|
||||
"""
|
||||
async with self.db.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
|
||||
doc_ids = [row['id'] for row in rows]
|
||||
|
||||
for doc_id in doc_ids:
|
||||
self.qdrant.delete_by_doc_id(doc_id)
|
||||
|
||||
if doc_ids:
|
||||
async with self.db.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||||
|
||||
logger.info("Глобальная БЗ очищена")
|
||||
356
rag/services/postgres_service.py
Normal file
356
rag/services/postgres_service.py
Normal file
@@ -0,0 +1,356 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис для работы с PostgreSQL (через asyncpg).
|
||||
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
|
||||
управление доступом к личным документам.
|
||||
|
||||
Полностью заменяет MySQLService. Реализованы все методы, необходимые для KBService и других компонентов.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
import asyncpg
|
||||
from asyncpg import Pool, Connection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostgresService:
|
||||
"""
|
||||
Асинхронный сервис для работы с PostgreSQL.
|
||||
Использует пул соединений asyncpg.
|
||||
"""
|
||||
|
||||
def __init__(self, host: str, port: int, user: str, password: str, db_name: str):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.db_name = db_name
|
||||
self.pool: Optional[Pool] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Подключение и инициализация таблиц
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""
|
||||
Создаёт пул соединений и инициализирует таблицы, если их нет.
|
||||
"""
|
||||
self.pool = await asyncpg.create_pool(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
user=self.user,
|
||||
password=self.password,
|
||||
database=self.db_name,
|
||||
min_size=1,
|
||||
max_size=10,
|
||||
command_timeout=60
|
||||
)
|
||||
await self._init_tables()
|
||||
logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")
|
||||
|
||||
async def _init_tables(self) -> None:
|
||||
"""
|
||||
Создаёт таблицы (если не существуют) с учётом всех необходимых колонок и индексов.
|
||||
Использует схему из schema_postgresql.sql (адаптированную для PostgreSQL).
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
# 1. Таблица documents
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id SERIAL PRIMARY KEY,
|
||||
source_name VARCHAR(255) NOT NULL,
|
||||
upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
owner_jid VARCHAR(255) NOT NULL,
|
||||
is_global BOOLEAN DEFAULT FALSE,
|
||||
collection_name VARCHAR(64) NOT NULL,
|
||||
metadata JSONB,
|
||||
room_jid VARCHAR(255) DEFAULT NULL,
|
||||
file_hash VARCHAR(64) DEFAULT NULL
|
||||
)
|
||||
""")
|
||||
# Индексы
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner ON documents(owner_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_global ON documents(is_global)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room ON documents(room_jid)")
|
||||
# Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS ft_source_name
|
||||
ON documents USING GIN (to_tsvector('russian', source_name))
|
||||
""")
|
||||
|
||||
# 2. Таблица document_access
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_access (
|
||||
doc_id INT,
|
||||
user_jid VARCHAR(255),
|
||||
PRIMARY KEY (doc_id, user_jid)
|
||||
)
|
||||
""")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_access_user ON document_access(user_jid)")
|
||||
|
||||
# 3. Тип ENUM для history.role (если не существует)
|
||||
await conn.execute("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE history_role AS ENUM ('user', 'assistant');
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
""")
|
||||
# Таблица history
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS history (
|
||||
id SERIAL PRIMARY KEY,
|
||||
user_jid VARCHAR(255) NOT NULL,
|
||||
role history_role NOT NULL,
|
||||
content TEXT,
|
||||
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
room_jid VARCHAR(255) DEFAULT NULL
|
||||
)
|
||||
""")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_user ON history(user_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_room ON history(room_jid)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_time ON history(timestamp)")
|
||||
|
||||
# 4. Таблица rooms
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS rooms (
|
||||
room_jid VARCHAR(255) PRIMARY KEY
|
||||
)
|
||||
""")
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Закрывает пул соединений."""
|
||||
if self.pool:
|
||||
await self.pool.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Документы
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_document(
|
||||
self,
|
||||
source_name: str,
|
||||
owner_jid: str,
|
||||
is_global: bool,
|
||||
collection_name: str,
|
||||
metadata: dict = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None
|
||||
) -> int:
|
||||
"""
|
||||
Добавляет запись о документе. Возвращает doc_id.
|
||||
|
||||
:param source_name: название документа (source_name)
|
||||
:param owner_jid: JID владельца
|
||||
:param is_global: глобальный ли документ
|
||||
:param collection_name: имя коллекции Qdrant
|
||||
:param metadata: произвольные метаданные (словарь)
|
||||
:param room_jid: JID комнаты (если документ комнатный)
|
||||
:param file_hash: SHA-256 хеш содержимого (для обновления версий)
|
||||
:return: id новой записи
|
||||
"""
|
||||
metadata_json = json.dumps(metadata, ensure_ascii=False) if metadata else None
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
INSERT INTO documents (source_name, owner_jid, is_global, collection_name, metadata, room_jid, file_hash)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id
|
||||
""",
|
||||
source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash
|
||||
)
|
||||
return row['id']
|
||||
|
||||
async def get_document(self, doc_id: int) -> Optional[Dict]:
|
||||
"""Возвращает запись документа по id."""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow("SELECT * FROM documents WHERE id = $1", doc_id)
|
||||
return dict(row) if row else None
|
||||
|
||||
async def add_access(self, doc_id: int, user_jid: str) -> None:
|
||||
"""
|
||||
Даёт пользователю доступ к личному документу (только для не-глобальных, не-комнатных).
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO document_access (doc_id, user_jid)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (doc_id, user_jid) DO NOTHING
|
||||
""",
|
||||
doc_id, user_jid
|
||||
)
|
||||
|
||||
async def can_access(self, doc_id: int, user_jid: str, is_global: bool) -> bool:
|
||||
"""Проверяет, имеет ли пользователь доступ к документу."""
|
||||
if is_global:
|
||||
return True
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT 1 FROM document_access WHERE doc_id = $1 AND user_jid = $2",
|
||||
doc_id, user_jid
|
||||
)
|
||||
return row is not None
|
||||
|
||||
async def find_document_by_name_and_room(self, source_name: str, room_jid: str) -> Optional[Dict]:
|
||||
"""
|
||||
Поиск документа в комнате по имени (source_name). Используется для обновления версий.
|
||||
Возвращает словарь с ключами id, file_hash или None.
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT id, file_hash FROM documents WHERE source_name = $1 AND room_jid = $2 LIMIT 1",
|
||||
source_name, room_jid
|
||||
)
|
||||
return dict(row) if row else None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# История диалогов
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_history(self, user_jid: str, role: str, content: str, room_jid: str = None) -> None:
|
||||
"""
|
||||
Сохраняет сообщение в историю.
|
||||
:param role: 'user' или 'assistant'
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"INSERT INTO history (user_jid, role, content, room_jid) VALUES ($1, $2, $3, $4)",
|
||||
user_jid, role, content, room_jid
|
||||
)
|
||||
|
||||
async def get_history(self, user_jid: str, room_jid: str = None, limit: int = 100) -> List[Dict]:
|
||||
"""
|
||||
Возвращает последние limit сообщений, отсортированных от старых к новым.
|
||||
Если room_jid указан – история комнаты, иначе – личная история пользователя.
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT role, content, timestamp
|
||||
FROM history
|
||||
WHERE room_jid = $1
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT $2
|
||||
""",
|
||||
room_jid, limit
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT role, content, timestamp
|
||||
FROM history
|
||||
WHERE user_jid = $1 AND room_jid IS NULL
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT $2
|
||||
""",
|
||||
user_jid, limit
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
async def clear_user_history(self, user_jid: str) -> None:
|
||||
"""Удаляет личную историю пользователя (не трогает комнаты)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM history WHERE user_jid = $1 AND room_jid IS NULL", user_jid)
|
||||
|
||||
async def clear_room_history(self, room_jid: str) -> None:
|
||||
"""Удаляет историю указанной комнаты."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM history WHERE room_jid = $1", room_jid)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Полнотекстовый поиск (fallback, если Qdrant недоступен)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
|
||||
"""
|
||||
Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
|
||||
Возвращает список документов (id, source_name, metadata).
|
||||
"""
|
||||
if not query or len(query.strip()) < 3:
|
||||
return []
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, source_name, metadata
|
||||
FROM documents
|
||||
WHERE room_jid = $1
|
||||
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||
LIMIT 10
|
||||
""",
|
||||
room_jid, query
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, source_name, metadata
|
||||
FROM documents
|
||||
WHERE (owner_jid = $1 OR is_global = TRUE)
|
||||
AND room_jid IS NULL
|
||||
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||
LIMIT 10
|
||||
""",
|
||||
user_jid, query
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Комнаты MUC
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def add_room(self, room_jid: str) -> None:
|
||||
"""Добавляет комнату в список (если ещё не добавлена)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"INSERT INTO rooms (room_jid) VALUES ($1) ON CONFLICT (room_jid) DO NOTHING",
|
||||
room_jid
|
||||
)
|
||||
|
||||
async def get_rooms(self) -> List[str]:
|
||||
"""Возвращает список JID комнат, в которых бот должен состоять."""
|
||||
async with self.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT room_jid FROM rooms")
|
||||
return [row['room_jid'] for row in rows]
|
||||
|
||||
async def remove_room(self, room_jid: str) -> None:
|
||||
"""Удаляет комнату из списка."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM rooms WHERE room_jid = $1", room_jid)
|
||||
|
||||
async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
|
||||
"""Сохраняет шаблон в БД (если с таким именем уже есть – перезаписывает)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
INSERT INTO room_templates (room_jid, name, file_path, file_hash)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (room_jid, name) DO UPDATE SET
|
||||
file_path = EXCLUDED.file_path,
|
||||
file_hash = EXCLUDED.file_hash,
|
||||
created_at = CURRENT_TIMESTAMP
|
||||
""", room_jid, name, file_path, file_hash)
|
||||
return True
|
||||
|
||||
async def get_template(self, room_jid: str, name: str) -> Optional[Dict]:
|
||||
"""Возвращает словарь {file_path, file_hash} или None."""
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow("SELECT file_path, file_hash FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||
return dict(row) if row else None
|
||||
|
||||
async def list_templates(self, room_jid: str) -> List[str]:
|
||||
"""Возвращает список названий шаблонов в комнате."""
|
||||
async with self.pool.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT name FROM room_templates WHERE room_jid = $1 ORDER BY name", room_jid)
|
||||
return [row['name'] for row in rows]
|
||||
|
||||
async def delete_template(self, room_jid: str, name: str) -> bool:
|
||||
"""Удаляет шаблон. Возвращает True, если удалён."""
|
||||
async with self.pool.acquire() as conn:
|
||||
result = await conn.execute("DELETE FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||
return result == "DELETE 1"
|
||||
452
rag/services/qdrant_service.py
Normal file
452
rag/services/qdrant_service.py
Normal file
@@ -0,0 +1,452 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис для взаимодействия с векторной базой данных Qdrant.
|
||||
Обеспечивает создание гибридной коллекции (dense + sparse), добавление точек,
|
||||
гибридный поиск с RRF, полнотекстовый поиск и удаление по doc_id.
|
||||
|
||||
ВНИМАНИЕ: Этот сервис НЕ использует генерацию эмбеддингов на лету через models.Document.
|
||||
Вместо этого он получает готовые dense-векторы из EmbeddingService (GigaChat)
|
||||
и sparse-векторы из FastEmbed (Qdrant/bm25).
|
||||
|
||||
ВАЖНО: В версиях qdrant-client 1.10+ метод search() заменён на query_points().
|
||||
В этом файле используется правильный API для вашей версии.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from qdrant_client import QdrantClient, models
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QdrantService:
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
grpc_port: int,
|
||||
collection_name: str,
|
||||
vector_size: int,
|
||||
distance: str,
|
||||
prefer_grpc: bool = False
|
||||
):
|
||||
"""
|
||||
Инициализация сервиса Qdrant.
|
||||
|
||||
Аргументы:
|
||||
host: хост Qdrant (обычно localhost)
|
||||
port: HTTP порт Qdrant (обычно 6333)
|
||||
grpc_port: gRPC порт Qdrant (обычно 6334)
|
||||
collection_name: имя коллекции в Qdrant
|
||||
vector_size: размерность плотных векторов (1024 для GigaChat Embeddings)
|
||||
distance: метрика расстояния (Cosine или Euclid)
|
||||
prefer_grpc: использовать ли gRPC вместо HTTP (по умолчанию False)
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.collection_name = collection_name
|
||||
self.vector_size = vector_size
|
||||
self.distance = models.Distance.COSINE if distance.lower() == "cosine" else models.Distance.EUCLID
|
||||
|
||||
# Имена векторных полей в коллекции
|
||||
self.dense_vector_name = "dense"
|
||||
self.sparse_vector_name = "sparse"
|
||||
|
||||
# Создаём клиент Qdrant
|
||||
self.client = QdrantClient(host=host, port=port, grpc_port=grpc_port, prefer_grpc=prefer_grpc)
|
||||
self._ensure_collection()
|
||||
logger.info(f"QdrantService инициализирован: коллекция {collection_name}, размер dense {vector_size}")
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""
|
||||
Проверяет существование коллекции. Если нет — создаёт её с поддержкой dense и sparse векторов.
|
||||
"""
|
||||
if not self.client.collection_exists(self.collection_name):
|
||||
logger.info(f"Коллекция {self.collection_name} не найдена. Создаю новую (dense + sparse)...")
|
||||
try:
|
||||
# Конфигурация dense-векторов (семантический поиск)
|
||||
vectors_config = {
|
||||
self.dense_vector_name: models.VectorParams(
|
||||
size=self.vector_size,
|
||||
distance=self.distance,
|
||||
)
|
||||
}
|
||||
# Конфигурация sparse-векторов (ключевой поиск BM25 с IDF)
|
||||
sparse_vectors_config = {
|
||||
self.sparse_vector_name: models.SparseVectorParams(
|
||||
modifier=models.Modifier.IDF
|
||||
)
|
||||
}
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=vectors_config,
|
||||
sparse_vectors_config=sparse_vectors_config
|
||||
)
|
||||
logger.info(f"Коллекция {self.collection_name} создана (dense size={self.vector_size}, sparse с IDF)")
|
||||
except Exception as e:
|
||||
# Если библиотека не поддерживает sparse (старая версия), создаём только dense
|
||||
logger.warning(f"Не удалось создать sparse-векторы: {e}. Создаю коллекцию только с dense.")
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config={
|
||||
self.dense_vector_name: models.VectorParams(
|
||||
size=self.vector_size,
|
||||
distance=self.distance,
|
||||
)
|
||||
}
|
||||
)
|
||||
logger.info(f"Коллекция {self.collection_name} создана (только dense size={self.vector_size})")
|
||||
# После создания коллекции добавляем payload-индексы
|
||||
self._create_payload_indexes()
|
||||
else:
|
||||
# Коллекция уже существует – проверяем, есть ли sparse-поле (логируем предупреждение, если нет)
|
||||
try:
|
||||
collection_info = self.client.get_collection(self.collection_name)
|
||||
if not hasattr(collection_info, 'config') or not hasattr(collection_info.config, 'params'):
|
||||
logger.debug("Не удалось проверить наличие sparse-поля – пропускаем")
|
||||
else:
|
||||
sparse_exists = (hasattr(collection_info.config.params, 'sparse_vectors') and
|
||||
self.sparse_vector_name in collection_info.config.params.sparse_vectors)
|
||||
if not sparse_exists:
|
||||
logger.warning(
|
||||
f"Коллекция {self.collection_name} не содержит sparse-векторного поля. "
|
||||
"Гибридный поиск будет работать только с dense. "
|
||||
"Рекомендуется пересоздать коллекцию с помощью --force или вручную."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Не удалось проверить конфигурацию коллекции: {e}")
|
||||
logger.debug(f"Коллекция {self.collection_name} уже существует")
|
||||
|
||||
def _create_payload_indexes(self) -> None:
|
||||
"""
|
||||
Создаёт индексы для полей payload для ускорения фильтрации и поиска.
|
||||
Индексы нужны для полей:
|
||||
- user_jid: для доступа к личным документам (keyword)
|
||||
- room_jid: для доступа к документам комнаты (keyword)
|
||||
- is_global: для доступа к глобальным документам (keyword)
|
||||
- doc_id: для быстрого удаления всех чанков документа (integer)
|
||||
- text: для полнотекстового поиска (text с настройками tokenizer=whitespace, lowercase, phrase_matching)
|
||||
"""
|
||||
logger.info("Создаю payload-индексы для коллекции...")
|
||||
|
||||
# Индексы для фильтрации (keyword)
|
||||
for field in ["user_jid", "room_jid", "is_global"]:
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name=field,
|
||||
field_schema="keyword"
|
||||
)
|
||||
logger.debug(f"Индекс keyword для поля '{field}' создан")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось создать индекс для '{field}': {e}")
|
||||
|
||||
# Индекс для doc_id (integer)
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="doc_id",
|
||||
field_schema="integer"
|
||||
)
|
||||
logger.debug("Индекс integer для поля 'doc_id' создан")
|
||||
except Exception as e:
|
||||
logger.warning(f"Не удалось создать индекс для 'doc_id': {e}")
|
||||
|
||||
# Полнотекстовый индекс для текста с детальными параметрами
|
||||
try:
|
||||
# Используем расширенную конфигурацию (tokenizer, lowercase, phrase_matching)
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="text",
|
||||
field_schema=models.TextIndexParams(
|
||||
tokenizer="whitespace",
|
||||
lowercase=True,
|
||||
min_token_len=1,
|
||||
max_token_len=0, # 0 = без ограничения
|
||||
phrase_matching=True
|
||||
)
|
||||
)
|
||||
logger.info("Полнотекстовый индекс для поля 'text' создан (tokenizer=whitespace, lowercase, phrase_matching)")
|
||||
except Exception as e:
|
||||
# Fallback для старых версий qdrant-client (без TextIndexParams)
|
||||
logger.warning(f"Не удалось создать расширенный индекс text: {e}. Пробую упрощённый.")
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
collection_name=self.collection_name,
|
||||
field_name="text",
|
||||
field_schema="text"
|
||||
)
|
||||
logger.info("Полнотекстовый индекс для поля 'text' создан (упрощённый)")
|
||||
except Exception as e2:
|
||||
logger.error(f"Не удалось создать индекс для 'text': {e2}")
|
||||
|
||||
logger.info(f"Payload-индексы для коллекции {self.collection_name} созданы")
|
||||
|
||||
def add_chunk(self, dense: List[float], sparse: Any, payload: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Добавляет точку (чанк) в коллекцию с двумя типами векторов: dense и sparse.
|
||||
|
||||
Аргументы:
|
||||
dense: плотный вектор (1024) от GigaChat
|
||||
sparse: разрежённый вектор от FastEmbed (SparseEmbedding) или None
|
||||
payload: метаданные чанка
|
||||
|
||||
Возвращает:
|
||||
str: уникальный идентификатор точки
|
||||
"""
|
||||
point_id = str(uuid4())
|
||||
|
||||
# Преобразуем sparse-вектор из формата FastEmbed в формат Qdrant
|
||||
sparse_vector = None
|
||||
if sparse is not None:
|
||||
if hasattr(sparse, 'indices') and hasattr(sparse, 'values'):
|
||||
# FastEmbed SparseEmbedding -> SparseVector
|
||||
sparse_vector = models.SparseVector(
|
||||
indices=sparse.indices.tolist(),
|
||||
values=sparse.values.tolist()
|
||||
)
|
||||
elif isinstance(sparse, dict) and 'indices' in sparse:
|
||||
# Уже словарь
|
||||
sparse_vector = models.SparseVector(**sparse)
|
||||
else:
|
||||
logger.warning(f"Неизвестный формат sparse: {type(sparse)}")
|
||||
|
||||
# Если sparse-вектор не задан, передаём пустой (чтобы не нарушить схему коллекции)
|
||||
if sparse_vector is None:
|
||||
sparse_vector = models.SparseVector(indices=[], values=[])
|
||||
|
||||
point = models.PointStruct(
|
||||
id=point_id,
|
||||
vector={
|
||||
self.dense_vector_name: dense,
|
||||
self.sparse_vector_name: sparse_vector
|
||||
},
|
||||
payload=payload
|
||||
)
|
||||
self.client.upsert(collection_name=self.collection_name, points=[point])
|
||||
logger.debug(f"Чанк добавлен: id={point_id}, doc_id={payload.get('doc_id')}")
|
||||
return point_id
|
||||
|
||||
def search_hybrid(
|
||||
self,
|
||||
dense_vector: List[float],
|
||||
sparse_vector: Any,
|
||||
user_jid: str,
|
||||
room_jid: str = None,
|
||||
top_k: int = 15,
|
||||
rrf_k: int = 60
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Гибридный поиск: объединяет результаты dense и sparse поиска через Reciprocal Rank Fusion (RRF).
|
||||
|
||||
Аргументы:
|
||||
dense_vector: плотный вектор запроса (1024) от GigaChat
|
||||
sparse_vector: разрежённый вектор запроса от FastEmbed (SparseEmbedding) или None
|
||||
user_jid: JID пользователя для фильтрации доступа
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
top_k: количество возвращаемых результатов
|
||||
rrf_k: константа сглаживания в RRF (обычно 60)
|
||||
|
||||
Возвращает:
|
||||
List[Dict]: список результатов с ключами score, payload, id
|
||||
"""
|
||||
# --- 1. Фильтр доступа ---
|
||||
if room_jid:
|
||||
# В комнате: документы комнаты ИЛИ глобальные
|
||||
filter_cond = models.Filter(
|
||||
should=[
|
||||
models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid)),
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
else:
|
||||
# В личном чате: личные документы пользователя ИЛИ глобальные
|
||||
filter_cond = models.Filter(
|
||||
should=[
|
||||
models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid)),
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
|
||||
# --- 2. Dense-поиск (семантический) ---
|
||||
dense_resp = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=dense_vector,
|
||||
using=self.dense_vector_name,
|
||||
limit=top_k,
|
||||
query_filter=filter_cond,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# --- 3. Sparse-поиск (ключевой) с преобразованием формата ---
|
||||
sparse_resp = None
|
||||
if sparse_vector is not None:
|
||||
# Преобразуем sparse-вектор аналогично методу add_chunk
|
||||
sparse_query = None
|
||||
if hasattr(sparse_vector, 'indices') and hasattr(sparse_vector, 'values'):
|
||||
# FastEmbed SparseEmbedding -> SparseVector
|
||||
sparse_query = models.SparseVector(
|
||||
indices=sparse_vector.indices.tolist(),
|
||||
values=sparse_vector.values.tolist()
|
||||
)
|
||||
elif isinstance(sparse_vector, dict) and 'indices' in sparse_vector:
|
||||
sparse_query = models.SparseVector(**sparse_vector)
|
||||
elif isinstance(sparse_vector, models.SparseVector):
|
||||
sparse_query = sparse_vector
|
||||
else:
|
||||
logger.warning(f"Неизвестный формат sparse_vector: {type(sparse_vector)}")
|
||||
|
||||
if sparse_query is not None:
|
||||
sparse_resp = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=sparse_query,
|
||||
using=self.sparse_vector_name,
|
||||
limit=top_k,
|
||||
query_filter=filter_cond,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# --- 4. Reciprocal Rank Fusion (RRF) ---
|
||||
scores = {} # point_id -> {score, payload, id}
|
||||
|
||||
for rank, hit in enumerate(dense_resp.points):
|
||||
rrf_score = 1.0 / (rrf_k + rank + 1)
|
||||
scores[hit.id] = {
|
||||
"score": rrf_score,
|
||||
"payload": hit.payload,
|
||||
"id": hit.id
|
||||
}
|
||||
|
||||
if sparse_resp:
|
||||
for rank, hit in enumerate(sparse_resp.points):
|
||||
rrf_score = 1.0 / (rrf_k + rank + 1)
|
||||
if hit.id in scores:
|
||||
scores[hit.id]["score"] += rrf_score
|
||||
else:
|
||||
scores[hit.id] = {
|
||||
"score": rrf_score,
|
||||
"payload": hit.payload,
|
||||
"id": hit.id
|
||||
}
|
||||
|
||||
# --- 5. Сортировка и возврат top_k ---
|
||||
sorted_items = sorted(scores.values(), key=lambda x: x["score"], reverse=True)
|
||||
return sorted_items[:top_k]
|
||||
|
||||
def search_keywords(self, query_text: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> List[Dict]:
|
||||
"""
|
||||
Выполняет полнотекстовый поиск по полю text (для точных запросов).
|
||||
|
||||
Аргументы:
|
||||
query_text: текст запроса (ищется как подстрока в text)
|
||||
user_jid: JID пользователя для фильтрации
|
||||
room_jid: JID комнаты (None для личного чата)
|
||||
top_k: количество возвращаемых результатов
|
||||
|
||||
Возвращает:
|
||||
List[Dict]: список результатов с ключами score, payload, id
|
||||
|
||||
Используется для точных запросов типа "выведи пункт 1.1",
|
||||
где семантический поиск может не справиться.
|
||||
"""
|
||||
# Формируем фильтр доступа
|
||||
if room_jid:
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid))
|
||||
],
|
||||
should=[
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
else:
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid))
|
||||
],
|
||||
should=[
|
||||
models.FieldCondition(key="is_global", match=models.MatchValue(value=True))
|
||||
]
|
||||
)
|
||||
|
||||
# Добавляем условие полнотекстового поиска (MatchText)
|
||||
filter_cond.must.append(
|
||||
models.FieldCondition(
|
||||
key="text",
|
||||
match=models.MatchText(text=query_text)
|
||||
)
|
||||
)
|
||||
|
||||
# Получаем точки через scroll
|
||||
points = []
|
||||
offset = None
|
||||
while len(points) < top_k:
|
||||
scroll_result = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=filter_cond,
|
||||
limit=top_k,
|
||||
offset=offset,
|
||||
with_payload=True
|
||||
)
|
||||
points.extend(scroll_result[0])
|
||||
offset = scroll_result[1]
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
# Простая сортировка по количеству совпавших ключевых слов (score не от Qdrant, а наша эвристика)
|
||||
query_words = set(query_text.lower().split())
|
||||
for point in points:
|
||||
text = point.payload.get('text', '').lower()
|
||||
matches = sum(1 for word in query_words if word in text)
|
||||
point._match_score = matches
|
||||
|
||||
points.sort(key=lambda p: getattr(p, '_match_score', 0), reverse=True)
|
||||
|
||||
return [
|
||||
{
|
||||
"score": getattr(p, '_match_score', 0),
|
||||
"payload": p.payload,
|
||||
"id": p.id
|
||||
}
|
||||
for p in points[:top_k]
|
||||
]
|
||||
|
||||
def delete_by_doc_id(self, doc_id: int) -> None:
|
||||
"""
|
||||
Удаляет все точки (чанки), принадлежащие указанному документу.
|
||||
|
||||
Аргументы:
|
||||
doc_id: идентификатор документа в PostgreSQL
|
||||
"""
|
||||
filter_cond = models.Filter(
|
||||
must=[
|
||||
models.FieldCondition(key="doc_id", match=models.MatchValue(value=doc_id))
|
||||
]
|
||||
)
|
||||
|
||||
points = []
|
||||
offset = None
|
||||
while True:
|
||||
scroll_result = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=filter_cond,
|
||||
limit=100,
|
||||
offset=offset
|
||||
)
|
||||
points.extend(scroll_result[0])
|
||||
offset = scroll_result[1]
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
ids = [p.id for p in points]
|
||||
if ids:
|
||||
self.client.delete(collection_name=self.collection_name, points_selector=ids)
|
||||
logger.info(f"Удалено {len(ids)} чанков для doc_id={doc_id}")
|
||||
|
||||
def delete_collection(self) -> None:
|
||||
"""Удаляет всю коллекцию (опасная операция, используется редко)."""
|
||||
if self.client.collection_exists(self.collection_name):
|
||||
self.client.delete_collection(self.collection_name)
|
||||
logger.warning(f"Коллекция {self.collection_name} удалена")
|
||||
137
rag/services/reranker_service.py
Normal file
137
rag/services/reranker_service.py
Normal file
@@ -0,0 +1,137 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Сервис переранжирования с использованием кросс-энкодеров.
|
||||
Загружает модель один раз при первом вызове и использует её для ранжирования фрагментов.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Tuple, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Импортируем sentence-transformers с обработкой ошибок
|
||||
try:
|
||||
from sentence_transformers import CrossEncoder
|
||||
HAS_CROSS_ENCODER = True
|
||||
except ImportError:
|
||||
HAS_CROSS_ENCODER = False
|
||||
logger.warning("sentence-transformers не установлен. Используйте LLM-переранжирование.")
|
||||
|
||||
|
||||
class RerankerService:
|
||||
"""
|
||||
Сервис переранжирования на основе кросс-энкодеров.
|
||||
Загружает модель при инициализации или при первом вызове.
|
||||
"""
|
||||
|
||||
# Рекомендуемая модель для русского языка
|
||||
DEFAULT_MODEL = "intfloat/multilingual-e5-reranker"
|
||||
|
||||
def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None):
|
||||
"""
|
||||
Инициализация сервиса.
|
||||
|
||||
Аргументы:
|
||||
model_name: имя модели (если None, используется DEFAULT_MODEL).
|
||||
device: устройство для вычислений ('cpu' или 'cuda'). Если None, автоопределение.
|
||||
"""
|
||||
self.model_name = model_name or self.DEFAULT_MODEL
|
||||
self.device = device or ("cuda" if self._cuda_available() else "cpu")
|
||||
self._model = None
|
||||
logger.info(f"RerankerService инициализирован с моделью {self.model_name} (device={self.device})")
|
||||
|
||||
@staticmethod
|
||||
def _cuda_available() -> bool:
|
||||
"""Проверяет доступность CUDA (для GPU)."""
|
||||
try:
|
||||
import torch
|
||||
return torch.cuda.is_available()
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def _load_model(self) -> None:
|
||||
"""Загружает модель кросс-энкодера (ленивая загрузка)."""
|
||||
if not HAS_CROSS_ENCODER:
|
||||
raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")
|
||||
|
||||
if self._model is None:
|
||||
logger.info(f"Загрузка модели {self.model_name} на устройство {self.device}...")
|
||||
self._model = CrossEncoder(self.model_name, device=self.device)
|
||||
logger.info(f"Модель {self.model_name} загружена")
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
fragments: List[str],
|
||||
top_k: Optional[int] = None
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""
|
||||
Ранжирует фрагменты по релевантности запросу.
|
||||
|
||||
Аргументы:
|
||||
query: текст запроса.
|
||||
fragments: список текстовых фрагментов.
|
||||
top_k: максимальное количество возвращаемых фрагментов (если None – все).
|
||||
|
||||
Возвращает:
|
||||
List[Tuple[str, float]]: список кортежей (фрагмент, оценка_релевантности),
|
||||
отсортированный по убыванию оценки.
|
||||
"""
|
||||
if not fragments:
|
||||
return []
|
||||
|
||||
self._load_model()
|
||||
|
||||
# Формируем пары (запрос, фрагмент) для оценки
|
||||
pairs = [(query, frag) for frag in fragments]
|
||||
|
||||
# Получаем оценки релевантности (числа от 0 до 1)
|
||||
try:
|
||||
scores = self._model.predict(pairs)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при оценке кросс-энкодером: {e}")
|
||||
# В случае ошибки возвращаем фрагменты в исходном порядке с оценкой 0.0
|
||||
return [(frag, 0.0) for frag in fragments]
|
||||
|
||||
# Собираем результаты и сортируем по убыванию оценки
|
||||
results = list(zip(fragments, scores))
|
||||
results.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Ограничиваем количество, если указано
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
logger.debug(f"Переранжирование: из {len(fragments)} фрагментов оставлено {len(results)}")
|
||||
return results
|
||||
|
||||
def rerank_context(
|
||||
self,
|
||||
query: str,
|
||||
context: str,
|
||||
top_k: Optional[int] = None
|
||||
) -> str:
|
||||
"""
|
||||
Удобная обёртка: принимает строку контекста с фрагментами, разделёнными '\n\n',
|
||||
возвращает отранжированную строку.
|
||||
|
||||
Аргументы:
|
||||
query: текст запроса.
|
||||
context: исходный контекст (фрагменты разделены '\n\n').
|
||||
top_k: максимальное количество возвращаемых фрагментов.
|
||||
|
||||
Возвращает:
|
||||
str: отранжированный контекст (фрагменты разделены '\n\n').
|
||||
"""
|
||||
if not context:
|
||||
return ""
|
||||
|
||||
# Разбиваем контекст на фрагменты
|
||||
fragments = [f.strip() for f in context.split("\n\n") if f.strip()]
|
||||
if not fragments:
|
||||
return ""
|
||||
|
||||
# Ранжируем
|
||||
ranked = self.rerank(query, fragments, top_k)
|
||||
|
||||
# Собираем обратно в строку (только текст, без оценок)
|
||||
return "\n\n".join([frag for frag, _ in ranked])
|
||||
Reference in New Issue
Block a user