v0.21.1-fastapi

rag/flow/splitter/schema.py

@@ -25,7 +25,7 @@ class SplitterFromUpstream(BaseModel):
     file: dict | None = Field(default=None)
     chunks: list[dict[str, Any]] | None = Field(default=None)
 
-    output_format: Literal["json", "markdown", "text", "html"] | None = Field(default=None)
+    output_format: Literal["json", "markdown", "text", "html", "chunks"] | None = Field(default=None)
 
     json_result: list[dict[str, Any]] | None = Field(default=None, alias="json")
     markdown_result: str | None = Field(default=None, alias="markdown")

rag/flow/splitter/splitter.py

@@ -19,7 +19,7 @@ import trio
 
 from api.utils import get_uuid
 from api.utils.base64_image import id2image, image2id
-from ocr.service import get_ocr_service
+from deepdoc.parser.pdf_parser import RAGFlowPdfParser
 from rag.flow.base import ProcessBase, ProcessParamBase
 from rag.flow.splitter.schema import SplitterFromUpstream
 from rag.nlp import naive_merge, naive_merge_with_images
@@ -87,7 +87,7 @@ class Splitter(ProcessBase):
         sections, section_images = [], []
         for o in from_upstream.json_result or []:
             sections.append((o.get("text", ""), o.get("position_tag", "")))
-            section_images.append(id2image(o.get("img_id"), partial(STORAGE_IMPL.get)))
+            section_images.append(id2image(o.get("img_id"), partial(STORAGE_IMPL.get, tenant_id=self._canvas._tenant_id)))
 
         chunks, images = naive_merge_with_images(
             sections,
@@ -96,20 +96,16 @@ class Splitter(ProcessBase):
             deli,
             self._param.overlapped_percent,
         )
-        ocr_service = get_ocr_service()
-        cks = []
-        for c, img in zip(chunks, images):
-            if not c.strip():
-                continue
-            cleaned_text = await ocr_service.remove_tag(c)
-            positions = await ocr_service.extract_positions(c)
-            cks.append({
-                "text": cleaned_text,
+        cks = [
+            {
+                "text": RAGFlowPdfParser.remove_tag(c),
                 "image": img,
-                "positions": [[pos[0][-1]+1, *pos[1:]] for pos in positions],
-            })
+                "positions": [[pos[0][-1]+1, *pos[1:]] for pos in RAGFlowPdfParser.extract_positions(c)],
+            }
+            for c, img in zip(chunks, images) if c.strip()
+        ]
         async with trio.open_nursery() as nursery:
             for d in cks:
-                nursery.start_soon(image2id, d, partial(STORAGE_IMPL.put), get_uuid())
+                nursery.start_soon(image2id, d, partial(STORAGE_IMPL.put, tenant_id=self._canvas._tenant_id), get_uuid())
         self.set_output("chunks",  cks)
         self.callback(1, "Done.")