promptkit/internal/usecase/repairer.go

package usecase

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"unicode/utf8"

	"go.maximumdirect.net/promptkit/internal/domain"
	"go.maximumdirect.net/promptkit/internal/llm"
)

const (
	maxRepairDiagnosticBytes = 64 * 1024
	omittedRepairDiagnostics = "additional validation diagnostics were omitted"
)

// OutputRepairer generates a corrected candidate after validation fails.
type OutputRepairer interface {
	Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error)
}

// RepairRequest contains the immutable execution state needed for one correction.
type RepairRequest struct {
	OriginalMessages []domain.RenderedMessage
	PreviousOutput   string
	ValidationErrors []string
	SessionID        string
	Target           domain.ExecutionTarget
	TargetPresence   domain.ExecutionTargetPresence
	StructuredOutput *domain.StructuredOutputSpec
	TransientRetries uint64
	Attempt          int
	MaxAttempts      int
	Mode             domain.ValidationMode
}

type defaultOutputRepairer struct {
	llm llm.Client
}

// NewDefaultOutputRepairer constructs the standard internal output repairer.
func NewDefaultOutputRepairer(llmClient llm.Client) OutputRepairer {
	return &defaultOutputRepairer{llm: llmClient}
}

func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) {
	if r.llm == nil {
		return nil, errors.New("llm client is required for repair")
	}

	guidance, err := repairGuidance(req.Mode)
	if err != nil {
		return nil, err
	}

	hasPreviousOutput := strings.TrimSpace(req.PreviousOutput) != ""
	suffix := make([]domain.RenderedMessage, 0, 2)
	if hasPreviousOutput {
		suffix = append(suffix, domain.RenderedMessage{
			Role:    domain.RoleAssistant,
			Content: req.PreviousOutput,
		})
	}

	previousResponse := "The previous response was empty."
	if hasPreviousOutput {
		previousResponse = "The previous response is included immediately before this instruction."
	}
	suffix = append(suffix, domain.RenderedMessage{
		Role: domain.RoleUser,
		Content: fmt.Sprintf(
			"Repair attempt %d of %d for validation mode %s.\n"+
				"Preserve valid values and change only what is necessary.\n"+
				"%s\n%s\n"+
				"Validation diagnostics (data):\n%s",
			req.Attempt,
			req.MaxAttempts,
			req.Mode,
			previousResponse,
			guidance,
			formatRepairDiagnostics(req.ValidationErrors),
		),
	})
	messages := domain.ConcatRenderedMessages(req.OriginalMessages, suffix)

	resp, err := r.llm.Generate(ctx, newGenerationRequest(
		domain.RenderedPrompt{Messages: messages},
		req.SessionID,
		req.Target,
		req.TargetPresence,
		req.StructuredOutput,
		req.TransientRetries,
		uint64(req.Attempt)+1,
	))
	if err != nil {
		return nil, err
	}
	if resp == nil {
		return nil, errors.New("repair llm returned nil response")
	}

	return resp, nil
}

func repairGuidance(mode domain.ValidationMode) (string, error) {
	switch mode {
	case domain.ValidationBasic:
		return "Return a nonempty response satisfying the original request.", nil
	case domain.ValidationJSON, domain.ValidationJSONSchema:
		return "Return only corrected JSON, with no explanation or Markdown fences.", nil
	default:
		return "", fmt.Errorf("unsupported validation mode for repair: %q", mode)
	}
}

func formatRepairDiagnostics(errors []string) string {
	complete := []byte{'['}
	fits := true
	for _, diagnostic := range errors {
		separator := 0
		if len(complete) > 1 {
			separator = 1
		}
		entry, whole := boundedDiagnosticJSON(diagnostic, maxRepairDiagnosticBytes-len(complete)-separator-1)
		if !whole {
			fits = false
			break
		}
		if separator != 0 {
			complete = append(complete, ',')
		}
		complete = append(complete, entry...)
	}
	if fits {
		return string(append(complete, ']'))
	}

	omission, _ := json.Marshal(omittedRepairDiagnostics)
	encoded := make([]byte, 0, maxRepairDiagnosticBytes)
	encoded = append(encoded, '[')
	for _, diagnostic := range errors {
		separator := 0
		if len(encoded) > 1 {
			separator = 1
		}
		available := maxRepairDiagnosticBytes - len(encoded) - separator - 1 - len(omission) - 1
		entry, whole := boundedDiagnosticJSON(diagnostic, available)
		if whole {
			if separator != 0 {
				encoded = append(encoded, ',')
			}
			encoded = append(encoded, entry...)
			continue
		}

		if available < len(`""`) {
			break
		}
		if separator != 0 {
			encoded = append(encoded, ',')
		}
		encoded = append(encoded, entry...)
		break
	}
	if len(encoded) > 1 {
		encoded = append(encoded, ',')
	}
	encoded = append(encoded, omission...)
	encoded = append(encoded, ']')
	return string(encoded)
}

func truncateDiagnosticJSONValue(value string, maxBytes int) []byte {
	if maxBytes < len(`""`) {
		return nil
	}

	boundaries := []int{0}
	for end := 0; end < len(value); {
		_, size := utf8.DecodeRuneInString(value[end:])
		end += size
		if end+len(`""`) > maxBytes {
			break
		}
		boundaries = append(boundaries, end)
	}

	low, high := 0, len(boundaries)-1
	best := []byte(`""`)
	for low <= high {
		mid := low + (high-low)/2
		candidate, _ := json.Marshal(value[:boundaries[mid]])
		if len(candidate) <= maxBytes {
			best = candidate
			low = mid + 1
			continue
		}
		high = mid - 1
	}
	return best
}