Expand module contracts for references and normalizer LLM access
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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...)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user