153 lines
3.6 KiB
Go
153 lines
3.6 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/gcfg"
|
||
"github.com/gogf/gf/v2/os/gctx"
|
||
)
|
||
|
||
type Dataset struct {
|
||
Id string `json:"id"`
|
||
Name string `json:"name"`
|
||
}
|
||
|
||
type ListDatasetsRes struct {
|
||
Code int `json:"code"`
|
||
Data []Dataset `json:"data"`
|
||
}
|
||
|
||
type DeleteReq struct {
|
||
Ids []string `json:"ids"`
|
||
}
|
||
|
||
type CommonResponse struct {
|
||
Code int `json:"code"`
|
||
Message string `json:"message"`
|
||
}
|
||
|
||
func main() {
|
||
ctx := gctx.New()
|
||
|
||
fmt.Println("🚀 RAGFlow知识库清理工具（简化版）")
|
||
fmt.Println("=" + strings.Repeat("=", 50))
|
||
|
||
// 设置配置文件路径
|
||
g.Cfg().GetAdapter().(*gcfg.AdapterFile).SetPath("../")
|
||
|
||
// 读取配置
|
||
baseURL := g.Cfg().MustGet(ctx, "ragflow.base_url").String()
|
||
apiKey := g.Cfg().MustGet(ctx, "ragflow.api_key").String()
|
||
|
||
if baseURL == "" || apiKey == "" {
|
||
fmt.Println("❌ RAGFlow配置缺失，请检查../config.yaml")
|
||
fmt.Println(" 需要配置: ragflow.base_url 和 ragflow.api_key")
|
||
os.Exit(1)
|
||
}
|
||
|
||
baseURL = strings.TrimSuffix(baseURL, "/")
|
||
fmt.Printf("📡 连接到RAGFlow: %s\n", baseURL)
|
||
|
||
// 创建HTTP客户端
|
||
client := &http.Client{
|
||
Timeout: 30 * time.Second,
|
||
}
|
||
|
||
// 1. 列出所有知识库
|
||
listURL := baseURL + "/api/v1/datasets?page=1&page_size=1000"
|
||
req, err := http.NewRequest("GET", listURL, nil)
|
||
if err != nil {
|
||
fmt.Printf("❌ 创建请求失败: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
fmt.Printf("❌ 查询知识库列表失败: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
body, _ := io.ReadAll(resp.Body)
|
||
var listRes ListDatasetsRes
|
||
if err := json.Unmarshal(body, &listRes); err != nil {
|
||
fmt.Printf("❌ 解析响应失败: %v\n", err)
|
||
fmt.Printf("响应内容: %s\n", string(body))
|
||
os.Exit(1)
|
||
}
|
||
|
||
if len(listRes.Data) == 0 {
|
||
fmt.Println("✅ 没有找到任何知识库，无需清理")
|
||
os.Exit(0)
|
||
}
|
||
|
||
// 显示知识库列表
|
||
fmt.Printf("\n📚 发现 %d 个知识库：\n", len(listRes.Data))
|
||
for i, dataset := range listRes.Data {
|
||
fmt.Printf(" %d. ID: %s, Name: %s\n", i+1, dataset.Id, dataset.Name)
|
||
}
|
||
|
||
// 二次确认
|
||
fmt.Printf("\n⚠️ 警告：即将删除所有 %d 个知识库！\n", len(listRes.Data))
|
||
fmt.Print("请输入 'YES' 确认删除: ")
|
||
|
||
var confirm string
|
||
fmt.Scanln(&confirm)
|
||
|
||
if confirm != "YES" {
|
||
fmt.Println("❌ 取消删除操作")
|
||
os.Exit(0)
|
||
}
|
||
|
||
// 2. 批量删除
|
||
var datasetIds []string
|
||
for _, dataset := range listRes.Data {
|
||
datasetIds = append(datasetIds, dataset.Id)
|
||
}
|
||
|
||
deleteReq := DeleteReq{Ids: datasetIds}
|
||
reqBody, _ := json.Marshal(deleteReq)
|
||
|
||
deleteURL := baseURL + "/api/v1/datasets"
|
||
delReq, err := http.NewRequest("DELETE", deleteURL, bytes.NewBuffer(reqBody))
|
||
if err != nil {
|
||
fmt.Printf("❌ 创建删除请求失败: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
delReq.Header.Set("Authorization", "Bearer "+apiKey)
|
||
delReq.Header.Set("Content-Type", "application/json")
|
||
|
||
fmt.Println("\n🗑️ 开始删除知识库...")
|
||
delResp, err := client.Do(delReq)
|
||
if err != nil {
|
||
fmt.Printf("❌ 删除失败: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
defer delResp.Body.Close()
|
||
|
||
delBody, _ := io.ReadAll(delResp.Body)
|
||
var commonRes CommonResponse
|
||
if err := json.Unmarshal(delBody, &commonRes); err != nil {
|
||
fmt.Printf("❌ 解析删除响应失败: %v\n", err)
|
||
fmt.Printf("响应内容: %s\n", string(delBody))
|
||
os.Exit(1)
|
||
}
|
||
|
||
if commonRes.Code != 0 {
|
||
fmt.Printf("❌ 删除失败: %s\n", commonRes.Message)
|
||
os.Exit(1)
|
||
}
|
||
|
||
fmt.Printf("✅ 成功删除 %d 个知识库！\n", len(datasetIds))
|
||
fmt.Println("=" + strings.Repeat("=", 50))
|
||
}
|