Agentic_Rag/scripts/prepare_data.py

from __future__ import annotations

import re
from pathlib import Path

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


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
    }
    policy_type = metadata_lines.get("政策类型", "")
    category = metadata_lines.get("适用类目", "")
    if not policy_type or not category:
        raise ValueError(f"文档缺少政策类型或适用类目：{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="commerce_policy",
            chunk_index=index,
            policy_type=policy_type,
            category=category,
            metadata={"title": title},
        )
        for index, section in enumerate(sections)
    ]


def prepare_sqlite(settings: Settings) -> int:
    return initialize_database(
        path=settings.sqlite_path,
        source_path=settings.orders_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 原始订单：{settings.orders_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()