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