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"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/stagename"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	stagewarnings "gitea.maximumdirect.net/eric/audita/internal/framework/warnings"
)

// 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"`
	PromptMetadata    map[string]any              `json:"prompt_metadata,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"`
	Warnings    []stagewarnings.StageWarning           `json:"warnings,omitempty"`
	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 = stagename.ProposalGeneration(req.ModuleInstance, sectionIndexPtr(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),
				llm.ConfiguredSecrets(req.Config),
			),
		}
	}

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

	requestMetadata := map[string]any{
		"module_key":         req.ModuleKey,
		"module_instance":    req.ModuleInstance,
		"replacement_policy": req.ReplacementPolicy,
		"section":            req.Section,
		"start_index":        req.StartIndex,
		"model":              model,
	}
	if len(req.PromptMetadata) > 0 {
		requestMetadata["prompt_metadata"] = req.PromptMetadata
	}
	requestMetadata["response_schema"] = responseSchema.DiagnosticsMap()

	if writer != nil {
		artifacts, _ = writer.WriteInteraction(
			stage,
			requestMetadata,
			map[string]any{
				"messages": messages,
			},
			response,
			errPayload(callErr),
		)
	}

	if callErr != nil {
		if structuredoutput.IsMalformedError(callErr) {
			return Result{
				Warnings:  []stagewarnings.StageWarning{newMalformedProposalWarning(req.Section, artifacts, callErr)},
				Artifacts: artifacts,
			}, 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,
		}

		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,
		Warnings:    nil,
		Artifacts:   artifacts,
	}, nil
}

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 errPayload(err error) any {
	if err == nil {
		return nil
	}
	return map[string]any{"error": err.Error()}
}

func newMalformedProposalWarning(section *contracts.SectionMetadata, artifacts InteractionArtifacts, err error) stagewarnings.StageWarning {
	warning := stagewarnings.StageWarning{
		Scope:                  stagewarnings.ScopeProposalGeneration,
		ReasonCode:             "proposal_response_malformed",
		Message:                strings.TrimSpace(err.Error()),
		DiagnosticArtifactPath: diagnosticArtifactPath(artifacts),
	}
	if section != nil {
		sectionIndex := section.Index
		warning.SectionIndex = &sectionIndex
	}
	return warning
}

func diagnosticArtifactPath(artifacts InteractionArtifacts) string {
	if artifacts.ErrorPayloadPath != "" {
		return artifacts.ErrorPayloadPath
	}
	return artifacts.ResponsePayloadPath
}

func sectionIndexPtr(section *contracts.SectionMetadata) *int {
	if section == nil {
		return nil
	}
	index := section.Index
	return &index
}

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
}