新增 rechunk任务接口

applications/async_task/__init__.py

@@ -1,5 +1,6 @@
 from .chunk_task import ChunkEmbeddingTask
 from .delete_task import DeleteTask
+from .auto_rechunk_task import AutoRechunkTask
 
 
-__all__ = ["ChunkEmbeddingTask", "DeleteTask"]
+__all__ = ["ChunkEmbeddingTask", "DeleteTask", "AutoRechunkTask"]

applications/async_task/auto_rechunk_task.py

@@ -0,0 +1,29 @@
+from applications.utils.http import AsyncHttpClient
+
+
+class AutoRechunkTask:
+    def __init__(self, mysql_client):
+        self.mysql_client = mysql_client
+
+    async def get_tasks(self):
+        query = """
+            SELECT doc_id, title, text, text_type, dataset_id
+            FROM contents
+            WHERE status = %s;
+        """
+        tasks = await self.mysql_client.async_fetch(query, params=(0,))
+        return tasks
+
+    @staticmethod
+    async def rechunk(task):
+        url = "http://192.168.100.31:8001/api/chunk"
+        async with AsyncHttpClient() as http_client:
+            response = await http_client.post(url, json=task, headers={"Content-Type": "application/json"})
+        return response
+
+    async def deal(self):
+        tasks = await self.get_tasks()
+        for task in tasks:
+            await self.rechunk(task)
+
+        return len(tasks)

routes/buleprint.py

@@ -9,6 +9,7 @@ from quart_cors import cors
 
 from applications.api import get_basic_embedding
 from applications.api import get_img_embedding
+from applications.async_task import AutoRechunkTask
 from applications.async_task import ChunkEmbeddingTask, DeleteTask
 from applications.config import (
     DEFAULT_MODEL,
@@ -469,3 +470,10 @@ async def chunk_list():
             },
         }
     )
+
+@server_bp.route("/auto_rechunk", methods=["GET"])
+async def auto_rechunk():
+    resource = get_resource_manager()
+    auto_rechunk_task = AutoRechunkTask(mysql_client=resource.mysql_client)
+    res = await auto_rechunk_task.deal()
+    return jsonify({"status_code": 200, "detail": "success", "cnt": len(res)})