Refactor document ID handling in API responses

- Updated API response structures across multiple endpoints to replace `document_id` with `doc_id` for consistency.
- Enhanced error handling in segment retrieval to ensure proper validation of document ownership.
- Improved response clarity by ensuring all relevant endpoints reflect the updated field naming convention.

openapi/kb/graph.go

@@ -74,7 +74,7 @@ func GetSegmentGraph(c *gin.Context) {
 	result := gin.H{
 		"entities":      []interface{}{},
 		"relationships": []interface{}{},
-		"document_id":   docID,
+		"doc_id":        docID,
 		"segment_id":    segmentID,
 		"options":       options,
 	}
@@ -130,7 +130,7 @@ func ExtractSegmentGraph(c *gin.Context) {
 	// Return mock response for now
 	result := gin.H{
 		"message":             "Entities and relationships extracted successfully",
-		"document_id":         docID,
+		"doc_id":              docID,
 		"segment_id":          segmentID,
 		"entities_count":      0,
 		"relationships_count": 0,
@@ -196,10 +196,10 @@ func ExtractSegmentGraphAsync(c *gin.Context) {
 
 	// Return job ID for status tracking
 	result := gin.H{
-		"job_id":      jobID,
-		"message":     "Graph extraction started",
-		"document_id": docID,
-		"segment_id":  segmentID,
+		"job_id":     jobID,
+		"message":    "Graph extraction started",
+		"doc_id":     docID,
+		"segment_id": segmentID,
 	}
 
 	response.RespondWithSuccess(c, response.StatusCreated, result)

openapi/kb/hit.go

@@ -37,9 +37,9 @@ func ScrollHits(c *gin.Context) {
 
 	// Parse query parameters for scroll options
 	options := map[string]interface{}{
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"limit":       100, // Default limit
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"limit":      100, // Default limit
 	}
 
 	// Parse limit (default: 100)
@@ -145,10 +145,10 @@ func GetHits(c *gin.Context) {
 
 	// Return mock response for now
 	result := gin.H{
-		"hits":        []interface{}{},
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"total":       0,
+		"hits":       []interface{}{},
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"total":      0,
 	}
 
 	if len(filter) > 0 {
@@ -199,10 +199,10 @@ func GetHit(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 	// TODO: Implement get hit detail logic
 	c.JSON(http.StatusOK, gin.H{
-		"hit":         nil,
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"hit_id":      hitID,
+		"hit":        nil,
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"hit_id":     hitID,
 	})
 }
 
@@ -233,10 +233,10 @@ func AddHits(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 	// TODO: Implement add hit logic
 	c.JSON(http.StatusOK, gin.H{
-		"message":     "Hit added successfully",
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"hit_id":      "placeholder-hit-id",
+		"message":    "Hit added successfully",
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"hit_id":     "placeholder-hit-id",
 	})
 }
 
@@ -281,7 +281,7 @@ func UpdateHits(c *gin.Context) {
 
 	result := gin.H{
 		"message":       "Hits updated successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"segment_id":    segmentID,
 		"updated_count": len(hitIDs),
 	}
@@ -352,7 +352,7 @@ func RemoveHits(c *gin.Context) {
 
 	result := gin.H{
 		"message":       "Hits removed successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"segment_id":    segmentID,
 		"hit_ids":       validHitIDs,
 		"removed_count": len(validHitIDs),

openapi/kb/score.go

@@ -70,7 +70,7 @@ func UpdateScores(c *gin.Context) {
 
 	result := gin.H{
 		"message":       "Scores updated successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"scores":        req.Scores,
 		"updated_count": len(req.Scores),
 	}

openapi/kb/segment.go

@@ -331,7 +331,7 @@ func RemoveSegments(c *gin.Context) {
 	}
 
 	// Perform remove segments operation
-	removedCount, err := kb.Instance.RemoveSegments(c.Request.Context(), validSegmentIDs)
+	removedCount, err := kb.Instance.RemoveSegments(c.Request.Context(), docID, validSegmentIDs)
 	if err != nil {
 		errorResp := &response.ErrorResponse{
 			Code:             response.ErrServerError.Code,
@@ -411,8 +411,8 @@ func GetSegments(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 	// TODO: Implement get segments logic
 	c.JSON(http.StatusOK, gin.H{
-		"segments":    []interface{}{},
-		"document_id": docID,
+		"segments": []interface{}{},
+		"doc_id":   docID,
 	})
 }
 
@@ -440,13 +440,55 @@ func GetSegment(c *gin.Context) {
 		return
 	}
 
+	// Check if KB instance exists
+	if kb.Instance == nil {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrServerError.Code,
+			ErrorDescription: "Knowledge base instance is not initialized",
+		}
+		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
+		return
+	}
+
 	// TODO: Implement document permission validation for docID
-	// TODO: Implement get single segment logic
-	c.JSON(http.StatusOK, gin.H{
-		"segment":     nil,
-		"document_id": docID,
-		"segment_id":  segmentID,
-	})
+
+	// Get the segment using KB interface
+	segment, err := kb.Instance.GetSegment(c.Request.Context(), docID, segmentID)
+	if err != nil {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrServerError.Code,
+			ErrorDescription: fmt.Sprintf("Failed to get segment: %v", err),
+		}
+		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
+		return
+	}
+
+	if segment == nil {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrInvalidRequest.Code,
+			ErrorDescription: "Segment not found",
+		}
+		response.RespondWithError(c, response.StatusNotFound, errorResp)
+		return
+	}
+
+	// Verify that the segment belongs to the specified document
+	if segment.DocumentID != docID {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrAccessDenied.Code,
+			ErrorDescription: "Segment does not belong to the specified document",
+		}
+		response.RespondWithError(c, response.StatusForbidden, errorResp)
+		return
+	}
+
+	result := gin.H{
+		"segment":    segment,
+		"doc_id":     docID,
+		"segment_id": segmentID,
+	}
+
+	response.RespondWithSuccess(c, response.StatusOK, result)
 }
 
 // ScrollSegments scrolls segments with iterator-style pagination
@@ -621,9 +663,9 @@ func AddSegmentsAsync(c *gin.Context) {
 
 	// Return job ID for status tracking
 	result := gin.H{
-		"job_id":      jobID,
-		"message":     "Segments addition started",
-		"document_id": docID,
+		"job_id":  jobID,
+		"message": "Segments addition started",
+		"doc_id":  docID,
 	}
 
 	response.RespondWithSuccess(c, response.StatusCreated, result)
@@ -680,9 +722,9 @@ func UpdateSegmentsAsync(c *gin.Context) {
 
 	// Return job ID for status tracking
 	result := gin.H{
-		"job_id":      jobID,
-		"message":     "Segments update started",
-		"document_id": docID,
+		"job_id":  jobID,
+		"message": "Segments update started",
+		"doc_id":  docID,
 	}
 
 	response.RespondWithSuccess(c, response.StatusCreated, result)
@@ -747,11 +789,11 @@ func GetSegmentParents(c *gin.Context) {
 
 	// Return mock response for now
 	result := gin.H{
-		"parents":     []interface{}{},
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"depth":       depth,
-		"total":       0,
+		"parents":    []interface{}{},
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"depth":      depth,
+		"total":      0,
 	}
 
 	response.RespondWithSuccess(c, response.StatusOK, result)

openapi/kb/vote.go

@@ -37,9 +37,9 @@ func ScrollVotes(c *gin.Context) {
 
 	// Parse query parameters for scroll options
 	options := map[string]interface{}{
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"limit":       100, // Default limit
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"limit":      100, // Default limit
 	}
 
 	// Parse limit (default: 100)
@@ -139,10 +139,10 @@ func GetVotes(c *gin.Context) {
 
 	// Return mock response for now
 	result := gin.H{
-		"votes":       []interface{}{},
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"total":       0,
+		"votes":      []interface{}{},
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"total":      0,
 	}
 
 	if len(filter) > 0 {
@@ -193,10 +193,10 @@ func GetVote(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 	// TODO: Implement get vote detail logic
 	c.JSON(http.StatusOK, gin.H{
-		"vote":        nil,
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"vote_id":     voteID,
+		"vote":       nil,
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"vote_id":    voteID,
 	})
 }
 
@@ -227,10 +227,10 @@ func AddVotes(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 	// TODO: Implement add vote logic
 	c.JSON(http.StatusOK, gin.H{
-		"message":     "Vote added successfully",
-		"document_id": docID,
-		"segment_id":  segmentID,
-		"vote_id":     "placeholder-vote-id",
+		"message":    "Vote added successfully",
+		"doc_id":     docID,
+		"segment_id": segmentID,
+		"vote_id":    "placeholder-vote-id",
 	})
 }
 
@@ -275,7 +275,7 @@ func UpdateVotes(c *gin.Context) {
 
 	result := gin.H{
 		"message":       "Votes updated successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"segment_id":    segmentID,
 		"updated_count": len(voteIDs),
 	}
@@ -346,7 +346,7 @@ func RemoveVotes(c *gin.Context) {
 
 	result := gin.H{
 		"message":       "Votes removed successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"segment_id":    segmentID,
 		"vote_ids":      validVoteIDs,
 		"removed_count": len(validVoteIDs),

openapi/kb/weight.go

@@ -56,7 +56,7 @@ func UpdateWeights(c *gin.Context) {
 	// TODO: Implement document permission validation for docID
 
 	// Perform batch update weight operation
-	updatedCount, err := kb.Instance.UpdateWeight(c.Request.Context(), req.Weights)
+	updatedCount, err := kb.Instance.UpdateWeight(c.Request.Context(), docID, req.Weights)
 	if err != nil {
 		errorResp := &response.ErrorResponse{
 			Code:             response.ErrServerError.Code,
@@ -69,7 +69,7 @@ func UpdateWeights(c *gin.Context) {
 	// Return success response
 	result := gin.H{
 		"message":       "Segment weights updated successfully",
-		"document_id":   docID,
+		"doc_id":        docID,
 		"weights":       req.Weights,
 		"updated_count": updatedCount,
 	}