feat: integrate flask-orjson for improved JSON serialization performance (#23935)

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -1,7 +1,7 @@
-import json
 from collections import defaultdict
 from typing import Any, Optional
 
+import orjson
 from pydantic import BaseModel
 
 from configs import dify_config
@@ -134,13 +134,13 @@ class Jieba(BaseKeyword):
         dataset_keyword_table = self.dataset.dataset_keyword_table
         keyword_data_source_type = dataset_keyword_table.data_source_type
         if keyword_data_source_type == "database":
-            dataset_keyword_table.keyword_table = json.dumps(keyword_table_dict, cls=SetEncoder)
+            dataset_keyword_table.keyword_table = dumps_with_sets(keyword_table_dict)
             db.session.commit()
         else:
             file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt"
             if storage.exists(file_key):
                 storage.delete(file_key)
-            storage.save(file_key, json.dumps(keyword_table_dict, cls=SetEncoder).encode("utf-8"))
+            storage.save(file_key, dumps_with_sets(keyword_table_dict).encode("utf-8"))
 
     def _get_dataset_keyword_table(self) -> Optional[dict]:
         dataset_keyword_table = self.dataset.dataset_keyword_table
@@ -156,12 +156,11 @@ class Jieba(BaseKeyword):
                 data_source_type=keyword_data_source_type,
             )
             if keyword_data_source_type == "database":
-                dataset_keyword_table.keyword_table = json.dumps(
+                dataset_keyword_table.keyword_table = dumps_with_sets(
                     {
                         "__type__": "keyword_table",
                         "__data__": {"index_id": self.dataset.id, "summary": None, "table": {}},
-                    },
-                    cls=SetEncoder,
+                    }
                 )
             db.session.add(dataset_keyword_table)
             db.session.commit()
@@ -252,8 +251,13 @@ class Jieba(BaseKeyword):
         self._save_dataset_keyword_table(keyword_table)
 
 
-class SetEncoder(json.JSONEncoder):
-    def default(self, obj):
-        if isinstance(obj, set):
-            return list(obj)
-        return super().default(obj)
+def set_orjson_default(obj: Any) -> Any:
+    """Default function for orjson serialization of set types"""
+    if isinstance(obj, set):
+        return list(obj)
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def dumps_with_sets(obj: Any) -> str:
+    """JSON dumps with set support using orjson"""
+    return orjson.dumps(obj, default=set_orjson_default).decode("utf-8")