yao/openapi/kb/segment.go

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)
}