moteva/server/internal/infrastructure/piagent/reference_context.go

package piagent

import (
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"unicode/utf8"

	"img_infinite_canvas/internal/domain/design"
)

type inlineReferenceMatch struct {
	start int
	end   int
	index int
	name  string
	url   string
}

var (
	inlineReferenceDirectivePattern = regexp.MustCompile(`(?i)(参考图|reference image)\s*(\d+)?\s*(?:[（(]([^）)\n]+)[）)])?\s*[:：]\s*(https?://[^\s\]，。]+|data:image/[^\s\]，。]+|/[^\s\]，。]+|attachment://reference-\d+)`)
	inlinePromptImageTokenPattern   = regexp.MustCompile(`\[@image:#(\d+):([^:\]]+):(?:\[([^\]]+)\]\(([^)]+)\)|([^\]]+))\]`)
)

func inlineReferencePlacementContext(prompt string, messages []design.AgentChatMessage) string {
	messageText := agentMessageInlineReferenceText(messages)
	source := prompt
	if len(collectInlineReferenceMatches(messageText)) > 0 && (len(collectInlineReferenceMatches(prompt)) == 0 || utf8.RuneCountInString(messageText) >= utf8.RuneCountInString(prompt)) {
		source = messageText
	}
	return inlineReferencePlacementContextFromText(source, prefersEnglish(prompt+"\n"+messageText))
}

func inlineReferencePlacementContextFromText(text string, english bool) string {
	matches := collectInlineReferenceMatches(text)
	if len(matches) == 0 {
		return "none"
	}
	if len(matches) > 16 {
		matches = matches[:16]
	}

	lines := make([]string, 0, len(matches)+1)
	if english {
		lines = append(lines, "Image input order follows the inline reference order. Nearby text before or after a reference is semantic role/context for that image.")
	} else {
		lines = append(lines, "图片输入顺序按参考图出现顺序；某张图前后的相邻文字是这张图的角色/用途线索。")
	}
	for order, match := range matches {
		beforeStart := 0
		if order > 0 {
			beforeStart = matches[order-1].end
		}
		afterEnd := len(text)
		if order+1 < len(matches) {
			afterEnd = matches[order+1].start
		}
		before := compactReferenceNeighbor(text[beforeStart:match.start], true)
		after := compactReferenceNeighbor(text[match.end:afterEnd], false)
		if english {
			lines = append(lines, englishReferencePlacementLine(order+1, match, before, after))
			continue
		}
		lines = append(lines, chineseReferencePlacementLine(order+1, match, before, after))
	}
	return strings.Join(lines, "\n")
}

func agentMessageInlineReferenceText(messages []design.AgentChatMessage) string {
	var builder strings.Builder
	imageIndex := 0
	writeText := func(text string) {
		text = strings.TrimSpace(text)
		if text == "" {
			return
		}
		if builder.Len() > 0 {
			builder.WriteString(" ")
		}
		builder.WriteString(text)
	}
	for _, message := range messages {
		if strings.TrimSpace(message.Role) != "" && strings.ToLower(strings.TrimSpace(message.Role)) != "user" {
			continue
		}
		for _, content := range message.Contents {
			switch strings.ToLower(strings.TrimSpace(content.Type)) {
			case "text":
				if strings.EqualFold(strings.TrimSpace(content.TextSource), "settings") {
					continue
				}
				writeText(stripInlineImageGeneratorDirectiveLines(content.Text))
			case "image", "mask":
				imageIndex++
				name := compactInlineReferenceName(content.ImageName)
				if name == "" {
					name = "image"
				}
				url := sanitizeReferenceURL(content.ImageURL)
				if url == "" {
					url = fmt.Sprintf("attachment://reference-%d", imageIndex)
				}
				writeText(fmt.Sprintf("参考图 %d（%s）：%s", imageIndex, name, url))
			}
		}
	}
	return builder.String()
}

func collectInlineReferenceMatches(text string) []inlineReferenceMatch {
	if strings.TrimSpace(text) == "" {
		return nil
	}
	matches := make([]inlineReferenceMatch, 0)

	for _, indexes := range inlinePromptImageTokenPattern.FindAllStringSubmatchIndex(text, -1) {
		match := inlineReferenceMatch{start: indexes[0], end: indexes[1]}
		match.index = atoiSubmatch(text, indexes, 2)
		match.name = submatch(text, indexes, 4)
		match.url = sanitizeReferenceURL(firstNonEmptySubmatch(text, indexes, 8, 10))
		if match.url != "" {
			matches = append(matches, match)
		}
	}

	for _, indexes := range inlineReferenceDirectivePattern.FindAllStringSubmatchIndex(text, -1) {
		match := inlineReferenceMatch{start: indexes[0], end: indexes[1]}
		match.index = atoiSubmatch(text, indexes, 4)
		match.name = submatch(text, indexes, 6)
		match.url = sanitizeReferenceURL(submatch(text, indexes, 8))
		if match.url != "" {
			matches = append(matches, match)
		}
	}

	sort.SliceStable(matches, func(i, j int) bool {
		if matches[i].start == matches[j].start {
			return matches[i].end > matches[j].end
		}
		return matches[i].start < matches[j].start
	})

	deduped := make([]inlineReferenceMatch, 0, len(matches))
	for _, match := range matches {
		overlaps := false
		for _, existing := range deduped {
			if match.start < existing.end && match.end > existing.start {
				overlaps = true
				break
			}
		}
		if !overlaps {
			deduped = append(deduped, match)
		}
	}
	return deduped
}

func chineseReferencePlacementLine(order int, match inlineReferenceMatch, before string, after string) string {
	parts := []string{fmt.Sprintf("- %s", referenceDisplayName(order, match, false))}
	if before != "" {
		parts = append(parts, fmt.Sprintf("前文“%s”", before))
	}
	if after != "" {
		parts = append(parts, fmt.Sprintf("后文“%s”", after))
	}
	if len(parts) == 1 {
		parts = append(parts, "无相邻文字")
	}
	return strings.Join(parts, "：")
}

func englishReferencePlacementLine(order int, match inlineReferenceMatch, before string, after string) string {
	parts := []string{fmt.Sprintf("- %s", referenceDisplayName(order, match, true))}
	if before != "" {
		parts = append(parts, fmt.Sprintf("previous text %q", before))
	}
	if after != "" {
		parts = append(parts, fmt.Sprintf("next text %q", after))
	}
	if len(parts) == 1 {
		parts = append(parts, "no adjacent text")
	}
	return strings.Join(parts, ": ")
}

func referenceDisplayName(order int, match inlineReferenceMatch, english bool) string {
	name := compactInlineReferenceName(match.name)
	if name == "" {
		name = "image"
	}
	if english {
		label := fmt.Sprintf("Reference image %d", order)
		if match.index > 0 && match.index != order {
			label += fmt.Sprintf(" / user label %d", match.index)
		}
		return fmt.Sprintf("%s (%s)", label, name)
	}
	label := fmt.Sprintf("参考图 %d", order)
	if match.index > 0 && match.index != order {
		label += fmt.Sprintf(" / 用户标注 %d", match.index)
	}
	return fmt.Sprintf("%s（%s）", label, name)
}

func compactReferenceNeighbor(value string, keepTail bool) string {
	value = strings.ReplaceAll(value, "\u00a0", " ")
	value = imagePromptURLPattern.ReplaceAllString(value, "[参考图]")
	value = strings.Join(strings.Fields(strings.TrimSpace(value)), " ")
	value = strings.Trim(value, " \t\r\n，。,.；;、")
	if value == "" {
		return ""
	}
	const maxRunes = 72
	if utf8.RuneCountInString(value) <= maxRunes {
		return value
	}
	runes := []rune(value)
	if keepTail {
		return "..." + string(runes[len(runes)-maxRunes:])
	}
	return string(runes[:maxRunes]) + "..."
}

func sanitizeReferenceURL(value string) string {
	return strings.TrimSpace(strings.TrimRight(value, "，。,.；;、"))
}

func compactInlineReferenceName(value string) string {
	return strings.Join(strings.Fields(value), " ")
}

func stripInlineImageGeneratorDirectiveLines(text string) string {
	lines := strings.Split(text, "\n")
	filtered := make([]string, 0, len(lines))
	for _, line := range lines {
		if strings.HasPrefix(strings.ToLower(strings.TrimSpace(line)), "image generator target node:") {
			continue
		}
		filtered = append(filtered, line)
	}
	return strings.TrimSpace(strings.Join(filtered, "\n"))
}

func submatch(text string, indexes []int, offset int) string {
	if offset+1 >= len(indexes) || indexes[offset] < 0 || indexes[offset+1] < 0 {
		return ""
	}
	return text[indexes[offset]:indexes[offset+1]]
}

func firstNonEmptySubmatch(text string, indexes []int, offsets ...int) string {
	for _, offset := range offsets {
		if value := submatch(text, indexes, offset); strings.TrimSpace(value) != "" {
			return value
		}
	}
	return ""
}

func atoiSubmatch(text string, indexes []int, offset int) int {
	value := strings.TrimSpace(submatch(text, indexes, offset))
	if value == "" {
		return 0
	}
	number, err := strconv.Atoi(value)
	if err != nil {
		return 0
	}
	return number
}