From e09b37e2855bf6af8fa4c8a0e3b3f45ae6fbc176 Mon Sep 17 00:00:00 2001
From: Liang Ding <d@b3log.org>
Date: Sun, 5 Mar 2023 12:16:47 +0800
Subject: [PATCH] =?UTF-8?q?:art:=20=E5=9D=97=E7=BA=A7=E5=8A=A0=E5=85=A5=20?=
 =?UTF-8?q?`=E4=BA=BA=E5=B7=A5=E6=99=BA=E8=83=BD`=20https://github.com/siy?=
 =?UTF-8?q?uan-note/siyuan/issues/7566?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 kernel/model/ai.go | 42 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/kernel/model/ai.go b/kernel/model/ai.go
index 401e69014..bc48dc5a9 100644
--- a/kernel/model/ai.go
+++ b/kernel/model/ai.go
@@ -19,7 +19,9 @@ package model
 import (
 	"bytes"
 
-	"github.com/siyuan-note/siyuan/kernel/sql"
+	"github.com/88250/lute/ast"
+	"github.com/88250/lute/parse"
+	"github.com/siyuan-note/siyuan/kernel/treenode"
 	"github.com/siyuan-note/siyuan/kernel/util"
 )
 
@@ -70,12 +72,40 @@ func isOpenAIAPIEnabled() bool {
 }
 
 func getBlocksContent(ids []string) string {
-	sqlBlocks := sql.GetBlocks(ids)
-	buf := bytes.Buffer{}
-	for _, sqlBlock := range sqlBlocks {
-		buf.WriteString(sqlBlock.Content)
-		buf.WriteString("\n\n")
+	var nodes []*ast.Node
+	trees := map[string]*parse.Tree{}
+	for _, id := range ids {
+		bt := treenode.GetBlockTree(id)
+		if nil == bt {
+			continue
+		}
+
+		var tree *parse.Tree
+		if tree = trees[bt.RootID]; nil == tree {
+			tree, _ = loadTreeByBlockID(bt.RootID)
+			if nil == tree {
+				continue
+			}
+
+			trees[bt.RootID] = tree
+		}
+
+		if node := treenode.GetNodeInTree(tree, id); nil != node {
+			if ast.NodeDocument == node.Type {
+				for child := node.FirstChild; nil != child; child = child.Next {
+					nodes = append(nodes, child)
+				}
+			} else {
+				nodes = append(nodes, node)
+			}
+		}
 	}
 
+	buf := bytes.Buffer{}
+	for _, node := range nodes {
+		content := treenode.NodeStaticContent(node, nil, true)
+		buf.WriteString(content)
+		buf.WriteString("\n\n")
+	}
 	return buf.String()
 }