60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
"""Development-only Mock knowledge documents and Milvus importer."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from uuid import uuid5, NAMESPACE_URL
|
|
|
|
from rag.chunk_config import resolve_chunk_config
|
|
from rag.chunking import chunk_document
|
|
from rag.cleaning import clean_document_text
|
|
from rag.document_parser import parse_document
|
|
from rag.embedding import embed_texts
|
|
|
|
|
|
_ROOT = Path(__file__).resolve().parent.parent
|
|
MOCK_DATA_DIR = _ROOT / "data" / "mock_knowledge"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MockDocument:
|
|
doc_id: str
|
|
title: str
|
|
collection: str
|
|
strategy: str
|
|
path: Path
|
|
|
|
|
|
MOCK_DOCUMENTS = (
|
|
MockDocument("mock-faq-001", "常见问题示例", "fin_faq", "qa_pair", MOCK_DATA_DIR / "faq.md"),
|
|
MockDocument("mock-fund-doc-001", "基金产品说明示例", "fin_fund_doc", "chapter_semantic", MOCK_DATA_DIR / "fund_product.md"),
|
|
MockDocument("mock-policy-001", "政策法规示例", "fin_policy", "chapter_semantic", MOCK_DATA_DIR / "policy.md"),
|
|
)
|
|
|
|
|
|
async def ingest_mock_documents(milvus_client, *, embedder=embed_texts) -> None:
|
|
"""Import all checked-in Markdown samples into their Milvus collections."""
|
|
config = resolve_chunk_config()
|
|
for document in MOCK_DOCUMENTS:
|
|
result = chunk_document(
|
|
clean_document_text(parse_document(document.path)).text,
|
|
document.strategy,
|
|
config=config,
|
|
)
|
|
vectors = await embedder([chunk.text for chunk in result.chunks])
|
|
rows = []
|
|
for index, (chunk, vector) in enumerate(zip(result.chunks, vectors)):
|
|
chunk_id = str(uuid5(NAMESPACE_URL, f"{document.doc_id}:{index}"))
|
|
rows.append(
|
|
{
|
|
"chunk_id": chunk_id,
|
|
"doc_id": document.doc_id,
|
|
"title": document.title,
|
|
"section_title": chunk.section_title or "",
|
|
"text": chunk.text,
|
|
"strategy": document.strategy,
|
|
"vector": vector,
|
|
}
|
|
)
|
|
await milvus_client.insert(collection_name=document.collection, data=rows)
|