mindmap update

main.py

@@ -23,6 +23,9 @@
 from fastapi import File, UploadFile
 from fastapi.responses import JSONResponse
 from graphviz import Digraph
+from concurrent.futures import ProcessPoolExecutor
+from fastapi import FastAPI, BackgroundTasks
+import asyncio
 
 load_dotenv()
 
@@ -72,6 +75,13 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
 
 app.add_middleware(LimitUploadSize, max_upload_size=1024 * 1024 * 10)
 
+executor = ProcessPoolExecutor()
+
+
+async def run_in_process(fn, *args):
+    loop = asyncio.get_running_loop()
+    return await loop.run_in_executor(executor, fn, *args)
+
 
 class TokenData(BaseModel):
     api_key: str
@@ -272,52 +282,130 @@ async def create_image_ocr(file: UploadFile = File(...), td: TokenData = Depends
         os.unlink(tmp_path)
 
 
-@app.post("/knowledge/mindmap/create", summary="Create a knowledge base mindmap from params",
-          description="Generating mind maps from given structured data", include_in_schema=False)
-async def create_mindmap(item: MindmapItem, td: bool = Depends(verify_api_key)):
+# @app.post("/knowledge/mindmap/create", summary="Create a knowledge base mindmap from params",
+#           description="Generating mind maps from given structured data", include_in_schema=False)
+# async def create_mindmap(item: MindmapItem, td: bool = Depends(verify_api_key)):
+#     try:
+#         log.info(f"create_mindmap: {item}")
+#         # 创建并构建思维导图
+#         graph = Digraph(comment=item.title, engine="sfdp")
+#         graph.attr(splines='curved', overlap='false', margin='0.4')  # 设置图的大小为A4纸尺寸
+#         build_mind_map(graph, item.title, None, structure=item.structure)
+#         fileuuid = str(uuid.uuid4())
+#         graph.render(os.path.join(DATA_DIR, fileuuid), format='png', cleanup=True)
+#         server_url = os.environ.get("GPTS_API_SERVER")
+#         if server_url.endswith("/"):
+#             server_url = server_url[:-1]
+#         return RestResult(code=0, msg="success", result=dict(data=f"{server_url}/assets/{fileuuid}.png"))
+#     except Exception as e:
+#         log.error(f"create_mindmap error: {e}")
+#         raise HTTPException(status_code=500, detail=str(e))
+
+
+def create_mindmap_task(task):
     try:
-        log.info(f"create_mindmap: {item}")
+        log.info(f"generate_mindmap params: {task['content']}")
+        airesp = create_mindma_data_by_openai(task['content'])
+        log.info(f"generate_mindmap result: {airesp}")
         # 创建并构建思维导图
+        # 将 JSON 字符串转换为 Python 字典
+        data = json.loads(airesp)
+        # 使用 model_validate 方法创建 MindmapItem 实例
+        item = MindmapItem.model_validate(data)
         graph = Digraph(comment=item.title, engine="sfdp")
-        graph.attr(splines='curved', overlap='false',  margin='0.4')  # 设置图的大小为A4纸尺寸
+        graph.attr(splines='curved')
         build_mind_map(graph, item.title, None, structure=item.structure)
-        fileuuid = str(uuid.uuid4())
-        graph.render(os.path.join(DATA_DIR, fileuuid), format='png', cleanup=True)
+        fileuuid = task["id"]
+        output_path = os.path.join(DATA_DIR, fileuuid)
+        graph.render(output_path, format='png', cleanup=True)
+        # 生成 DOT 文件
+        dot_path = output_path + ".dot"
+        with open(dot_path, "w") as dot_file:
+            dot_file.write(graph.source)
+
         server_url = os.environ.get("GPTS_API_SERVER")
         if server_url.endswith("/"):
             server_url = server_url[:-1]
-        return RestResult(code=0, msg="success", result=dict(data=f"{server_url}/assets/{fileuuid}.png"))
+        task["status"] = "done"
+        task["remark"] = "The task is completed, please access the URL information!"
+        task["image_url"] = f"{server_url}/assets/{fileuuid}.png"
+        task["dot_url"] = f"{server_url}/assets/{fileuuid}.dot"
+        with open(os.path.join(DATA_DIR, f"{fileuuid}.json"), "w") as f:
+            data = json.dumps(task)
+            f.write(data)
     except Exception as e:
-        log.error(f"create_mindmap error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        import traceback
+        traceback.print_exc()
 
 
-@app.get("/knowledge/mindmap/generate", summary="Create a knowledge base mindmap from query content",
-         description="Generating mind maps from given content")
-async def generate_mindmap(content: str = Query(...), td: bool = Depends(verify_api_key)):
+# @app.get("/knowledge/mindmap/generate", summary="Create a knowledge base mindmap from query content",
+#          description="Generating mind maps from given content")
+# async def generate_mindmap(content: str = Query(...), td: bool = Depends(verify_api_key)):
+#     try:
+#         log.info(f"generate_mindmap params: {content}")
+#         airesp = await create_mindma_data_by_openai(content)
+#         log.info(f"generate_mindmap result: {airesp}")
+#         # 创建并构建思维导图
+#         # 将 JSON 字符串转换为 Python 字典
+#         data = json.loads(airesp)
+#         # 使用 model_validate 方法创建 MindmapItem 实例
+#         item = MindmapItem.model_validate(data)
+#         graph = Digraph(comment=item.title, engine="sfdp")
+#         graph.attr(splines='curved')
+#         build_mind_map(graph, item.title, None, structure=item.structure)
+#         fileuuid = str(uuid.uuid4())
+#         graph.render(os.path.join(DATA_DIR, fileuuid), format='png', cleanup=True)
+#         server_url = os.environ.get("GPTS_API_SERVER")
+#         if server_url.endswith("/"):
+#             server_url = server_url[:-1]
+#         return RestResult(code=0, msg="success", result=dict(data=f"{server_url}/assets/{fileuuid}.png"))
+#     except Exception as e:
+#         log.error(f"generate_mindmap error: {e}")
+#         raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/knowledge/mindmap/task/add", summary="Create a mindmap generate task from content",
+         description="Create a mindmap generate task from content")
+async def generate_mindmap_task_add(background_tasks: BackgroundTasks, content: str = Query(...),
+                                    td: bool = Depends(verify_api_key)):
     try:
-        log.info(f"generate_mindmap params: {content}")
-        airesp = create_mindma_data_by_openai(content)
-        log.info(f"generate_mindmap result: {airesp}")
-        # 创建并构建思维导图
-        # 将 JSON 字符串转换为 Python 字典
-        data = json.loads(airesp)
-        # 使用 model_validate 方法创建 MindmapItem 实例
-        item = MindmapItem.model_validate(data)
-        graph = Digraph(comment=item.title, engine="sfdp")
-        graph.attr(splines='curved')
-        build_mind_map(graph, item.title, None, structure=item.structure)
-        fileuuid = str(uuid.uuid4())
-        graph.render(os.path.join(DATA_DIR, fileuuid), format='png', cleanup=True)
+        log.info(f"generate_mindmap_task_add params: {content}")
         server_url = os.environ.get("GPTS_API_SERVER")
         if server_url.endswith("/"):
             server_url = server_url[:-1]
-        return RestResult(code=0, msg="success", result=dict(data=f"{server_url}/assets/{fileuuid}.png"))
+        taskid = str(uuid.uuid4())
+        task = dict(
+            id=taskid,
+            content=content,
+            status="pending",
+            image_url=f"{server_url}/assets/{taskid}.png",
+            dot_url=f"{server_url}/assets/{taskid}.dot",
+            status_url=f"{server_url}/knowledge/mindmap/task/result/{taskid}",
+            remark="The task is being processed, please remember to save the URL information and revisit it later!",
+        )
+        with open(os.path.join(DATA_DIR, f"{taskid}.json"), "w") as f:
+            data = json.dumps(task)
+            f.write(data)
+        background_tasks.add_task(create_mindmap_task, task)
+        log.info(f"generate_mindmap_task_add result: {task}")
+        return RestResult(code=0, msg="task add success", result=dict(data=task))
     except Exception as e:
         log.error(f"generate_mindmap error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@app.get("/knowledge/mindmap/task/result/{taskid}", summary="Get the mindmap generate task result",
+         description="Get the mindmap generate task result")
+async def generate_mindmap_task_result(taskid: str):
+    if not re.match(r'^[\w-]+$', taskid):
+        raise HTTPException(status_code=400, detail="Invalid task ID format")
+
+    file_path = os.path.join(DATA_DIR, f"{taskid}.json")
+    if not os.path.isfile(file_path):
+        raise HTTPException(status_code=404, detail="File not found")
+    return FileResponse(file_path)
+
+
 if __name__ == "__main__":
     import uvicorn
 

tests/test.http

@@ -66,7 +66,7 @@ Authorization: Bearer a99e05501a0405531caf783eef419b56a5a32f57b64ae3b89587b3a0d5
 
 ###
 
-GET http://127.0.0.1:8700/knowledge/mindmap/generate?content=根据Python知识创建一个思维导图，最多60节点
+GET http://127.0.0.1:8700/knowledge/mindmap/task/add?content=根据C++基础创建一个思维导图
 Accept: application/json
 Content-Type: application/json
 Authorization: Bearer a99e05501a0405531caf783eef419b56a5a32f57b64ae3b89587b3a0d5202ee167d80d727a1b8181

tools/generate_mindmap.py

@@ -6,9 +6,9 @@
 from main import MindmapItem
 
 
-def generate_mindmap():
+async def generate_mindmap():
     try:
-        airesp = create_mindma_data_by_openai("根据微积分基础整理一个学习计划思维导图")
+        airesp = await create_mindma_data_by_openai("根据微积分基础整理一个学习计划思维导图")
         # 创建并构建思维导图
         data = json.loads(airesp)
         item = MindmapItem.model_validate(data)