Merge pull request #1123 from trheyi/main

Update segment graph API to include entity-based relationship retrieval

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

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.14 h1:uv/0Bq533iFdnMHZdRBTOlaNMdb1+ZxXIlHDZHIHcvg=
-github.com/ulikunitz/xz v0.5.12/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
+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=

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
+	// Prepare the response
-	segmentGraph, err := kb.Instance.GetSegmentGraph(c.Request.Context(), docID, segmentID)
+	result := gin.H{
+		"doc_id":     docID,
+		"segment_id": segmentID,
+	}
+
+	// Get entities if requested
+	if includeEntities {
+		entities, err := kb.Instance.GetSegmentEntities(c.Request.Context(), docID, segmentID)
 		if err != nil {
 			errorResp := &response.ErrorResponse{
-			Code:             "segment_not_found",
+				Code:             "segment_entities_error",
-			ErrorDescription: fmt.Sprintf("Failed to get segment graph: %v", err),
+				ErrorDescription: fmt.Sprintf("Failed to get segment entities: %v", err),
 			}
 			response.RespondWithError(c, response.StatusNotFound, errorResp)
 			return
 		}
-
+		result["entities"] = entities
-	// Prepare the response based on query parameters
+		result["entities_count"] = len(entities)
-	result := gin.H{
-		"doc_id":     segmentGraph.DocID,
-		"segment_id": segmentGraph.SegmentID,
 	}
 
-	// Add entities if requested
+	// Get relationships if requested (using entity-based query for better results)
-	if includeEntities {
-		result["entities"] = segmentGraph.Entities
-		result["entities_count"] = len(segmentGraph.Entities)
-	}
-
-	// Add relationships if requested
 	if includeRelationships {
-		result["relationships"] = segmentGraph.Relationships
+		relationships, err := kb.Instance.GetSegmentRelationshipsByEntities(c.Request.Context(), docID, segmentID)
-		result["relationships_count"] = len(segmentGraph.Relationships)
+		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),
+		"entities_count":      extractionResult.EntitiesCount,      // Use count from structure
-		"relationships_count": len(extractionResult.ExtractedRelationships),
+		"relationships_count": extractionResult.RelationshipsCount, // Use count from structure
-		"entities":            extractionResult.ExtractedEntities,
+		"extraction_model":    extractionResult.ExtractionModel,
-		"relationships":       extractionResult.ExtractedRelationships,
 		"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)

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)