Skip to content

fix: preserve full document chunks #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
suguanYang merged 1 commit into from
May 12, 2026
Merged

fix: preserve full document chunks #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 0 additions & 29 deletions apps/worker/app/core/tasks/kb_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down
16 changes: 9 additions & 7 deletions apps/worker/tests/contract/test_parse_task_contract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")

Expand All @@ -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 = (
Expand Down Expand Up @@ -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(
Expand Down
23 changes: 3 additions & 20 deletions packages/shared-python/shared/services/job_lifecycle_sync.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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"):
Expand Down Expand Up @@ -313,32 +305,23 @@ 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))
existing = result.scalar_one_or_none()

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,
Expand Down
134 changes: 0 additions & 134 deletions packages/shared-python/shared/services/retrieval/publication_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
Loading