From c46d08c6130ecfe24b2e452fa6bdbc7b40a845e9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 12 May 2026 14:13:02 +0800 Subject: [PATCH] fix: preserve full document chunks --- apps/worker/app/core/tasks/kb_tasks.py | 29 ---- .../contract/test_parse_task_contract.py | 16 ++- .../shared/services/job_lifecycle_sync.py | 23 +-- .../services/retrieval/publication_service.py | 134 ------------------ 4 files changed, 12 insertions(+), 190 deletions(-) diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index e614e9c1b..19acc3605 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/kb_tasks.py @@ -676,33 +676,6 @@ def _parse(job_id: str, user_id: str | None): ) result_s3_key = result_bundle.zip_key - publication_chunks = chunks - dedup_stats = None - try: - from shared.services.retrieval.publication_service import RetrievalPublicationService - - with get_sync_db_context() as db: - job_record = db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - if job_record: - gc_namespace = ( - JobMetadataHelper.get_field(job_metadata, "namespace") - or "default" - ) - publication_chunks, dedup_stats = ( - RetrievalPublicationService.garbage_collect_and_dedup_local_media( - db, - job_id=job_id, - user_id=str(job_record.user_id), - namespace=gc_namespace, - add_dir=str(add_dir) if add_dir else "", - chunks=chunks, - ) - ) - except Exception as e: - logger.error(f"[{job_id}] GC failed (non-fatal): {e}") - stored_count = 0 lifecycle_service.update_progress( @@ -713,14 +686,12 @@ def _parse(job_id: str, user_id: str | None): lifecycle_service.finalize_job_success( job_id=job_id, chunks=chunks, - publication_chunks=publication_chunks, result_s3_key=result_s3_key, checksum=checksum_value, zip_size=zip_size, stored_count=stored_count, delivery_mode="url", section_summaries=section_summaries, - chunk_dedup_stats=dedup_stats, ) logger.info( diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 9c0edfad7..12f672267 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -768,6 +768,10 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", False) + def fake_cleanup_task_workspace(workspace_dir: str | None) -> bool: + captured_artifacts["workspace_dir"] = workspace_dir + return True + def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { "exists": storage_key == s3_key, @@ -878,6 +882,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(kb_tasks, "cleanup_task_workspace", fake_cleanup_task_workspace) result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") @@ -896,6 +901,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: "table_chunks": 0, "total_pages": None, } + workspace_dir = Path(str(captured_artifacts["workspace_dir"])) + assert list(workspace_dir.rglob("images/duplicate.png")) with engine.begin() as connection: job_result_row = ( @@ -946,13 +953,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: ) assert job_chunk_ids == ["duplicate-text", "duplicate-image", "new-text"] - assert document_chunk_ids == ["new-text"] - assert dict(job_result_row["document_metadata"])["chunk_dedup"] == { - "total_incoming": 3, - "duplicates_skipped": 2, - "new_chunks_inserted": 1, - "overlap_by_document": {existing_document_id: 2}, - } + assert document_chunk_ids == ["duplicate-text", "duplicate-image", "new-text"] + assert "chunk_overlap" not in dict(job_result_row["document_metadata"] or {}) def test_should_initialize_billing_once_for_concurrent_parse_tasks( diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index c43e6635d..cd9dcc8ef 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -57,17 +57,15 @@ def finalize_job_success( checksum: str, zip_size: int, chunks: Optional[List[Dict[str, Any]]] = None, - publication_chunks: Optional[List[Dict[str, Any]]] = None, stored_count: int = 0, delivery_mode: str = "url", section_summaries: Optional[Dict[str, str]] = None, - chunk_dedup_stats: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Finalize a successful job in a single atomic transaction. Steps (all within one DB transaction): 1. Upsert JobResult + replace full result chunks - 2. Publish retrieval state from publication chunks + 2. Publish document state from full result chunks 3. Mark job as DONE via state machine (CAS) 4. Create WebhookEvent if webhook_enabled 5. COMMIT @@ -85,15 +83,9 @@ def finalize_job_success( inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=zip_size, - chunk_dedup_stats=chunk_dedup_stats, ) normalized_chunks = chunks or [] - normalized_publication_chunks = ( - publication_chunks - if publication_chunks is not None - else normalized_chunks - ) self._replace_chunks(db, job_result.id, normalized_chunks) previous_document_scope = ( self._retrieval_publication.get_existing_document_scope( @@ -106,7 +98,7 @@ def finalize_job_success( db, job_id=job_id, job_result_id=job_result.id, - chunks=normalized_publication_chunks, + chunks=normalized_chunks, ) ) if published_document_state is not None and not published_document_state.get("skipped_all_duplicate"): @@ -313,7 +305,6 @@ def _upsert_job_result( inline_payload: Optional[Dict[str, Any]] = None, result_s3_key: Optional[str] = None, result_size: Optional[int] = None, - chunk_dedup_stats: Optional[Dict[str, Any]] = None, ) -> JobResult: """Create or update JobResult row.""" result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) @@ -321,24 +312,16 @@ def _upsert_job_result( if existing: existing.delivery_mode = delivery_mode - doc_meta = existing.document_metadata or {} - if chunk_dedup_stats: - doc_meta["chunk_dedup"] = chunk_dedup_stats - existing.document_metadata = doc_meta existing.inline_payload = inline_payload existing.result_s3_key = result_s3_key existing.result_size = result_size db.flush() return existing - doc_meta = {} - if chunk_dedup_stats: - doc_meta["chunk_dedup"] = chunk_dedup_stats - job_result = JobResult( job_id=job_id, delivery_mode=delivery_mode, - document_metadata=doc_meta, + document_metadata={}, inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=result_size, diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 5322fa809..80aaa4ae0 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -8,7 +8,6 @@ from __future__ import annotations -from collections import defaultdict from datetime import datetime, timezone from typing import Any, Dict, List, Optional from uuid import uuid4 @@ -36,139 +35,6 @@ def utc_now_naive() -> datetime: class RetrievalPublicationService: - - # ── Chunk-level content-hash dedup ────────────────────────────────── - # Mirrors graph_builder._dedup_chunks_by_content but operates on - # the DB document_chunks table instead of local knowledge_graph.json. - - @staticmethod - def _collect_existing_chunk_id_map( - db: Session, - *, - user_id: str, - namespace: str, - ) -> Dict[str, str]: - """Return {chunk_id -> document_id} for all active document chunks - in the given (user_id, namespace) scope. - - Only considers chunks belonging to the *current* revision of each - active document (Document.current_job_result_id == DocumentChunk.job_result_id). - """ - rows = db.execute( - select(DocumentChunk.chunk_id, DocumentChunk.document_id) - .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), - ) - .where( - Document.user_id == user_id, - Document.namespace == namespace, - Document.status == "active", - ) - ).all() - return {row[0]: row[1] for row in rows} - - @staticmethod - def _dedup_chunks_by_content( - new_chunks: List[Dict[str, Any]], - existing_chunk_map: Dict[str, str], - ) -> tuple[List[Dict[str, Any]], Dict[str, int]]: - """Filter new_chunks: discard any whose chunk_id already exists. - - Uses the same deterministic know_id (content-hash) comparison as - graph_builder._dedup_chunks_by_content. - - Returns: - (deduped_chunks, overlap_by_document) - - deduped_chunks: chunks whose chunk_id is NOT in existing_chunk_map - - overlap_by_document: {document_id: count} of skipped chunks per - existing document (for observability logging) - """ - overlap_by_document: Dict[str, int] = defaultdict(int) - deduped: List[Dict[str, Any]] = [] - skipped = 0 - - for chunk in new_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if cid and cid in existing_chunk_map: - skipped += 1 - overlap_by_document[existing_chunk_map[cid]] += 1 - else: - deduped.append(chunk) - - if skipped > 0: - logger.warning( - f"📊 DB chunk dedup: {skipped}/{len(new_chunks)} duplicate chunks " - f"skipped (by chunk_id), {len(deduped)} new chunks to insert. " - f"Overlap by document: {dict(overlap_by_document)}" - ) - return deduped, dict(overlap_by_document) - - @classmethod - def garbage_collect_and_dedup_local_media( - cls, - db: Session, - *, - job_id: str, - user_id: str, - namespace: str, - add_dir: str, - chunks: List[Dict[str, Any]], - ) -> tuple[List[Dict[str, Any]], Dict[str, Any]]: - """ - Deduplicates chunks against the DB and physically deletes associated redundant - media files (images/tables) from the local add_dir before ZIP packaging. - Returns the deduplicated chunks. - """ - import os - - logger.info(f"[{job_id}] Starting local GC for redundant media files in namespace: {namespace}...") - try: - existing_map = cls._collect_existing_chunk_id_map( - db, user_id=user_id, namespace=namespace - ) - deduped_chunks, overlap = cls._dedup_chunks_by_content(chunks, existing_map) - - stats = { - "total_incoming": len(chunks), - "duplicates_skipped": len(chunks) - len(deduped_chunks), - "new_chunks_inserted": len(deduped_chunks), - "overlap_by_document": overlap, - } - - if len(deduped_chunks) < len(chunks): - active_paths = set() - for c in deduped_chunks: - fp = c.get("metadata", {}).get("file_path") or c.get("file_path") - if fp: - active_paths.add(fp) - - deleted_count = 0 - if add_dir and os.path.exists(add_dir): - for c in chunks: - fp = c.get("metadata", {}).get("file_path") or c.get("file_path") - if fp and fp not in active_paths: - abs_path = os.path.join(add_dir, fp) - if os.path.exists(abs_path): - os.remove(abs_path) - deleted_count += 1 - - logger.info(f"[{job_id}] GC complete: permanently removed {deleted_count} redundant local media files.") - return deduped_chunks, stats - else: - logger.info(f"[{job_id}] GC complete: no redundant chunks found.") - return chunks, stats - except Exception as e: - logger.error(f"[{job_id}] GC failed (non-fatal): {e}") - stats = { - "total_incoming": len(chunks), - "duplicates_skipped": 0, - "new_chunks_inserted": len(chunks), - "overlap_by_document": {}, - } - return chunks, stats - # ── Public API ────────────────────────────────────────────────────── def get_existing_document_scope(