diff --git a/internal/handler/datasets.go b/internal/handler/datasets.go index aa0d896cb0..b856509511 100644 --- a/internal/handler/datasets.go +++ b/internal/handler/datasets.go @@ -18,7 +18,11 @@ package handler import ( "encoding/json" + "fmt" "net/http" + "ragflow/internal/engine" + "ragflow/internal/engine/types" + "sort" "strconv" "strings" @@ -198,3 +202,243 @@ func (h *DatasetsHandler) DeleteDatasets(c *gin.Context) { "data": result, }) } + +// GetKnowledgeGraph handles GET /api/v1/datasets/:dataset_id/graph. +func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + datasetID := strings.TrimSpace(c.Param("dataset_id")) + if datasetID == "" { + jsonError(c, common.CodeDataError, "dataset_id is required") + return + } + + dataset, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + tenantID, _ := dataset["tenant_id"].(string) + if tenantID == "" { + jsonError(c, common.CodeDataError, "tenant_id is required") + return + } + + docEngine := engine.Get() + if docEngine == nil { + jsonError(c, common.CodeServerError, "Document engine is not initialized") + return + } + + indexName := fmt.Sprintf("ragflow_%s", tenantID) + exists, err := docEngine.TableExists(c.Request.Context(), indexName) + if err != nil { + jsonError(c, common.CodeServerError, err.Error()) + return + } + + result := gin.H{ + "graph": map[string]interface{}{}, + "mind_map": map[string]interface{}{}, + } + if !exists { + jsonResponse(c, common.CodeSuccess, result, "success") + return + } + + searchResult, err := docEngine.Search(c.Request.Context(), &types.SearchRequest{ + IndexNames: []string{indexName}, + KbIDs: []string{datasetID}, + Offset: 0, + Limit: 1, + SelectFields: []string{"content_with_weight", "knowledge_graph_kwd"}, + Filter: map[string]interface{}{ + "kb_id": []string{datasetID}, + "knowledge_graph_kwd": []string{"graph"}, + }, + }) + if err != nil { + jsonError(c, common.CodeServerError, err.Error()) + return + } + if searchResult == nil || len(searchResult.Chunks) == 0 { + jsonResponse(c, common.CodeSuccess, result, "success") + return + } + + chunk := searchResult.Chunks[0] + graphType := firstStringValue(chunk["knowledge_graph_kwd"]) + contentWithWeight, _ := chunk["content_with_weight"].(string) + if strings.TrimSpace(contentWithWeight) == "" { + jsonResponse(c, common.CodeSuccess, result, "success") + return + } + + var graphData map[string]interface{} + if err := json.Unmarshal([]byte(contentWithWeight), &graphData); err != nil { + jsonResponse(c, common.CodeSuccess, result, "success") + return + } + if len(graphData) == 0 { + jsonResponse(c, common.CodeSuccess, result, "success") + return + } + + if graphType == "" { + graphType = "graph" + } + if graphType == "graph" { + sortKnowledgeGraph(graphData) + result["graph"] = graphData + } else { + result[graphType] = graphData + } + + jsonResponse(c, common.CodeSuccess, result, "success") +} + +// DeleteKnowledgeGraph handles DELETE /api/v1/datasets/:dataset_id/graph. +func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + datasetID := strings.TrimSpace(c.Param("dataset_id")) + if datasetID == "" { + jsonError(c, common.CodeDataError, "dataset_id is required") + return + } + + dataset, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + tenantID, _ := dataset["tenant_id"].(string) + if tenantID == "" { + jsonError(c, common.CodeDataError, "tenant_id is required") + return + } + + docEngine := engine.Get() + if docEngine == nil { + jsonError(c, common.CodeServerError, "Document engine is not initialized") + return + } + + indexName := fmt.Sprintf("ragflow_%s", tenantID) + if _, err := docEngine.Delete(c.Request.Context(), map[string]interface{}{ + "knowledge_graph_kwd": []string{"graph", "subgraph", "entity", "relation", "community_report"}, + }, indexName, datasetID); err != nil { + jsonError(c, common.CodeServerError, err.Error()) + return + } + + jsonResponse(c, common.CodeSuccess, true, "success") +} + +func firstStringValue(value interface{}) string { + switch v := value.(type) { + case string: + return strings.TrimSpace(v) + case []string: + if len(v) > 0 { + return strings.TrimSpace(v[0]) + } + case []interface{}: + for _, item := range v { + if s, ok := item.(string); ok { + s = strings.TrimSpace(s) + if s != "" { + return s + } + } + } + } + return "" +} + +func sortKnowledgeGraph(graphData map[string]interface{}) { + nodes := mapSlice(graphData["nodes"]) + if len(nodes) > 0 { + sort.Slice(nodes, func(i, j int) bool { + return numericValue(nodes[i]["pagerank"]) > numericValue(nodes[j]["pagerank"]) + }) + if len(nodes) > 256 { + nodes = nodes[:256] + } + graphData["nodes"] = nodes + } + + edges := mapSlice(graphData["edges"]) + if len(edges) > 0 { + nodeIDSet := make(map[string]struct{}, len(nodes)) + for _, node := range nodes { + if id, ok := node["id"].(string); ok { + nodeIDSet[id] = struct{}{} + } + } + filteredEdges := make([]map[string]interface{}, 0, len(edges)) + for _, edge := range edges { + source, _ := edge["source"].(string) + target, _ := edge["target"].(string) + if source == "" || target == "" || source == target { + continue + } + if _, ok := nodeIDSet[source]; !ok { + continue + } + if _, ok := nodeIDSet[target]; !ok { + continue + } + filteredEdges = append(filteredEdges, edge) + } + sort.Slice(filteredEdges, func(i, j int) bool { + return numericValue(filteredEdges[i]["weight"]) > numericValue(filteredEdges[j]["weight"]) + }) + if len(filteredEdges) > 128 { + filteredEdges = filteredEdges[:128] + } + graphData["edges"] = filteredEdges + } +} + +func mapSlice(value interface{}) []map[string]interface{} { + raw, ok := value.([]interface{}) + if !ok { + return nil + } + result := make([]map[string]interface{}, 0, len(raw)) + for _, item := range raw { + if m, ok := item.(map[string]interface{}); ok { + result = append(result, m) + } + } + return result +} + +func numericValue(value interface{}) float64 { + switch v := value.(type) { + case float64: + return v + case float32: + return float64(v) + case int: + return float64(v) + case int64: + return float64(v) + case json.Number: + f, _ := v.Float64() + return f + default: + return 0 + } +} diff --git a/internal/router/router.go b/internal/router/router.go index ddb13e0ac9..5b8a840e7a 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -179,6 +179,8 @@ func (r *Router) Setup(engine *gin.Engine) { { datasets.GET("", r.datasetsHandler.ListDatasets) datasets.GET("/:dataset_id", r.datasetsHandler.GetDataset) + datasets.GET("/:dataset_id/graph", r.datasetsHandler.GetKnowledgeGraph) + datasets.DELETE("/:dataset_id/graph", r.datasetsHandler.DeleteKnowledgeGraph) datasets.POST("", r.datasetsHandler.CreateDataset) datasets.DELETE("", r.datasetsHandler.DeleteDatasets) datasets.POST("/search", r.chunkHandler.RetrievalTest)