Enable correction-aware downstream D&D extractors

This commit is contained in:
2026-08-26 23:59:05 +00:00
parent 759d32403f
commit a26d6ed042
12 changed files with 118 additions and 55 deletions

View File

@@ -347,7 +347,7 @@ All named producers truthfully advertise and implement the exact-response
protocol, with unchanged ordinary extraction behavior. This stage is one
Terra prompt.
## Stage 7 — Migrate Downstream D&D Extractors
## Stage 7 — Migrate Downstream D&D Extractors
### Goal

View File

@@ -208,19 +208,25 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
var response extractionResponse
inputs := shared.PromptInputs(sourceInput, req.References)
inputs[NPCRegistryReferenceSlot] = npcRegistry.PromptInput()
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key,
PromptID: PromptID,
PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile,
SessionID: req.SessionID,
StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
Correction: req.Correction,
Inputs: inputs,
}, &response); err != nil {
}, &response)
if err != nil {
return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("complete structured output: %w", err)
}
candidate, err := shared.ModelCandidateFromResponse(completion)
if err != nil {
return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("capture model candidate: %w", err)
}
canonicalizeResponse(&response, order, req.Source.ID)
return contracts.TypedExtractionResult[dnd.CombatTurnList]{Value: canonicalCombatTurnList(response, req.Source.ID)}, nil
return contracts.TypedExtractionResult[dnd.CombatTurnList]{Value: canonicalCombatTurnList(response, req.Source.ID), ModelCandidate: candidate}, nil
}
func ModuleSpec() pipeline.ModuleSpec {
@@ -228,6 +234,7 @@ func ModuleSpec() pipeline.ModuleSpec {
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.CombatTurnListKind,

View File

@@ -448,7 +448,7 @@ func TestExtractorManifestMetadataAndFingerprints(t *testing.T) {
func TestModuleSpecAndRegistration(t *testing.T) {
wantSlots := referenceSlots()
got := ModuleSpec()
if got.Key != Key || got.Stage != pipeline.StageExtract || got.ArtifactKind != dnd.CombatTurnListKind || !reflect.DeepEqual(got.Requires, []string{"chunks", "source.transcript"}) || !reflect.DeepEqual(got.Provides, []string{"dnd.combat_turns"}) || !reflect.DeepEqual(got.ReferenceSlots, wantSlots) {
if got.Key != Key || got.Stage != pipeline.StageExtract || got.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 || got.ArtifactKind != dnd.CombatTurnListKind || !reflect.DeepEqual(got.Requires, []string{"chunks", "source.transcript"}) || !reflect.DeepEqual(got.Provides, []string{"dnd.combat_turns"}) || !reflect.DeepEqual(got.ReferenceSlots, wantSlots) {
t.Fatalf("ModuleSpec() = %#v, want combat extractor contract", got)
}
got.Requires[0] = "changed"

View File

@@ -142,19 +142,26 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
inputs[name] = input
}
var response extractionResponse
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key,
PromptID: PromptID,
PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile,
SessionID: req.SessionID,
StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
Correction: req.Correction,
Inputs: inputs,
}, &response); err != nil {
}, &response)
if err != nil {
return contracts.TypedExtractionResult[dnd.EnemyEventList]{}, extractorErrorf("complete structured output: %w", err)
}
candidate, err := shared.ModelCandidateFromResponse(completion)
if err != nil {
return contracts.TypedExtractionResult[dnd.EnemyEventList]{}, extractorErrorf("capture model candidate: %w", err)
}
return contracts.TypedExtractionResult[dnd.EnemyEventList]{
Value: canonicalEnemyEventList(response, shared.NewSourceRefOrder(req.Source), req.Source.ID),
ModelCandidate: candidate,
}, nil
}
@@ -177,6 +184,7 @@ func ModuleSpec() pipeline.ModuleSpec {
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.EnemyEventListKind,

View File

@@ -137,7 +137,7 @@ func TestConstructorSpecOptionsAndSafeMetadata(t *testing.T) {
first.Requires[0] = "changed"
first.ReferenceSlots[0].AcceptedMediaTypes[0] = "changed"
second := ModuleSpec()
if second.Requires[0] != "chunks" || second.ArtifactKind != dnd.EnemyEventListKind || second.ExecutionClass != contracts.ExecutionClassLLMBacked || second.ReferenceSlots[0].AcceptedMediaTypes[0] == "changed" {
if second.Requires[0] != "chunks" || second.ArtifactKind != dnd.EnemyEventListKind || second.ExecutionClass != contracts.ExecutionClassLLMBacked || second.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 || second.ReferenceSlots[0].AcceptedMediaTypes[0] == "changed" {
t.Fatalf("ModuleSpec() reused mutable state: %#v", second)
}
registry := pipeline.NewExtractorRegistry()

View File

@@ -155,17 +155,23 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
var response extractionResponse
inputs := shared.PromptInputs(sourceInput, req.References)
inputs[ItemRegistryReferenceSlot] = registry.PromptInput()
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key, PromptID: PromptID, PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile, SessionID: req.SessionID, StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts, Inputs: inputs,
}, &response); err != nil {
ProfileID: req.LLMProfile, SessionID: req.SessionID, StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
Correction: req.Correction, Inputs: inputs,
}, &response)
if err != nil {
return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
}
candidate, err := shared.ModelCandidateFromResponse(completion)
if err != nil {
return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("capture model candidate: %w", err)
}
value, err := canonicalItemOccurrenceList(response, order, req.Source.ID, registry)
if err != nil {
return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("map item occurrence response: %w", err)
}
return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{Value: value}, nil
return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{Value: value, ModelCandidate: candidate}, nil
}
func ModuleSpec() pipeline.ModuleSpec {
@@ -173,6 +179,7 @@ func ModuleSpec() pipeline.ModuleSpec {
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.ItemOccurrenceListKind,

View File

@@ -11,11 +11,15 @@ import (
func TestExtractGroundsOccurrencesInRequiredRegistry(t *testing.T) {
id := itemidentity.DeriveID("Torch")
client := &fakeItemOccurrencesLLMClient{response: extractionResponse{Occurrences: []itemOccurrenceResponse{
{Name: "Torch", Kind: "lost", From: "party", SourceRefs: responseRefs(1, 1)},
}}}
rawResponse := []byte(`{"occurrences":[{"name":"Torch","kind":"lost","from":"party","source_refs":[{"start_unit_id":1,"end_unit_id":1}]}]}`)
client := &fakeItemOccurrencesLLMClient{content: append([]byte(nil), rawResponse...)}
req := extractionRequest()
req.References = itemRegistryReferences(t)
correction, err := contracts.NewSemanticCorrection([]byte(`{"occurrences":[]}`), "Keep the transcript-grounded item occurrence.")
if err != nil {
t.Fatalf("NewSemanticCorrection() error = %v", err)
}
req.Correction = correction
result, err := newExtractor(t, client, req.References).Extract(context.Background(), req)
if err != nil {
t.Fatal(err)
@@ -26,10 +30,30 @@ func TestExtractGroundsOccurrencesInRequiredRegistry(t *testing.T) {
if refs := result.Value.Occurrences[0].SourceRefs; len(refs) != 1 || refs[0].SourceID != req.Source.ID || refs[0].StartUnitID != 1 || refs[0].EndUnitID != 1 {
t.Fatalf("occurrence evidence = %#v, want current-source unit range", refs)
}
input := client.requests[0].Inputs[ItemRegistryReferenceSlot]
request := client.requests[0]
input := request.Inputs[ItemRegistryReferenceSlot]
if input.Name != ItemRegistryReferenceSlot || string(input.Content) != `{"items":[{"name":"Torch"}]}` || strings.Contains(string(input.Content), "item:sha256:") {
t.Fatalf("registry prompt input = %#v, want names-only projection", input)
}
if request.Correction == nil || string(request.Correction.AssistantResponse) != `{"occurrences":[]}` || request.Correction.UserGuidance != "Keep the transcript-grounded item occurrence." {
t.Fatalf("correction = %#v, want exact request correction", request.Correction)
}
for _, material := range []string{string(request.Correction.AssistantResponse), request.Correction.UserGuidance} {
if strings.Contains(material, id) || strings.Contains(material, "item:sha256:") {
t.Fatalf("correction material leaked opaque item identity: %q", material)
}
}
correction.AssistantResponse[0] = '['
if got := string(request.Correction.AssistantResponse); got != `{"occurrences":[]}` {
t.Fatalf("captured correction changed after caller mutation: %q", got)
}
if result.ModelCandidate == nil || result.ModelCandidate.Protocol != contracts.CorrectionProtocolSingleResponseV1 || string(result.ModelCandidate.Response) != string(rawResponse) {
t.Fatalf("model candidate = %#v, want exact validated response", result.ModelCandidate)
}
client.content[0] = '['
if got := string(result.ModelCandidate.Response); got != string(rawResponse) {
t.Fatalf("model candidate changed after provider buffer mutation: %q", got)
}
}
func TestExtractCanonicalizesComparisonEquivalentNames(t *testing.T) {

View File

@@ -36,6 +36,9 @@ func testSourceRefs() []source.SourceRef {
func TestModuleSpecDeclaresRequiredRegistry(t *testing.T) {
spec := ModuleSpec()
if spec.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 {
t.Fatalf("correction protocol = %q, want %q", spec.CorrectionProtocol, contracts.CorrectionProtocolSingleResponseV1)
}
var slot contracts.ReferenceSlot
for _, candidate := range spec.ReferenceSlots {
if candidate.Name == ItemRegistryReferenceSlot {

View File

@@ -151,22 +151,29 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
var response extractionResponse
inputs := shared.PromptInputs(sourceInput, req.References)
inputs[LocationRegistryReferenceSlot] = grounding.PromptInput()
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key, PromptID: PromptID, PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile, SessionID: req.SessionID, StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts, Inputs: inputs,
}, &response); err != nil {
ProfileID: req.LLMProfile, SessionID: req.SessionID, StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
Correction: req.Correction, Inputs: inputs,
}, &response)
if err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
}
candidate, err := shared.ModelCandidateFromResponse(completion)
if err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("capture model candidate: %w", err)
}
occurrences, err := canonicalOccurrenceList(response, shared.NewSourceRefOrder(req.Source), req.Source.ID, grounding)
if err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("resolve location grounding: %w", err)
}
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{Value: occurrences}, nil
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{Value: occurrences, ModelCandidate: candidate}, nil
}
func ModuleSpec() pipeline.ModuleSpec {
return pipeline.ModuleSpec{
Key: Key, Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...), Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.LocationOccurrenceListKind, ReferenceSlots: referenceSlots(),
}

View File

@@ -212,7 +212,7 @@ func TestExtractorContractsMetadataAndFailures(t *testing.T) {
}
spec := ModuleSpec()
if spec.Key != Key || spec.Stage != pipeline.StageExtract || spec.ExecutionClass != contracts.ExecutionClassLLMBacked || spec.ArtifactKind != dnd.LocationOccurrenceListKind {
if spec.Key != Key || spec.Stage != pipeline.StageExtract || spec.ExecutionClass != contracts.ExecutionClassLLMBacked || spec.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 || spec.ArtifactKind != dnd.LocationOccurrenceListKind {
t.Fatalf("ModuleSpec() = %#v", spec)
}
var slot contracts.ReferenceSlot

View File

@@ -159,23 +159,29 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
var response extractionResponse
inputs := shared.PromptInputs(sourceInput, req.References)
inputs[NPCRegistryReferenceSlot] = npcRegistry.PromptInput()
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key,
PromptID: PromptID,
PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile,
SessionID: req.SessionID,
StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
Correction: req.Correction,
Inputs: inputs,
}, &response); err != nil {
}, &response)
if err != nil {
return contracts.TypedExtractionResult[dnd.NPCOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
}
candidate, err := shared.ModelCandidateFromResponse(completion)
if err != nil {
return contracts.TypedExtractionResult[dnd.NPCOccurrenceList]{}, extractorErrorf("capture model candidate: %w", err)
}
canonicalizeResponse(&response, order, req.Source.ID)
value, err := canonicalOccurrenceList(response, req.Source.ID, npcRegistry)
if err != nil {
return contracts.TypedExtractionResult[dnd.NPCOccurrenceList]{}, extractorErrorf("resolve NPC names against registry: %w", err)
}
return contracts.TypedExtractionResult[dnd.NPCOccurrenceList]{Value: value}, nil
return contracts.TypedExtractionResult[dnd.NPCOccurrenceList]{Value: value, ModelCandidate: candidate}, nil
}
func ModuleSpec() pipeline.ModuleSpec {
@@ -183,6 +189,7 @@ func ModuleSpec() pipeline.ModuleSpec {
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.NPCOccurrenceListKind,

View File

@@ -285,7 +285,7 @@ func TestExtractRejectsInvalidRequestsAndProviderFailures(t *testing.T) {
func TestModuleSpecRegistrationMetadataAndFingerprints(t *testing.T) {
got := ModuleSpec()
if got.Key != Key || got.Stage != pipeline.StageExtract || got.ArtifactKind != dnd.NPCOccurrenceListKind || !reflect.DeepEqual(got.Requires, []string{"chunks", "source.transcript"}) || !reflect.DeepEqual(got.Provides, []string{"dnd.npc_occurrences"}) {
if got.Key != Key || got.Stage != pipeline.StageExtract || got.CorrectionProtocol != contracts.CorrectionProtocolSingleResponseV1 || got.ArtifactKind != dnd.NPCOccurrenceListKind || !reflect.DeepEqual(got.Requires, []string{"chunks", "source.transcript"}) || !reflect.DeepEqual(got.Provides, []string{"dnd.npc_occurrences"}) {
t.Fatalf("ModuleSpec() = %#v", got)
}
var registrySlot contracts.ReferenceSlot