add mindmap gen

common.py

@@ -1,8 +1,8 @@
 import tiktoken
-
 import hashlib
 import secrets
-
+from graphviz import Digraph
+import random
 from PIL import ImageFilter, ImageEnhance
 
 
@@ -67,9 +67,49 @@ def optimize_text_by_openai(content):
     return response.choices[0].message.content
 
 
-def build_mind_map(graph, node, parent, structure, color='black'):
+def generate_light_color(pcolor: str):
+    """生成比给定颜色浅一半的颜色"""
+    r, g, b = int(pcolor[1:3], 16), int(pcolor[3:5], 16), int(pcolor[5:7], 16)
+    r = int(r + (255 - r) / 2)
+    g = int(g + (255 - g) / 2)
+    b = int(b + (255 - b) / 2)
+    return '#%02x%02x%02x' % (r, g, b)
+
+
+def generate_random_dark_color():
+    """
+    生成随机深色的函数。
+    通过确保RGB值不会太高，从而生成深色调。
+    """
+    r = random.randint(0, 100)
+    g = random.randint(0, 100)
+    b = random.randint(0, 100)
+    return f'#{r:02x}{g:02x}{b:02x}'
+
+
+# 改进的思维导图构建函数
+def build_mind_map(graph, node, parent, structure, level=0, parent_color=None):
+    # 根据层级设置样式
+    if level == 0:  # 根节点
+        node_color = generate_random_dark_color()
+        graph.node(node, style='filled', color=node_color, fontsize="21", fontname='SimHei', fontcolor='white',
+                   shape='tripleoctagon',peripheries="2", label=node)
+    elif level == 1:  # 第二层节点
+        node_color = generate_random_dark_color()
+        graph.node(node, style='filled', color=node_color, fontsize="18", fontname='SimHei', fontcolor='white',
+                   shape='hexagon',peripheries="2", label=node)
+    elif level == 2:  # 第三层节点
+        node_color = generate_light_color(parent_color)
+        graph.node(node, style='filled', color=node_color, fontsize="16", shape='note', fontname='SimHei', label=node)
+    else:  # 其他层级
+        node_color = generate_light_color(parent_color)
+        graph.node(node, style='solid', color=node_color,fontsize="14", shape='egg', fontname='SimHei', label=node)
+
+    # 连接节点
     if parent:
-        graph.edge(parent, node, color=color)
+        graph.edge(parent, node,  penwidth='2.0', color=parent_color if level == 1 else node_color)
+
+    # 递归构建子节点
     for child in structure.get(node, []):
-        graph.node(child)
-        build_mind_map(graph, child, node, structure, color=color)
+        build_mind_map(graph, child, node, structure, level=level + 1,
+                       parent_color=node_color if level == 1 else parent_color)

main.py

@@ -91,10 +91,15 @@ class IndexItem(BaseModel):
 
 
 class MindmapItem(BaseModel):
-    title: str = Field("Mindmap", title="Mindmap Title", description="Mindmap Title")
-    structure: Dict[str, List[str]] = Field({}, title="Mindmap Structure data",
-                                            description="Mindmap Structure data")
-
+    title: str = Field(default="Mindmap", title="Mindmap Title", description="Mindmap Title",
+                       example="Python 学习")
+    structure: Dict[str, List[str]] = Field(default={}, title="Mindmap Structure data",
+                                            description="Mindmap Structure data",
+                                            example={
+                                                "Python 学习": ["基础知识", "高级主题"],
+                                                "基础知识": ["变量", "数据类型", "控制流"],
+                                                "高级主题": ["面向对象", "装饰器", "迭代器"]
+                                            })
 
 class IndexSearchItem(BaseModel):
     collection: str = Field("default", title="Collection name",
@@ -266,14 +271,15 @@ async def create_mindmap(item: MindmapItem, td: bool = Depends(verify_api_key)):
     try:
         log.info(f"create_mindmap: {item}")
         # 创建并构建思维导图
-        graph = Digraph(comment=item.title)
+        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='pdf', cleanup=True)
+        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}.pdf"))
+        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))

tests/test.http

@@ -44,22 +44,29 @@ Content-Type: application/json
 Authorization: Bearer a99e05501a0405531caf783eef419b56a5a32f57b64ae3b89587b3a0d5202ee167d80d727a1b8181
 
 {
-    "title": "Python学习",
+    "title": "对数",
     "structure": {
-        "Python学习": ["基础知识", "高级主题"],
-        "基础知识": ["变量", "数据类型", "控制流"],
-        "高级主题": ["面向对象", "装饰器", "迭代器"],
-        "变量": ["定义", "赋值"],
-        "数据类型": ["整数", "浮点数", "字符串"],
-        "控制流": ["if语句", "for循环", "while循环"],
-        "面向对象": ["类", "继承", "多态"],
-        "装饰器": ["定义", "应用"],
-        "迭代器": ["创建", "使用"]
+        "对数": ["定义", "性质", "类型", "应用"],
+        "定义": ["对数的概念", "对数的表示"],
+        "性质": ["对数定律", "换底公式"],
+        "类型": ["常用对数", "自然对数"],
+        "应用": ["科学计算", "复利计算", "音乐理论", "信息熵"],
+        "对数的概念": ["对数与指数的关系"],
+        "对数的表示": ["底数", "真数"],
+        "对数定律": ["乘法定律", "除法定律", "幂的定律"],
+        "换底公式": ["公式描述"],
+        "常用对数": ["十进制", "应用举例"],
+        "自然对数": ["以e为底", "特性"],
+        "科学计算": ["对数尺", "数据压缩"],
+        "复利计算": ["金融领域应用"],
+        "音乐理论": ["音程与频率"],
+        "信息熵": ["信息论中的应用"]
     }
 }
 
 
 
+
 ###
 
 POST http://127.0.0.1:8700/knowledge/search