xueyiming 1 тиждень тому
батько
коміт
e8be966dfa

+ 2 - 1
applications/async_task/chunk_task.py

@@ -35,6 +35,7 @@ class ChunkEmbeddingTask(TopicAwarePackerV2):
         dataset_id = data.get("dataset_id", 0)  # 默认知识库 id 为 0
         re_chunk = data.get("re_chunk", False)
         dont_chunk = data.get("dont_chunk", False)
+        ext = data.get("ext", None)
         if re_chunk:
             await self.content_manager.update_content_info(
                 doc_id=doc_id,
@@ -46,7 +47,7 @@ class ChunkEmbeddingTask(TopicAwarePackerV2):
             flag = True
         else:
             flag = await self.content_manager.insert_content(
-                doc_id, text, text_type, title, dataset_id
+                doc_id, text, text_type, title, dataset_id, ext
             )
         if not flag:
             return []

+ 4 - 4
applications/utils/mysql/contents.py

@@ -2,14 +2,14 @@ from .base import BaseMySQLClient
 
 
 class Contents(BaseMySQLClient):
-    async def insert_content(self, doc_id, text, text_type, title, dataset_id):
+    async def insert_content(self, doc_id, text, text_type, title, dataset_id, ext):
         query = """
             INSERT IGNORE INTO contents
-                (doc_id, text, text_type, title, dataset_id)
-            VALUES (%s, %s, %s, %s, %s);
+                (doc_id, text, text_type, title, dataset_id, ext)
+            VALUES (%s, %s, %s, %s, %s, %s);
         """
         return await self.pool.async_save(
-            query=query, params=(doc_id, text, text_type, title, dataset_id)
+            query=query, params=(doc_id, text, text_type, title, dataset_id, ext)
         )
 
     async def update_content_info(self, doc_id, text, text_type, title, dataset_id):