package kb import ( "net/http" "strconv" "strings" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi/response" ) // Segment Management Handlers // AddSegments adds segments to a document func AddSegments(c *gin.Context) { var req AddSegmentsRequest // Parse and bind JSON request if err := c.ShouldBindJSON(&req); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid request format: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate request if err := req.Validate(); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), } 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 } // Convert request to UpsertOptions upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions() if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Failed to convert request to upsert options: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Perform add segments operation segmentIDs, err := kb.Instance.AddSegments(c.Request.Context(), req.DocID, req.SegmentTexts, upsertOptions) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to add segments: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response result := gin.H{ "message": "Segments added successfully", "collection_id": req.CollectionID, "doc_id": req.DocID, "segment_ids": segmentIDs, "segments_count": len(segmentIDs), } response.RespondWithSuccess(c, response.StatusCreated, result) } // UpdateSegments updates segments manually func UpdateSegments(c *gin.Context) { var req UpdateSegmentsRequest // Parse and bind JSON request if err := c.ShouldBindJSON(&req); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid request format: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate request if err := req.Validate(); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), } 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 } // Convert request to UpsertOptions upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions() if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Failed to convert request to upsert options: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Perform update segments operation updatedCount, err := kb.Instance.UpdateSegments(c.Request.Context(), req.SegmentTexts, upsertOptions) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to update segments: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response result := gin.H{ "message": "Segments updated successfully", "collection_id": req.CollectionID, "updated_count": updatedCount, "segments_count": len(req.SegmentTexts), } response.RespondWithSuccess(c, response.StatusOK, result) } // RemoveSegments removes segments by IDs func RemoveSegments(c *gin.Context) { // TODO: Implement remove segments logic c.JSON(http.StatusOK, gin.H{"message": "Segments removed"}) } // RemoveSegmentsByDocID removes all segments of a document func RemoveSegmentsByDocID(c *gin.Context) { // TODO: Implement remove segments by document ID logic c.JSON(http.StatusOK, gin.H{"message": "Segments removed by document ID"}) } // GetSegments gets segments by IDs func GetSegments(c *gin.Context) { // TODO: Implement get segments logic c.JSON(http.StatusOK, gin.H{"segments": []interface{}{}}) } // GetSegment gets a single segment by ID func GetSegment(c *gin.Context) { // TODO: Implement get single segment logic c.JSON(http.StatusOK, gin.H{"segment": nil}) } // ScrollSegments scrolls segments with iterator-style pagination func ScrollSegments(c *gin.Context) { // Parse docID from URL path parameter docID := c.Param("docID") if docID == "" { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "docID 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 } // Parse query parameters for scroll options options := &types.ScrollSegmentsOptions{ IncludeMetadata: true, // Default to include metadata } // Parse limit (default: 100) if limitStr := c.Query("limit"); limitStr != "" { if limit, err := strconv.Atoi(limitStr); err == nil && limit > 0 { options.Limit = limit } } if options.Limit == 0 { options.Limit = 100 // Default limit } // Parse scroll_id parameter for continuing pagination if scrollID := strings.TrimSpace(c.Query("scroll_id")); scrollID != "" { options.ScrollID = scrollID } // Parse order_by parameter if orderBy := strings.TrimSpace(c.Query("order_by")); orderBy != "" { options.OrderBy = strings.Split(orderBy, ",") // Trim spaces from each field for i, field := range options.OrderBy { options.OrderBy[i] = strings.TrimSpace(field) } } // Parse fields parameter if fields := strings.TrimSpace(c.Query("fields")); fields != "" { options.Fields = strings.Split(fields, ",") // Trim spaces from each field for i, field := range options.Fields { options.Fields[i] = strings.TrimSpace(field) } } // Parse include options if includeNodes := c.Query("include_nodes"); includeNodes == "true" { options.IncludeNodes = true } if includeRelationships := c.Query("include_relationships"); includeRelationships == "true" { options.IncludeRelationships = true } if includeMetadata := c.Query("include_metadata"); includeMetadata == "false" { options.IncludeMetadata = false } // Parse filter parameters (basic implementation for common filters) filter := make(map[string]interface{}) if score := c.Query("score"); score != "" { if scoreVal, err := strconv.ParseFloat(score, 64); err == nil { filter["score"] = scoreVal } } if weight := c.Query("weight"); weight != "" { if weightVal, err := strconv.ParseFloat(weight, 64); err == nil { filter["weight"] = weightVal } } if vote := c.Query("vote"); vote != "" { if voteVal, err := strconv.Atoi(vote); err == nil { filter["vote"] = voteVal } } if len(filter) > 0 { options.Filter = filter } // Call GraphRag ScrollSegments method result, err := kb.Instance.ScrollSegments(c.Request.Context(), docID, options) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to scroll segments: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response response.RespondWithSuccess(c, response.StatusOK, result) }