geo/server/internal/tenant/app/kol_assist_runtime.go

package app

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

	"github.com/google/uuid"

	"github.com/geo-platform/tenant-api/internal/shared/auth"
	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/shared/response"
)

const kolAssistSystemPrompt = "You are a prompt designer for marketing content. Respond ONLY with JSON. Output the `content` field first. Do not wrap in markdown fences."

var kolAssistResponseSchema = []byte(`{
  "type": "object",
  "additionalProperties": false,
  "required": ["content"],
  "properties": {
    "content": {
      "type": "string"
    }
  }
}`)

func (s *KolAssistService) Run(
	ctx context.Context,
	actor auth.Actor,
	req AssistRequest,
	onDelta func(string),
) (*KolAssistResult, string, error) {
	if _, err := s.profileSvc.RequireActiveKol(ctx, actor); err != nil {
		return nil, "", err
	}
	if err := s.llm.Validate(); err != nil {
		return nil, "", response.ErrServiceUnavailable(50341, "llm_unavailable", err.Error())
	}

	req, _, err := normalizeAssistRequest(req)
	if err != nil {
		return nil, "", err
	}
	if req.Mode == kolAssistModeGenerate {
		req.Description = s.enrichGenerateDescription(ctx, req.Description)
	}

	reservation, err := s.reserveStreamingAIPoints(ctx, actor, req)
	if err != nil {
		return nil, "", err
	}

	generateReq, err := BuildKolAssistGenerateRequest(req)
	if err != nil {
		s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
		return nil, "", response.ErrInternal(50087, "kol_assist_prompt_build_failed", err.Error())
	}

	streamAccumulator := newKolAssistContentStream(onDelta)
	generated, err := s.llm.Generate(ctx, generateReq, streamAccumulator.ConsumeRawDelta)
	if err != nil {
		s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
		return nil, "", err
	}

	result, err := ParseKolAssistResponse(generated.Content)
	if err != nil {
		s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
		return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error())
	}
	result, err = FinalizeKolAssistResult(req, result)
	if err != nil {
		s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
		return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error())
	}

	completeCtx, cancel := newGenerationCleanupContext()
	if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, generated.Model); err != nil {
		cancel()
		return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", "failed to confirm ai points usage")
	}
	cancel()

	return result, generated.Model, nil
}

func BuildKolAssistGenerateRequest(req AssistRequest) (llm.GenerateRequest, error) {
	userPrompt, err := BuildKolAssistUserPrompt(req)
	if err != nil {
		return llm.GenerateRequest{}, err
	}

	return llm.GenerateRequest{
		Prompt: fmt.Sprintf("SYSTEM:\n%s\n\nUSER:\n%s", kolAssistSystemPrompt, userPrompt),
		ResponseFormat: &llm.ResponseFormat{
			Type:        llm.ResponseFormatTypeJSONSchema,
			Name:        "kol_assist_response",
			Description: "JSON object with content:string",
			SchemaJSON:  kolAssistResponseSchema,
			Strict:      true,
		},
	}, nil
}

func BuildKolAssistUserPrompt(req AssistRequest) (string, error) {
	switch strings.ToLower(strings.TrimSpace(req.Mode)) {
	case kolAssistModeGenerate:
		return strings.TrimSpace(fmt.Sprintf(`
你的任务是把“用户需求 + 参考材料”转换成“最终可直接给另一个 AI 使用的提示词模板”。

严格要求：
1. 只输出最终提示词模板，不要输出分析、总结、复盘、学习笔记、解释说明。
2. 如果提供了参考文章或参考网页，只学习其选题切入、标题策略、结构框架、叙事节奏、风格语气、互动与转化方式，然后抽象成可复用规则。
3. 任何会变化的信息都必须参数化成 {{变量名}} 占位符，不要把具体的人名、品牌、平台、标题、数字、时间、地区、产品、案例、链接、经历写死。
4. 若用户基于具体文章/网页生成模板，结果中至少包含 3 个 {{}} 占位符。
5. 结果必须是“生成同类内容”的模板，不是对参考文章的分析。
6. 模板至少覆盖：创作目标、选题要求、标题要求、结构要求、风格要求、互动/转化要求、限制项。
7. 输出中文，直接给出最终模板正文。

用户需求与参考材料：
%s
`, req.Description)), nil
	case kolAssistModeOptimize:
		schemaJSON := "null"
		if req.Schema != nil {
			payload, err := json.Marshal(req.Schema)
			if err != nil {
				return "", fmt.Errorf("encode schema: %w", err)
			}
			schemaJSON = string(payload)
		}
		instruction := strings.TrimSpace(req.Description)
		if instruction == "" {
			instruction = "Make the prompt clearer, more complete, and easier for another AI to follow."
		}
		return fmt.Sprintf(
			"Optimize the following prompt according to the instruction. Keep variable placeholders intact.\n\nInstruction:\n%s\n\nCurrent schema: %s\n\nPROMPT:\n%s",
			instruction,
			schemaJSON,
			req.CurrentContent,
		), nil
	default:
		return "", fmt.Errorf("unsupported assist mode %q", req.Mode)
	}
}

func ParseKolAssistResponse(content string) (*KolAssistResult, error) {
	type rawPayload struct {
		Content string          `json:"content"`
		Schema  json.RawMessage `json:"schema"`
	}

	var lastErr error
	for _, candidate := range extractJSONCandidates(content) {
		var payload rawPayload
		if err := json.Unmarshal([]byte(candidate), &payload); err != nil {
			lastErr = err
			continue
		}

		result := &KolAssistResult{
			Content: strings.TrimSpace(payload.Content),
			Schema:  KolSchemaJSON{Variables: []KolVariableDefinition{}},
		}
		if result.Content == "" {
			lastErr = fmt.Errorf("kol assist result content is empty")
			continue
		}

		schema, err := decodeKolAssistSchema(payload.Schema)
		if err != nil {
			lastErr = err
			continue
		}
		result.Schema = schema
		return result, nil
	}

	if lastErr == nil {
		lastErr = fmt.Errorf("empty content")
	}
	return nil, fmt.Errorf("decode kol assist result: %w", lastErr)
}

func FinalizeKolAssistResult(req AssistRequest, result *KolAssistResult) (*KolAssistResult, error) {
	if result == nil {
		return nil, fmt.Errorf("kol assist result is nil")
	}

	baseSchema := result.Schema
	if req.Mode == kolAssistModeOptimize && req.Schema != nil {
		merged := make([]KolVariableDefinition, 0, len(req.Schema.Variables)+len(result.Schema.Variables))
		merged = append(merged, req.Schema.Variables...)
		merged = append(merged, result.Schema.Variables...)
		baseSchema = KolSchemaJSON{Variables: merged}
	}

	result.Schema = SyncKolSchemaWithContent(result.Content, baseSchema)
	if err := ValidateSchema(result.Schema); err != nil {
		return nil, fmt.Errorf("invalid kol assist schema: %w", err)
	}

	return result, nil
}

func SyncKolSchemaWithContent(content string, schema KolSchemaJSON) KolSchemaJSON {
	keys := extractKolPlaceholderKeys(content)
	if len(keys) == 0 {
		return KolSchemaJSON{Variables: []KolVariableDefinition{}}
	}

	schema = normalizeKolSchema(schema)

	byKey := make(map[string][]KolVariableDefinition, len(schema.Variables))
	for _, variable := range schema.Variables {
		key := strings.TrimSpace(variable.Key)
		if key == "" {
			key = strings.TrimSpace(variable.Label)
		}
		if key == "" {
			continue
		}
		byKey[key] = append(byKey[key], variable)
	}

	usedIDs := make(map[string]struct{}, len(keys))
	variables := make([]KolVariableDefinition, 0, len(keys))

	for _, key := range keys {
		var variable KolVariableDefinition
		if candidates := byKey[key]; len(candidates) > 0 {
			variable = candidates[0]
			byKey[key] = candidates[1:]
		} else {
			variable = KolVariableDefinition{
				Key:      key,
				Label:    key,
				Type:     KolVariableTypeInput,
				Required: true,
			}
		}

		if strings.TrimSpace(variable.Key) == "" {
			variable.Key = key
		}
		if strings.TrimSpace(variable.Label) == "" {
			variable.Label = key
		}

		variable = normalizeKolVariableDefinition(variable)
		id := strings.TrimSpace(variable.ID)
		if id == "" || !strings.HasPrefix(id, "var_") {
			id = newKolVariableID()
		}
		if _, exists := usedIDs[id]; exists {
			id = newKolVariableID()
		}
		variable.ID = id
		usedIDs[id] = struct{}{}
		variables = append(variables, variable)
	}

	return normalizeKolSchema(KolSchemaJSON{Variables: variables})
}

func extractKolPlaceholderKeys(content string) []string {
	text := strings.TrimSpace(content)
	if text == "" {
		return []string{}
	}

	seen := make(map[string]struct{})
	keys := make([]string, 0)
	matches := kolPlaceholderRE.FindAllStringSubmatch(text, -1)
	for _, match := range matches {
		if len(match) < 2 {
			continue
		}
		key := strings.TrimSpace(match[1])
		if key == "" {
			continue
		}
		if _, exists := seen[key]; exists {
			continue
		}
		seen[key] = struct{}{}
		keys = append(keys, key)
	}

	return keys
}

func newKolVariableID() string {
	return "var_" + strings.ReplaceAll(uuid.NewString(), "-", "")[:8]
}

type kolAssistContentStream struct {
	raw     strings.Builder
	content string
	onDelta func(string)
}

func newKolAssistContentStream(onDelta func(string)) *kolAssistContentStream {
	return &kolAssistContentStream{onDelta: onDelta}
}

func (s *kolAssistContentStream) ConsumeRawDelta(chunk string) {
	if s == nil || s.onDelta == nil || chunk == "" {
		return
	}

	s.raw.WriteString(chunk)
	current := extractKolAssistContentFromPartialJSON(s.raw.String())
	if current == "" || current == s.content {
		return
	}
	if strings.HasPrefix(current, s.content) {
		delta := current[len(s.content):]
		s.content = current
		if delta != "" {
			s.onDelta(delta)
		}
	}
}

func extractKolAssistContentFromPartialJSON(raw string) string {
	keyIndex := strings.Index(raw, `"content"`)
	if keyIndex < 0 {
		return ""
	}

	index := keyIndex + len(`"content"`)
	for index < len(raw) && isKolAssistJSONWhitespace(raw[index]) {
		index++
	}
	if index >= len(raw) || raw[index] != ':' {
		return ""
	}
	index++
	for index < len(raw) && isKolAssistJSONWhitespace(raw[index]) {
		index++
	}
	if index >= len(raw) || raw[index] != '"' {
		return ""
	}
	index++

	var builder strings.Builder
	for index < len(raw) {
		ch := raw[index]
		if ch == '"' {
			return builder.String()
		}
		if ch != '\\' {
			builder.WriteByte(ch)
			index++
			continue
		}

		if index+1 >= len(raw) {
			return builder.String()
		}
		escaped := raw[index+1]
		switch escaped {
		case '"', '\\', '/':
			builder.WriteByte(escaped)
			index += 2
		case 'b':
			builder.WriteByte('\b')
			index += 2
		case 'f':
			builder.WriteByte('\f')
			index += 2
		case 'n':
			builder.WriteByte('\n')
			index += 2
		case 'r':
			builder.WriteByte('\r')
			index += 2
		case 't':
			builder.WriteByte('\t')
			index += 2
		case 'u':
			if index+6 > len(raw) {
				return builder.String()
			}
			var codePoint rune
			for _, hexChar := range raw[index+2 : index+6] {
				codePoint <<= 4
				switch {
				case hexChar >= '0' && hexChar <= '9':
					codePoint += rune(hexChar - '0')
				case hexChar >= 'a' && hexChar <= 'f':
					codePoint += rune(hexChar-'a') + 10
				case hexChar >= 'A' && hexChar <= 'F':
					codePoint += rune(hexChar-'A') + 10
				default:
					return builder.String()
				}
			}
			builder.WriteRune(codePoint)
			index += 6
		default:
			return builder.String()
		}
	}

	return builder.String()
}

func isKolAssistJSONWhitespace(value byte) bool {
	return value == ' ' || value == '\n' || value == '\r' || value == '\t'
}

func decodeKolAssistSchema(raw json.RawMessage) (KolSchemaJSON, error) {
	cleaned := strings.TrimSpace(string(raw))
	if cleaned == "" || cleaned == "null" {
		return KolSchemaJSON{Variables: []KolVariableDefinition{}}, nil
	}

	var payload struct {
		Variables []json.RawMessage `json:"variables"`
	}
	if err := json.Unmarshal([]byte(cleaned), &payload); err != nil {
		return KolSchemaJSON{}, fmt.Errorf("decode kol assist schema: %w", err)
	}

	variables := make([]KolVariableDefinition, 0, len(payload.Variables))
	for _, item := range payload.Variables {
		if len(item) == 0 {
			continue
		}

		var variable KolVariableDefinition
		if err := json.Unmarshal(item, &variable); err == nil {
			variables = append(variables, variable)
			continue
		}

		var key string
		if err := json.Unmarshal(item, &key); err == nil {
			key = strings.TrimSpace(key)
			if key == "" {
				continue
			}
			variables = append(variables, KolVariableDefinition{
				Key:      key,
				Label:    key,
				Type:     KolVariableTypeInput,
				Required: true,
			})
			continue
		}

		return KolSchemaJSON{}, fmt.Errorf("decode kol assist schema variable: unsupported item %s", string(item))
	}

	return normalizeKolSchema(KolSchemaJSON{Variables: variables}), nil
}