support-bot-saas/prepare_data.py

import json
import pickle
import re
from pathlib import Path
from typing import List

import faiss
import numpy as np
from langchain_core.embeddings import Embeddings
from sentence_transformers import SentenceTransformer

EMBED_MODEL = "intfloat/multilingual-e5-large"
DATA_DIR = Path(__file__).parent
TICKETS_PATTERN = "tickets*.json"
VECTOR_DIR = DATA_DIR / "vector_store"
VECTOR_INDEX = VECTOR_DIR / "index.faiss.bin"
VECTOR_META = VECTOR_DIR / "metadata.pkl"

NON_SOLUTION = re.compile(
    r"(уточнит|приложите скрин|какая ошибка|с какой проблемой"
    r"|нет обратной связи|запрос завершу|открыть его снова"
    r"|откройте новый|обратиться в службу поддержки"
    r"|свяжитесь с технической"
    r"|напишите нам|позвоните нам)", re.I
)


class LocalEmbeddings(Embeddings):
    def __init__(self, model_name: str):
        self.model = SentenceTransformer(model_name, device="cuda")

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        emb = self.model.encode(texts, normalize_embeddings=True, show_progress_bar=True)
        return emb.tolist()

    def embed_query(self, text: str) -> List[float]:
        emb = self.model.encode([text], normalize_embeddings=True)
        return emb[0].tolist()


def clean_text(text: str) -> str:
    text = re.sub(r'\[/?q\].*?\[/q\]', '', text, flags=re.DOTALL | re.IGNORECASE)
    text = re.sub(r'https?://\S+', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def has_real_solution(messages: list[dict]) -> str | None:
    support_texts = []
    for m in messages:
        if m["role"] == "support":
            cleaned = clean_text(m["text"])
            if cleaned:
                support_texts.append(cleaned)
    if not support_texts:
        return None
    combined = "\n".join(support_texts)
    if len(combined) < 100:
        return None
    last = support_texts[-1]
    if NON_SOLUTION.search(last):
        return None
    if last.strip().endswith("?") and len(last) < 150:
        return None
    return combined


def get_client_dialogue(messages: list[dict]) -> str:
    texts = []
    for m in messages:
        if m["role"] == "client":
            cleaned = clean_text(m["text"])
            if cleaned:
                texts.append(cleaned)
    return "\n".join(texts)


def build_chunk(ticket: dict) -> dict | None:
    msgs = ticket.get("messages") or []
    if not msgs:
        return None
    description = clean_text(ticket.get("description", ""))
    solution = has_real_solution(msgs)
    if not solution:
        return None

    client_msgs = get_client_dialogue(msgs)

    # Текст для поиска — только описание + переписка клиента
    search_parts = []
    if description:
        search_parts.append(description)
    if client_msgs:
        search_parts.append(client_msgs)
    search_text = "\n".join(search_parts).strip()
    if not search_text:
        return None

    # Полный текст для показа и LLM
    full_text = (
        f"Категория: {ticket['category']}\n"
        f"Проблема: {description}\n"
        f"Решение: {solution}"
    )

    metadata = {
        "ticket_id": ticket["ticket_id"],
        "client": ticket["client"],
        "category": ticket["category"],
        "full_text": full_text,
    }
    return {"search_text": search_text, "metadata": metadata}


def load_all_tickets() -> list:
    all_tickets = []
    for f in sorted(DATA_DIR.glob(TICKETS_PATTERN)):
        tickets = json.loads(f.read_text(encoding="utf-8"))
        print(f"  {f.name}: {len(tickets)} тикетов")
        all_tickets.extend(tickets)
    return all_tickets


def save_faiss_index(index, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(faiss.serialize_index(index))


def load_faiss_index(path: Path):
    return faiss.deserialize_index(np.frombuffer(path.read_bytes(), dtype=np.uint8))


def main():
    print("Загрузка файлов tickets*.json ...")
    tickets = load_all_tickets()
    print(f"Всего тикетов: {len(tickets)}")

    chunks = []
    skipped = 0
    for t in tickets:
        chunk = build_chunk(t)
        if chunk is not None:
            chunks.append(chunk)
        else:
            skipped += 1
    print(f"Сформировано чанков: {len(chunks)}")
    print(f"Пропущено (нет решения): {skipped}")

    print(f"Загрузка модели эмбеддингов: {EMBED_MODEL} ...")
    embeddings = LocalEmbeddings(EMBED_MODEL)

    search_texts = [c["search_text"] for c in chunks]
    metadatas = [c["metadata"] for c in chunks]

    print("Создание эмбеддингов и FAISS индекса...")
    vectors = embeddings.embed_documents(search_texts)
    dim = len(vectors[0])
    index = faiss.IndexFlatIP(dim)
    index.add(np.array(vectors, dtype=np.float32))

    save_faiss_index(index, VECTOR_INDEX)
    with open(VECTOR_META, "wb") as f:
        pickle.dump({"metadatas": metadatas}, f)

    print(f"FAISS индекс сохранён: {VECTOR_DIR}")
    print(f"Размерность: {dim}")
    print(f"Векторов в индексе: {index.ntotal}")


if __name__ == "__main__":
    main()