diff --git a/minirag/minirag.py b/minirag/minirag.py index fed6b19..396d403 100644 --- a/minirag/minirag.py +++ b/minirag/minirag.py @@ -33,6 +33,8 @@ StorageNameSpace, QueryParam, DocStatus, + DocProcessingStatus, + TextChunkSchema, ) @@ -353,40 +355,55 @@ async def ainsert( if isinstance(ids, str): ids = [ids] - await self.apipeline_enqueue_documents(input, ids) - await self.apipeline_process_enqueue_documents( - split_by_character, split_by_character_only - ) - - # Perform additional entity extraction as per original ainsert logic - inserting_chunks = { - compute_mdhash_id(dp["content"], prefix="chunk-"): { - **dp, - "full_doc_id": doc_id, - } - for doc_id, status_doc in ( - await self.doc_status.get_docs_by_status(DocStatus.PROCESSED) - ).items() - for dp in self.chunking_func( - status_doc.content, - self.chunk_overlap_token_size, - self.chunk_token_size, - self.tiktoken_model_name, + extraction_failures = [] + try: + await self.apipeline_enqueue_documents(input, ids) + documents_to_extract = await self.apipeline_process_enqueue_documents( + split_by_character, split_by_character_only ) - } - if inserting_chunks: - logger.info("Performing entity extraction on newly processed chunks") - await extract_entities( - inserting_chunks, - knowledge_graph_inst=self.chunk_entity_relation_graph, - entity_vdb=self.entities_vdb, - entity_name_vdb=self.entity_name_vdb, - relationships_vdb=self.relationships_vdb, - global_config=asdict(self), - ) - - await self._insert_done() + for doc_id, (status_doc, inserting_chunks) in ( + documents_to_extract.items() + ): + try: + logger.info( + "Performing entity extraction for %s newly processed chunks", + len(inserting_chunks), + ) + await extract_entities( + inserting_chunks, + knowledge_graph_inst=self.chunk_entity_relation_graph, + entity_vdb=self.entities_vdb, + entity_name_vdb=self.entity_name_vdb, + relationships_vdb=self.relationships_vdb, + global_config=asdict(self), + ) + except Exception as error: + await self._update_document_status( + doc_id, + status_doc, + DocStatus.FAILED, + len(inserting_chunks), + error=str(error), + ) + extraction_failures.append((doc_id, error)) + logger.exception("Entity extraction failed for document %s", doc_id) + continue + + await self._update_document_status( + doc_id, + status_doc, + DocStatus.PROCESSED, + len(inserting_chunks), + ) + finally: + await self._insert_done() + + if extraction_failures: + failed_doc_ids = ", ".join(doc_id for doc_id, _ in extraction_failures) + raise RuntimeError( + f"Entity extraction failed for document(s): {failed_doc_ids}" + ) from extraction_failures[0][1] async def apipeline_enqueue_documents( self, input: str | list[str], ids: list[str] | None = None @@ -448,15 +465,39 @@ async def apipeline_enqueue_documents( await self.doc_status.upsert(new_docs) logger.info(f"Stored {len(new_docs)} new unique documents") + async def _update_document_status( + self, + doc_id: str, + status_doc: DocProcessingStatus, + status: DocStatus, + chunks_count: int, + error: str | None = None, + ) -> None: + status_data = { + "status": status, + "chunks_count": chunks_count, + "content": status_doc.content, + "content_summary": status_doc.content_summary, + "content_length": status_doc.content_length, + "created_at": status_doc.created_at, + "updated_at": datetime.now().isoformat(), + } + if error is not None: + status_data["error"] = error + await self.doc_status.upsert({doc_id: status_data}) + async def apipeline_process_enqueue_documents( self, split_by_character: str | None = None, split_by_character_only: bool = False, - ) -> None: + ) -> dict[str, tuple[DocProcessingStatus, dict[str, TextChunkSchema]]]: """ - Process pending documents by splitting them into chunks, processing - each chunk for entity and relation extraction, and updating the - document status. + Stage pending documents for graph extraction. + + Documents remain in PROCESSING until entity and relation extraction + completes in ``ainsert``. The returned mapping contains only chunks + staged during this invocation, so already processed documents are not + extracted again. """ processing_docs, failed_docs, pending_docs = await asyncio.gather( self.doc_status.get_docs_by_status(DocStatus.PROCESSING), @@ -471,13 +512,14 @@ async def apipeline_process_enqueue_documents( } if not to_process_docs: logger.info("No documents to process") - return + return {} docs_batches = [ list(to_process_docs.items())[i : i + self.max_parallel_insert] for i in range(0, len(to_process_docs), self.max_parallel_insert) ] logger.info(f"Number of batches to process: {len(docs_batches)}") + documents_to_extract = {} for batch_idx, docs_batch in enumerate(docs_batches): for doc_id, status_doc in docs_batch: @@ -493,25 +535,33 @@ async def apipeline_process_enqueue_documents( self.tiktoken_model_name, ) } - await asyncio.gather( - self.chunks_vdb.upsert(chunks), - self.full_docs.upsert({doc_id: {"content": status_doc.content}}), - self.text_chunks.upsert(chunks), + await self._update_document_status( + doc_id, + status_doc, + DocStatus.PROCESSING, + len(chunks), ) - await self.doc_status.upsert( - { - doc_id: { - "status": DocStatus.PROCESSED, - "chunks_count": len(chunks), - "content": status_doc.content, - "content_summary": status_doc.content_summary, - "content_length": status_doc.content_length, - "created_at": status_doc.created_at, - "updated_at": datetime.now().isoformat(), - } - } - ) - logger.info("Document processing pipeline completed") + try: + await asyncio.gather( + self.chunks_vdb.upsert(chunks), + self.full_docs.upsert( + {doc_id: {"content": status_doc.content}} + ), + self.text_chunks.upsert(chunks), + ) + except Exception as error: + await self._update_document_status( + doc_id, + status_doc, + DocStatus.FAILED, + len(chunks), + error=str(error), + ) + raise + documents_to_extract[doc_id] = (status_doc, chunks) + + logger.info("Document staging pipeline completed") + return documents_to_extract async def _insert_done(self): tasks = [] @@ -608,4 +658,4 @@ async def _delete_by_entity_done(self): if storage_inst is None: continue tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback()) - await asyncio.gather(*tasks) \ No newline at end of file + await asyncio.gather(*tasks) diff --git a/minirag/utils.py b/minirag/utils.py index a4f0e60..ac0fffa 100644 --- a/minirag/utils.py +++ b/minirag/utils.py @@ -6,6 +6,7 @@ import logging import os import re +import unicodedata from dataclasses import dataclass from functools import wraps from hashlib import md5 @@ -64,8 +65,15 @@ def compute_args_hash(*args, cache_type: str | None = None) -> str: def clean_text(text: str) -> str: - """Clean text by removing null bytes (0x00) and whitespace""" - return text.strip().replace("\x00", "") + """Clean text by removing null bytes (0x00) and whitespace, and normalising + its Unicode composition. + + See :func:`clean_str` for why the NFC pass matters. Here it also decides + document identity: ``doc_id`` is an md5 of this function's output, so + without it the same document typed on one keyboard and copied from another + hashes differently and is indexed twice. + """ + return unicodedata.normalize("NFC", text.strip().replace("\x00", "")) def get_content_summary(content: str, max_length: int = 100) -> str: @@ -179,7 +187,15 @@ def clean_str(input: Any) -> str: result = html.unescape(input.strip()) # https://stackoverflow.com/questions/4324790/removing-control-characters-from-a-string-in-python - return re.sub(r"[\x00-\x1f\x7f-\x9f]", "", result) + result = re.sub(r"[\x00-\x1f\x7f-\x9f]", "", result) + # Accented scripts have more than one valid encoding of the same text: "ệ" + # is either U+1EC7 or "e" + U+0323 + U+0302. Both render identically, so the + # difference is invisible, but the strings compare unequal and hash + # differently — and these strings become identity keys (graph node ids and + # the md5 behind every `ent-` vector id). Vietnamese hits this constantly + # because sources disagree: IMEs emit NFC, the macOS filesystem emits NFD. + # Pure-ASCII text is unaffected, which is why this went unnoticed. + return unicodedata.normalize("NFC", result) def is_float_regex(value): diff --git a/tests/test_incremental_indexing.py b/tests/test_incremental_indexing.py new file mode 100644 index 0000000..6eafe54 --- /dev/null +++ b/tests/test_incremental_indexing.py @@ -0,0 +1,143 @@ +import unittest +from datetime import datetime +from unittest.mock import AsyncMock, patch + +from minirag.base import DocProcessingStatus, DocStatus +from minirag.minirag import MiniRAG + + +class FakeDocStatusStorage: + def __init__(self, docs_by_status=None): + self.docs_by_status = docs_by_status or {} + self.upserts = [] + + async def get_docs_by_status(self, status): + return self.docs_by_status.get(status, {}) + + async def upsert(self, data): + self.upserts.append(data) + + +class FakeStorage: + def __init__(self): + self.upserts = [] + + async def upsert(self, data): + self.upserts.append(data) + + +def make_status_doc(content="A document"): + now = datetime.now().isoformat() + return DocProcessingStatus( + content=content, + content_summary=content, + content_length=len(content), + status=DocStatus.PENDING, + created_at=now, + updated_at=now, + ) + + +def make_chunk(content="A document"): + return { + "tokens": 2, + "content": content, + "chunk_order_index": 0, + "full_doc_id": "doc-1", + } + + +def make_bare_rag(doc_status): + rag = object.__new__(MiniRAG) + rag.doc_status = doc_status + rag.chunk_entity_relation_graph = object() + rag.entities_vdb = object() + rag.entity_name_vdb = object() + rag.relationships_vdb = object() + rag.chunk_overlap_token_size = 0 + rag.chunk_token_size = 100 + rag.tiktoken_model_name = "gpt-4o-mini" + rag.max_parallel_insert = 2 + rag.chunks_vdb = FakeStorage() + rag.full_docs = FakeStorage() + rag.text_chunks = FakeStorage() + rag._insert_done = AsyncMock() + return rag + + +class IncrementalIndexingTests(unittest.IsolatedAsyncioTestCase): + async def test_duplicate_insert_does_not_extract_processed_documents(self): + existing_doc = make_status_doc("Already indexed") + doc_status = FakeDocStatusStorage( + {DocStatus.PROCESSED: {"doc-existing": existing_doc}} + ) + rag = make_bare_rag(doc_status) + rag.chunking_func = lambda *_: [make_chunk("Already indexed")] + rag.apipeline_enqueue_documents = AsyncMock() + rag.apipeline_process_enqueue_documents = AsyncMock(return_value={}) + + with patch("minirag.minirag.extract_entities", new_callable=AsyncMock) as extract: + await rag.ainsert("Already indexed") + + extract.assert_not_awaited() + rag._insert_done.assert_awaited_once() + + async def test_successful_extraction_marks_only_staged_document_processed(self): + status_doc = make_status_doc() + doc_status = FakeDocStatusStorage() + rag = make_bare_rag(doc_status) + staged_chunks = {"chunk-new": make_chunk()} + rag.apipeline_enqueue_documents = AsyncMock() + rag.apipeline_process_enqueue_documents = AsyncMock( + return_value={"doc-new": (status_doc, staged_chunks)} + ) + + with ( + patch("minirag.minirag.asdict", return_value={}), + patch("minirag.minirag.extract_entities", new_callable=AsyncMock) as extract, + ): + await rag.ainsert("A document") + + extract.assert_awaited_once() + self.assertEqual(doc_status.upserts[-1]["doc-new"]["status"], DocStatus.PROCESSED) + + async def test_failed_extraction_marks_document_failed(self): + status_doc = make_status_doc() + doc_status = FakeDocStatusStorage() + rag = make_bare_rag(doc_status) + staged_chunks = {"chunk-new": make_chunk()} + rag.apipeline_enqueue_documents = AsyncMock() + rag.apipeline_process_enqueue_documents = AsyncMock( + return_value={"doc-new": (status_doc, staged_chunks)} + ) + + with ( + patch("minirag.minirag.asdict", return_value={}), + patch( + "minirag.minirag.extract_entities", + new_callable=AsyncMock, + side_effect=RuntimeError("LLM unavailable"), + ), + self.assertRaisesRegex(RuntimeError, "doc-new"), + ): + await rag.ainsert("A document") + + self.assertEqual(doc_status.upserts[-1]["doc-new"]["status"], DocStatus.FAILED) + self.assertEqual(doc_status.upserts[-1]["doc-new"]["error"], "LLM unavailable") + rag._insert_done.assert_awaited_once() + + async def test_staging_keeps_document_processing_until_extraction(self): + status_doc = make_status_doc() + doc_status = FakeDocStatusStorage({DocStatus.PENDING: {"doc-new": status_doc}}) + rag = make_bare_rag(doc_status) + rag.chunking_func = lambda *_: [make_chunk()] + + staged = await rag.apipeline_process_enqueue_documents() + + self.assertIn("doc-new", staged) + self.assertEqual(doc_status.upserts[0]["doc-new"]["status"], DocStatus.PROCESSING) + self.assertEqual(len(rag.chunks_vdb.upserts), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_unicode_normalization.py b/tests/test_unicode_normalization.py new file mode 100644 index 0000000..b9ff950 --- /dev/null +++ b/tests/test_unicode_normalization.py @@ -0,0 +1,67 @@ +import unicodedata +import unittest + +from minirag.utils import clean_str, clean_text, compute_mdhash_id + +# Same text, both valid Unicode encodings of it. They render identically. +VI_TEXT = "Nguyễn Văn An đi đánh cầu lông ở Nhà thi đấu Phú Thọ." +VI_ENTITY = "PHÚ THỌ" + + +def nfc(s): + return unicodedata.normalize("NFC", s) + + +def nfd(s): + return unicodedata.normalize("NFD", s) + + +class TestUnicodeNormalization(unittest.TestCase): + def test_the_two_forms_really_do_differ(self): + """Guard the premise: without normalisation these are distinct strings.""" + self.assertNotEqual(nfc(VI_TEXT), nfd(VI_TEXT)) + self.assertNotEqual(len(nfc(VI_TEXT).encode()), len(nfd(VI_TEXT).encode())) + + def test_clean_text_collapses_both_forms(self): + self.assertEqual(clean_text(nfc(VI_TEXT)), clean_text(nfd(VI_TEXT))) + + def test_clean_str_collapses_both_forms(self): + self.assertEqual(clean_str(nfc(VI_ENTITY)), clean_str(nfd(VI_ENTITY))) + + def test_document_id_is_stable_across_forms(self): + """doc_id drives dedup; an unstable one re-indexes the same document.""" + a = compute_mdhash_id(clean_text(nfc(VI_TEXT)), prefix="doc-") + b = compute_mdhash_id(clean_text(nfd(VI_TEXT)), prefix="doc-") + self.assertEqual(a, b) + + def test_entity_vector_id_is_stable_across_forms(self): + a = compute_mdhash_id(clean_str(nfc(VI_ENTITY)), prefix="ent-") + b = compute_mdhash_id(clean_str(nfd(VI_ENTITY)), prefix="ent-") + self.assertEqual(a, b) + + def test_output_is_nfc(self): + for form in (nfc, nfd): + self.assertEqual(clean_text(form(VI_TEXT)), nfc(clean_text(form(VI_TEXT)))) + self.assertEqual(clean_str(form(VI_ENTITY)), nfc(clean_str(form(VI_ENTITY)))) + + def test_upper_before_clean_str_still_converges(self): + """operate.py upper-cases entity names before cleaning them.""" + self.assertEqual( + clean_str(nfd("Nguyễn Văn An").upper()), + clean_str(nfc("Nguyễn Văn An").upper()), + ) + + def test_existing_behaviour_is_preserved(self): + self.assertEqual(clean_text(" hello\x00 world "), "hello world") + self.assertEqual(clean_str(" &text\x07 "), "&text") + self.assertEqual(clean_str(123), 123) # non-str passes through untouched + + def test_ascii_is_untouched(self): + """The change must be a no-op for text that has only one encoding.""" + for s in ("Wolfgang Schulz", "the badminton court", "LIHUA"): + self.assertEqual(clean_str(s), s) + self.assertEqual(clean_text(s), s) + + +if __name__ == "__main__": + unittest.main()