audita/internal/framework/proposal_generation/generate.go

// Package proposal_generation provides shared, deterministic LLM-backed
// proposal-generation helpers. It only produces candidate proposals; validation
// and application remain runner responsibilities.
package proposal_generation

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"

	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
)

// InteractionDiagnosticsWriter writes machine-readable prompt/response artifacts.
type InteractionDiagnosticsWriter interface {
	WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error)
}

// InteractionArtifacts contains written diagnostics paths.
type InteractionArtifacts struct {
	RequestMetadataPath string `json:"request_metadata_path,omitempty"`
	RequestPayloadPath  string `json:"request_payload_path,omitempty"`
	ResponsePayloadPath string `json:"response_payload_path,omitempty"`
	ErrorPayloadPath    string `json:"error_payload_path,omitempty"`
}

// StructuredCorrectionProposal is one LLM response correction payload.
type StructuredCorrectionProposal struct {
	TargetSegmentID int     `json:"id"`
	OriginalText    string  `json:"original_text"`
	CorrectedText   string  `json:"corrected_text"`
	Confidence      float64 `json:"confidence"`
}

// StructuredCorrectionSet is the reusable structured LLM response model for
// candidate correction proposals.
type StructuredCorrectionSet struct {
	Corrections []StructuredCorrectionProposal `json:"corrections"`
}

// Request captures reusable proposal-generation inputs for future modules.
type Request struct {
	ModuleKey         string                      `json:"module_key"`
	ModuleInstance    string                      `json:"module_instance"`
	ReplacementPolicy proposals.ReplacementPolicy `json:"replacement_policy"`
	WorkingTranscript *schema.Transcript          `json:"-"`
	Section           *contracts.SectionMetadata  `json:"section,omitempty"`
	Glossary          *schema.Glossary            `json:"-"`
	Config            *config.Config              `json:"-"`
	Messages          []contracts.LLMMessage      `json:"messages"`
	Model             string                      `json:"model,omitempty"`
	StageName         string                      `json:"stage_name,omitempty"`
	StartIndex        int                         `json:"start_index"`
	LLMClient         contracts.StructuredLLMClient
	Scheduler         contracts.LLMScheduler
	DiagnosticsDir    string
	DiagnosticsWriter InteractionDiagnosticsWriter
}

// Result contains generated candidate proposals and optional diagnostics paths.
type Result struct {
	Corrections []proposals.CorrectionProposal         `json:"corrections"`
	Enriched    []proposals.EnrichedCorrectionProposal `json:"enriched"`
	Artifacts   InteractionArtifacts                   `json:"artifacts,omitempty"`
}

// GenerateCandidates executes one structured LLM call and deterministically maps
// its correction-set response into framework proposal types.
func GenerateCandidates(ctx context.Context, req Request) (Result, error) {
	if strings.TrimSpace(req.ModuleKey) == "" {
		return Result{}, fmt.Errorf("module key must not be empty")
	}
	if strings.TrimSpace(req.ModuleInstance) == "" {
		return Result{}, fmt.Errorf("module instance must not be empty")
	}
	if req.StartIndex < 0 {
		return Result{}, fmt.Errorf("start index must be non-negative")
	}
	if req.LLMClient == nil {
		return Result{}, fmt.Errorf("structured LLM client is required")
	}
	if len(req.Messages) == 0 {
		return Result{}, fmt.Errorf("messages must not be empty")
	}

	stage := strings.TrimSpace(req.StageName)
	if stage == "" {
		stage = buildStageName(req.ModuleInstance, req.Section)
	}
	model := resolveModel(req.Config, req.Model)
	messages := append([]contracts.LLMMessage(nil), req.Messages...)

	var writer InteractionDiagnosticsWriter
	if req.DiagnosticsWriter != nil {
		writer = req.DiagnosticsWriter
	} else if strings.TrimSpace(req.DiagnosticsDir) != "" {
		writer = diagnosticsWriterAdapter{
			writer: llm.NewDiagnosticsWriter(
				filepath.Join(req.DiagnosticsDir, req.ModuleInstance),
				proposalGenerationSecrets(req.Config),
			),
		}
	}

	var (
		response  StructuredCorrectionSet
		callErr   error
		artifacts InteractionArtifacts
	)
	call := func(callCtx context.Context) error {
		_, callErr = req.LLMClient.CompleteStructured(callCtx, contracts.StructuredCompletionRequest{
			StageName: stage,
			Messages:  messages,
			Model:     model,
		}, &response)
		return callErr
	}
	if req.Scheduler != nil {
		callErr = req.Scheduler.Run(ctx, call)
	} else {
		callErr = call(ctx)
	}

	if writer != nil {
		artifacts, _ = writer.WriteInteraction(
			stage,
			map[string]any{
				"module_key":         req.ModuleKey,
				"module_instance":    req.ModuleInstance,
				"replacement_policy": req.ReplacementPolicy,
				"section":            req.Section,
				"start_index":        req.StartIndex,
				"model":              model,
			},
			map[string]any{
				"messages": messages,
			},
			response,
			errPayload(callErr),
		)
	}

	if callErr != nil {
		return Result{}, fmt.Errorf("proposal generation completion failed: %w", callErr)
	}

	corrections := make([]proposals.CorrectionProposal, 0, len(response.Corrections))
	enriched := make([]proposals.EnrichedCorrectionProposal, 0, len(response.Corrections))
	for i, raw := range response.Corrections {
		candidate := proposals.CorrectionProposal{
			TargetSegmentID: raw.TargetSegmentID,
			OriginalText:    raw.OriginalText,
			CorrectedText:   raw.CorrectedText,
			Confidence:      raw.Confidence,
		}
		if err := candidate.Validate(); err != nil {
			return Result{}, fmt.Errorf("invalid structured correction at index %d: %w", i, err)
		}

		corrections = append(corrections, candidate)
		enrichedCandidate := proposals.EnrichedCorrectionProposal{
			CorrectionProposal: candidate,
			ProposalMetadata: proposals.ProposalMetadata{
				ProposalIndex:  req.StartIndex + i,
				ModuleKey:      req.ModuleKey,
				ModuleInstance: req.ModuleInstance,
			},
		}
		if req.Section != nil {
			sectionIndex := req.Section.Index
			enrichedCandidate.SectionIndex = &sectionIndex
		}
		enriched = append(enriched, enrichedCandidate)
	}

	return Result{
		Corrections: corrections,
		Enriched:    enriched,
		Artifacts:   artifacts,
	}, nil
}

func buildStageName(moduleInstance string, section *contracts.SectionMetadata) string {
	base := fmt.Sprintf("%s:proposal-generation", moduleInstance)
	if section == nil {
		return base
	}
	return fmt.Sprintf("%s:section-%04d", base, section.Index)
}

func resolveModel(cfg *config.Config, override string) string {
	if strings.TrimSpace(override) != "" {
		return strings.TrimSpace(override)
	}
	if cfg == nil {
		return ""
	}
	return llm.ResolvePrimaryConfig(*cfg).Model
}

func proposalGenerationSecrets(cfg *config.Config) []string {
	if cfg == nil {
		return nil
	}
	return []string{
		cfg.PrimaryLLM.APIKey,
		cfg.ValidationLLM.APIKey,
		cfg.EffectiveValidationLLMConfig().APIKey,
	}
}

func errPayload(err error) any {
	if err == nil {
		return nil
	}
	return map[string]any{"error": err.Error()}
}

type diagnosticsWriterAdapter struct {
	writer *llm.DiagnosticsWriter
}

func (a diagnosticsWriterAdapter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
	artifacts, err := a.writer.WriteInteraction(stage, requestMetadata, requestPayload, responsePayload, errorPayload)
	if err != nil {
		return InteractionArtifacts{}, err
	}
	return InteractionArtifacts{
		RequestMetadataPath: artifacts.RequestMetadataPath,
		RequestPayloadPath:  artifacts.RequestPayloadPath,
		ResponsePayloadPath: artifacts.ResponsePayloadPath,
		ErrorPayloadPath:    artifacts.ErrorPayloadPath,
	}, nil
}