ZhouYI/06_agentic_rag/VPN_agentic_rag/scripts/prepare_data.py

from __future__ import annotations

import re
from pathlib import Path

from app.config import PROJECT_ROOT, Settings, get_settings
from app.embeddings import create_embedding_provider
from app.milvus_store import IndexedDocument, MilvusVectorStore
from app.sql_store import initialize_database


INCIDENTS_SOURCE_PATH = PROJECT_ROOT / "data" / "source" / "incidents.csv"


def split_document(path: Path) -> list[IndexedDocument]:
    text = path.read_text(encoding="utf-8").strip()
    if not text:
        raise ValueError(f"Milvus 原始文档为空：{path}")

    metadata_lines = {
        line.split("：", 1)[0]: line.split("：", 1)[1].strip()
        for line in text.splitlines()
        if "：" in line
    }
    service = metadata_lines.get("服务", "")
    fault_type = metadata_lines.get("故障类型", "")
    client_os = metadata_lines.get("适用系统", "")
    client_version = metadata_lines.get("客户端版本", "")
    document_type = metadata_lines.get("文档类别", "")
    required_metadata = {
        "服务": service,
        "故障类型": fault_type,
        "适用系统": client_os,
        "客户端版本": client_version,
        "文档类别": document_type,
    }
    missing_metadata = [
        field_name
        for field_name, value in required_metadata.items()
        if not value
    ]
    if missing_metadata:
        raise ValueError(
            f"文档缺少元数据 {', '.join(missing_metadata)}：{path.name}"
        )

    # 课程文档以二级标题作为稳定 Chunk 边界，保留章节语义和来源定位。
    sections = [
        section.strip()
        for section in re.split(r"(?=^## )", text, flags=re.MULTILINE)
        if section.strip()
    ]
    title = text.splitlines()[0].lstrip("# ")
    return [
        IndexedDocument(
            content=section,
            source=path.name,
            doc_type=document_type,
            chunk_index=index,
            service=service,
            fault_type=fault_type,
            client_os=client_os,
            client_version=client_version,
            metadata={
                "title": title,
                "document_type": document_type,
            },
        )
        for index, section in enumerate(sections)
    ]


def prepare_sqlite(settings: Settings) -> int:
    return initialize_database(
        path=settings.sqlite_path,
        source_path=INCIDENTS_SOURCE_PATH,
        reset=True,
    )


def prepare_milvus(settings: Settings) -> tuple[int, int, int]:
    document_paths = sorted(settings.documents_path.glob("*.md"))
    if not document_paths:
        raise FileNotFoundError(f"Milvus 原始文档不存在：{settings.documents_path}")

    # 先完成所有文档校验和切分，再连接 Milvus，避免半批脏数据。
    documents = [
        document
        for path in document_paths
        for document in split_document(path)
    ]
    embeddings = create_embedding_provider(settings)
    store = MilvusVectorStore(
        uri=settings.milvus_uri,
        token=settings.milvus_token.get_secret_value(),
        collection_name=settings.milvus_collection,
        embeddings=embeddings,
    )
    # 演示数据采用全量重建，确保 Schema、索引和向量维度保持一致。
    store.create_collection(recreate=True)
    inserted = store.insert_documents(documents)
    return len(document_paths), inserted, embeddings.dimension


def main() -> None:
    settings = get_settings()
    print("数据准备开始")
    print(f"- SQLite 原始工单：{INCIDENTS_SOURCE_PATH}")
    print(f"- Milvus 原始文档：{settings.documents_path}")

    sqlite_rows = prepare_sqlite(settings)
    document_count, chunk_count, dimension = prepare_milvus(settings)

    print(f"SQLite 完成：database={settings.sqlite_path}，rows={sqlite_rows}")
    print(
        "Milvus 完成："
        f"collection={settings.milvus_collection}，"
        f"documents={document_count}，chunks={chunk_count}，dimension={dimension}"
    )


if __name__ == "__main__":
    main()