thread-master/backend/internal/library/matrix/generate.go

package matrix

import (
	"context"
	"encoding/json"
	"fmt"
	"regexp"
	"strings"

	libprompt "haixun-backend/internal/library/prompt"
	"haixun-backend/internal/library/threadspost"
	domai "haixun-backend/internal/model/ai/domain/usecase"
)

type Row struct {
	SortOrder        int      `json:"sort_order"`
	SearchTag        string   `json:"search_tag"`
	Angle            string   `json:"angle"`
	Hook             string   `json:"hook"`
	Text             string   `json:"text"`
	ReferenceNotes   string   `json:"reference_notes"`
	SourcePermalinks []string `json:"source_permalinks"`
	Rationale        string   `json:"rationale"`
}

type GenerateResult struct {
	Rows []Row `json:"rows"`
}

type MaterialPost struct {
	SearchTag string
	Author    string
	Text      string
	Permalink string
	Priority  string
}

type GenerateInput struct {
	Persona       string
	TopicLabel    string
	AudienceBrief string
	ProductBrief  string
	Posts         []MaterialPost
	Count         int
}

var codeFenceRE = regexp.MustCompile("(?s)```(?:json)?\\s*(.*?)\\s*```")

func BuildUserPrompt(in GenerateInput) (string, error) {
	count := in.Count
	if count <= 0 {
		count = 5
	}
	personaBlock := ""
	if strings.TrimSpace(in.Persona) != "" {
		personaBlock = "人設與語氣：\n" + strings.TrimSpace(in.Persona) + "\n"
	}
	audience := strings.TrimSpace(in.AudienceBrief)
	if audience == "" {
		audience = "（未指定）"
	}
	product := strings.TrimSpace(in.ProductBrief)
	if product == "" {
		product = "（尚未填寫）"
	}
	topic := strings.TrimSpace(in.TopicLabel)
	if topic == "" {
		topic = "未指定"
	}
	return libprompt.MatrixPlacementUser(map[string]string{
		"persona_block":   personaBlock,
		"topic_label":     topic,
		"audience_line":   audience,
		"product_brief":   product,
		"post_count":      fmt.Sprintf("%d", len(in.Posts)),
		"materials_block": buildMaterialsBlock(in.Posts),
		"count":           fmt.Sprintf("%d", count),
	})
}

func buildMaterialsBlock(posts []MaterialPost) string {
	if len(posts) == 0 {
		return "（無素材）"
	}
	lines := make([]string, 0, len(posts))
	for i, post := range posts {
		lines = append(lines, fmt.Sprintf(
			"%d. [%s/%s] @%s\n%s\n連結：%s",
			i+1,
			strings.TrimSpace(post.Priority),
			strings.TrimSpace(post.SearchTag),
			strings.TrimSpace(post.Author),
			strings.TrimSpace(post.Text),
			strings.TrimSpace(post.Permalink),
		))
	}
	return strings.Join(lines, "\n\n")
}

type TextGenerator interface {
	GenerateText(ctx context.Context, req domai.GenerateRequest) (*domai.GenerateResult, error)
}

// CopyMatrixGenerateRequest tunes token budget for multi-row JSON matrix output.
func CopyMatrixGenerateRequest(base domai.GenerateRequest, count int) domai.GenerateRequest {
	if count <= 0 {
		count = 5
	}
	temp := 0.45
	tokens := 8192 + count*1024
	if tokens > 16384 {
		tokens = 16384
	}
	base.Temperature = &temp
	base.MaxTokens = &tokens
	return base
}

func MatrixRetryUserPrompt(count int) string {
	if count <= 0 {
		count = 5
	}
	return fmt.Sprintf(
		"上一則回覆的 JSON 不完整或無法解析。請只回傳單一 JSON 物件 {\"rows\":[...]}，共 %d 篇，不要用 markdown 程式碼區塊、不要加說明文字。\n"+
			"每筆含 sort_order、search_tag、angle、hook、text、reference_notes、source_permalinks、rationale。\n"+
			"reference_notes 與 rationale 各不超過 40 字；text 建議 80～220 字。",
		count,
	)
}

// GenerateCopyOutput calls the LLM and parses matrix JSON, with one compact-json retry.
func GenerateCopyOutput(
	ctx context.Context,
	ai TextGenerator,
	req domai.GenerateRequest,
	count int,
) (GenerateResult, error) {
	genReq := CopyMatrixGenerateRequest(req, count)
	result, err := ai.GenerateText(ctx, genReq)
	if err != nil {
		return GenerateResult{}, err
	}
	parsed, parseErr := ParseGenerateOutput(result.Text)
	if parseErr == nil {
		if count > 0 && len(parsed.Rows) < count {
			parseErr = fmt.Errorf("matrix truncated: got %d rows, want %d", len(parsed.Rows), count)
		} else {
			return parsed, nil
		}
	}
	retryReq := CopyMatrixGenerateRequest(domai.GenerateRequest{
		Provider:   req.Provider,
		Model:      req.Model,
		Credential: req.Credential,
		System:     req.System,
		Messages: append(
			append([]domai.Message{}, req.Messages...),
			domai.Message{Role: "assistant", Content: result.Text},
			domai.Message{Role: "user", Content: MatrixRetryUserPrompt(count)},
		),
	}, count)
	retryResult, retryErr := ai.GenerateText(ctx, retryReq)
	if retryErr != nil {
		return GenerateResult{}, fmt.Errorf("matrix retry failed: %w (first parse: %v)", retryErr, parseErr)
	}
	retryParsed, retryParseErr := ParseGenerateOutput(retryResult.Text)
	if retryParseErr != nil {
		return GenerateResult{}, retryParseErr
	}
	if count > 0 && len(retryParsed.Rows) < count {
		return GenerateResult{}, fmt.Errorf("matrix truncated after retry: got %d rows, want %d", len(retryParsed.Rows), count)
	}
	return retryParsed, nil
}

func ParseGenerateOutput(raw string) (GenerateResult, error) {
	payload, err := extractJSONObject(raw)
	if err != nil {
		return GenerateResult{}, err
	}
	out, err := decodeGenerateResult(payload)
	if err == nil {
		return normalizeGenerateResult(out)
	}
	repaired, repairErr := repairTruncatedJSONObject(payload)
	if repairErr != nil {
		return GenerateResult{}, err
	}
	out, err = decodeGenerateResult(repaired)
	if err != nil {
		return GenerateResult{}, fmt.Errorf("parse matrix json: %w", err)
	}
	return normalizeGenerateResult(out)
}

func decodeGenerateResult(payload []byte) (GenerateResult, error) {
	var out GenerateResult
	if err := json.Unmarshal(payload, &out); err != nil {
		return GenerateResult{}, err
	}
	return out, nil
}

func normalizeGenerateResult(out GenerateResult) (GenerateResult, error) {
	if len(out.Rows) == 0 {
		return GenerateResult{}, fmt.Errorf("matrix rows missing")
	}
	for i := range out.Rows {
		out.Rows[i].Text = trimText(out.Rows[i].Text)
		if out.Rows[i].Text == "" {
			return GenerateResult{}, fmt.Errorf("matrix row %d empty", i+1)
		}
		if out.Rows[i].SortOrder <= 0 {
			out.Rows[i].SortOrder = i + 1
		}
	}
	return out, nil
}

func trimText(text string) string {
	text = strings.TrimSpace(text)
	runes := []rune(text)
	if len(runes) > threadspost.MaxPublishRunes {
		return string(runes[:threadspost.MaxPublishRunes])
	}
	return text
}

func extractJSONObject(raw string) ([]byte, error) {
	raw = strings.TrimSpace(raw)
	if m := codeFenceRE.FindStringSubmatch(raw); len(m) == 2 {
		raw = strings.TrimSpace(m[1])
	}
	start := strings.Index(raw, "{")
	if start < 0 {
		return nil, fmt.Errorf("matrix output missing json object")
	}
	slice := raw[start:]
	end := strings.LastIndex(slice, "}")
	if end <= 0 {
		return []byte(slice), nil
	}
	return []byte(slice[:end+1]), nil
}

func repairTruncatedJSONObject(payload []byte) ([]byte, error) {
	if len(payload) == 0 || payload[0] != '{' {
		return nil, fmt.Errorf("matrix output missing json object")
	}
	stack := make([]byte, 0, 8)
	inString := false
	escaped := false
	for _, b := range payload {
		if inString {
			if escaped {
				escaped = false
				continue
			}
			if b == '\\' {
				escaped = true
				continue
			}
			if b == '"' {
				inString = false
			}
			continue
		}
		switch b {
		case '"':
			inString = true
		case '{':
			stack = append(stack, '}')
		case '[':
			stack = append(stack, ']')
		case '}', ']':
			if len(stack) > 0 && stack[len(stack)-1] == b {
				stack = stack[:len(stack)-1]
			}
		}
	}
	repaired := append([]byte(nil), payload...)
	if inString {
		repaired = append(repaired, '"')
	}
	for i := len(stack) - 1; i >= 0; i-- {
		repaired = append(repaired, stack[i])
	}
	return repaired, nil
}