Merge branch 'main' into feature/websocket-streaming-api

Author: SmartDever02

Dockerfile

@@ -10,11 +10,10 @@ WORKDIR /ragflow
 # Copy models downloaded via download_deps.py
 RUN mkdir -p /ragflow/rag/res/deepdoc /root/.ragflow
 RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/huggingface.co,target=/huggingface.co \
-    cp /huggingface.co/InfiniFlow/huqie/huqie.txt.trie /ragflow/rag/res/ && \
     tar --exclude='.*' -cf - \
         /huggingface.co/InfiniFlow/text_concat_xgb_v1.0 \
         /huggingface.co/InfiniFlow/deepdoc \
-        | tar -xf - --strip-components=3 -C /ragflow/rag/res/deepdoc 
+        | tar -xf - --strip-components=3 -C /ragflow/rag/res/deepdoc
 
 # https://github.com/chrismattmann/tika-python
 # This is the only way to run python-tika without internet access. Without this set, the default is to check the tika version and pull latest every time from Apache.

agent/canvas.py

@@ -91,9 +91,6 @@ def __init__(self, dsl: str, tenant_id=None, task_id=None):
     def load(self):
         self.components = self.dsl["components"]
         cpn_nms = set([])
-        for k, cpn in self.components.items():
-            cpn_nms.add(cpn["obj"]["component_name"])
-
         for k, cpn in self.components.items():
             cpn_nms.add(cpn["obj"]["component_name"])
             param = component_class(cpn["obj"]["component_name"] + "Param")()

agent/component/agent_with_tools.py

@@ -327,7 +327,7 @@ def delta(txt):
         self.set_output("content", answer)
 
     @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
-    def _invoke(self, **kwargs):
+    async def _invoke_async(self, **kwargs):
         if self.check_if_canceled("LLM processing"):
             return
 
@@ -338,22 +338,25 @@ def clean_formated_answer(ans: str) -> str:
 
         prompt, msg, _ = self._prepare_prompt_variables()
         error: str = ""
-        output_structure=None
+        output_structure = None
         try:
-            output_structure = self._param.outputs['structured']
+            output_structure = self._param.outputs["structured"]
         except Exception:
             pass
         if output_structure and isinstance(output_structure, dict) and output_structure.get("properties") and len(output_structure["properties"]) > 0:
-            schema=json.dumps(output_structure, ensure_ascii=False, indent=2)
-            prompt += structured_output_prompt(schema)
-            for _ in range(self._param.max_retries+1):
+            schema = json.dumps(output_structure, ensure_ascii=False, indent=2)
+            prompt_with_schema = prompt + structured_output_prompt(schema)
+            for _ in range(self._param.max_retries + 1):
                 if self.check_if_canceled("LLM processing"):
                     return
 
-                _, msg = message_fit_in([{"role": "system", "content": prompt}, *msg], int(self.chat_mdl.max_length * 0.97))
+                _, msg_fit = message_fit_in(
+                    [{"role": "system", "content": prompt_with_schema}, *deepcopy(msg)],
+                    int(self.chat_mdl.max_length * 0.97),
+                )
                 error = ""
-                ans = self._generate(msg)
-                msg.pop(0)
+                ans = await self._generate_async(msg_fit)
+                msg_fit.pop(0)
                 if ans.find("**ERROR**") >= 0:
                     logging.error(f"LLM response error: {ans}")
                     error = ans
@@ -362,26 +365,31 @@ def clean_formated_answer(ans: str) -> str:
                     self.set_output("structured", json_repair.loads(clean_formated_answer(ans)))
                     return
                 except Exception:
-                    msg.append({"role": "user", "content": "The answer can't not be parsed as JSON"})
+                    msg_fit.append({"role": "user", "content": "The answer can't not be parsed as JSON"})
                     error = "The answer can't not be parsed as JSON"
             if error:
                 self.set_output("_ERROR", error)
             return
 
         downstreams = self._canvas.get_component(self._id)["downstream"] if self._canvas.get_component(self._id) else []
         ex = self.exception_handler()
-        if any([self._canvas.get_component_obj(cid).component_name.lower()=="message" for cid in downstreams]) and not (ex and ex["goto"]):
-            self.set_output("content", partial(self._stream_output_async, prompt, msg))
+        if any([self._canvas.get_component_obj(cid).component_name.lower() == "message" for cid in downstreams]) and not (
+            ex and ex["goto"]
+        ):
+            self.set_output("content", partial(self._stream_output_async, prompt, deepcopy(msg)))
             return
 
-        for _ in range(self._param.max_retries+1):
+        error = ""
+        for _ in range(self._param.max_retries + 1):
             if self.check_if_canceled("LLM processing"):
                 return
 
-            _, msg = message_fit_in([{"role": "system", "content": prompt}, *msg], int(self.chat_mdl.max_length * 0.97))
+            _, msg_fit = message_fit_in(
+                [{"role": "system", "content": prompt}, *deepcopy(msg)], int(self.chat_mdl.max_length * 0.97)
+            )
             error = ""
-            ans = self._generate(msg)
-            msg.pop(0)
+            ans = await self._generate_async(msg_fit)
+            msg_fit.pop(0)
             if ans.find("**ERROR**") >= 0:
                 logging.error(f"LLM response error: {ans}")
                 error = ans
@@ -395,23 +403,9 @@ def clean_formated_answer(ans: str) -> str:
             else:
                 self.set_output("_ERROR", error)
 
-    def _stream_output(self, prompt, msg):
-        _, msg = message_fit_in([{"role": "system", "content": prompt}, *msg], int(self.chat_mdl.max_length * 0.97))
-        answer = ""
-        for ans in self._generate_streamly(msg):
-            if self.check_if_canceled("LLM streaming"):
-                return
-
-            if ans.find("**ERROR**") >= 0:
-                if self.get_exception_default_value():
-                    self.set_output("content", self.get_exception_default_value())
-                    yield self.get_exception_default_value()
-                else:
-                    self.set_output("_ERROR", ans)
-                return
-            yield ans
-            answer += ans
-        self.set_output("content", answer)
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
+    def _invoke(self, **kwargs):
+        return asyncio.run(self._invoke_async(**kwargs))
 
     def add_memory(self, user:str, assist:str, func_name: str, params: dict, results: str, user_defined_prompt:dict={}):
         summ = tool_call_summary(self.chat_mdl, func_name, params, results, user_defined_prompt)

agent/component/llm.py

@@ -49,16 +49,19 @@ def __init__(self, tools_map: dict[str, object], callback: partial):
         self.callback = callback
 
     def tool_call(self, name: str, arguments: dict[str, Any]) -> Any:
+        return asyncio.run(self.tool_call_async(name, arguments))
+
+    async def tool_call_async(self, name: str, arguments: dict[str, Any]) -> Any:
         assert name in self.tools_map, f"LLM tool {name} does not exist"
         st = timer()
         tool_obj = self.tools_map[name]
         if isinstance(tool_obj, MCPToolCallSession):
-            resp = tool_obj.tool_call(name, arguments, 60)
+            resp = await asyncio.to_thread(tool_obj.tool_call, name, arguments, 60)
         else:
             if hasattr(tool_obj, "invoke_async") and asyncio.iscoroutinefunction(tool_obj.invoke_async):
-                resp = asyncio.run(tool_obj.invoke_async(**arguments))
+                resp = await tool_obj.invoke_async(**arguments)
             else:
-                resp = asyncio.run(asyncio.to_thread(tool_obj.invoke, **arguments))
+                resp = await asyncio.to_thread(tool_obj.invoke, **arguments)
 
         self.callback(name, arguments, resp, elapsed_time=timer()-st)
         return resp

agent/tools/base.py

@@ -33,7 +33,7 @@
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api.db.services.tenant_llm_service import TenantLLMService
-from api.db.services.task_service import TaskService, queue_tasks
+from api.db.services.task_service import TaskService, queue_tasks, cancel_all_task_of
 from api.db.services.dialog_service import meta_filter, convert_conditions
 from api.utils.api_utils import check_duplicate_ids, construct_json_result, get_error_data_result, get_parser_config, get_result, server_error_response, token_required, \
     get_request_json
@@ -321,9 +321,7 @@ async def update_doc(tenant_id, dataset_id, document_id):
             try:
                 if not DocumentService.update_by_id(doc.id, {"status": str(status)}):
                     return get_error_data_result(message="Database error (Document update)!")
-
                 settings.docStoreConn.update({"doc_id": doc.id}, {"available_int": status}, search.index_name(kb.tenant_id), doc.kb_id)
-                return get_result(data=True)
             except Exception as e:
                 return server_error_response(e)
 
@@ -350,12 +348,10 @@ async def update_doc(tenant_id, dataset_id, document_id):
     }
     renamed_doc = {}
     for key, value in doc.to_dict().items():
-        if key == "run":
-            renamed_doc["run"] = run_mapping.get(str(value))
         new_key = key_mapping.get(key, key)
         renamed_doc[new_key] = value
         if key == "run":
-            renamed_doc["run"] = run_mapping.get(value)
+            renamed_doc["run"] = run_mapping.get(str(value))
 
     return get_result(data=renamed_doc)
 
@@ -839,6 +835,8 @@ async def stop_parsing(tenant_id, dataset_id):
             return get_error_data_result(message=f"You don't own the document {id}.")
         if int(doc[0].progress) == 1 or doc[0].progress == 0:
             return get_error_data_result("Can't stop parsing document with progress at 0 or 1")
+        # Send cancellation signal via Redis to stop background task
+        cancel_all_task_of(id)
         info = {"run": "2", "progress": 0, "chunk_num": 0}
         DocumentService.update_by_id(id, info)
         settings.docStoreConn.delete({"doc_id": doc[0].id}, search.index_name(tenant_id), dataset_id)

api/apps/sdk/doc.py

@@ -148,6 +148,7 @@ class Storage(Enum):
     AWS_S3 = 4
     OSS = 5
     OPENDAL = 6
+    GCS = 7
 
 # environment
 # ENV_STRONG_TEST_COUNT = "STRONG_TEST_COUNT"

common/constants.py

@@ -31,6 +31,7 @@
 import rag.utils.opensearch_conn
 from rag.utils.azure_sas_conn import RAGFlowAzureSasBlob
 from rag.utils.azure_spn_conn import RAGFlowAzureSpnBlob
+from rag.utils.gcs_conn import RAGFlowGCS
 from rag.utils.minio_conn import RAGFlowMinio
 from rag.utils.opendal_conn import OpenDALStorage
 from rag.utils.s3_conn import RAGFlowS3
@@ -109,6 +110,7 @@
 OB = {}
 OSS = {}
 OS = {}
+GCS = {}
 
 DOC_MAXIMUM_SIZE: int = 128 * 1024 * 1024
 DOC_BULK_SIZE: int = 4
@@ -151,7 +153,8 @@ class StorageFactory:
         Storage.AZURE_SAS: RAGFlowAzureSasBlob,
         Storage.AWS_S3: RAGFlowS3,
         Storage.OSS: RAGFlowOSS,
-        Storage.OPENDAL: OpenDALStorage
+        Storage.OPENDAL: OpenDALStorage,
+        Storage.GCS: RAGFlowGCS,
     }
 
     @classmethod
@@ -250,7 +253,7 @@ def init_settings():
     else:
         raise Exception(f"Not supported doc engine: {DOC_ENGINE}")
 
-    global AZURE, S3, MINIO, OSS
+    global AZURE, S3, MINIO, OSS, GCS
     if STORAGE_IMPL_TYPE in ['AZURE_SPN', 'AZURE_SAS']:
         AZURE = get_base_config("azure", {})
     elif STORAGE_IMPL_TYPE == 'AWS_S3':
@@ -259,6 +262,8 @@ def init_settings():
         MINIO = decrypt_database_config(name="minio")
     elif STORAGE_IMPL_TYPE == 'OSS':
         OSS = get_base_config("oss", {})
+    elif STORAGE_IMPL_TYPE == 'GCS':
+        GCS = get_base_config("gcs", {})
 
     global STORAGE_IMPL
     STORAGE_IMPL = StorageFactory.create(Storage[STORAGE_IMPL_TYPE])

common/settings.py

@@ -60,6 +60,8 @@ user_default_llm:
 #   access_key: 'access_key'
 #   secret_key: 'secret_key'
 #   region: 'region'
+#gcs:
+#  bucket: 'bridgtl-edm-d-bucket-ragflow'
 # oss:
 #   access_key: 'access_key'
 #   secret_key: 'secret_key'

conf/service_conf.yaml

@@ -25,6 +25,8 @@
 
 
 def vision_figure_parser_figure_data_wrapper(figures_data_without_positions):
+    if not figures_data_without_positions:
+        return []
     return [
         (
             (figure_data[1], [figure_data[0]]),
@@ -35,7 +37,9 @@ def vision_figure_parser_figure_data_wrapper(figures_data_without_positions):
     ]
 
 
-def vision_figure_parser_docx_wrapper(sections,tbls,callback=None,**kwargs):
+def vision_figure_parser_docx_wrapper(sections, tbls, callback=None,**kwargs):
+    if not tbls:
+        return []
     try:
         vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT)
         callback(0.7, "Visual model detected. Attempting to enhance figure extraction...")
@@ -53,6 +57,8 @@ def vision_figure_parser_docx_wrapper(sections,tbls,callback=None,**kwargs):
 
 
 def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
+    if not tbls:
+        return []
     try:
         vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT)
         callback(0.7, "Visual model detected. Attempting to enhance figure extraction...")