全面更换orjson

src/chat/knowledge/embedding_store.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-import json
+import orjson
 import os
 import math
 import asyncio
@@ -277,8 +277,11 @@ class EmbeddingStore:
                 test_vectors[str(idx)] = self._get_embedding(s)
         
         with open(self.get_test_file_path(), "w", encoding="utf-8") as f:
-            json.dump(test_vectors, f, ensure_ascii=False, indent=2)
-        
+            f.write(orjson.dumps(
+                test_vectors,
+                option=orjson.OPT_INDENT_2
+            ).decode('utf-8'))
+
         logger.info("测试字符串嵌入向量保存完成")
 
     def load_embedding_test_vectors(self):
@@ -287,7 +290,7 @@ class EmbeddingStore:
         if not os.path.exists(path):
             return None
         with open(path, "r", encoding="utf-8") as f:
-            return json.load(f)
+            return orjson.loads(f.read())
 
     def check_embedding_model_consistency(self):
         """校验当前模型与本地嵌入模型是否一致（使用多线程优化）"""
@@ -416,7 +419,9 @@ class EmbeddingStore:
             logger.info(f"{self.namespace}嵌入库的FaissIndex保存成功")
             logger.info(f"正在保存{self.namespace}嵌入库的idx2hash映射到文件{self.idx2hash_file_path}")
             with open(self.idx2hash_file_path, "w", encoding="utf-8") as f:
-                f.write(json.dumps(self.idx2hash, ensure_ascii=False, indent=4))
+                f.write(orjson.dumps(
+                    self.idx2hash,  option=orjson.OPT_INDENT_2
+                ).decode('utf-8'))
             logger.info(f"{self.namespace}嵌入库的idx2hash映射保存成功")
 
     def load_from_file(self) -> None:
@@ -457,7 +462,7 @@ class EmbeddingStore:
                 logger.info(f"正在加载{self.namespace}嵌入库的idx2hash映射...")
                 logger.debug(f"正在从文件{self.idx2hash_file_path}中加载{self.namespace}嵌入库的idx2hash映射")
                 with open(self.idx2hash_file_path, "r") as f:
-                    self.idx2hash = json.load(f)
+                    self.idx2hash = orjson.loads(f.read())
                 logger.info(f"{self.namespace}嵌入库的idx2hash映射加载成功")
             else:
                 raise Exception(f"文件{self.idx2hash_file_path}不存在")

src/chat/knowledge/ie_process.py

@@ -1,5 +1,5 @@
 import asyncio
-import json
+import orjson
 import time
 from typing import List, Union
 
@@ -20,7 +20,7 @@ def _extract_json_from_text(text: str):
     try:
         fixed_json = repair_json(text)
         if isinstance(fixed_json, str):
-            parsed_json = json.loads(fixed_json)
+            parsed_json = orjson.loads(fixed_json)
         else:
             parsed_json = fixed_json
 
@@ -95,9 +95,10 @@ def _entity_extract(llm_req: LLMRequest, paragraph: str) -> List[str]:
 def _rdf_triple_extract(llm_req: LLMRequest, paragraph: str, entities: list) -> List[List[str]]:
     """对段落进行实体提取，返回提取出的实体列表（JSON格式）"""
     rdf_extract_context = prompt_template.build_rdf_triple_extract_context(
-        paragraph, entities=json.dumps(entities, ensure_ascii=False)
+        paragraph, entities=orjson.dumps(entities).decode('utf-8')
     )
 
+
     # 使用 asyncio.run 来运行异步方法
     try:
         # 如果当前已有事件循环在运行，使用它

src/chat/knowledge/kg_manager.py

@@ -1,4 +1,4 @@
-import json
+import orjson
 import os
 import time
 from typing import Dict, List, Tuple
@@ -74,7 +74,7 @@ class KGManager:
         # 保存段落hash到文件
         with open(self.pg_hash_file_path, "w", encoding="utf-8") as f:
             data = {"stored_paragraph_hashes": list(self.stored_paragraph_hashes)}
-            f.write(json.dumps(data, ensure_ascii=False, indent=4))
+            f.write(orjson.dumps(data, option=orjson.OPT_INDENT_2).decode('utf-8'))
 
     def load_from_file(self):
         """从文件加载KG数据"""
@@ -88,7 +88,7 @@ class KGManager:
 
         # 加载段落hash
         with open(self.pg_hash_file_path, "r", encoding="utf-8") as f:
-            data = json.load(f)
+            data = orjson.loads(f.read())
             self.stored_paragraph_hashes = set(data["stored_paragraph_hashes"])
 
         # 加载实体计数

src/chat/knowledge/open_ie.py

@@ -1,4 +1,4 @@
-import json
+import orjson
 import os
 import glob
 from typing import Any, Dict, List
@@ -113,7 +113,7 @@ class OpenIE:
         data_list = []
         for file in json_files:
             with open(file, "r", encoding="utf-8") as f:
-                data = json.load(f)
+                data = orjson.loads(f.read())
                 data_list.append(data)
         if not data_list:
             # print(f"111111111111111111111Root Path : \n{ROOT_PATH}")

src/chat/knowledge/utils/json_fix.py

@@ -1,4 +1,4 @@
-import json
+import orjson
 from json_repair import repair_json
 
 
@@ -45,7 +45,7 @@ def fix_broken_generated_json(json_str: str) -> str:
     - Iterating over the JSON string once to determine and fix unclosed braces or brackets.
     - Ensuring braces and brackets inside string literals are not considered.
 
-    If the original json_str string can be successfully loaded by json.loads(), will directly return it without any modification.
+    If the original json_str string can be successfully loaded by orjson.loads(), will directly return it without any modification.
 
     Args:
         json_str (str): The malformed JSON string to be fixed.
@@ -56,9 +56,9 @@ def fix_broken_generated_json(json_str: str) -> str:
 
     try:
         # Try to load the JSON to see if it is valid
-        json.loads(json_str)
+        orjson.loads(json_str)
         return json_str  # Return as-is if valid
-    except json.JSONDecodeError: ...
+    except orjson.JSONDecodeError: ...
 
     # Step 1: Remove trailing content after the last comma.
     last_comma_index = json_str.rfind(",")
@@ -80,7 +80,7 @@ def new_fix_broken_generated_json(json_str: str) -> str:
     """
     使用 json-repair 库修复格式错误的 JSON 字符串。
 
-    如果原始 json_str 字符串可以被 json.loads() 成功加载，则直接返回而不进行任何修改。
+    如果原始 json_str 字符串可以被 orjson.loads() 成功加载，则直接返回而不进行任何修改。
 
     参数:
         json_str (str): 需要修复的格式错误的 JSON 字符串。
@@ -90,8 +90,8 @@ def new_fix_broken_generated_json(json_str: str) -> str:
     """
     try:
         # 尝试加载 JSON 以查看其是否有效
-        json.loads(json_str)
+        orjson.loads(json_str)
         return json_str  # 如果有效则按原样返回
-    except json.JSONDecodeError:
+    except orjson.JSONDecodeError:
         # 如果无效，则尝试修复它
         return repair_json(json_str)