From 0252e58b04190b2bea049a7ea1f1f1be43cc14e8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 29 Aug 2025 11:50:14 +0800 Subject: [PATCH] Update segment graph API to include entity-based relationship retrieval - Enhanced GetSegmentGraph function to retrieve segment entities and relationships based on query parameters, improving data access. - Introduced new endpoint GetSegmentRelationshipsByEntities for fetching relationships connected to entities, with appropriate error handling. - Updated response structures to reflect counts and types of queries used, ensuring clarity in API responses. - Modified go.mod and go.sum to update the dependency for github.com/ulikunitz/xz to version 0.5.14. --- go.mod | 2 +- go.sum | 4 +- openapi/kb/graph.go | 116 ++++++++++++++++++++++++++++++++++---------- openapi/kb/kb.go | 1 + 4 files changed, 95 insertions(+), 28 deletions(-) diff --git a/go.mod b/go.mod index 8e52b915..2a00b0da 100644 --- a/go.mod +++ b/go.mod @@ -138,7 +138,7 @@ require ( github.com/tiendc/go-deepcopy v1.6.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect - github.com/ulikunitz/xz v0.5.12 // indirect + github.com/ulikunitz/xz v0.5.14 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect diff --git a/go.sum b/go.sum index 04f91aec..efd6fcad 100644 --- a/go.sum +++ b/go.sum @@ -311,8 +311,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= github.com/ulikunitz/xz v0.5.9/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= -github.com/ulikunitz/xz v0.5.12 h1:37Nm15o69RwBkXM0J6A5OlE67RZTfzUxTj8fB3dfcsc= -github.com/ulikunitz/xz v0.5.12/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= +github.com/ulikunitz/xz v0.5.14 h1:uv/0Bq533iFdnMHZdRBTOlaNMdb1+ZxXIlHDZHIHcvg= +github.com/ulikunitz/xz v0.5.14/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c= github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI= github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY= diff --git a/openapi/kb/graph.go b/openapi/kb/graph.go index 0e41bea4..5a571d12 100644 --- a/openapi/kb/graph.go +++ b/openapi/kb/graph.go @@ -53,33 +53,41 @@ func GetSegmentGraph(c *gin.Context) { includeEntities := c.DefaultQuery("include_entities", "true") != "false" includeRelationships := c.DefaultQuery("include_relationships", "true") != "false" - // Call the GraphRag instance to get segment graph - segmentGraph, err := kb.Instance.GetSegmentGraph(c.Request.Context(), docID, segmentID) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: "segment_not_found", - ErrorDescription: fmt.Sprintf("Failed to get segment graph: %v", err), - } - response.RespondWithError(c, response.StatusNotFound, errorResp) - return - } - - // Prepare the response based on query parameters + // Prepare the response result := gin.H{ - "doc_id": segmentGraph.DocID, - "segment_id": segmentGraph.SegmentID, + "doc_id": docID, + "segment_id": segmentID, } - // Add entities if requested + // Get entities if requested if includeEntities { - result["entities"] = segmentGraph.Entities - result["entities_count"] = len(segmentGraph.Entities) + entities, err := kb.Instance.GetSegmentEntities(c.Request.Context(), docID, segmentID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: "segment_entities_error", + ErrorDescription: fmt.Sprintf("Failed to get segment entities: %v", err), + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return + } + result["entities"] = entities + result["entities_count"] = len(entities) } - // Add relationships if requested + // Get relationships if requested (using entity-based query for better results) if includeRelationships { - result["relationships"] = segmentGraph.Relationships - result["relationships_count"] = len(segmentGraph.Relationships) + relationships, err := kb.Instance.GetSegmentRelationshipsByEntities(c.Request.Context(), docID, segmentID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: "segment_relationships_error", + ErrorDescription: fmt.Sprintf("Failed to get segment relationships: %v", err), + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return + } + result["relationships"] = relationships + result["relationships_count"] = len(relationships) + result["query_type"] = "by_entities" // Indicate we're using entity-based relationship query } response.RespondWithSuccess(c, response.StatusOK, result) @@ -197,6 +205,63 @@ func GetSegmentRelationships(c *gin.Context) { response.RespondWithSuccess(c, response.StatusOK, result) } +// GetSegmentRelationshipsByEntities gets all relationships connected to entities in this segment +func GetSegmentRelationshipsByEntities(c *gin.Context) { + // Extract docID from URL path + docID := c.Param("docID") + if docID == "" { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Document ID is required", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Extract segmentID from URL path + segmentID := c.Param("segmentID") + if segmentID == "" { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Segment ID is required", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Check if kb.Instance is available + if kb.Instance == nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Knowledge base not initialized", + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Call the GraphRag instance to get segment relationships by entities + relationships, err := kb.Instance.GetSegmentRelationshipsByEntities(c.Request.Context(), docID, segmentID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: "segment_relationships_by_entities_error", + ErrorDescription: fmt.Sprintf("Failed to get segment relationships by entities: %v", err), + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return + } + + // Prepare the response + result := gin.H{ + "doc_id": docID, + "segment_id": segmentID, + "relationships": relationships, + "relationships_count": len(relationships), + "query_type": "by_entities", // Indicate this is entity-based query + } + + response.RespondWithSuccess(c, response.StatusOK, result) +} + // ExtractSegmentGraph re-extracts entities and relationships for a specific segment (synchronous) func ExtractSegmentGraph(c *gin.Context) { // Extract docID from URL path @@ -347,16 +412,17 @@ func ExtractSegmentGraph(c *gin.Context) { return } - // Build response + // Build response using the simplified SegmentExtractionResult structure result := map[string]interface{}{ "message": "Entities and relationships extracted successfully", "doc_id": extractionResult.DocID, "segment_id": extractionResult.SegmentID, - "entities_count": len(extractionResult.ExtractedEntities), - "relationships_count": len(extractionResult.ExtractedRelationships), - "entities": extractionResult.ExtractedEntities, - "relationships": extractionResult.ExtractedRelationships, + "entities_count": extractionResult.EntitiesCount, // Use count from structure + "relationships_count": extractionResult.RelationshipsCount, // Use count from structure + "extraction_model": extractionResult.ExtractionModel, "extraction_options": extractOptions, + // Note: Detailed entities and relationships are no longer returned + // Frontend should use separate APIs (GetSegmentEntities/GetSegmentRelationships) if needed } response.RespondWithSuccess(c, response.StatusOK, result) diff --git a/openapi/kb/kb.go b/openapi/kb/kb.go index de254484..812f7a1a 100644 --- a/openapi/kb/kb.go +++ b/openapi/kb/kb.go @@ -54,6 +54,7 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) { group.GET("/documents/:docID/segments/:segmentID/graph", GetSegmentGraph) group.GET("/documents/:docID/segments/:segmentID/entities", GetSegmentEntities) group.GET("/documents/:docID/segments/:segmentID/relationships", GetSegmentRelationships) + group.GET("/documents/:docID/segments/:segmentID/relationships/by-entities", GetSegmentRelationshipsByEntities) group.POST("/documents/:docID/segments/:segmentID/extract", ExtractSegmentGraph) group.POST("/documents/:docID/segments/:segmentID/extract/async", ExtractSegmentGraphAsync)