优化切片，增加章节图片信息

chunk_text.py

@@ -342,6 +342,15 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str:
             ins.get("image_footnote"),
         )
 
+    if record_type == "chart":
+        return join_nonempty_parts(
+            ocr_text,
+            ins.get("chart_caption"),
+            ins.get("content"),
+            ins.get("img_path"),
+            ins.get("chart_footnote"),
+        )
+
     if record_type == "equation":
         return join_nonempty_parts(
             ocr_text,
@@ -371,6 +380,7 @@ def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
         "text",
         "table",
         "image",
+        "chart",
         "equation",
         "code",
         "list",
@@ -385,6 +395,7 @@ def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
     caption_field = {
         "image": "image_caption",
         "table": "table_caption",
+        "chart": "chart_caption",
         "code": "code_caption",
     }.get(ins.get("type"))
     if not caption_field:
@@ -810,85 +821,67 @@ def split_content_bbox(data, max_length=8000):
 
 
 
-# def group_records(records):
+def _group_records_without_parent_context(records):
-#     result = []
-#     current_group = []
-
-#     for i, ins in enumerate(records):
-#         if 'text_level' in ins and ins['text_level']:
-#             # 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
-#             if current_group:
-#                 result.append(current_group)
-#                 current_group = []  # 开始新一组
-#             current_group.append(ins)
-#         else:
-#             # 如果当前记录没有text_level或其值为空，则继续添加到当前组
-#             current_group.append(ins)
-
-#     # 别忘了将最后一个组添加到结果中
-#     if current_group:
-#         result.append(current_group)
-
-#     return result
-
-def group_records(records):
-    """
-    按 text_level 分组，并为每个组合适的父级标题上下文。
-    
-    规则：
-    - text_level=N (N>1) 的组 → 在组开头注入最近的 text_level=1 到 N-1 的标题
-    """
     result = []
     current_group = []
 
-    # 维护标题栈：level -> {text, page_idx, bbox}
-    heading_stack = {}
-    
     for i, ins in enumerate(records):
-        text_level = ins.get('text_level')
+        if 'text_level' in ins and ins['text_level']:
-        
+            # 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
-        if text_level is not None and isinstance(text_level, int) and text_level > 0:
-            # 更新标题栈
-            heading_stack[text_level] = {
-                'text': ins.get('text', ''),
-                'page_idx': ins.get('page_idx', -1),
-                'bbox': ins.get('bbox', []),
-                'level': text_level
-            }
-            # 清除更低层级的标题（高层级变化后，低层级失效）
-            levels_to_remove = [l for l in list(heading_stack.keys()) if l > text_level]
-            for l in levels_to_remove:
-                del heading_stack[l]
-            
-            # 遇到新标题，先保存当前组
             if current_group:
                 result.append(current_group)
-                current_group = []
+                current_group = []  # 开始新一组
-            
+            current_group.append(ins)
-            # 注入父级标题（如果当前层级 > 1）
-            if text_level > 1:
-                prefix_records = []
-                for level in sorted(heading_stack.keys()):
-                    if level < text_level:
-                        h = heading_stack[level]
-                        prefix_records.append({
-                            'type': 'text',
-                            'text': h['text'],
-                            'text_level': level,
-                            'page_idx': h['page_idx'],
-                            'bbox': h['bbox']
-                        })
-                current_group = prefix_records + [ins]
         else:
-                current_group = [ins]
+            # 如果当前记录没有text_level或其值为空，则继续添加到当前组
-        else:
-            # 非标题记录，继续添加到当前组
             current_group.append(ins)
 
+    # 别忘了将最后一个组添加到结果中
     if current_group:
         result.append(current_group)
 
     return result
+
+
+def _get_group_level(group):
+    if not group:
+        return None
+
+    level = group[0].get("text_level")
+    if isinstance(level, int) and level > 0:
+        return level
+    return None
+
+
+def group_records(records):
+    """
+    Split records by text_level, then prepend ancestor section content.
+
+    For example, a text_level=2 section will include the nearest text_level=1
+    group content before its own content. A text_level=3 section will include
+    the nearest level 1 and level 2 group content before its own content.
+    """
+    base_groups = _group_records_without_parent_context(records)
+    result = []
+    section_stack = []  # [(level, merged_group)]
+
+    for group in base_groups:
+        level = _get_group_level(group)
+        if level is None:
+            result.append(group)
+            continue
+
+        while section_stack and section_stack[-1][0] >= level:
+            section_stack.pop()
+
+        parent_group = section_stack[-1][1] if section_stack else []
+        merged_group = parent_group + group
+        result.append(merged_group)
+        section_stack.append((level, merged_group))
+
+    return result
+
+
 def get_chunk_bbox_otherfile(data, max_length=8000):
     """
     处理没有 bbox 信息的文件数据。
@@ -1262,7 +1255,7 @@ def find_ship_info_by_hull(json_file_path, data):
     except json.JSONDecodeError:
 	        print("错误：JSON 文件格式不正确")
         return None
-# def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
+def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
     """ 
     合并过短的切片：
     - 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
@@ -1358,110 +1351,6 @@ def find_ship_info_by_hull(json_file_path, data):
             ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
 
     return result
-
-def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
-    """ 
-    合并过短的切片：
-    - 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
-    - 若处于末尾无下一个切片，则反向合并到上一个切片末尾
-    - content 用换行拼接，positions 顺序拼接
-    
-    Args:
-        slices: [{"content": str, "positions": [...]}]
-        min_length: 短切片的字符长度阈值（含）
-        filename: 用于实体提取的文件名
-    
-    Returns:
-        合并后的 slices 列表
-    """
-    if not slices:
-        return slices
-
-    result = []
-    pending_contents = []  # 缓存等待合并到"下一个"的短切片 content
-    pending_positions = []  # 缓存对应的 positions
-
-    for ins in slices:
-        content = ins.get("content", "") or ""
-        positions = ins.get("positions", []) or []
-
-        if len(content) <= min_length:
-            # 暂存，等到下一个正常长度的切片再合并
-            pending_contents.append(content)
-            pending_positions.extend(positions)
-        else:
-            # 正常切片：把暂存的短切片合并到它的开头
-            if pending_contents:
-                merged_prefix = "\n".join(pending_contents)
-                content = merged_prefix + ("\n" if merged_prefix else "") + content
-                positions = pending_positions + positions
-                pending_contents = []
-                pending_positions = []
-            
-            result.append({
-                "content": content,
-                "positions": positions
-            })
-
-    # 收尾：如果末尾还有未合并的短切片（后面没有正常切片可合并）
-    # 则反向合并到上一个切片末尾
-    if pending_contents:
-        merged_suffix = "\n".join(pending_contents)
-        if result:
-            last = result[-1]
-            last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
-            last["positions"] = (last["positions"] or []) + pending_positions
-        else:
-            # 极端情况：所有切片都很短，整体作为一个切片返回
-            result.append({
-                "content": merged_suffix,
-                "positions": pending_positions
-            })
-
-    try:
-        final_result = get_entity(filename)
-        xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
-        if xinghao is None:  # 确保 xinghao 为 None 时不会报错
-            xinghao = {"model_name": "", "ship_name": ""}
-            print("未找到匹配的舰船信息")
-
-        # 安全处理 xinghao 为 None 的情况
-        model_name = ""
-        if xinghao and 'model_name' in xinghao:
-            model_name = xinghao['model_name']
-            
-        other_data = format_entity_text(final_result)
-        logger.info(f"文件名实体提取成功: {other_data}")
-    except Exception as e:
-        logger.warning(f"文件名实体提取失败: {e}")
-        other_data = ""
-        model_name = ""  # 确保异常时 model_name 有定义
-
-    # 将提取的信息（如舰艇名、型号）注入到每个切片的开头
-    for ins in result:
-        content = ins['content']
-        # 如果没有提取到有效数据，则跳过注入
-            
-        newline_idx = content.find('\n')
-        info = other_data
-        if model_name:
-            info += f'，型号为{model_name}'
-        
-        # ===== 修改部分：拼接文件名 + 实体信息 =====
-        # 构建前缀：#文件名为：xxx.pdf(实体信息）
-        filename_prefix = f"#文件名为：{filename}"
-        if info:
-            filename_prefix += f"({info}）"
-        # 注意：这里使用全角右括号"）"以匹配你的示例，如需半角请改为 ")"
-        
-        # 将前缀插入到 content 的最开头
-        if content:
-            ins['content'] = filename_prefix + "\n" + content
-        else:
-            ins['content'] = filename_prefix
-        # ==========================================
-
-    return result
 def find_and_read_content_list(directory, original_filename, encoding='utf-8'):
     """
     根据原始文件名，在指定目录及其子目录下查找对应的 _content_list.json 文件并返回内容。
@@ -1527,6 +1416,8 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
         slices = get_chunk_bbox(replaced_content_list)
         slices_check = chunk_check(slices, 8000)
         slices_check = merge_short_slices(slices_check, min_length=30,filename=filename) 
+        print(11111111111111111111111111111111111111)
+        print(len(images))
         return {"slices": slices_check, "images": images}
 
     except json.JSONDecodeError:
@@ -1596,7 +1487,7 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -
         
 if __name__ == "__main__":
 
-    filepath ="/app/mineru_output/舰船抗沉损管训练仿真系统研究_content_list.json"
+    filepath = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\舰船抗沉损管训练仿真系统研究_content_list.json"
     try:
         with open(filepath, 'r', encoding='utf-8') as file:
             data = json.load(file)