wehub-resource-sync
456 lines
17 KiB
Python
456 lines
17 KiB
Python
"""Tests for directive-only metadata cleanup when (re)queuing to PENDING.
|
|
|
|
Covers:
|
|
* ``doc_status_reset_metadata`` — keeps long-lived directives
|
|
(``process_options`` / ``source_file``, legacy ``source_file_name``
|
|
tolerant) and drops every per-attempt timing/result field.
|
|
* ``doc_status_metadata_has_attempt_fields`` — the precise trigger used to
|
|
decide whether an already-PENDING doc needs normalising.
|
|
* ``_validate_and_fix_document_consistency`` — interrupted docs reset to a
|
|
clean PENDING (storage AND in-memory), already-PENDING docs carrying stale
|
|
per-attempt fields normalised in place, and the locked semantic for
|
|
docs that hold non-attempt custom metadata.
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from lightrag.base import DocProcessingStatus, DocStatus
|
|
from lightrag.lightrag import LightRAG
|
|
from lightrag.utils import EmbeddingFunc, Tokenizer
|
|
from lightrag.utils_pipeline import (
|
|
doc_status_metadata_has_attempt_fields,
|
|
doc_status_reset_metadata,
|
|
)
|
|
|
|
pytestmark = pytest.mark.offline
|
|
|
|
|
|
# A metadata blob carrying every per-attempt field a document could have
|
|
# accumulated across parse/analyze/process plus the two long-lived directives.
|
|
_FULL_ATTEMPT_METADATA = {
|
|
"process_options": "iF",
|
|
"source_file": "report.pdf",
|
|
"parse_start_time": 100,
|
|
"parse_end_time": 200,
|
|
"parse_stage_skipped": True,
|
|
"parse_warnings": {"docx": "missing paraId"},
|
|
"analyzing_start_time": 300,
|
|
"analyzing_end_time": 400,
|
|
"analyzing_stage_skipped": True,
|
|
"chunk_opts": "size=512, overlap=256",
|
|
"process_start_time": 500,
|
|
"process_end_time": 600,
|
|
"parse_format": "raw",
|
|
"parse_engine": "native",
|
|
"chunk_method": "fixed_token",
|
|
"skip_kg": True,
|
|
"mm_chunks": 3,
|
|
"hard_fallback_split": "256 -> 512",
|
|
}
|
|
|
|
|
|
class _SimpleTokenizerImpl:
|
|
def encode(self, content: str) -> list[int]:
|
|
return [ord(ch) for ch in content]
|
|
|
|
def decode(self, tokens: list[int]) -> str:
|
|
return "".join(chr(t) for t in tokens)
|
|
|
|
|
|
async def _dummy_embedding(texts: list[str]) -> np.ndarray:
|
|
return np.ones((len(texts), 8), dtype=float)
|
|
|
|
|
|
async def _dummy_llm(*args, **kwargs) -> str:
|
|
return "ok"
|
|
|
|
|
|
def _deterministic_chunking(
|
|
tokenizer,
|
|
content: str,
|
|
split_by_character,
|
|
split_by_character_only: bool,
|
|
chunk_overlap_token_size: int,
|
|
chunk_token_size: int,
|
|
) -> list[dict]:
|
|
return [{"tokens": 1, "content": content, "chunk_order_index": 0}]
|
|
|
|
|
|
def _status_to_text(status: object) -> str:
|
|
if isinstance(status, DocStatus):
|
|
return status.value
|
|
return str(status).replace("DocStatus.", "").lower()
|
|
|
|
|
|
async def _build_rag(tmp_path, test_name: str) -> LightRAG:
|
|
workspace = f"{test_name}_{uuid4().hex[:8]}"
|
|
rag = LightRAG(
|
|
working_dir=str(tmp_path / test_name),
|
|
workspace=workspace,
|
|
llm_model_func=_dummy_llm,
|
|
embedding_func=EmbeddingFunc(
|
|
embedding_dim=8,
|
|
max_token_size=8192,
|
|
func=_dummy_embedding,
|
|
),
|
|
tokenizer=Tokenizer("test-tokenizer", _SimpleTokenizerImpl()),
|
|
chunking_func=_deterministic_chunking,
|
|
max_parallel_insert=1,
|
|
)
|
|
await rag.initialize_storages()
|
|
return rag
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# doc_status_reset_metadata unit tests
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
def test_reset_metadata_keeps_directives_drops_attempt_fields():
|
|
result = doc_status_reset_metadata({"metadata": dict(_FULL_ATTEMPT_METADATA)})
|
|
assert result == {"process_options": "iF", "source_file": "report.pdf"}
|
|
|
|
|
|
def test_reset_metadata_normalizes_legacy_source_file_name():
|
|
result = doc_status_reset_metadata(
|
|
{"metadata": {"source_file_name": "legacy.docx", "process_options": "t"}}
|
|
)
|
|
assert result == {"process_options": "t", "source_file": "legacy.docx"}
|
|
|
|
|
|
def test_reset_metadata_prefers_new_source_file_over_legacy():
|
|
result = doc_status_reset_metadata(
|
|
{"metadata": {"source_file": "new.pdf", "source_file_name": "old.pdf"}}
|
|
)
|
|
assert result == {"source_file": "new.pdf"}
|
|
|
|
|
|
def test_reset_metadata_handles_empty_or_invalid_metadata():
|
|
assert doc_status_reset_metadata({"metadata": {}}) == {}
|
|
assert doc_status_reset_metadata({"metadata": None}) == {}
|
|
assert doc_status_reset_metadata({}) == {}
|
|
assert doc_status_reset_metadata(None) == {}
|
|
# blank directive values are not carried
|
|
assert (
|
|
doc_status_reset_metadata(
|
|
{"metadata": {"process_options": "", "source_file": ""}}
|
|
)
|
|
== {}
|
|
)
|
|
|
|
|
|
def test_has_attempt_fields_trigger():
|
|
assert doc_status_metadata_has_attempt_fields({"metadata": {"parse_end_time": 1}})
|
|
assert doc_status_metadata_has_attempt_fields(
|
|
{"metadata": {"process_options": "F", "chunk_opts": "x"}}
|
|
)
|
|
# only directives / custom non-attempt fields -> no trigger
|
|
assert not doc_status_metadata_has_attempt_fields(
|
|
{"metadata": {"process_options": "F", "source_file": "a.pdf"}}
|
|
)
|
|
assert not doc_status_metadata_has_attempt_fields(
|
|
{"metadata": {"custom_field": "x"}}
|
|
)
|
|
assert not doc_status_metadata_has_attempt_fields({"metadata": {}})
|
|
assert not doc_status_metadata_has_attempt_fields({"metadata": None})
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# _validate_and_fix_document_consistency integration tests
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_interrupted_reset_clears_attempt_metadata_in_storage_and_memory(
|
|
tmp_path,
|
|
):
|
|
rag = await _build_rag(tmp_path, "reset_clears_attempt_metadata")
|
|
try:
|
|
doc_id = "doc-interrupted"
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
await rag.full_docs.upsert(
|
|
{doc_id: {"content": "interrupted doc", "file_path": "report.pdf"}}
|
|
)
|
|
await rag.doc_status.upsert(
|
|
{
|
|
doc_id: {
|
|
"status": DocStatus.PROCESSING,
|
|
"content_summary": "interrupted",
|
|
"content_length": 15,
|
|
"chunks_count": 1,
|
|
"chunks_list": ["c-1"],
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"file_path": "report.pdf",
|
|
"track_id": "track-1",
|
|
"error_msg": "old error",
|
|
"metadata": dict(_FULL_ATTEMPT_METADATA),
|
|
}
|
|
}
|
|
)
|
|
|
|
to_process_docs = await rag.doc_status.get_docs_by_status(DocStatus.PROCESSING)
|
|
pipeline_status = {"latest_message": "", "history_messages": []}
|
|
await rag._validate_and_fix_document_consistency(
|
|
to_process_docs=to_process_docs,
|
|
pipeline_status=pipeline_status,
|
|
pipeline_status_lock=asyncio.Lock(),
|
|
)
|
|
|
|
# Storage: directives only, error cleared, chunks preserved.
|
|
stored = await rag.doc_status.get_by_id(doc_id)
|
|
assert stored is not None
|
|
assert _status_to_text(stored["status"]) == "pending"
|
|
assert stored["metadata"] == {
|
|
"process_options": "iF",
|
|
"source_file": "report.pdf",
|
|
}
|
|
assert stored["error_msg"] == ""
|
|
assert stored["chunks_list"] == ["c-1"]
|
|
|
|
# In-memory status_doc carried forward by workers is cleaned too.
|
|
assert to_process_docs[doc_id].metadata == {
|
|
"process_options": "iF",
|
|
"source_file": "report.pdf",
|
|
}
|
|
assert _status_to_text(to_process_docs[doc_id].status) == "pending"
|
|
finally:
|
|
await rag.finalize_storages()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stale_pending_normalized_clean_pending_untouched(tmp_path):
|
|
rag = await _build_rag(tmp_path, "stale_pending_normalized")
|
|
try:
|
|
stale_id = "doc-stale-pending"
|
|
clean_id = "doc-clean-pending"
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
await rag.full_docs.upsert(
|
|
{
|
|
stale_id: {"content": "stale", "file_path": "stale.pdf"},
|
|
clean_id: {"content": "clean", "file_path": "clean.pdf"},
|
|
}
|
|
)
|
|
await rag.doc_status.upsert(
|
|
{
|
|
stale_id: {
|
|
"status": DocStatus.PENDING,
|
|
"content_summary": "stale",
|
|
"content_length": 5,
|
|
"chunks_count": 0,
|
|
"chunks_list": [],
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"file_path": "stale.pdf",
|
|
"track_id": "track-stale",
|
|
"error_msg": "",
|
|
"metadata": {
|
|
"process_options": "F",
|
|
"parse_end_time": 999,
|
|
"analyzing_start_time": 111,
|
|
},
|
|
},
|
|
clean_id: {
|
|
"status": DocStatus.PENDING,
|
|
"content_summary": "clean",
|
|
"content_length": 5,
|
|
"chunks_count": 0,
|
|
"chunks_list": [],
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"file_path": "clean.pdf",
|
|
"track_id": "track-clean",
|
|
"error_msg": "",
|
|
"metadata": {"process_options": "F", "source_file": "clean.pdf"},
|
|
},
|
|
}
|
|
)
|
|
|
|
to_process_docs = await rag.doc_status.get_docs_by_status(DocStatus.PENDING)
|
|
pipeline_status = {"latest_message": "", "history_messages": []}
|
|
await rag._validate_and_fix_document_consistency(
|
|
to_process_docs=to_process_docs,
|
|
pipeline_status=pipeline_status,
|
|
pipeline_status_lock=asyncio.Lock(),
|
|
)
|
|
|
|
# Stale PENDING normalized to directives-only (storage + memory).
|
|
stale_stored = await rag.doc_status.get_by_id(stale_id)
|
|
assert stale_stored["metadata"] == {"process_options": "F"}
|
|
assert to_process_docs[stale_id].metadata == {"process_options": "F"}
|
|
|
|
# Clean PENDING left untouched: no rewrite (updated_at unchanged).
|
|
clean_stored = await rag.doc_status.get_by_id(clean_id)
|
|
assert clean_stored["metadata"] == {
|
|
"process_options": "F",
|
|
"source_file": "clean.pdf",
|
|
}
|
|
assert clean_stored["updated_at"] == now
|
|
finally:
|
|
await rag.finalize_storages()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_custom_metadata_semantic_lock(tmp_path):
|
|
"""Locks the chosen action semantic for already-PENDING docs:
|
|
|
|
* only custom (non-attempt) metadata, no stale key -> left untouched;
|
|
* stale key + custom field -> collapsed to directives-only (custom
|
|
dropped), matching the interrupted-reset path.
|
|
"""
|
|
rag = await _build_rag(tmp_path, "pending_custom_metadata_lock")
|
|
try:
|
|
custom_only_id = "doc-custom-only"
|
|
stale_plus_custom_id = "doc-stale-plus-custom"
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
await rag.full_docs.upsert(
|
|
{
|
|
custom_only_id: {"content": "c", "file_path": "custom.pdf"},
|
|
stale_plus_custom_id: {"content": "s", "file_path": "mixed.pdf"},
|
|
}
|
|
)
|
|
await rag.doc_status.upsert(
|
|
{
|
|
custom_only_id: {
|
|
"status": DocStatus.PENDING,
|
|
"content_summary": "c",
|
|
"content_length": 1,
|
|
"chunks_count": 0,
|
|
"chunks_list": [],
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"file_path": "custom.pdf",
|
|
"track_id": "track-custom",
|
|
"error_msg": "",
|
|
"metadata": {"custom_field": "keep-me"},
|
|
},
|
|
stale_plus_custom_id: {
|
|
"status": DocStatus.PENDING,
|
|
"content_summary": "s",
|
|
"content_length": 1,
|
|
"chunks_count": 0,
|
|
"chunks_list": [],
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"file_path": "mixed.pdf",
|
|
"track_id": "track-mixed",
|
|
"error_msg": "",
|
|
"metadata": {
|
|
"process_options": "F",
|
|
"custom_field": "drop-me",
|
|
"parse_start_time": 42,
|
|
},
|
|
},
|
|
}
|
|
)
|
|
|
|
to_process_docs = await rag.doc_status.get_docs_by_status(DocStatus.PENDING)
|
|
pipeline_status = {"latest_message": "", "history_messages": []}
|
|
await rag._validate_and_fix_document_consistency(
|
|
to_process_docs=to_process_docs,
|
|
pipeline_status=pipeline_status,
|
|
pipeline_status_lock=asyncio.Lock(),
|
|
)
|
|
|
|
custom_stored = await rag.doc_status.get_by_id(custom_only_id)
|
|
assert custom_stored["metadata"] == {"custom_field": "keep-me"}
|
|
assert custom_stored["updated_at"] == now
|
|
|
|
mixed_stored = await rag.doc_status.get_by_id(stale_plus_custom_id)
|
|
assert mixed_stored["metadata"] == {"process_options": "F"}
|
|
assert "custom_field" not in mixed_stored["metadata"]
|
|
finally:
|
|
await rag.finalize_storages()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_and_normalize_with_string_status(tmp_path):
|
|
"""Non-JSON backends (Redis/Mongo/OpenSearch/Postgres) hydrate
|
|
``DocProcessingStatus.status`` as a raw string (``"processing"`` /
|
|
``"pending"``) rather than a ``DocStatus`` member. Because ``DocStatus``
|
|
subclasses ``str``, the enum-membership checks in
|
|
``_validate_and_fix_document_consistency`` still match — this test pins
|
|
that behaviour so a future change to ``DocStatus`` (e.g. dropping the
|
|
``str`` base) cannot silently stop string-status docs from being cleaned.
|
|
|
|
Builds the ``to_process_docs`` objects directly with string statuses to
|
|
exercise the string path (``get_docs_by_status`` would return enum-valued
|
|
members from the in-memory store).
|
|
"""
|
|
rag = await _build_rag(tmp_path, "reset_string_status")
|
|
try:
|
|
processing_id = "doc-string-processing"
|
|
pending_id = "doc-string-pending"
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
await rag.full_docs.upsert(
|
|
{
|
|
processing_id: {"content": "p", "file_path": "proc.pdf"},
|
|
pending_id: {"content": "q", "file_path": "pend.pdf"},
|
|
}
|
|
)
|
|
|
|
# Status is a plain string, exactly as a non-JSON backend hydrates it.
|
|
to_process_docs = {
|
|
processing_id: DocProcessingStatus(
|
|
content_summary="p",
|
|
content_length=1,
|
|
file_path="proc.pdf",
|
|
status="processing",
|
|
created_at=now,
|
|
updated_at=now,
|
|
track_id="track-proc",
|
|
chunks_count=1,
|
|
chunks_list=["c-1"],
|
|
error_msg="old error",
|
|
metadata=dict(_FULL_ATTEMPT_METADATA),
|
|
),
|
|
pending_id: DocProcessingStatus(
|
|
content_summary="q",
|
|
content_length=1,
|
|
file_path="pend.pdf",
|
|
status="pending",
|
|
created_at=now,
|
|
updated_at=now,
|
|
track_id="track-pend",
|
|
chunks_count=0,
|
|
chunks_list=[],
|
|
error_msg="",
|
|
metadata={"process_options": "F", "analyzing_start_time": 7},
|
|
),
|
|
}
|
|
assert isinstance(to_process_docs[processing_id].status, str)
|
|
assert not isinstance(to_process_docs[processing_id].status, DocStatus)
|
|
|
|
pipeline_status = {"latest_message": "", "history_messages": []}
|
|
await rag._validate_and_fix_document_consistency(
|
|
to_process_docs=to_process_docs,
|
|
pipeline_status=pipeline_status,
|
|
pipeline_status_lock=asyncio.Lock(),
|
|
)
|
|
|
|
# String "processing" is reset to a clean PENDING.
|
|
proc_stored = await rag.doc_status.get_by_id(processing_id)
|
|
assert proc_stored is not None
|
|
assert _status_to_text(proc_stored["status"]) == "pending"
|
|
assert proc_stored["metadata"] == {
|
|
"process_options": "iF",
|
|
"source_file": "report.pdf",
|
|
}
|
|
assert to_process_docs[processing_id].metadata == {
|
|
"process_options": "iF",
|
|
"source_file": "report.pdf",
|
|
}
|
|
|
|
# String "pending" carrying a stale field is normalized in place.
|
|
pend_stored = await rag.doc_status.get_by_id(pending_id)
|
|
assert pend_stored is not None
|
|
assert pend_stored["metadata"] == {"process_options": "F"}
|
|
assert to_process_docs[pending_id].metadata == {"process_options": "F"}
|
|
finally:
|
|
await rag.finalize_storages()
|