83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
|
"""Validated, all-or-nothing ingestion of a document into Milvus."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from uuid import NAMESPACE_URL, uuid5
|
||
|
|
|
||
|
|
from rag.chunk_config import ChunkConfig, resolve_chunk_config
|
||
|
|
from rag.chunking import chunk_document
|
||
|
|
from rag.cleaning import clean_document_text
|
||
|
|
from rag.embedding import EMBEDDING_DIMENSION, embed_texts
|
||
|
|
from rag.milvus_delete import _escape_filter_value
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger("rag.ingestion")
|
||
|
|
|
||
|
|
|
||
|
|
async def ingest_document_atomic(
|
||
|
|
document_text: str,
|
||
|
|
doc_id: str,
|
||
|
|
title: str,
|
||
|
|
collection_name: str,
|
||
|
|
strategy: str,
|
||
|
|
*,
|
||
|
|
milvus_client,
|
||
|
|
embedder=embed_texts,
|
||
|
|
config: ChunkConfig | None = None,
|
||
|
|
) -> dict:
|
||
|
|
"""Validate, embed, and insert one document without leaving partial rows."""
|
||
|
|
if not doc_id or not doc_id.strip():
|
||
|
|
raise ValueError("doc_id must not be empty")
|
||
|
|
if not title or not title.strip():
|
||
|
|
raise ValueError("title must not be empty")
|
||
|
|
if not collection_name or not collection_name.strip():
|
||
|
|
raise ValueError("collection_name must not be empty")
|
||
|
|
|
||
|
|
resolved_config = config or resolve_chunk_config()
|
||
|
|
cleaned = clean_document_text(document_text)
|
||
|
|
chunking = chunk_document(cleaned.text, strategy, config=resolved_config)
|
||
|
|
if not chunking.chunks:
|
||
|
|
raise ValueError("document produced no chunks")
|
||
|
|
|
||
|
|
try:
|
||
|
|
vectors = await embedder([chunk.text for chunk in chunking.chunks])
|
||
|
|
if len(vectors) != len(chunking.chunks) or any(
|
||
|
|
len(vector) != EMBEDDING_DIMENSION for vector in vectors
|
||
|
|
):
|
||
|
|
raise ValueError(f"Embedding dimension must be {EMBEDDING_DIMENSION}")
|
||
|
|
rows = [
|
||
|
|
{
|
||
|
|
"chunk_id": str(uuid5(NAMESPACE_URL, f"{doc_id}:{index}")),
|
||
|
|
"doc_id": doc_id,
|
||
|
|
"title": title,
|
||
|
|
"section_title": chunk.section_title or "",
|
||
|
|
"text": chunk.text,
|
||
|
|
"strategy": strategy,
|
||
|
|
"vector": vector,
|
||
|
|
}
|
||
|
|
for index, (chunk, vector) in enumerate(zip(chunking.chunks, vectors))
|
||
|
|
]
|
||
|
|
await milvus_client.insert(collection_name=collection_name, data=rows)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("atomic Milvus ingestion failed for doc_id=%s", doc_id)
|
||
|
|
try:
|
||
|
|
await milvus_client.delete(
|
||
|
|
collection_name=collection_name,
|
||
|
|
filter=f'doc_id == "{_escape_filter_value(doc_id)}"',
|
||
|
|
)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("failed to clean up document rows for doc_id=%s", doc_id)
|
||
|
|
raise
|
||
|
|
|
||
|
|
return {
|
||
|
|
"doc_id": doc_id,
|
||
|
|
"collection_name": collection_name,
|
||
|
|
"strategy": strategy,
|
||
|
|
"actual_strategy": chunking.actual_strategy,
|
||
|
|
"degraded": chunking.degraded,
|
||
|
|
"warning": chunking.warning,
|
||
|
|
"cleaning_changed": cleaned.changed,
|
||
|
|
"cleaning_warnings": cleaned.warnings,
|
||
|
|
"chunk_count": len(rows),
|
||
|
|
}
|