Expand module contracts for references and normalizer LLM access

This commit is contained in:
2026-07-05 16:06:47 +00:00
parent 84c4c06712
commit 39e49d7f77
23 changed files with 248 additions and 36 deletions

View File

@@ -162,6 +162,10 @@ func (chunker compositionChunker) Key() string {
return "generic-chunker"
}
func (chunker compositionChunker) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (chunker compositionChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
if req.Source == nil {
return contracts.ChunkResult{}, errors.New("source document is required")
@@ -264,6 +268,10 @@ func (normalizer compositionNormalizer) Key() string {
return "generic-normalizer"
}
func (normalizer compositionNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (normalizer compositionNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
return contracts.NormalizeResult{Candidates: req.Candidates}, nil
}

View File

@@ -58,6 +58,7 @@ type SourceChunk struct {
type ChunkRequest struct {
Source *source.SourceDocument `json:"-"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
@@ -71,6 +72,7 @@ type ChunkResult struct {
type Chunker interface {
Key() string
ReferenceSlots() []ReferenceSlot
Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error)
}
@@ -165,6 +167,8 @@ type NormalizeRequest struct {
Source *source.SourceDocument `json:"-"`
LaneID string `json:"lane_id"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
@@ -177,6 +181,7 @@ type NormalizeResult struct {
type Normalizer interface {
Key() string
ReferenceSlots() []ReferenceSlot
Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error)
}

View File

@@ -378,6 +378,10 @@ func (chunker fakeChunker) Key() string {
return chunker.key
}
func (chunker fakeChunker) ReferenceSlots() []ReferenceSlot {
return nil
}
func (chunker fakeChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
return ChunkResult{
Chunks: []SourceChunk{
@@ -400,6 +404,10 @@ func (chunker *recordingChunker) Key() string {
return chunker.key
}
func (chunker *recordingChunker) ReferenceSlots() []ReferenceSlot {
return nil
}
func (chunker *recordingChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
chunker.request = req
return fakeChunker{key: chunker.key}.Chunk(ctx, req)
@@ -487,6 +495,10 @@ func (normalizer fakeNormalizer) Key() string {
return normalizer.key
}
func (normalizer fakeNormalizer) ReferenceSlots() []ReferenceSlot {
return nil
}
func (normalizer fakeNormalizer) Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error) {
return NormalizeResult{Candidates: req.Candidates}, nil
}

View File

@@ -330,6 +330,10 @@ func (chunker registryChunker) Key() string {
return chunker.key
}
func (chunker registryChunker) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (chunker registryChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
return contracts.ChunkResult{}, nil
}
@@ -354,6 +358,10 @@ func (normalizer registryNormalizer) Key() string {
return normalizer.key
}
func (normalizer registryNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (normalizer registryNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
return contracts.NormalizeResult{}, nil
}

View File

@@ -87,7 +87,7 @@ func validateModuleSpec(kind string, expectedStage ModuleStage, spec ModuleSpec)
if spec.Stage != expectedStage {
return fmt.Errorf("%s %q must use %q stage, got %q", kind, spec.Key, expectedStage, spec.Stage)
}
if spec.Stage != StageExtract && len(spec.ReferenceSlots) > 0 {
if !referenceSlotStage(spec.Stage) && len(spec.ReferenceSlots) > 0 {
return fmt.Errorf("%s %q must not declare reference slots", kind, spec.Key)
}
if err := validateReferenceSlots(spec.ReferenceSlots); err != nil {
@@ -96,6 +96,10 @@ func validateModuleSpec(kind string, expectedStage ModuleStage, spec ModuleSpec)
return nil
}
func referenceSlotStage(stage ModuleStage) bool {
return stage == StageChunk || stage == StageExtract || stage == StageNormalize
}
func sortedRegistryKeys[C any](constructors map[string]C) []string {
if len(constructors) == 0 {
return nil

View File

@@ -7,19 +7,117 @@ import (
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)
func TestValidateModuleSpecRejectsReferenceSlotsForNonExtractors(t *testing.T) {
err := validateModuleSpec("chunker", StageChunk, ModuleSpec{
Key: "generic",
Stage: StageChunk,
ReferenceSlots: []contracts.ReferenceSlot{
{Name: "roster"},
},
})
if err == nil {
t.Fatal("validateModuleSpec() error = nil, want error")
func TestValidateModuleSpecAllowsReferenceSlotsForEligibleStages(t *testing.T) {
tests := []struct {
name string
kind string
stage ModuleStage
}{
{name: "chunker", kind: "chunker", stage: StageChunk},
{name: "extractor", kind: "extractor", stage: StageExtract},
{name: "normalizer", kind: "normalizer", stage: StageNormalize},
}
if !strings.Contains(err.Error(), "reference slots") {
t.Fatalf("validateModuleSpec() error = %q, want reference slots context", err.Error())
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
spec := normalizeModuleSpec(ModuleSpec{
Key: "module",
Stage: test.stage,
ReferenceSlots: []contracts.ReferenceSlot{
{Name: "roster", Description: "Character roster", MaxBytes: 1024},
},
})
err := validateModuleSpec(test.kind, test.stage, spec)
if err != nil {
t.Fatalf("validateModuleSpec() error = %v, want nil", err)
}
})
}
}
func TestValidateModuleSpecRejectsReferenceSlotsForIneligibleStages(t *testing.T) {
tests := []struct {
name string
kind string
stage ModuleStage
}{
{name: "input", kind: "input adapter", stage: StageInput},
{name: "merge", kind: "merger", stage: StageMerge},
{name: "validate", kind: "validator", stage: StageValidate},
{name: "output", kind: "output encoder", stage: StageOutput},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
spec := normalizeModuleSpec(ModuleSpec{
Key: "module",
Stage: test.stage,
ReferenceSlots: []contracts.ReferenceSlot{
{Name: "roster"},
},
})
err := validateModuleSpec(test.kind, test.stage, spec)
if err == nil {
t.Fatal("validateModuleSpec() error = nil, want error")
}
if !strings.Contains(err.Error(), "reference slots") {
t.Fatalf("validateModuleSpec() error = %q, want reference slots context", err.Error())
}
})
}
}
func TestValidateModuleSpecRejectsInvalidReferenceSlotsForEligibleStages(t *testing.T) {
invalidSlots := []struct {
name string
slots []contracts.ReferenceSlot
want string
}{
{
name: "empty name",
slots: []contracts.ReferenceSlot{{Name: " "}},
want: "name",
},
{
name: "duplicate name after trim",
slots: []contracts.ReferenceSlot{
{Name: "roster"},
{Name: " roster "},
},
want: "duplicated",
},
{
name: "negative max bytes",
slots: []contracts.ReferenceSlot{{Name: "roster", MaxBytes: -1}},
want: "max_bytes",
},
}
eligibleStages := []struct {
name string
kind string
stage ModuleStage
}{
{name: "chunk", kind: "chunker", stage: StageChunk},
{name: "extract", kind: "extractor", stage: StageExtract},
{name: "normalize", kind: "normalizer", stage: StageNormalize},
}
for _, stage := range eligibleStages {
for _, invalid := range invalidSlots {
t.Run(stage.name+"/"+invalid.name, func(t *testing.T) {
spec := normalizeModuleSpec(ModuleSpec{
Key: "module",
Stage: stage.stage,
ReferenceSlots: invalid.slots,
})
err := validateModuleSpec(stage.kind, stage.stage, spec)
if err == nil {
t.Fatal("validateModuleSpec() error = nil, want error")
}
if !strings.Contains(err.Error(), invalid.want) {
t.Fatalf("validateModuleSpec() error = %q, want %q", err.Error(), invalid.want)
}
})
}
}
}

View File

@@ -118,6 +118,10 @@ func (chunker integrationChunker) Key() string {
return "chunk"
}
func (chunker integrationChunker) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (chunker integrationChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
return contracts.ChunkResult{
Chunks: []contracts.SourceChunk{
@@ -186,6 +190,10 @@ func (normalizer integrationNormalizer) Key() string {
return "normalize"
}
func (normalizer integrationNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (normalizer integrationNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
return contracts.NormalizeResult{Candidates: req.Candidates}, nil
}

View File

@@ -224,6 +224,7 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
Source: doc,
LaneID: lane.ID,
Candidates: mergeResult.Candidates,
LLMClient: input.LLMClient,
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,

View File

@@ -464,6 +464,10 @@ func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) {
if len(extractor.seenLLMClients) != 2 || extractor.seenLLMClients[0] == nil || extractor.seenLLMClients[1] == nil {
t.Fatalf("seen LLM clients = %#v, want client for each chunk", extractor.seenLLMClients)
}
normalizer := modules.normalizers["normalize"]
if len(normalizer.requests) != 1 || normalizer.requests[0].LLMClient == nil {
t.Fatalf("normalizer LLM client = %#v, want client on normalize request", normalizer.requests)
}
if extractor.seenMetadata[0]["request"] != "test" {
t.Fatalf("seen metadata = %#v, want request metadata", extractor.seenMetadata)
}
@@ -1314,6 +1318,10 @@ func (chunker *runnerChunker) Key() string {
return chunker.key
}
func (chunker *runnerChunker) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (chunker *runnerChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
chunker.requests = append(chunker.requests, req)
return contracts.ChunkResult{
@@ -1421,6 +1429,10 @@ func (normalizer *runnerNormalizer) Key() string {
return normalizer.key
}
func (normalizer *runnerNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (normalizer *runnerNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
normalizer.requests = append(normalizer.requests, req)
candidates := append([]artifacts.ArtifactCandidate(nil), normalizer.result...)

View File

@@ -216,6 +216,10 @@ func (chunker walkingSkeletonChunker) Key() string {
return "fake/chunk"
}
func (chunker walkingSkeletonChunker) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (chunker walkingSkeletonChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
if len(req.Source.Units) < 3 {
return contracts.ChunkResult{}, fmt.Errorf("fixture source must contain at least three units")
@@ -336,6 +340,10 @@ func (normalizer walkingSkeletonNormalizer) Key() string {
return DefaultNormalizeModule
}
func (normalizer walkingSkeletonNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (normalizer walkingSkeletonNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
return contracts.NormalizeResult{Candidates: req.Candidates}, nil
}