xinmi/internal/workbench/categories.go

package workbench

import (
	"context"
	"fmt"
	"slices"
	"sort"
	"strings"

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

type SegmentCategory struct {
	ID              string `json:"id"`
	KnowledgeBaseID string `json:"knowledge_base_id"`
	Name            string `json:"name"`
}

// Resolve the permitted task scope before enumerating its categories. Explicit
// files remain explicit: selecting a category must never open the entire base.
func (s *Service) segmentCategories(ctx context.Context, d TaskData) ([]SegmentCategory, error) {
	result := []SegmentCategory{}
	if s.knowledgeBases == nil || s.knowledge == nil {
		return result, nil
	}
	if err := s.validateKnowledgeScope(ctx, d.KBIDs, d.KnowledgeIDs, d.KnowledgeTags); err != nil {
		return nil, err
	}
	tenant, _, _, err := identity(ctx)
	if err != nil {
		return nil, err
	}
	baseIDs := slices.Clone(d.KBIDs)
	for id := range d.KnowledgeTags {
		baseIDs = append(baseIDs, id)
	}
	files, err := s.knowledge.GetKnowledgeBatchWithSharedAccess(ctx, tenant, d.KnowledgeIDs)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if file != nil {
			baseIDs = append(baseIDs, file.KnowledgeBaseID)
		}
	}
	sort.Strings(baseIDs)
	for _, id := range slices.Compact(baseIDs) {
		base, err := s.knowledgeBases.GetKnowledgeBaseByID(ctx, id)
		if err != nil || base == nil {
			return nil, ErrForbidden
		}
		if err := s.validateKnowledgeScope(ctx, []string{id}, nil, nil); err != nil {
			return nil, err
		}
		q := s.db.WithContext(ctx).Model(&types.KnowledgeTag{}).Where("tenant_id=? AND knowledge_base_id=?", base.TenantID, id)
		if tags := d.KnowledgeTags[id]; len(tags) > 0 && len(d.KnowledgeFolders[id]) == 0 {
			q = q.Where("id IN ?", tags)
		}
		allowedFiles := []string{}
		for _, file := range files {
			if file != nil && file.KnowledgeBaseID == id {
				allowedFiles = append(allowedFiles, file.ID)
			}
		}
		wholeBase := slices.Contains(d.KBIDs, id) && len(d.KnowledgeTags[id]) == 0 && len(d.KnowledgeFolders[id]) == 0
		if !wholeBase && len(allowedFiles) > 0 {
			q = q.Where("id IN (SELECT tag_id FROM knowledge_tag_relations WHERE knowledge_id IN ?)", allowedFiles)
		}
		var rows []types.KnowledgeTag
		if err := q.Order("sort_order, name, id").Find(&rows).Error; err != nil {
			return nil, err
		}
		for _, row := range rows {
			result = append(result, SegmentCategory{row.ID, row.KnowledgeBaseID, row.Name})
		}
	}
	return result, nil
}

func normalizeCategoryRecommendations(segments []Segment, options []SegmentCategory) error {
	allowed := map[string]bool{}
	for _, option := range options {
		allowed[option.ID] = true
	}
	for i := range segments {
		seg := &segments[i]
		if seg.RecommendedCategoryID != "" && !allowed[seg.RecommendedCategoryID] {
			return fmt.Errorf("%w: AI 推荐了当前 Agent 资料范围以外的分类，请重新拆解", ErrInvalid)
		}
		seg.CategoryID = seg.RecommendedCategoryID
		if seg.CategoryID == "" {
			seg.CategoryReason = "当前范围没有合适分类，沿用本段 Agent 的资料范围。"
		} else if strings.TrimSpace(seg.CategoryReason) == "" {
			return fmt.Errorf("%w: AI 未说明分类推荐理由，请重新拆解", ErrInvalid)
		}
	}
	return nil
}

func (s *Service) agentCategoryOptions(ctx context.Context, d TaskData, flow FlowData) (map[string][]SegmentCategory, error) {
	result := map[string][]SegmentCategory{}
	for _, node := range flow.Nodes {
		if _, exists := result[node.AgentID]; exists {
			continue
		}
		scope := d
		selected := flow
		selected.Nodes = []FlowNode{{ID: node.ID, AgentID: node.AgentID}}
		if err := s.inheritFlow(ctx, &scope, selected); err != nil {
			return nil, err
		}
		options, err := s.segmentCategories(ctx, scope)
		if err != nil {
			return nil, err
		}
		result[node.AgentID] = options
	}
	return result, nil
}

func validateAgentCategoryRecommendations(segments []Segment, options map[string][]SegmentCategory) error {
	for _, segment := range segments {
		if segment.RecommendedCategoryID != "" && !slices.ContainsFunc(options[segment.AgentID], func(option SegmentCategory) bool {
			return option.ID == segment.RecommendedCategoryID
		}) {
			return fmt.Errorf("%w: 本段推荐分类不属于负责 Agent 的资料范围，请重新拆解", ErrInvalid)
		}
	}
	return nil
}

// A recommendation is optional metadata. Discard invalid AI suggestions without
// widening the assigned Agent's scope or losing an otherwise valid decomposition.
func reconcileCategoryRecommendations(segments []Segment, options []SegmentCategory, byAgent map[string][]SegmentCategory) error {
	for i := range segments {
		segment := &segments[i]
		if segment.RecommendedCategoryID == "" {
			continue
		}
		allowed := func(option SegmentCategory) bool { return option.ID == segment.RecommendedCategoryID }
		if !slices.ContainsFunc(options, allowed) || !slices.ContainsFunc(byAgent[segment.AgentID], allowed) || strings.TrimSpace(segment.CategoryReason) == "" {
			segment.RecommendedCategoryID = ""
			segment.CategoryReason = ""
		}
	}
	if err := normalizeCategoryRecommendations(segments, options); err != nil {
		return err
	}
	return validateAgentCategoryRecommendations(segments, byAgent)
}

func (s *Service) applySegmentCategory(ctx context.Context, scope *TaskData, categoryID string) error {
	if categoryID == "" {
		return nil
	}
	options, err := s.segmentCategories(ctx, *scope)
	if err != nil {
		return err
	}
	var selected *SegmentCategory
	for i := range options {
		if options[i].ID == categoryID {
			selected = &options[i]
			break
		}
	}
	if selected == nil {
		return fmt.Errorf("%w: 本段分类不在当前 Agent 的知识范围内，请重新选择", ErrInvalid)
	}
	baseID := selected.KnowledgeBaseID
	tenant, _, _, err := identity(ctx)
	if err != nil {
		return err
	}
	files, err := s.knowledge.GetKnowledgeBatchWithSharedAccess(ctx, tenant, scope.KnowledgeIDs)
	if err != nil {
		return err
	}
	fileIDs := []string{}
	for _, file := range files {
		if file != nil && file.KnowledgeBaseID == baseID {
			fileIDs = append(fileIDs, file.ID)
		}
	}
	wholeBase := slices.Contains(scope.KBIDs, baseID) && len(scope.KnowledgeTags[baseID]) == 0 && len(scope.KnowledgeFolders[baseID]) == 0
	if wholeBase || len(fileIDs) == 0 {
		scope.KBIDs, scope.KnowledgeIDs = nil, nil
		scope.KnowledgeFolders = nil
		scope.KnowledgeTags = map[string][]string{baseID: {categoryID}}
		return nil
	}
	var ids []string
	if err := s.db.WithContext(ctx).Table("knowledge_tag_relations").Where("tag_id=? AND knowledge_id IN ?", categoryID, fileIDs).Pluck("knowledge_id", &ids).Error; err != nil {
		return err
	}
	if len(ids) == 0 {
		return fmt.Errorf("%w: 所选分类在当前资料范围内没有文件", ErrInvalid)
	}
	scope.KBIDs, scope.KnowledgeFolders, scope.KnowledgeTags = nil, nil, nil
	scope.KnowledgeIDs = ids
	return nil
}