refactor: 重构simple

llm_web_kit/extractor/html/recognizer/text.py

@@ -195,7 +195,7 @@ def __combine_text(self, text1:str, text2:str, lang='en') -> str:
             return self.replace_entities(txt.strip(), entities_map)
         else:
             # 根据text1的最后一个字符和text2的第一个字符判断两个text之间的连接
-            if (text2[0] in string.punctuation) or (text2[0] in special_symbols) or (text2[0] in other_symbols) or (text1 and text1[-1] in other_symbols):
+            if (text2 and text2[0] in string.punctuation) or (text2 and text2[0] in special_symbols) or (text2 and text2[0] in other_symbols) or (text1 and text1[-1] in other_symbols):
                 words_sep = ''
             else:
                 words_sep = ' '

llm_web_kit/simple.py

@@ -5,12 +5,14 @@
 
 from llm_web_kit.config.cfg_reader import load_pipe_tpl
 from llm_web_kit.extractor.extractor_chain import ExtractSimpleFactory
-from llm_web_kit.extractor.html.extractor import (
-    HTMLPageLayoutType, MagicHTMLFIleFormatorExtractor,
-    NoClipHTMLFIleFormatorExtractor)
 from llm_web_kit.input.datajson import DataJson
 
 
+class PipeType:
+    HTML = 'html'
+    NOCLIP = 'noclip_html'
+
+
 class ExtractorType:
     HTML = 'html'
     PDF = 'pdf'
@@ -19,24 +21,32 @@ class ExtractorType:
 
 class ExtractorFactory:
     """factory class for extractor."""
-    html_extractor = None
+    magic_html_extractor = None
+    noclip_html_extractor = None
     pdf_extractor = None
     ebook_extractor = None
 
     @staticmethod
-    def get_extractor(extractor_type: str):
+    def get_extractor(extractor_type: str, pipe_tpl_name: str):
         if extractor_type == ExtractorType.HTML:
-            if ExtractorFactory.html_extractor is None:
-                extractor_cfg = load_pipe_tpl('html')
-                chain = ExtractSimpleFactory.create(extractor_cfg)
-                ExtractorFactory.html_extractor = chain
-            return ExtractorFactory.html_extractor
+            if pipe_tpl_name == PipeType.HTML:
+                if ExtractorFactory.magic_html_extractor is None:
+                    extractor_cfg = load_pipe_tpl(pipe_tpl_name)
+                    chain = ExtractSimpleFactory.create(extractor_cfg)
+                    ExtractorFactory.magic_html_extractor = chain
+                return ExtractorFactory.magic_html_extractor
+            if pipe_tpl_name == PipeType.NOCLIP:
+                if ExtractorFactory.noclip_html_extractor is None:
+                    extractor_cfg = load_pipe_tpl(pipe_tpl_name)
+                    chain = ExtractSimpleFactory.create(extractor_cfg)
+                    ExtractorFactory.noclip_html_extractor = chain
+                return ExtractorFactory.noclip_html_extractor
         else:
             raise ValueError(f'Invalid extractor type: {extractor_type}')
 
 
 def __extract_main_html_by_no_clip_html(url:str, html_content: str, raw_html:str) -> DataJson:
-    extractor = NoClipHTMLFIleFormatorExtractor(load_pipe_tpl('noclip_html'))
+    extractor = ExtractorFactory.get_extractor(ExtractorType.HTML, PipeType.NOCLIP)
     if raw_html == '':
         raw_html = html_content
     input_data_dict = {
@@ -54,14 +64,8 @@ def __extract_main_html_by_no_clip_html(url:str, html_content: str, raw_html:str
     return result
 
 
-def __extract_main_html_by_maigic_html(url:str, html_str: str, page_layout_type:str) -> DataJson:
-    magic_html_extractor = MagicHTMLFIleFormatorExtractor(load_pipe_tpl('html'))
-    main_html, method, title = magic_html_extractor._extract_main_html(html_str, url, page_layout_type)
-    return main_html, title
-
-
 def __extract_html(url:str, html_content: str) -> DataJson:
-    extractor = ExtractorFactory.get_extractor(ExtractorType.HTML)
+    extractor = ExtractorFactory.get_extractor(ExtractorType.HTML, PipeType.HTML)
     input_data_dict = {
         'track_id': str(uuid.uuid4()),
         'url': url,
@@ -94,7 +98,10 @@ def extract_html_to_mm_md(url:str, html_content: str, clip_html=True, raw_html='
     return result.get_content_list().to_mm_md()
 
 
-def extract_main_html_by_maigic_html(url:str, html_str: str, page_layout_type:str = HTMLPageLayoutType.LAYOUT_ARTICLE) -> str:
-    """extract main html."""
-    result = __extract_main_html_by_maigic_html(url, html_str, page_layout_type)
-    return result[0], result[1]
+def extract_main_html(url:str, html_content: str, clip_html=True, raw_html='') -> str:
+    if clip_html:
+        result = __extract_html(url, html_content)
+    else:
+        result = __extract_main_html_by_no_clip_html(url, html_content, raw_html)
+    main_html = result.get('main_html')
+    return main_html

tests/llm_web_kit/simple/test_simple.py

@@ -2,7 +2,7 @@
 import unittest
 
 from llm_web_kit.simple import (extract_html_to_md, extract_html_to_mm_md,
-                                extract_main_html_by_maigic_html)
+                                extract_main_html)
 
 
 class TestSimple(unittest.TestCase):
@@ -136,9 +136,13 @@ def test_extract_pure_html_to_mm_md(self):
         mm_md = extract_html_to_mm_md(self.url, self.html_content, clip_html=True)
         self.assertEqual(mm_md, '# Test Content\n\nThis is a test paragraph.\n\n![Test Image](e5db82b5bf63d49d80c5533616892d3386f43955369520986d67653c700fc53c)\n')
 
-    def test_extract_magic_html(self):
-        magic_html, title = extract_main_html_by_maigic_html(self.url, self.html_content)
-        self.assertEqual(magic_html, '<div><body><h1>Test Content</h1><p>This is a test paragraph.</p><img src="https://example.com/image.jpg" alt="Test Image"></body></div>')
+    def test_extract_magic_main_html(self):
+        magic_main_html = extract_main_html(self.url, self.html_content, clip_html=True)
+        self.assertEqual(magic_main_html, '<div><body><h1>Test Content</h1><p>This is a test paragraph.</p><img src="https://example.com/image.jpg" alt="Test Image"></body></div>')
+
+    def test_extract_noclip_main_html(self):
+        magic_main_html = extract_main_html(self.url, self.html_content, clip_html=False, raw_html=self.html_content)
+        self.assertEqual(magic_main_html, '<html><body><h1>Test Content</h1><p>This is a test paragraph.</p><img src="https://example.com/image.jpg" alt="Test Image"></body></html>')
 
     def test_extract_real_html_to_md(self):
         md = extract_html_to_md(self.url, self.real_html_content, clip_html=False)