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()