From 72b3ef630cff13be376d186f927df0670baa012b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 7 May 2026 17:43:41 +0800 Subject: [PATCH] refactor: consolidate chunk connection logic and improve parsing in services --- .../contract/test_parse_task_contract.py | 10 +- .../services/chunks/chunk_connections.py | 222 ++++++++++++++++++ .../chunks/dataframe_chunk_converter.py | 134 ++--------- .../services/storage/zip_result_service.py | 182 ++------------ 4 files changed, 266 insertions(+), 282 deletions(-) create mode 100644 packages/shared-python/shared/services/chunks/chunk_connections.py diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index bf02fe4d5..7fa0ed28e 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -223,7 +223,15 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "summary": "", "know_id": "kid-1", "tokens": "", - "connectto": "", + "connectto": json.dumps( + [ + { + "target": "table-1", + "relation": "embeds", + "ref": "[tables/table-1.html]", + } + ] + ), "addtime": "now", "page_nums": "1", }, diff --git a/packages/shared-python/shared/services/chunks/chunk_connections.py b/packages/shared-python/shared/services/chunks/chunk_connections.py new file mode 100644 index 000000000..efb0804f1 --- /dev/null +++ b/packages/shared-python/shared/services/chunks/chunk_connections.py @@ -0,0 +1,222 @@ +"""Build canonical chunk connection metadata.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any, TypeAlias, TypedDict + +from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans + + +class PositionPayload(TypedDict): + start: int + end: int + + +class ConnectionPayload(TypedDict, total=False): + target: str + relation: str + ref: str + position: PositionPayload + score: float + keywords: list[str] + + +RelationshipRef: TypeAlias = str | ChunkRefSpan +ConnectionValue: TypeAlias = str | ConnectionPayload +ConnectionKey: TypeAlias = tuple[str, str, str] +PositionKey: TypeAlias = tuple[str, str] + + +def parse_relationship_refs(type_value: object, content: str) -> list[RelationshipRef]: + parsed_relationships = _parse_type_relationship_refs(type_value) + if parsed_relationships: + return [relationship for relationship in parsed_relationships] + return [span for span in extract_chunk_ref_spans(content)] + + +def build_resource_target_map( + chunks: Sequence[Mapping[str, Any]], + *, + image_files_map: Mapping[str, Mapping[str, Any]] | None = None, + table_files_map: Mapping[str, Mapping[str, Any]] | None = None, +) -> dict[str, str]: + target_map: dict[str, str] = {} + for chunk in chunks: + chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id") or "").strip() + if not chunk_id: + continue + + chunk_type = str(chunk.get("type", "")).strip().split("\n", 1)[0].lower() + if chunk_type not in {"image", "table"}: + continue + + metadata = chunk.get("metadata", {}) + file_path = "" + if isinstance(metadata, dict): + file_path = str(metadata.get("file_path") or "").strip() + if not file_path: + file_map = image_files_map if chunk_type == "image" else table_files_map + file_info = file_map.get(chunk_id) if file_map else None + if file_info: + file_path = str(file_info.get("file_path") or "").strip() + + path_alias = str(chunk.get("path") or "").strip() + aliases = {file_path, path_alias} + for alias in list(aliases): + if alias: + aliases.add(f"[{alias}]") + for alias in aliases: + if alias: + target_map[alias] = chunk_id + return target_map + + +def convert_refs_to_embed_connections( + refs: Sequence[RelationshipRef], target_map: Mapping[str, str] +) -> list[ConnectionPayload]: + connections: list[ConnectionPayload] = [] + for ref in refs: + if isinstance(ref, dict): + ref_text = str(ref.get("ref") or "").strip() + start = ref.get("start") + end = ref.get("end") + else: + ref_text = str(ref or "").strip() + start = None + end = None + if not ref_text: + continue + + target_id = target_map.get(ref_text) + if not target_id and ref_text.startswith("[") and ref_text.endswith("]"): + target_id = target_map.get(ref_text[1:-1].strip()) + if not target_id: + continue + + connection: ConnectionPayload = { + "target": target_id, + "relation": "embeds", + "ref": ref_text, + } + if isinstance(start, int) and isinstance(end, int): + connection["position"] = { + "start": start, + "end": end, + } + connections.append(connection) + return connections + + +def normalize_connect_to_targets( + connects: object, target_map: Mapping[str, str] +) -> list[ConnectionPayload]: + if connects is None or connects == "": + return [] + + raw_items = connects if isinstance(connects, list) else [connects] + normalized: list[ConnectionPayload] = [] + for item in raw_items: + if item is None or item == "": + continue + + if isinstance(item, dict): + target = str(item.get("target") or "").strip() + normalized_target = target_map.get(target, target) + if not normalized_target: + continue + + normalized_item: ConnectionPayload = { + "target": normalized_target, + "relation": str(item.get("relation") or "related"), + } + score = item.get("score") + if isinstance(score, (int, float)): + normalized_item["score"] = float(score) + keywords = item.get("keywords") + if isinstance(keywords, list): + normalized_item["keywords"] = [str(keyword) for keyword in keywords] + ref = item.get("ref") + if ref: + normalized_item["ref"] = str(ref) + position = item.get("position") + if isinstance(position, dict): + start = position.get("start") + end = position.get("end") + if isinstance(start, int) and isinstance(end, int): + normalized_item["position"] = {"start": start, "end": end} + normalized.append(normalized_item) + continue + + target = str(item or "").strip() + normalized_target = target_map.get(target, target) + if normalized_target: + normalized.append( + { + "target": normalized_target, + "relation": "related", + "score": 1.0, + "keywords": [], + } + ) + return normalized + + +def merge_connections( + *connection_lists: Sequence[ConnectionValue], +) -> list[ConnectionValue]: + merged: list[ConnectionValue] = [] + unpositioned_indexes: dict[ConnectionKey, int] = {} + positioned_keys: dict[ConnectionKey, set[PositionKey]] = {} + for connection_list in connection_lists: + for item in connection_list or []: + if not isinstance(item, dict): + continue + key = _get_connection_key(item) + position_key = _get_connection_position_key(item) + if position_key is None: + if key in unpositioned_indexes or key in positioned_keys: + continue + unpositioned_indexes[key] = len(merged) + merged.append(item) + continue + + key_positions = positioned_keys.setdefault(key, set()) + if position_key in key_positions: + continue + key_positions.add(position_key) + unpositioned_index = unpositioned_indexes.pop(key, None) + if unpositioned_index is None: + merged.append(item) + else: + merged[unpositioned_index] = item + return merged + + +def _parse_type_relationship_refs(type_value: object) -> list[str]: + if not isinstance(type_value, str) or "\n" not in type_value: + return [] + lines = [line.strip() for line in type_value.split("\n") if line.strip()] + return [line for line in lines[1:] if line.upper() != "PTXT"] + + +def _get_connection_key(item: ConnectionValue) -> ConnectionKey: + if not isinstance(item, dict): + return ("", "related", "") + return ( + str(item.get("target") or ""), + str(item.get("relation") or "related"), + str(item.get("ref") or ""), + ) + + +def _get_connection_position_key(item: ConnectionValue) -> PositionKey | None: + if not isinstance(item, dict): + return None + position = item.get("position") + if not isinstance(position, dict): + return None + return ( + str(position.get("start", "")), + str(position.get("end", "")), + ) diff --git a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py index 6b2cb7b37..9332c044f 100644 --- a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py +++ b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py @@ -5,13 +5,21 @@ import json import os import uuid -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Sequence from typing import Dict, Literal, Protocol, TypeAlias, TypedDict, Union, cast import pandas as pd from loguru import logger -from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans +from shared.services.chunks.chunk_connections import ( + ConnectionValue, + RelationshipRef, + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) class _ParserRow(Protocol): @@ -24,28 +32,12 @@ def __len__(self) -> int: ... def iterrows(self) -> Iterable[tuple[object, _ParserRow]]: ... -class PositionPayload(TypedDict): - start: int - end: int - - -class ConnectionPayload(TypedDict, total=False): - target: str - relation: str - ref: str - position: PositionPayload - score: float - keywords: list[str] - - JsonPrimitive: TypeAlias = str | int | float | bool | None JsonValue: TypeAlias = Union[ JsonPrimitive, list["JsonValue"], dict[str, "JsonValue"], ] -RelationshipRef: TypeAlias = str | ChunkRefSpan -ConnectionValue: TypeAlias = str | ConnectionPayload ChunkType: TypeAlias = Literal["text", "image", "table"] @@ -134,15 +126,6 @@ def _safe_parse_tokens(value: object) -> list[str]: return [] -def _safe_parse_relationships(value: object) -> list[str]: - if _is_missing(value): - return [] - if not isinstance(value, str) or "\n" not in value: - return [] - lines = [line.strip() for line in value.split("\n") if line.strip()] - return [line for line in lines[1:] if line.upper() != "PTXT"] - - def _normalize_resource_ref(ref: RelationshipRef) -> str: if isinstance(ref, dict): ref_text = str(ref.get("ref") or "").strip() @@ -197,87 +180,6 @@ def _parse_connect_to(value: object) -> list[ConnectionValue]: ] -def _build_resource_target_map(chunks: Sequence[ChunkPayload]) -> dict[str, str]: - target_map: dict[str, str] = {} - for chunk in chunks: - if chunk["type"] not in {"image", "table"}: - continue - chunk_id = str(chunk["chunk_id"] or chunk["know_id"]).strip() - if not chunk_id: - continue - metadata = chunk["metadata"] - file_path = "" - file_path = str(metadata.get("file_path") or "").strip() - path_alias = chunk["path"].strip() - aliases = {file_path, path_alias} - for alias in list(aliases): - if alias: - aliases.add(f"[{alias}]") - for alias in aliases: - if alias: - target_map[alias] = chunk_id - return target_map - - -def _refs_to_embed_connections( - refs: Sequence[RelationshipRef], target_map: Mapping[str, str] -) -> list[ConnectionPayload]: - connections: list[ConnectionPayload] = [] - for ref in refs: - if isinstance(ref, dict): - ref_text = str(ref.get("ref") or "").strip() - start = ref.get("start") - end = ref.get("end") - else: - ref_text = str(ref or "").strip() - start = None - end = None - if not ref_text: - continue - target_id = target_map.get(ref_text) - if not target_id and ref_text.startswith("[") and ref_text.endswith("]"): - target_id = target_map.get(ref_text[1:-1].strip()) - if not target_id: - continue - connection: ConnectionPayload = { - "target": target_id, - "relation": "embeds", - "ref": ref_text, - } - if isinstance(start, int) and isinstance(end, int): - connection["position"] = { - "start": start, - "end": end, - } - connections.append(connection) - return connections - - -def _merge_connections( - *connection_lists: Sequence[ConnectionValue], -) -> list[ConnectionValue]: - merged: list[ConnectionValue] = [] - seen: set[tuple[str, str, str, str, str]] = set() - for connection_list in connection_lists: - for item in connection_list or []: - if not isinstance(item, dict): - continue - position = item.get("position") - position_data = position if isinstance(position, dict) else {} - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - str(position_data.get("start", "")), - str(position_data.get("end", "")), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - return merged - - def _parse_page_numbers(value: object) -> list[int]: if _is_missing(value): return [] @@ -305,10 +207,7 @@ def _get_chunk_type(value: object) -> ChunkType: def _get_relationship_refs(type_value: object, content: str) -> list[RelationshipRef]: - parsed_relationships = _safe_parse_relationships(type_value) - if parsed_relationships: - return [relationship for relationship in parsed_relationships] - return [span for span in extract_chunk_ref_spans(content)] + return parse_relationship_refs(type_value, content) def _get_connect_to(metadata: ChunkMetadata) -> list[ConnectionValue]: @@ -399,18 +298,21 @@ def dataframe_to_chunks(df: _ParserDataFrame | None) -> list[Dict[str, JsonValue } ) - resource_target_map = _build_resource_target_map(chunks) + resource_target_map = build_resource_target_map(chunks) for chunk in chunks: metadata = chunk["metadata"] relationship_refs = metadata.pop("_relationship_refs", []) if chunk["type"] != "text": continue - embed_connections = _refs_to_embed_connections( + embed_connections = convert_refs_to_embed_connections( relationship_refs, resource_target_map ) - metadata["connect_to"] = _merge_connections( + metadata["connect_to"] = merge_connections( embed_connections, - _get_connect_to(metadata), + normalize_connect_to_targets( + _get_connect_to(metadata), + resource_target_map, + ), ) logger.debug(f"DataFrame conversion completed: chunk count={len(chunks)}") diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index 2ae93b1f3..bfd1513f3 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -13,7 +13,13 @@ from loguru import logger from PIL import Image -from shared.utils.chunk_refs import extract_chunk_ref_spans +from shared.services.chunks.chunk_connections import ( + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) from shared.utils.text_utils import truncate_content_preview import pandas as pd @@ -201,163 +207,11 @@ def _format_chunks( table_files_map: Dict[str, Dict[str, Any]], ) -> List[Dict[str, Any]]: """Convert chunks data to ZIP specification format""" - - def safe_parse_rels(type_val): - """Safely parse in-document relationship fields from type metadata.""" - rels = [] - if type_val and isinstance(type_val, str): - if "\n" in type_val: - lines = [ - line.strip() for line in type_val.split("\n") if line.strip() - ] - rels.extend([line for line in lines[1:] if line.upper() != "PTXT"]) - return rels if rels else [] - - def build_resource_target_map() -> Dict[str, str]: - """Build ref/path -> chunk_id aliases for image and table chunks.""" - target_map: Dict[str, str] = {} - for chunk in chunks: - chunk_id = str( - chunk.get("chunk_id") or chunk.get("know_id") or "" - ).strip() - if not chunk_id: - continue - chunk_type_str = ( - str(chunk.get("type", "")).strip().split("\n", 1)[0].lower() - ) - if chunk_type_str not in {"image", "table"}: - continue - metadata = chunk.get("metadata", {}) - file_path = "" - if isinstance(metadata, dict): - file_path = str(metadata.get("file_path") or "").strip() - if not file_path: - file_info = ( - image_files_map.get(chunk_id) - if chunk_type_str == "image" - else table_files_map.get(chunk_id) - ) - if file_info: - file_path = str(file_info.get("file_path") or "").strip() - path_alias = str(chunk.get("path") or "").strip() - aliases = {file_path, path_alias} - for alias in list(aliases): - if alias: - aliases.add(f"[{alias}]") - for alias in aliases: - if alias: - target_map[alias] = chunk_id - return target_map - - def normalize_connect_to( - connects, target_map: Dict[str, str] - ) -> List[Dict[str, Any]]: - """Normalize connect_to to chunk_id-based entries.""" - if not connects: - return [] - - raw_items = connects if isinstance(connects, list) else [connects] - normalized = [] - for item in raw_items: - if not item: - continue - - if isinstance(item, dict): - target = str(item.get("target") or "").strip() - normalized_target = target_map.get(target, target) - if not normalized_target: - continue - - normalized_item = { - "target": normalized_target, - "relation": item.get("relation", "related"), - } - if "score" in item: - normalized_item["score"] = item.get("score", 1.0) - if "keywords" in item: - normalized_item["keywords"] = item.get("keywords", []) - if "ref" in item and item.get("ref"): - normalized_item["ref"] = item.get("ref") - if "position" in item and isinstance(item.get("position"), dict): - normalized_item["position"] = item.get("position") - normalized.append(normalized_item) - continue - - item_str = str(item).strip() - if not item_str: - continue - normalized_target = target_map.get(item_str, item_str) - normalized.append( - { - "target": normalized_target, - "relation": "related", - "score": 1.0, - "keywords": [], - } - ) - - return normalized - - def refs_to_embed_connections( - refs: List[Any], target_map: Dict[str, str] - ) -> List[Dict[str, Any]]: - """Convert resource refs to connect_to embeds entries.""" - normalized = [] - for ref in refs: - if isinstance(ref, dict): - ref_str = str(ref.get("ref") or "").strip() - start = ref.get("start") - end = ref.get("end") - else: - ref_str = str(ref or "").strip() - start = None - end = None - if not ref_str: - continue - target_id = target_map.get(ref_str) - if not target_id and ref_str.startswith("[") and ref_str.endswith("]"): - target_id = target_map.get(ref_str[1:-1].strip()) - if not target_id: - continue - connection: Dict[str, Any] = { - "target": target_id, - "relation": "embeds", - "ref": ref_str, - } - if isinstance(start, int) and isinstance(end, int): - connection["position"] = { - "start": start, - "end": end, - } - normalized.append(connection) - return normalized - - def merge_connections( - *connection_lists: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Merge connect_to entries while keeping stable order.""" - merged: List[Dict[str, Any]] = [] - seen = set() - for connection_list in connection_lists: - for item in connection_list or []: - if not isinstance(item, dict): - continue - position = item.get("position") - position_data = position if isinstance(position, dict) else {} - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - str(position_data.get("start", "")), - str(position_data.get("end", "")), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - return merged - - resource_target_map = build_resource_target_map() + resource_target_map = build_resource_target_map( + chunks, + image_files_map=image_files_map, + table_files_map=table_files_map, + ) formatted = [] for chunk in chunks: @@ -404,16 +258,14 @@ def merge_connections( ) # Convert in-text resource refs into embeds edges. - relationship_refs = safe_parse_rels( - chunk.get("type_raw") or chunk_type_str + relationship_refs = parse_relationship_refs( + chunk.get("type_raw") or chunk_type_str, + str(content), ) - if not relationship_refs: - relationship_refs = extract_chunk_ref_spans(content) - - embed_connections = refs_to_embed_connections( + embed_connections = convert_refs_to_embed_connections( relationship_refs, resource_target_map ) - related_connections = normalize_connect_to( + related_connections = normalize_connect_to_targets( existing_metadata.get("connect_to") or chunk.get("connect_to") or chunk.get("connectto"),