xinmi/internal/application/service/chat_pipeline/extract_entity_restrictions.go

package chatpipeline

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"unicode"

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

// These are labelled source lists, not medical classifications. All offsets
// refer to the unchanged chunk; normalized text is used only for comparison.
var graphListHeadingPattern = regexp.MustCompile(`(?m)[〔〖【\[][^〔〖【\[〕〗】\]\n]{1,40}[〕〗】\]]|^[ \t]*[^\n:：〔〖【\[]{1,30}[:：]`)
var graphListPredicatePattern = regexp.MustCompile(`谨\s*慎\s*使\s*用|慎\s*[用甩](?:\s*于)?|不\s*适\s*[宜用合](?:\s*于)?|不\s*宜(?:\s*于)?|禁\s*止\s*使\s*用|禁\s*用(?:\s*于)?|不\s*[得可]\s*使\s*用|不\s*能\s*使\s*用`)
var graphListAmbiguityPattern = regexp.MustCompile(`[和及与或但若如除仅]|以及|同时|慎|禁|不适|使用|服用|适用|适宜|用于`)

type graphRestrictionMember struct {
	start, end   int
	target       string
	sourceTarget string
	predicate    string
	label        string
	explicit     bool
	ambiguous    bool
}

type graphRestrictionScope struct {
	start, end, headingEnd int
	section                int
	ownerText              string
	heading                string
	label                  string
	members                []graphRestrictionMember
	ambiguous              bool
}

func graphRestrictionLabel(predicate string) string {
	switch strings.TrimSuffix(graphComparisonText(predicate), "于") {
	case "慎用", "慎甩", "谨慎使用":
		return "慎用于"
	case "不适宜", "不适用", "不适合", "不宜":
		return "不适用于"
	case "禁用", "禁止使用", "不得使用", "不可使用", "不能使用":
		return "禁用于"
	}
	return ""
}

func graphRestrictionHeadingLabel(heading string) (string, bool) {
	name := strings.Trim(graphComparisonText(heading), "〔〖【[〕〗】]:：")
	for _, suffix := range []string{"人群", "对象"} {
		if !strings.HasSuffix(name, suffix) {
			continue
		}
		predicate := strings.TrimSuffix(name, suffix)
		label := graphRestrictionLabel(predicate)
		// An OCR-damaged restrictive heading stays a literal block. It must not
		// lend a guessed typed predicate to its members.
		return label, label != "" || strings.HasPrefix(predicate, "不") ||
			strings.HasPrefix(predicate, "慎") || strings.HasPrefix(predicate, "禁")
	}
	return "", false
}

func graphTrimSpan(content string, start, end int) (int, int) {
	text := content[start:end]
	left := strings.TrimLeftFunc(text, unicode.IsSpace)
	start += len(text) - len(left)
	end = start + len(strings.TrimRightFunc(left, unicode.IsSpace))
	return start, end
}

func graphRestrictionScopes(content string) []graphRestrictionScope {
	if !graphHasProductDocumentation(content) {
		return nil
	}
	var scopes []graphRestrictionScope
	offset := 0
	for sectionIndex, section := range graphProductSections(content) {
		headings := graphListHeadingPattern.FindAllStringIndex(section, -1)
		for index, heading := range headings {
			label, recognized := graphRestrictionHeadingLabel(section[heading[0]:heading[1]])
			if !recognized {
				continue
			}
			end := len(section)
			if index+1 < len(headings) {
				end = headings[index+1][0]
			}
			body := section[heading[1]:end]
			if boundary := strings.IndexAny(body, "。！？!?"); boundary >= 0 {
				end = heading[1] + boundary
			}
			if boundary := strings.Index(section[heading[1]:end], "\n\n"); boundary >= 0 {
				end = heading[1] + boundary
			}
			start, end := graphTrimSpan(content, offset+heading[0], offset+end)
			scope := graphRestrictionScope{
				start: start, end: end, headingEnd: offset + heading[1],
				section: sectionIndex + 1, ownerText: section[:heading[0]],
				heading: section[heading[0]:heading[1]], label: label, ambiguous: label == "",
			}
			memberStart := scope.headingEnd
			for i, r := range content[scope.headingEnd:scope.end] {
				if !strings.ContainsRune("、，,；;", r) {
					continue
				}
				memberEnd := scope.headingEnd + i
				scope.members = append(scope.members, graphParseRestrictionMember(content, memberStart, memberEnd, label))
				memberStart = memberEnd + len(string(r))
			}
			scope.members = append(scope.members, graphParseRestrictionMember(content, memberStart, scope.end, label))
			for i, member := range scope.members {
				scope.ambiguous = scope.ambiguous || member.ambiguous
				// A suffix predicate after bare coordinated items may govern all of
				// them. Only a heading can supply the omitted predicate here.
				if i > 0 && member.explicit && !scope.members[i-1].explicit {
					scope.ambiguous = true
				}
				if member.target != "" && strings.Count(content[scope.headingEnd:scope.end], member.target) != 1 {
					scope.ambiguous = true
				}
			}
			scopes = append(scopes, scope)
		}
		offset += len(section)
	}
	return scopes
}

func graphParseRestrictionMember(content string, start, end int, headingLabel string) graphRestrictionMember {
	start, end = graphTrimSpan(content, start, end)
	member := graphRestrictionMember{start: start, end: end, label: headingLabel}
	raw := content[start:end]
	predicates := graphListPredicatePattern.FindAllStringIndex(raw, -1)
	if len(predicates) == 0 {
		member.target = raw
		member.sourceTarget = raw
	} else if len(predicates) == 1 {
		predicate := predicates[0]
		member.explicit = true
		member.predicate = raw[predicate[0]:predicate[1]]
		member.label = graphRestrictionLabel(member.predicate)
		before, after := strings.TrimSpace(raw[:predicate[0]]), strings.TrimSpace(raw[predicate[1]:])
		switch {
		case before != "" && after == "":
			member.sourceTarget = before
			member.target = strings.TrimSpace(strings.TrimPrefix(before, "对"))
		case before == "" && after != "":
			member.target = after
			member.sourceTarget = after
		default:
			member.ambiguous = true
		}
	} else {
		member.ambiguous = true
	}
	member.ambiguous = member.ambiguous || member.target == "" || member.label == "" ||
		graphListAmbiguityPattern.MatchString(graphComparisonText(member.target))
	return member
}

func graphRestrictionTypeMatches(expected, actual string) bool {
	return expected == actual || (expected == "禁用于" && actual == "不能用于")
}

func graphRestrictionTargetMatches(member graphRestrictionMember, target string) bool {
	return target != "" && (target == member.target || target == member.sourceTarget)
}

// Display comparisons remove whitespace only, never case or OCR distinctions.
func graphRestrictionDisplayText(value string) string {
	return strings.Join(strings.Fields(value), "")
}

func graphRestrictionHeadingName(heading string) string {
	return strings.Trim(graphRestrictionDisplayText(heading), "〔〖【[〕〗】]:：")
}

func graphRestrictionDisplayParts(value string) (heading, body string, ok bool) {
	value = graphRestrictionDisplayText(value)
	span := graphListHeadingPattern.FindStringIndex(value)
	if span == nil || span[0] != 0 {
		return "", "", false
	}
	heading, body = value[:span[1]], value[span[1]:]
	_, ok = graphRestrictionHeadingLabel(heading)
	return heading, body, ok
}

func graphRestrictionCompleteTargetMatches(scope graphRestrictionScope, target, content string) bool {
	heading, body, ok := graphRestrictionDisplayParts(target)
	if !ok || body == "" || body != graphRestrictionDisplayText(content[scope.headingEnd:scope.end]) {
		return false
	}
	sourceHeading := graphRestrictionDisplayText(scope.heading)
	name := graphRestrictionHeadingName(scope.heading)
	return heading == sourceHeading || heading == name+":" || heading == name+"："
}

func graphRestrictionSubjectMatches(scope graphRestrictionScope, subject string) bool {
	return graphGenericProduct(subject) || (graphProductIdentity(subject) != "" &&
		strings.Contains(graphProductIdentity(scope.ownerText), graphProductIdentity(subject)))
}

func graphRestrictionEvidenceTouches(scope graphRestrictionScope, evidence, content string) bool {
	if evidence == "" {
		return false
	}
	for offset := 0; offset < len(content); {
		index := strings.Index(content[offset:], evidence)
		if index < 0 {
			break
		}
		start := offset + index
		end := start + len(evidence)
		if start < scope.end && end > scope.headingEnd {
			return true
		}
		offset = end
	}
	return false
}

// Independent source/target anchors allow semantic diagnostics even when the
// model changed its quotation. They never repair that quotation or output edge.
func graphRestrictionRelation(content string, relation *types.GraphRelation, scopes []graphRestrictionScope) ([]string, bool, error) {
	if relation == nil || len(scopes) == 0 {
		return nil, false, nil
	}
	target := strings.TrimSpace(relation.Node2)
	type match struct {
		scope  graphRestrictionScope
		member int
		whole  bool
	}
	var candidates []match
	related := false
	for _, scope := range scopes {
		whole := graphRestrictionCompleteTargetMatches(scope, target, content)
		memberIndex := -1
		for i, member := range scope.members {
			if graphRestrictionTargetMatches(member, target) {
				memberIndex = i
			}
		}
		touches := graphRestrictionEvidenceTouches(scope, relation.Evidence, content)
		useLabel := graphPositiveUseLabel(relation.Type) || relation.Type == "慎用于" ||
			relation.Type == "禁用于" || relation.Type == "不适用于" || relation.Type == "不能用于"
		if !whole && memberIndex < 0 && !(touches && (useLabel || relation.Type == "关联" || relation.Type == "约束")) {
			continue
		}
		related = true
		if graphRestrictionSubjectMatches(scope, relation.Node1) {
			candidates = append(candidates, match{scope: scope, member: memberIndex, whole: whole})
		}
	}
	if !related {
		return nil, false, nil
	}
	if len(candidates) > 1 && strings.Contains(content, relation.Evidence) && relation.Evidence != "" {
		var anchored []match
		for _, candidate := range candidates {
			if graphRestrictionEvidenceTouches(candidate.scope, relation.Evidence, content) {
				anchored = append(anchored, candidate)
			}
		}
		candidates = anchored
	}
	if len(candidates) != 1 {
		return nil, true, fmt.Errorf("%w: restriction list has ambiguous target or product scope; preserve a complete literal list constraint with source-specific evidence", ErrGraphNeedsReview)
	}
	candidate := candidates[0]
	scope := candidate.scope
	if candidate.whole && relation.Type == "约束" {
		if !strings.Contains(relation.Evidence, content[scope.start:scope.end]) {
			return nil, true, fmt.Errorf("%w: complete restriction constraint requires unchanged heading and list evidence in source section %d", ErrGraphNeedsReview, scope.section)
		}
		return []string{graphComparisonText(content[scope.start:scope.end])}, true, nil
	}
	if scope.ambiguous || candidate.member < 0 {
		return nil, true, fmt.Errorf("%w: ambiguous restriction list in source section %d requires its complete literal block as a constraint, not an inferred population edge", ErrGraphNeedsReview, scope.section)
	}
	member := scope.members[candidate.member]
	var failures []error
	if !graphRestrictionTypeMatches(member.label, relation.Type) {
		failures = append(failures, fmt.Errorf("%w: restriction strength mismatch in source section %d item %d: expected %s from its own predicate, not %s from a neighboring item", ErrGraphNeedsReview, scope.section, candidate.member+1, member.label, relation.Type))
	}
	claim := content[scope.start:member.end]
	clause := scope.heading + content[member.start:member.end]
	if member.explicit {
		claim = content[member.start:member.end]
		clause = claim
	}
	if !strings.Contains(relation.Evidence, claim) {
		failures = append(failures, fmt.Errorf("%w: restriction evidence must include the unchanged governing predicate and target in source section %d item %d; a heading-governed item needs its heading and intervening text", ErrGraphNeedsReview, scope.section, candidate.member+1))
	}
	return []string{graphComparisonText(clause)}, true, errors.Join(failures...)
}

func validateGraphRestrictionCoverage(graph *types.GraphData, content string) error {
	scopes := graphRestrictionScopes(content)
	var failures []error
	for _, scope := range scopes {
		covered := make([]bool, len(scope.members))
		complete := false
		for _, relation := range graph.Relation {
			if relation == nil || relation.Evidence == "" || !strings.Contains(content, relation.Evidence) ||
				!graphRestrictionEvidenceTouches(scope, relation.Evidence, content) {
				continue
			}
			_, scoped, err := graphRestrictionRelation(content, relation, scopes)
			if !scoped || err != nil || !graphRestrictionSubjectMatches(scope, relation.Node1) {
				continue
			}
			if relation.Type == "约束" && graphRestrictionCompleteTargetMatches(scope, relation.Node2, content) &&
				strings.Contains(relation.Evidence, content[scope.start:scope.end]) {
				complete = true
				break
			}
			for i, member := range scope.members {
				if graphRestrictionTargetMatches(member, strings.TrimSpace(relation.Node2)) && graphRestrictionTypeMatches(member.label, relation.Type) {
					covered[i] = true
				}
			}
		}
		if !complete {
			complete = !scope.ambiguous
			for _, present := range covered {
				complete = complete && present
			}
		}
		if !complete {
			failures = append(failures, fmt.Errorf("%w: restriction list in source section %d is not fully covered; preserve each unambiguous member with its own strength or the complete literal heading and list as one constraint", ErrGraphNeedsReview, scope.section))
		}
	}
	return errors.Join(failures...)
}

func validateGraphRestrictionAttributes(graph *types.GraphData, content string) error {
	scopes := graphRestrictionScopes(content)
	if len(scopes) == 0 {
		return nil
	}
	sectionCount := len(graphProductSections(content))
	var failures []error
	for i, node := range graph.Node {
		if node == nil {
			continue
		}
		_, headingNode := graphRestrictionHeadingLabel(node.Name)
		for j, attribute := range node.Attributes {
			heading, _, labelled := graphRestrictionDisplayParts(attribute)
			if !headingNode && !labelled {
				continue
			}
			name := graphRestrictionHeadingName(heading)
			if headingNode {
				name = graphRestrictionHeadingName(node.Name)
			}
			var candidates []graphRestrictionScope
			for _, scope := range scopes {
				if graphRestrictionHeadingName(scope.heading) != name {
					continue
				}
				// Attributes have no evidence anchor to disambiguate a generic
				// product name across sections or a repeated standalone heading.
				if headingNode || (graphRestrictionSubjectMatches(scope, node.Name) &&
					(!graphGenericProduct(node.Name) || sectionCount == 1)) {
					candidates = append(candidates, scope)
				}
			}
			if len(candidates) != 1 {
				failures = append(failures, fmt.Errorf("%w: entity %d attribute %d has an ambiguous or unsupported restriction product/heading scope; require one literal local source block", ErrGraphNeedsReview, i+1, j+1))
				continue
			}
			scope := candidates[0]
			supported := graphRestrictionCompleteTargetMatches(scope, attribute, content)
			if headingNode && !labelled {
				matches := 0
				for _, member := range scope.members {
					if graphRestrictionDisplayText(attribute) != "" &&
						graphRestrictionDisplayText(attribute) == graphRestrictionDisplayText(content[member.start:member.end]) {
						matches++
					}
				}
				supported = matches == 1
			}
			if !supported {
				failures = append(failures, fmt.Errorf("%w: entity %d attribute %d has unsupported restriction text in source section %d; preserve literal OCR spelling, each local predicate and the complete labelled list", ErrGraphNeedsReview, i+1, j+1, scope.section))
			}
		}
	}
	return errors.Join(failures...)
}