Enterprise-Knowledge-Base/internal/agent/tools/list_knowledge_chunks.go

package tools

import (
	"context"
	"fmt"
	"strings"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

// ListKnowledgeChunksTool retrieves chunk snapshots for a specific knowledge document.
type ListKnowledgeChunksTool struct {
	BaseTool
	tenantID         uint
	chunkService     interfaces.ChunkService
	knowledgeService interfaces.KnowledgeService
}

// NewListKnowledgeChunksTool creates a new tool instance.
func NewListKnowledgeChunksTool(
	tenantID uint,
	knowledgeService interfaces.KnowledgeService,
	chunkService interfaces.ChunkService,
) *ListKnowledgeChunksTool {
	description := `Retrieve paged chunks for a document (knowledge) by knowledge_id.

## When to Use

- Need deterministic chunk previews for a known document
- Want to quickly confirm how many chunks a document contains
- Require surrounding context around a chunk_index returned by search results
- Need content snippets without running an additional search query

Avoid when:
- You don't know the knowledge_id (use knowledge_search first)

## Parameters

- knowledge_id (required): Target document/knowledge ID
- limit (optional): Number of chunks to fetch (default 20, max 100).
- offset (optional): Offset to start fetching chunks from (default 0).`

	return &ListKnowledgeChunksTool{
		BaseTool:         NewBaseTool("list_knowledge_chunks", description),
		tenantID:         tenantID,
		chunkService:     chunkService,
		knowledgeService: knowledgeService,
	}
}

// Parameters returns the JSON schema describing accepted arguments.
func (t *ListKnowledgeChunksTool) Parameters() map[string]interface{} {
	return map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"knowledge_id": map[string]interface{}{
				"type":        "string",
				"description": "Knowledge/document ID to inspect",
			},
			"limit": map[string]interface{}{
				"type":        "integer",
				"description": "Number of chunks to fetch (default 20, max 100)",
				"default":     20,
				"minimum":     1,
				"maximum":     100,
			},
			"offset": map[string]interface{}{
				"type":        "integer",
				"description": "Offset to start fetching chunks from (default 0)",
				"default":     0,
				"minimum":     0,
			},
		},
		"required": []string{"knowledge_id", "limit", "offset"},
	}
}

// Execute performs the chunk fetch against the chunk service.
func (t *ListKnowledgeChunksTool) Execute(ctx context.Context, args map[string]interface{}) (*types.ToolResult, error) {
	knowledgeID, ok := args["knowledge_id"].(string)
	if !ok || strings.TrimSpace(knowledgeID) == "" {
		return &types.ToolResult{
			Success: false,
			Error:   "knowledge_id is required",
		}, fmt.Errorf("knowledge_id is required")
	}
	knowledgeID = strings.TrimSpace(knowledgeID)

	chunkLimit := 20
	offset := 0
	if rawLimit, exists := args["limit"]; exists {
		switch v := rawLimit.(type) {
		case float64:
			chunkLimit = int(v)
		case int:
			chunkLimit = v
		}
	}
	if rawOffset, exists := args["offset"]; exists {
		switch v := rawOffset.(type) {
		case float64:
			offset = int(v)
		case int:
			offset = v
		}
	}
	if offset < 0 {
		offset = 0
	}

	pagination := &types.Pagination{
		Page:     offset/chunkLimit + 1,
		PageSize: chunkLimit,
	}

	chunks, total, err := t.chunkService.GetRepository().ListPagedChunksByKnowledgeID(ctx,
		t.tenantID, knowledgeID, pagination, []types.ChunkType{types.ChunkTypeText}, "")
	if err != nil {
		return &types.ToolResult{
			Success: false,
			Error:   fmt.Sprintf("failed to list chunks: %v", err),
		}, err
	}
	if chunks == nil {
		return &types.ToolResult{
			Success: false,
			Error:   "chunk query returned no data",
		}, fmt.Errorf("chunk query returned no data")
	}

	totalChunks := total
	fetched := len(chunks)

	knowledgeTitle := t.lookupKnowledgeTitle(ctx, knowledgeID)

	output := t.buildOutput(knowledgeID, knowledgeTitle, totalChunks, fetched, chunkLimit, chunks)

	formattedChunks := make([]map[string]interface{}, 0, len(chunks))
	for idx, c := range chunks {
		formattedChunks = append(formattedChunks, map[string]interface{}{
			"seq":             idx + 1,
			"chunk_id":        c.ID,
			"chunk_index":     c.ChunkIndex,
			"content":         c.Content,
			"chunk_type":      c.ChunkType,
			"knowledge_id":    c.KnowledgeID,
			"knowledge_base":  c.KnowledgeBaseID,
			"start_at":        c.StartAt,
			"end_at":          c.EndAt,
			"parent_chunk_id": c.ParentChunkID,
		})
	}

	return &types.ToolResult{
		Success: true,
		Output:  output,
		Data: map[string]interface{}{
			"knowledge_id":    knowledgeID,
			"knowledge_title": knowledgeTitle,
			"total_chunks":    totalChunks,
			"fetched_chunks":  fetched,
			"page":            pagination.Page,
			"page_size":       pagination.PageSize,
			"chunks":          formattedChunks,
		},
	}, nil
}

func (t *ListKnowledgeChunksTool) lookupKnowledgeTitle(ctx context.Context, knowledgeID string) string {
	if t.knowledgeService == nil {
		return ""
	}
	knowledge, err := t.knowledgeService.GetKnowledgeByID(ctx, knowledgeID)
	if err != nil || knowledge == nil {
		return ""
	}
	return strings.TrimSpace(knowledge.Title)
}

func (t *ListKnowledgeChunksTool) buildOutput(
	knowledgeID string,
	knowledgeTitle string,
	total int64,
	fetched int,
	chunkLimit int,
	chunks []*types.Chunk,
) string {
	builder := &strings.Builder{}
	builder.WriteString("=== 知识文档分块 ===\n\n")

	if knowledgeTitle != "" {
		builder.WriteString(fmt.Sprintf("文档: %s (%s)\n", knowledgeTitle, knowledgeID))
	} else {
		builder.WriteString(fmt.Sprintf("文档 ID: %s\n", knowledgeID))
	}
	builder.WriteString(fmt.Sprintf("总分块数: %d\n", total))
	builder.WriteString(fmt.Sprintf("本次拉取: %d 条（offset=%d）\n\n", fetched, chunkLimit))

	if fetched == 0 {
		builder.WriteString("未找到任何分块，请确认文档是否已完成解析。\n")
		if total > 0 {
			builder.WriteString("文档存在但当前页数据为空，请检查分页参数。\n")
		}
		return builder.String()
	}

	builder.WriteString("=== 分块内容预览 ===\n\n")
	for idx, c := range chunks {
		builder.WriteString(fmt.Sprintf("Chunk #%d (Index %d)\n", idx+1, c.ChunkIndex+1))
		builder.WriteString(fmt.Sprintf("  chunk_id: %s\n", c.ID))
		builder.WriteString(fmt.Sprintf("  类型: %s\n", c.ChunkType))
		builder.WriteString(fmt.Sprintf("  内容: %s\n\n", summarizeContent(c.Content)))
	}

	if int64(fetched) < total {
		builder.WriteString("提示：文档仍有更多分块，可调整 offset 或多次调用以获取全部内容。\n")
	}

	return builder.String()
}

func summarizeContent(content string) string {
	cleaned := strings.TrimSpace(content)
	if cleaned == "" {
		return "(空内容)"
	}

	return strings.TrimSpace(string(cleaned))
}