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.
This commit is contained in:
Max
parent
bc85518bc3
commit
0252e58b04
4 changed files with 95 additions and 28 deletions
2
go.mod
2
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
|
||||
|
|
|
|||
4
go.sum
4
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=
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue