Integrate extraction validation retries

internal/cli/assembled_spell_pipeline_contract_test.go

@@ -21,6 +21,10 @@ import (
 
 const assembledSpellExtractorKey = "test/dnd/spell-casts"
 
+const assembledCorrectingSpellExtractorKey = "test/dnd/correcting-spell-casts"
+
+const assembledDirectSpellValidatorKey = "test/dnd/direct-spell-correction"
+
 func TestAssembledSpellPipelineNormalizesMergedCasts(t *testing.T) {
 	registries, resolved, extractor := assembledSpellPipeline(t, assembledSpellPipelineOptions{})
 	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
@@ -101,6 +105,30 @@ func TestAssembledSpellPipelineNormalizesMergedCasts(t *testing.T) {
 	}
 }
 
+func TestAssembledSpellPipelineCorrectsRejectedDirectExtraction(t *testing.T) {
+	registries, resolved, extractor := assembledCorrectingSpellPipeline(t)
+	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{})
+	if err != nil {
+		t.Fatalf("Prepare() error = %v, want nil", err)
+	}
+
+	output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{
+		Prepared:       prepared,
+		RawInput:       readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"),
+		ChunkCacheMode: pipeline.ChunkCacheBypass,
+	})
+	if err != nil {
+		t.Fatalf("Run() error = %v, want nil", err)
+	}
+	if len(output.Rejected) != 0 || output.Manifest.ValidationStatus != "approved" || len(output.NormalizeOutputs) != 1 {
+		t.Fatalf("run output = %#v, want corrected accepted spell output", output)
+	}
+	correction := extractor.correctionSnapshot()
+	if correction == nil || string(correction.AssistantResponse) != `{"spell":"Mysterious Burst"}` || correction.UserGuidance != "use a known spell name" {
+		t.Fatalf("extract correction = %#v, want exact rejected model response and validator guidance", correction)
+	}
+}
+
 func TestAssembledSpellPipelineHonorsNormalizeValidatorOverride(t *testing.T) {
 	registries, resolved, _ := assembledSpellPipeline(t, assembledSpellPipelineOptions{normalizeValidatorOverride: true})
 	var normalizeChain *pipeline.ResolvedValidatorChain
@@ -161,8 +189,8 @@ func TestAssembledSpellPipelinePromotesTerminalUnknownSpellWarning(t *testing.T)
 	if !reflect.DeepEqual(rejectedFile.Rejected, output.Rejected) {
 		t.Fatalf("rejected file = %#v, run rejections = %#v, want durable rejection diagnostic", rejectedFile.Rejected, output.Rejected)
 	}
-	if len(output.Warnings) != 1 || output.Warnings[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Warnings[0].Scope != "spell_casts[0]" {
-		t.Fatalf("warnings = %#v, want terminal normalize catalog warning", output.Warnings)
+	if len(output.Warnings) != 2 || output.Warnings[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Warnings[0].Scope != "spell_casts[0]" || output.Warnings[1].ReasonCode != "spell_not_near_source" {
+		t.Fatalf("warnings = %#v, want complete terminal normalize validation warnings", output.Warnings)
 	}
 }
 
@@ -206,6 +234,47 @@ type assembledSpellPipelineOptions struct {
 	unknownSpell               bool
 }
 
+func assembledCorrectingSpellPipeline(t *testing.T) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledCorrectingSpellExtractor) {
+	t.Helper()
+	components := productionTestComponents(t)
+	extractor := &assembledCorrectingSpellExtractor{}
+	if err := pipeline.RegisterExtractor[dnd.SpellList](components.registries.Extractors, pipeline.ModuleSpec{
+		Key:                assembledCorrectingSpellExtractorKey,
+		Stage:              pipeline.StageExtract,
+		ExecutionClass:     contracts.ExecutionClassLLMBacked,
+		CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
+		Requires:           []string{"chunks", "source.transcript"},
+		Provides:           []string{"dnd.spell_casts"},
+		ArtifactKind:       dnd.SpellListKind,
+	}, func() (contracts.Extractor[dnd.SpellList], error) {
+		return extractor, nil
+	}); err != nil {
+		t.Fatalf("register correcting extractor: %v", err)
+	}
+	if err := pipeline.RegisterTypedValidator[dnd.SpellList](components.registries.Validators, dnd.SpellListKind, pipeline.ValidatorSpec{Key: assembledDirectSpellValidatorKey, ExecutionClass: contracts.ExecutionClassDeterministic}, func() (contracts.TypedValidator[dnd.SpellList], error) {
+		return assembledDirectSpellValidator{}, nil
+	}); err != nil {
+		t.Fatalf("register direct spell validator: %v", err)
+	}
+
+	extract := pipeline.Binding(assembledCorrectingSpellExtractorKey)
+	extract.Retries = 1
+	extract.Validators = pipeline.ValidatorOverride{Set: true, Validators: []pipeline.ModuleBinding{{Module: assembledDirectSpellValidatorKey}}}
+	resolved, err := pipeline.ResolvePipeline(pipeline.PipelineProfile{
+		ID:    "assembled-dnd-correcting-spells",
+		Input: pipeline.Binding("seriatim"),
+		Chunk: pipeline.ModuleBinding{Module: "generic", Options: map[string]any{"max_units": 1}},
+		Artifacts: map[string]pipeline.ArtifactLaneProfile{
+			"spells": {Extract: extract, Normalize: pipeline.Binding(spellnormalize.Key)},
+		},
+		Output: pipeline.Binding("json"),
+	}, pipeline.ResolveOptions{}, catalogFromRegistries(components.registries))
+	if err != nil {
+		t.Fatalf("ResolvePipeline() error = %v, want nil", err)
+	}
+	return components.registries, resolved, extractor
+}
+
 func assembledSpellPipeline(t *testing.T, options assembledSpellPipelineOptions) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledSpellExtractor) {
 	t.Helper()
 	components := productionTestComponents(t)
@@ -251,6 +320,69 @@ type assembledSpellExtractor struct {
 	unknownSpell bool
 }
 
+type assembledCorrectingSpellExtractor struct {
+	mu         sync.Mutex
+	correction *contracts.SemanticCorrection
+}
+
+func (*assembledCorrectingSpellExtractor) Key() string { return assembledCorrectingSpellExtractorKey }
+
+func (*assembledCorrectingSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }
+
+func (e *assembledCorrectingSpellExtractor) Extract(_ context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
+	if req.Source == nil || req.Chunk == nil {
+		return contracts.TypedExtractionResult[dnd.SpellList]{}, fmt.Errorf("correcting assembled extractor requires source and chunk")
+	}
+	response := `{"spell":"accepted"}`
+	value := dnd.SpellList{SpellCasts: []dnd.SpellCast{}}
+	if req.Chunk.Index == 0 && req.Correction == nil {
+		response = `{"spell":"Mysterious Burst"}`
+		value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Mysterious Burst", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}}
+	}
+	if req.Chunk.Index == 0 && req.Correction != nil {
+		correction, err := contracts.CloneSemanticCorrection(req.Correction)
+		if err != nil {
+			return contracts.TypedExtractionResult[dnd.SpellList]{}, err
+		}
+		e.mu.Lock()
+		e.correction = correction
+		e.mu.Unlock()
+		value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Cure Wounds", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}}
+	}
+	candidate, err := contracts.NewModelCandidate([]byte(response), contracts.CorrectionProtocolSingleResponseV1)
+	if err != nil {
+		return contracts.TypedExtractionResult[dnd.SpellList]{}, err
+	}
+	return contracts.TypedExtractionResult[dnd.SpellList]{Value: value, ModelCandidate: candidate}, nil
+}
+
+func (e *assembledCorrectingSpellExtractor) correctionSnapshot() *contracts.SemanticCorrection {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	correction, err := contracts.CloneSemanticCorrection(e.correction)
+	if err != nil {
+		return nil
+	}
+	return correction
+}
+
+type assembledDirectSpellValidator struct{}
+
+func (assembledDirectSpellValidator) Name() string { return assembledDirectSpellValidatorKey }
+
+func (assembledDirectSpellValidator) ExecutionClass() contracts.ExecutionClass {
+	return contracts.ExecutionClassDeterministic
+}
+
+func (assembledDirectSpellValidator) Validate(_ context.Context, req contracts.TypedValidationRequest[dnd.SpellList]) (contracts.ValidationResult, error) {
+	for _, cast := range req.Value.SpellCasts {
+		if cast.Spell == "Mysterious Burst" {
+			return contracts.ValidationResult{Approved: false, ReasonCode: "unknown_spell", Message: "spell is not in the catalog", CorrectionGuidance: "use a known spell name"}, nil
+		}
+	}
+	return contracts.ValidationResult{Approved: true}, nil
+}
+
 func (e *assembledSpellExtractor) Key() string { return assembledSpellExtractorKey }
 
 func (*assembledSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }