Initial commit: TrueConf Chatbot КЛЕВЕР
This commit is contained in:
@@ -0,0 +1,650 @@
|
||||
import os
|
||||
import sys
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
import json
|
||||
import hashlib
|
||||
import requests
|
||||
|
||||
from glob import glob
|
||||
from datetime import datetime
|
||||
|
||||
# Импорт SentenceTransformer УДАЛЕН для разгрузки ОЗУ сервера
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
VectorParams,
|
||||
PointStruct,
|
||||
PayloadSchemaType,
|
||||
Filter,
|
||||
FieldCondition,
|
||||
MatchValue
|
||||
)
|
||||
|
||||
# =========================================================
|
||||
# ИМПОРТ ЦЕНТРАЛЬНОГО КОНФИГА БОТА
|
||||
# =========================================================
|
||||
# Добавляем корневую папку бота в sys.path, чтобы скрипт увидел модуль config,
|
||||
# даже если скрипт запускается из подпапки.
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
import config.config as config
|
||||
|
||||
# =========================================================
|
||||
# CONFIG (Берем данные из центрального config.py)
|
||||
# =========================================================
|
||||
DOCS_DIR = config.SEARCH_DOCS_PATH
|
||||
QDRANT_HOST = config.QDRANT_HOST
|
||||
QDRANT_PORT = config.QDRANT_PORT
|
||||
COLLECTION_NAME = config.QDRANT_COLLECTION
|
||||
|
||||
# Формируем путь к реестру внутри папки логов, указанной в конфиге
|
||||
REGISTRY_PATH = os.path.join(config.LOGS_PATH, "indexing_registry.json")
|
||||
|
||||
# =========================================================
|
||||
# ЛОКАЛЬНЫЕ НАСТРОЙКИ ИНДЕКСАТОРА
|
||||
# =========================================================
|
||||
UNSTRUCTURED_API_URL = "http://172.16.80.250:9000/general/v0/general"
|
||||
BATCH_SIZE = 32
|
||||
RECREATE_COLLECTION = False
|
||||
MIN_TEXT_LENGTH = 25
|
||||
|
||||
# =========================================================
|
||||
# INIT
|
||||
# =========================================================
|
||||
|
||||
# Локальная инициализация нейросети убрана. Работаем через API.
|
||||
|
||||
print("🔄 Подключение к Qdrant...")
|
||||
|
||||
client = QdrantClient(
|
||||
host=QDRANT_HOST,
|
||||
port=QDRANT_PORT
|
||||
)
|
||||
|
||||
# =========================================================
|
||||
# REGISTRY & HASHING
|
||||
# =========================================================
|
||||
|
||||
def get_file_hash(filepath):
|
||||
"""Вычисляет MD5 хэш физического файла"""
|
||||
hasher = hashlib.md5()
|
||||
with open(filepath, 'rb') as f:
|
||||
buf = f.read()
|
||||
hasher.update(buf)
|
||||
return hasher.hexdigest()
|
||||
|
||||
def load_registry():
|
||||
"""Загружает текущее состояние проиндексированных файлов"""
|
||||
if os.path.exists(REGISTRY_PATH):
|
||||
with open(REGISTRY_PATH, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
return {}
|
||||
|
||||
def save_registry(registry):
|
||||
"""Сохраняет реестр на диск"""
|
||||
os.makedirs(os.path.dirname(REGISTRY_PATH), exist_ok=True)
|
||||
with open(REGISTRY_PATH, 'w', encoding='utf-8') as f:
|
||||
json.dump(registry, f, ensure_ascii=False, indent=4)
|
||||
|
||||
# =========================================================
|
||||
# HELPERS
|
||||
# =========================================================
|
||||
|
||||
def sha256_text(text):
|
||||
return hashlib.sha256(
|
||||
text.encode("utf-8")
|
||||
).hexdigest()
|
||||
|
||||
def clean_text(text):
|
||||
if not text:
|
||||
return ""
|
||||
|
||||
text = re.sub(
|
||||
r"\s+",
|
||||
" ",
|
||||
text
|
||||
)
|
||||
|
||||
return text.strip()
|
||||
|
||||
def is_page_number(text):
|
||||
return bool(
|
||||
re.match(
|
||||
r"^Стр\.\s*\d+\s*из\s*\d+$",
|
||||
text
|
||||
)
|
||||
)
|
||||
|
||||
def is_header_noise(text):
|
||||
patterns = [
|
||||
"АО «ХК «Сибцем»",
|
||||
"Тип документа:",
|
||||
"Ведущее подразделение:",
|
||||
"Дата утверждения:",
|
||||
"Редакция 1",
|
||||
"Оглавление"
|
||||
]
|
||||
|
||||
return any(
|
||||
p in text
|
||||
for p in patterns
|
||||
)
|
||||
|
||||
def is_junk(item, text):
|
||||
if not text:
|
||||
return True
|
||||
|
||||
if len(text) < MIN_TEXT_LENGTH:
|
||||
return True
|
||||
|
||||
if item["type"] in [
|
||||
"Header",
|
||||
"Footer",
|
||||
"Image"
|
||||
]:
|
||||
return True
|
||||
|
||||
if is_page_number(text):
|
||||
return True
|
||||
|
||||
if is_header_noise(text):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
# =========================================================
|
||||
# TABLE PARSER
|
||||
# =========================================================
|
||||
|
||||
def html_table_to_text(html):
|
||||
html = re.sub(
|
||||
r"\n+",
|
||||
" ",
|
||||
html
|
||||
)
|
||||
|
||||
rows = re.findall(
|
||||
r"<tr.*?>(.*?)</tr>",
|
||||
html,
|
||||
re.I | re.S
|
||||
)
|
||||
|
||||
parsed_rows = []
|
||||
|
||||
for row in rows:
|
||||
cells = re.findall(
|
||||
r"<t[hd].*?>(.*?)</t[hd]>",
|
||||
row,
|
||||
re.I | re.S
|
||||
)
|
||||
|
||||
clean_cells = []
|
||||
|
||||
for cell in cells:
|
||||
box = re.sub(
|
||||
r"<[^>]+>",
|
||||
" ",
|
||||
cell
|
||||
)
|
||||
box = clean_text(box)
|
||||
clean_cells.append(box)
|
||||
|
||||
if clean_cells:
|
||||
parsed_rows.append(
|
||||
clean_cells
|
||||
)
|
||||
|
||||
if not parsed_rows:
|
||||
return ""
|
||||
|
||||
headers = parsed_rows[0]
|
||||
|
||||
result_lines = []
|
||||
|
||||
for row_idx, row in enumerate(parsed_rows[1:], 1):
|
||||
result_lines.append(
|
||||
f"Запись {row_idx}:"
|
||||
)
|
||||
|
||||
for col_idx, cell in enumerate(row):
|
||||
if not cell:
|
||||
continue
|
||||
|
||||
if col_idx < len(headers):
|
||||
header = headers[col_idx]
|
||||
else:
|
||||
header = f"Колонка {col_idx+1}"
|
||||
|
||||
result_lines.append(
|
||||
f"{header}: {cell}"
|
||||
)
|
||||
|
||||
result_lines.append("")
|
||||
|
||||
return "\n".join(result_lines).strip()
|
||||
|
||||
# =========================================================
|
||||
# MERGE FRAGMENTS
|
||||
# =========================================================
|
||||
|
||||
def merge_fragments(elements):
|
||||
merged = []
|
||||
buffer = []
|
||||
buffer_page = None
|
||||
buffer_section = None
|
||||
|
||||
def flush_buffer():
|
||||
nonlocal buffer
|
||||
|
||||
if not buffer:
|
||||
return
|
||||
|
||||
text = clean_text(
|
||||
" ".join(buffer)
|
||||
)
|
||||
|
||||
if text:
|
||||
merged.append({
|
||||
"type": "MergedText",
|
||||
"text": text,
|
||||
"page_number": buffer_page,
|
||||
"section_title": buffer_section
|
||||
})
|
||||
|
||||
buffer = []
|
||||
|
||||
for el in elements:
|
||||
text = clean_text(
|
||||
el["text"]
|
||||
)
|
||||
|
||||
if not text:
|
||||
continue
|
||||
|
||||
el_type = el["type"]
|
||||
page = el["metadata"].get("page_number", 0)
|
||||
section = el.get("section_title", "")
|
||||
|
||||
# =========================================
|
||||
# TABLE
|
||||
# =========================================
|
||||
if el_type == "Table":
|
||||
flush_buffer()
|
||||
|
||||
merged.append({
|
||||
"type": "Table",
|
||||
"text": text,
|
||||
"html": el["metadata"].get("text_as_html", ""),
|
||||
"page_number": page,
|
||||
"section_title": section
|
||||
})
|
||||
continue
|
||||
|
||||
# =========================================
|
||||
# SHORT TEXT
|
||||
# =========================================
|
||||
if el_type in [
|
||||
"NarrativeText",
|
||||
"ListItem",
|
||||
"UncategorizedText"
|
||||
]:
|
||||
if len(text) < 80:
|
||||
buffer.append(text)
|
||||
buffer_page = page
|
||||
buffer_section = section
|
||||
else:
|
||||
flush_buffer()
|
||||
merged.append({
|
||||
"type": "Text",
|
||||
"text": text,
|
||||
"page_number": page,
|
||||
"section_title": section
|
||||
})
|
||||
else:
|
||||
flush_buffer()
|
||||
|
||||
flush_buffer()
|
||||
|
||||
return merged
|
||||
|
||||
# =========================================================
|
||||
# QDRANT INIT & DELETE
|
||||
# =========================================================
|
||||
|
||||
def init_collection():
|
||||
collections = [
|
||||
c.name
|
||||
for c in client.get_collections().collections
|
||||
]
|
||||
|
||||
if (
|
||||
RECREATE_COLLECTION
|
||||
and COLLECTION_NAME in collections
|
||||
):
|
||||
print("🗑 Удаление коллекции...")
|
||||
client.delete_collection(COLLECTION_NAME)
|
||||
collections = [
|
||||
c.name
|
||||
for c in client.get_collections().collections
|
||||
]
|
||||
|
||||
if COLLECTION_NAME not in collections:
|
||||
print("⚙️ Создание коллекции...")
|
||||
client.create_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(
|
||||
size=1024,
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
|
||||
indexes = [
|
||||
("metadata.filename", PayloadSchemaType.KEYWORD),
|
||||
("metadata.section_title", PayloadSchemaType.KEYWORD),
|
||||
("metadata.page_number", PayloadSchemaType.INTEGER),
|
||||
("metadata.content_type", PayloadSchemaType.KEYWORD)
|
||||
]
|
||||
|
||||
for field_name, field_schema in indexes:
|
||||
client.create_payload_index(
|
||||
collection_name=COLLECTION_NAME,
|
||||
field_name=field_name,
|
||||
field_schema=field_schema
|
||||
)
|
||||
|
||||
def delete_file_from_qdrant(filename):
|
||||
"""Удаляет все векторы, привязанные к конкретному файлу"""
|
||||
client.delete(
|
||||
collection_name=COLLECTION_NAME,
|
||||
points_selector=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="metadata.filename",
|
||||
match=MatchValue(value=filename)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
print(f" 🗑 Удалены старые данные файла {filename} из базы.")
|
||||
|
||||
# =========================================================
|
||||
# PARSE DOCUMENT
|
||||
# =========================================================
|
||||
|
||||
def parse_document(file_path):
|
||||
filename = os.path.basename(
|
||||
file_path
|
||||
)
|
||||
|
||||
print(f"\n📄 Парсинг: {filename}")
|
||||
|
||||
payload_data = {
|
||||
"strategy": "hi_res",
|
||||
"languages": ["rus"],
|
||||
"pdf_infer_table_structure": "true"
|
||||
}
|
||||
|
||||
ext = filename.lower().split(".")[-1]
|
||||
|
||||
if ext == "pdf":
|
||||
mime_type = "application/pdf"
|
||||
else:
|
||||
mime_type = (
|
||||
"application/vnd.openxmlformats-officedocument."
|
||||
"wordprocessingml.document"
|
||||
)
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
files = {
|
||||
"files": (
|
||||
filename,
|
||||
f,
|
||||
mime_type
|
||||
)
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
UNSTRUCTURED_API_URL,
|
||||
files=files,
|
||||
data=payload_data,
|
||||
timeout=900
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
return response.json()
|
||||
|
||||
# =========================================================
|
||||
# INDEX DOCUMENT
|
||||
# =========================================================
|
||||
|
||||
def index_document(file_path, data):
|
||||
filename = os.path.basename(file_path)
|
||||
|
||||
cleaned = []
|
||||
current_title = "Общая информация"
|
||||
|
||||
# =====================================================
|
||||
# CLEAN
|
||||
# =====================================================
|
||||
for item in data:
|
||||
raw_text = item.get("text", "")
|
||||
text = clean_text(raw_text)
|
||||
|
||||
if is_junk(item, text):
|
||||
continue
|
||||
|
||||
if item["type"] == "Title":
|
||||
current_title = text
|
||||
continue
|
||||
|
||||
item["section_title"] = current_title
|
||||
cleaned.append(item)
|
||||
|
||||
# =====================================================
|
||||
# MERGE
|
||||
# =====================================================
|
||||
merged = merge_fragments(cleaned)
|
||||
points = []
|
||||
|
||||
# =====================================================
|
||||
# BUILD CHUNKS
|
||||
# =====================================================
|
||||
for idx, chunk in enumerate(merged):
|
||||
chunk_type = chunk["type"]
|
||||
section_title = chunk["section_title"]
|
||||
page_number = chunk["page_number"]
|
||||
|
||||
# =================================================
|
||||
# TABLE / TEXT
|
||||
# =================================================
|
||||
if chunk_type == "Table":
|
||||
table_text = html_table_to_text(chunk["html"])
|
||||
content = (
|
||||
f"Раздел: {section_title}\n"
|
||||
f"Тип: Таблица\n\n"
|
||||
f"{table_text}"
|
||||
)
|
||||
else:
|
||||
content = (
|
||||
f"Раздел: {section_title}\n\n"
|
||||
f"{chunk['text']}"
|
||||
)
|
||||
|
||||
content = clean_text(content)
|
||||
|
||||
if len(content) < 40:
|
||||
continue
|
||||
|
||||
# =================================================
|
||||
# VECTOR (Переведено на центральный Embedding API)
|
||||
# =================================================
|
||||
try:
|
||||
payload_api = {
|
||||
"model": config.EMBEDDING_MODEL_NAME, # Отправляем "bge-m3"
|
||||
"input_type": "passage", # Для сохранения документов используем тип passage
|
||||
"texts": [content]
|
||||
}
|
||||
headers_api = {"Content-Type": "application/json"}
|
||||
if hasattr(config, 'EMBEDDING_API_TOKEN') and config.EMBEDDING_API_TOKEN:
|
||||
headers_api["X-Embedding-Token"] = config.EMBEDDING_API_TOKEN
|
||||
|
||||
api_resp = requests.post(config.EMBEDDING_API_URL, json=payload_api, headers=headers_api)
|
||||
api_resp.raise_for_status()
|
||||
vector = api_resp.json()["embeddings"][0]
|
||||
except Exception as e:
|
||||
print(f"❌ Ошибка векторизации чанка {idx} через API: {e}")
|
||||
continue # Пропускаем чанк, если центральный сервис недоступен
|
||||
|
||||
chunk_hash = sha256_text(
|
||||
filename + str(idx) + content
|
||||
)
|
||||
|
||||
point_id = str(
|
||||
uuid.uuid5(
|
||||
uuid.NAMESPACE_DNS,
|
||||
chunk_hash
|
||||
)
|
||||
)
|
||||
|
||||
payload = {
|
||||
"page_content": content,
|
||||
"metadata": {
|
||||
"filename": filename,
|
||||
"source": file_path,
|
||||
"chunk_index": idx,
|
||||
"section_title": section_title,
|
||||
"page_number": page_number,
|
||||
"content_type": chunk_type,
|
||||
"chunk_hash": chunk_hash
|
||||
}
|
||||
}
|
||||
|
||||
points.append(
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=vector,
|
||||
payload=payload
|
||||
)
|
||||
)
|
||||
|
||||
# =================================================
|
||||
# UPSERT BATCH
|
||||
# =================================================
|
||||
if len(points) >= BATCH_SIZE:
|
||||
client.upsert(
|
||||
collection_name=COLLECTION_NAME,
|
||||
points=points
|
||||
)
|
||||
print(f" ⬆️ batch={len(points)}")
|
||||
points = []
|
||||
|
||||
# =====================================================
|
||||
# FINAL UPSERT
|
||||
# =====================================================
|
||||
if points:
|
||||
client.upsert(
|
||||
collection_name=COLLECTION_NAME,
|
||||
points=points
|
||||
)
|
||||
|
||||
print(f" ✅ Загружено chunks: {len(merged)}")
|
||||
return len(merged)
|
||||
|
||||
# =========================================================
|
||||
# MAIN
|
||||
# =========================================================
|
||||
|
||||
def main():
|
||||
init_collection()
|
||||
|
||||
target_files = []
|
||||
|
||||
for ext in ["*.pdf", "*.docx"]:
|
||||
target_files.extend(
|
||||
glob(
|
||||
os.path.join(DOCS_DIR, "**", ext),
|
||||
recursive=True
|
||||
)
|
||||
)
|
||||
|
||||
# ФИЛЬТРАЦИЯ СКРЫТЫХ ВРЕМЕННЫХ ФАЙЛОВ WORD (~$)
|
||||
target_files = [f for f in target_files if not os.path.basename(f).startswith("~$")]
|
||||
|
||||
target_files.sort()
|
||||
|
||||
if not target_files:
|
||||
print("❌ PDF/DOCX файлы не найдены")
|
||||
return
|
||||
|
||||
print(f"📂 Найдено файлов для индексации: {len(target_files)}")
|
||||
print(f"📡 Использование удаленной векторизации: {config.EMBEDDING_API_URL} ({config.EMBEDDING_MODEL_NAME})")
|
||||
|
||||
# Загружаем реестр
|
||||
registry = load_registry()
|
||||
|
||||
# =====================================================
|
||||
# PROCESS
|
||||
# =====================================================
|
||||
for file_path in target_files:
|
||||
filename = os.path.basename(file_path)
|
||||
current_hash = get_file_hash(file_path)
|
||||
|
||||
file_info = registry.get(filename)
|
||||
|
||||
# Проверка: Файл уже есть, хэш совпадает, статус был успешным
|
||||
if file_info and file_info.get("hash") == current_hash and file_info.get("status") == "SUCCESS":
|
||||
print(f"\n⏭️ Пропуск (не изменился): {filename}")
|
||||
continue
|
||||
|
||||
print(f"\n🔄 Обработка: {filename}")
|
||||
|
||||
# Если файл есть в реестре со статусом SUCCESS, но хэш другой (файл обновили)
|
||||
if file_info and file_info.get("status") == "SUCCESS":
|
||||
print(f" ⚠️ Обнаружена новая версия файла! Обновляем индекс...")
|
||||
delete_file_from_qdrant(filename)
|
||||
|
||||
try:
|
||||
data = parse_document(file_path)
|
||||
chunks_count = index_document(file_path, data)
|
||||
|
||||
# Запись успешного результата
|
||||
registry[filename] = {
|
||||
"hash": current_hash,
|
||||
"status": "SUCCESS",
|
||||
"chunks": chunks_count,
|
||||
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
save_registry(registry)
|
||||
|
||||
# ПАУЗА ПРИ НЕДОСТУПНОСТИ СЕРВЕРА ПАРСИНГА
|
||||
except requests.exceptions.ConnectionError:
|
||||
print(f" ⚠️ Сервер парсинга документов недоступен. Ждем 60 секунд...")
|
||||
time.sleep(60)
|
||||
registry[filename] = {
|
||||
"hash": current_hash,
|
||||
"status": "ERROR",
|
||||
"error_details": "Connection refused",
|
||||
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
save_registry(registry)
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ Ошибка обработки файла {filename}: {e}")
|
||||
|
||||
# Запись ошибки в реестр
|
||||
registry[filename] = {
|
||||
"hash": current_hash,
|
||||
"status": "ERROR",
|
||||
"error_details": str(e),
|
||||
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
save_registry(registry)
|
||||
|
||||
print("\n🎉 ГОТОВО. База Qdrant полностью синхронизирована через API.")
|
||||
|
||||
# =========================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user