descrybe/apps/api/internal/processing/llm_json.go

package processing

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

// Local / weak-model defaults (8k-class context). See docs/local-llm-tuning.md.
const (
	DefaultStructuredTemp = 0.2
	MaxTokensEnhance      = 350
	MaxTokensSEO          = 180
	MaxTokensCampaign     = 650
	MaxProductDescRunes   = 400
	MaxAttrKeys           = 10
	MaxAttrValueRunes     = 60
	MaxBrandInjectRunes   = 500
	MaxCampaignProducts   = 8
	MaxCampaignNameRunes  = 80
)

// CompleteOptions tunes a single chat completion for structured tasks.
type CompleteOptions struct {
	MaxTokens   int
	Temperature float64 // 0 → client default (≤0.3 for structured)
}

// CompleterWithOptions is optional; OpenAIClient implements it.
type CompleterWithOptions interface {
	CompleteWithOptions(ctx context.Context, system, user string, opts CompleteOptions) (Completion, error)
}

// CompleteOnce calls CompleterWithOptions when available, else Complete.
func CompleteOnce(ctx context.Context, c Completer, system, user string, opts CompleteOptions) (Completion, error) {
	if c == nil {
		return Completion{}, fmt.Errorf("completer not configured")
	}
	if co, ok := c.(CompleterWithOptions); ok {
		return co.CompleteWithOptions(ctx, system, user, opts)
	}
	return c.Complete(ctx, system, user)
}

// StripJSONFences removes markdown code fences and isolates the outermost JSON object/array.
func StripJSONFences(text string) string {
	text = strings.TrimSpace(text)
	if text == "" {
		return ""
	}
	text = strings.TrimPrefix(text, "```json")
	text = strings.TrimPrefix(text, "```JSON")
	text = strings.TrimPrefix(text, "```")
	text = strings.TrimSuffix(text, "```")
	text = strings.TrimSpace(text)
	objAt := strings.Index(text, "{")
	arrAt := strings.Index(text, "[")
	// Prefer whichever structure appears first so array-of-objects is not sliced mid-stream.
	if arrAt >= 0 && (objAt < 0 || arrAt < objAt) {
		if j := strings.LastIndex(text, "]"); j > arrAt {
			return strings.TrimSpace(text[arrAt : j+1])
		}
	}
	if objAt >= 0 {
		if j := strings.LastIndex(text, "}"); j > objAt {
			return strings.TrimSpace(text[objAt : j+1])
		}
	}
	return text
}

// ParseJSONObject parses a model reply into a JSON object (fence-tolerant).
// Some local/weak models wrap the payload in a one-element array; accept that
// by promoting the first object element (preferring name/description keys).
func ParseJSONObject(text string) (map[string]any, error) {
	text = StripJSONFences(text)
	if text == "" {
		return nil, fmt.Errorf("empty json")
	}
	var obj map[string]any
	objErr := json.Unmarshal([]byte(text), &obj)
	if objErr == nil {
		return obj, nil
	}
	var arr []any
	if err := json.Unmarshal([]byte(text), &arr); err != nil {
		return nil, objErr
	}
	if len(arr) == 0 {
		return nil, fmt.Errorf("empty json array")
	}
	var fallback map[string]any
	for _, el := range arr {
		m, ok := el.(map[string]any)
		if !ok || m == nil {
			continue
		}
		if fallback == nil {
			fallback = m
		}
		if _, hasName := m["name"]; hasName {
			return m, nil
		}
		if _, hasDesc := m["description"]; hasDesc {
			return m, nil
		}
	}
	if fallback != nil {
		return fallback, nil
	}
	return nil, fmt.Errorf("json array has no object elements")
}

// CompleteJSON runs a structured completion and retries once if JSON parse fails.
func CompleteJSON(ctx context.Context, c Completer, system, user string, opts CompleteOptions) (Completion, map[string]any, error) {
	comp, err := CompleteOnce(ctx, c, system, user, opts)
	if err != nil {
		return Completion{}, nil, err
	}
	obj, err := ParseJSONObject(comp.Text)
	if err == nil {
		return comp, obj, nil
	}
	retryUser := user + "\n\nINVALID. Reply with ONLY one JSON object. No markdown, no prose."
	comp2, err2 := CompleteOnce(ctx, c, system, retryUser, opts)
	if err2 != nil {
		return comp, nil, err2
	}
	obj2, err3 := ParseJSONObject(comp2.Text)
	if err3 != nil {
		comp2.PromptTokens += comp.PromptTokens
		comp2.OutputTokens += comp.OutputTokens
		comp2.TotalTokens += comp.TotalTokens
		return comp2, nil, err3
	}
	comp2.PromptTokens += comp.PromptTokens
	comp2.OutputTokens += comp.OutputTokens
	comp2.TotalTokens += comp.TotalTokens
	return comp2, obj2, nil
}

// CompactAttrs keeps title-relevant key attributes only (sorted keys, capped).
func CompactAttrs(attrs map[string]any, maxKeys int) map[string]any {
	if len(attrs) == 0 {
		return map[string]any{}
	}
	if maxKeys <= 0 {
		maxKeys = MaxAttrKeys
	}
	keys := make([]string, 0, len(attrs))
	for k := range attrs {
		k = strings.TrimSpace(k)
		if k == "" {
			continue
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	// Prefer common retail keys first.
	priority := []string{"brand", "Brand", "color", "Color", "material", "Material", "size", "Size", "model", "Model", "gtin", "GTIN", "ean", "EAN"}
	ordered := make([]string, 0, len(keys))
	seen := map[string]bool{}
	for _, p := range priority {
		for _, k := range keys {
			if strings.EqualFold(k, p) && !seen[k] {
				ordered = append(ordered, k)
				seen[k] = true
			}
		}
	}
	for _, k := range keys {
		if !seen[k] {
			ordered = append(ordered, k)
		}
	}
	if len(ordered) > maxKeys {
		ordered = ordered[:maxKeys]
	}
	out := make(map[string]any, len(ordered))
	for _, k := range ordered {
		v := stringFromAny(attrs[k])
		if v == "" {
			continue
		}
		out[k] = truncateRunes(v, MaxAttrValueRunes)
	}
	return out
}

// CompactBrandPrompt caps brand-kit injection for small context windows.
func CompactBrandPrompt(block string) string {
	block = strings.TrimSpace(block)
	if block == "" {
		return ""
	}
	return truncateRunes(SanitizeText(block), MaxBrandInjectRunes)
}

// ProductEnhanceUser builds a short user prompt for title/description enhance.
func ProductEnhanceUser(category, name, description string, attrs map[string]any) string {
	var b strings.Builder
	b.WriteString("Category: ")
	b.WriteString(SanitizeText(category))
	b.WriteString("\nName: ")
	b.WriteString(SanitizeText(truncateRunes(name, 200)))
	b.WriteString("\nDesc: ")
	b.WriteString(SanitizeText(truncateRunes(description, MaxProductDescRunes)))
	compact := CompactAttrs(attrs, MaxAttrKeys)
	if len(compact) > 0 {
		b.WriteString("\nAttrs: ")
		b.WriteString(sanitizeJSON(compact))
	}
	return b.String()
}