From b83bfb6e7c2dcd669beaf8bc50742df019a74d0c Mon Sep 17 00:00:00 2001
From: liukaiwen <kl3157@columbia.edu>
Date: Fri, 13 Jun 2025 16:22:01 +0800
Subject: [PATCH] <feat>: add parse_single in MapItemToHtmlTagsParser for
 single html extraction

---
 .../main_html_parser/parser/tag_mapping.py    | 50 +++++++++++++++++++
 .../processor/test_tag_mapping.py             | 15 ++++++
 2 files changed, 65 insertions(+)

diff --git a/llm_web_kit/main_html_parser/parser/tag_mapping.py b/llm_web_kit/main_html_parser/parser/tag_mapping.py
index 8d525281..8878cb4b 100644
--- a/llm_web_kit/main_html_parser/parser/tag_mapping.py
+++ b/llm_web_kit/main_html_parser/parser/tag_mapping.py
@@ -72,6 +72,28 @@ def parse(self, pre_data: PreDataJson) -> PreDataJson:
             raise TagMappingParserException(e)
         return pre_data
 
+    def parse_single(self, pre_data: PreDataJson) -> PreDataJson:
+        """
+            skip element dict construct step, remove all non-main tags in template tagged html directly
+            for single-html extraction plan
+            Args:
+                pre_root:
+            Returns:
+                PreDataJson: 包含映射结果的PreDataJson对象
+        """
+        try:
+            template_tag_html = pre_data[PreDataJsonKey.TYPICAL_RAW_TAG_HTML]
+            response_json = pre_data[PreDataJsonKey.LLM_RESPONSE]
+            root = html.fromstring(template_tag_html)
+            # 直接抽取正文
+            content_list = self.tag_main_html(response_json, root)
+            template_extract_html = self.__extract_main_directly(root)
+            pre_data[PreDataJsonKey.TYPICAL_MAIN_HTML] = template_extract_html
+            pre_data[PreDataJsonKey.HTML_TARGET_LIST] = content_list
+        except Exception as e:
+            raise TagMappingParserException(e)
+        return pre_data
+
     def __get_max_width_layer(self, element_dict):
         max_length = 0
         max_width_layer = 0
@@ -139,6 +161,34 @@ def tag_parent(self, pre_root):
                 parent.set('magic_main_html', 'True')
                 cur = parent
 
+    def __extract_main_directly(self, pre_root):
+        def iter_process(elem):
+            if isinstance(elem, etree._Comment):
+                return
+            magic_main_html = elem.get('magic_main_html', None)
+            if magic_main_html:
+                # 查找所有子孙节点中 magic_main_html='True' 的元素
+                matching_elements = elem.xpath(
+                    './/*[@magic_main_html="True"]'
+                )
+                # 给正文最小单元节点的子孙节点补上正文标识，避免被删除
+                if len(matching_elements) == 0:
+                    for child in elem.iterdescendants():  # 仅遍历子孙节点（不包括自身）
+                        child.set('magic_main_html', 'True')
+            else:
+                # 非正文节点直接删除
+                parent = elem.getparent()
+                if parent is None:
+                    return
+                parent.remove(elem)
+            for elem_child in elem:
+                iter_process(elem_child)
+
+        if pre_root is None:
+            return None
+        iter_process(pre_root)
+        return html.tostring(pre_root, encoding='utf-8').decode()
+
     def tag_main_html(self, response, pre_root):
         content_list = []
         for elem in pre_root.iter():
diff --git a/tests/llm_web_kit/main_html_parser/processor/test_tag_mapping.py b/tests/llm_web_kit/main_html_parser/processor/test_tag_mapping.py
index ac41a902..fe90fed7 100644
--- a/tests/llm_web_kit/main_html_parser/processor/test_tag_mapping.py
+++ b/tests/llm_web_kit/main_html_parser/processor/test_tag_mapping.py
@@ -62,3 +62,18 @@ def test_construct_main_tree_fail_by_similarity(self):
         pre_data = parser.parse(pre_data)
         construct_success = pre_data.get(PreDataJsonKey.TYPICAL_MAIN_HTML_SUCCESS)
         self.assertEqual(False, construct_success)
+
+    def test_parse_single(self):
+        data = []
+        raw_html_path = base_dir.joinpath('assets/test_tag_mapping_web.jsonl')
+        with open(raw_html_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                data.append(json.loads(line.strip()))  # 解析每行 JSON
+        mock_dict = data[0]
+        pre_data = PreDataJson(mock_dict['pre_data'])
+        pre_data[PreDataJsonKey.TYPICAL_RAW_HTML] = pre_data[PreDataJsonKey.TYPICAL_RAW_TAG_HTML]
+        parser = MapItemToHtmlTagsParser({})
+        pre_data = parser.parse_single(pre_data)
+        content_list = pre_data[PreDataJsonKey.HTML_TARGET_LIST]
+        self.assertEqual(content_list, mock_dict['expected_content_list'])
+        self.assertEqual(len(pre_data['typical_main_html']), 2269)