eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enable correction-aware downstream D&D extractors

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-26 23:59:05 +00:00
parent 759d32403f
commit a26d6ed042
12 changed files with 118 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
internal/modules/dnd/extract/itemoccurrences/extractor.go
View File

@@ -155,28 +155,35 @@ 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 {
return pipeline.ModuleSpec{
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.ItemOccurrenceListKind,
ReferenceSlots: referenceSlots(),
Key: Key,
Stage: pipeline.StageExtract,
ExecutionClass: contracts.ExecutionClassLLMBacked,
CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
Requires: append([]string(nil), requiredCapabilities...),
Provides: append([]string(nil), providedCapabilities...),
ArtifactKind: dnd.ItemOccurrenceListKind,
ReferenceSlots: referenceSlots(),
}
}

32
internal/modules/dnd/extract/itemoccurrences/extractor_test.go
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) {

3
internal/modules/dnd/extract/itemoccurrences/registry_test.go
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 {