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

Add pipeline stage contracts

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-03 15:24:07 +00:00
parent 3cf2ac577f
commit 75a0a9fa79
3 changed files with 480 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
internal/framework/contracts/composition_test.go
View File

@@ -12,14 +12,22 @@ import (
)
var _ contracts.InputAdapter = compositionAdapter{}
var _ contracts.Chunker = compositionChunker{}
var _ contracts.Extractor = compositionExtractor{}
var _ contracts.Merger = compositionMerger{}
var _ contracts.Normalizer = compositionNormalizer{}
var _ contracts.Validator = compositionValidator{}
var _ contracts.OutputEncoder = compositionOutputEncoder{}
func TestContractsComposeAcrossPackages(t *testing.T) {
ctx := context.Background()
adapter := compositionAdapter{}
chunker := compositionChunker{}
extractor := compositionExtractor{}
merger := compositionMerger{}
normalizer := compositionNormalizer{}
validator := compositionValidator{}
encoder := compositionOutputEncoder{}
doc, err := adapter.Parse(ctx, contracts.ParseRequest{SourceID: "source-1"})
if err != nil {
@@ -29,7 +37,22 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
t.Fatalf("ValidateDocument() error = %v, want nil", err)
}
extraction, err := extractor.Extract(ctx, contracts.ExtractionRequest{Source: doc})
chunking, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
Metadata: map[string]any{"max_units": 2},
})
if err != nil {
t.Fatalf("Chunk() error = %v, want nil", err)
}
if len(chunking.Chunks) != 1 {
t.Fatalf("len(Chunks) = %d, want 1", len(chunking.Chunks))
}
extraction, err := extractor.Extract(ctx, contracts.ExtractionRequest{
Source: doc,
Chunk: &chunking.Chunks[0],
AmbientContext: map[string]any{"synopsis": "example synopsis"},
})
if err != nil {
t.Fatalf("Extract() error = %v, want nil", err)
}
@@ -44,9 +67,38 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
}
}
merge, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
LaneID: candidate.ArtifactType,
ChunkArtifacts: []contracts.ChunkArtifacts{
{
Chunk: chunking.Chunks[0],
Candidates: extraction.Candidates,
},
},
})
if err != nil {
t.Fatalf("Merge() error = %v, want nil", err)
}
if len(merge.Candidates) != 1 {
t.Fatalf("len(merge.Candidates) = %d, want 1", len(merge.Candidates))
}
normalize, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
LaneID: candidate.ArtifactType,
Candidates: merge.Candidates,
})
if err != nil {
t.Fatalf("Normalize() error = %v, want nil", err)
}
if len(normalize.Candidates) != 1 {
t.Fatalf("len(normalize.Candidates) = %d, want 1", len(normalize.Candidates))
}
validation, err := validator.Validate(ctx, contracts.ValidationRequest{
Source: doc,
Candidates: extraction.Candidates,
Candidates: normalize.Candidates,
})
if err != nil {
t.Fatalf("Validate() error = %v, want nil", err)
@@ -62,6 +114,22 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
if decision.CandidateIndex != candidate.Index {
t.Fatalf("CandidateIndex = %d, want %d", decision.CandidateIndex, candidate.Index)
}
output, err := encoder.Encode(ctx, contracts.OutputRequest{
Manifest: artifacts.RunManifest{RunID: "run-1"},
Approved: []artifacts.Artifact{
artifacts.ArtifactFromCandidate(normalize.Candidates[0]),
},
})
if err != nil {
t.Fatalf("Encode() error = %v, want nil", err)
}
if output.ContentType != "application/json" {
t.Fatalf("ContentType = %q, want application/json", output.ContentType)
}
if len(output.Bytes) == 0 {
t.Fatal("len(Bytes) = 0, want encoded bytes")
}
}
type compositionAdapter struct{}
@@ -83,6 +151,30 @@ func (adapter compositionAdapter) Parse(ctx context.Context, req contracts.Parse
}, nil
}
type compositionChunker struct{}
func (chunker compositionChunker) Key() string {
return "generic-chunker"
}
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")
}
return contracts.ChunkResult{
Chunks: []contracts.SourceChunk{
{
ID: req.Source.ID + ":chunk:0",
SourceID: req.Source.ID,
Index: 0,
Units: append([]source.SourceUnit(nil), req.Source.Units...),
Metadata: map[string]any{"strategy": "whole-document"},
},
},
}, nil
}
type compositionExtractor struct{}
func (extractor compositionExtractor) Key() string {
@@ -105,6 +197,13 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
if req.Source == nil {
return contracts.ExtractionResult{}, errors.New("source document is required")
}
units := req.Source.Units
if req.Chunk != nil {
units = req.Chunk.Units
}
if req.AmbientContext["synopsis"] == "" {
return contracts.ExtractionResult{}, errors.New("ambient synopsis is required")
}
return contracts.ExtractionResult{
Candidates: []artifacts.ArtifactCandidate{
@@ -117,8 +216,8 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
SourceRefs: []source.SourceRef{
{
SourceID: req.Source.ID,
StartUnitID: req.Source.Units[0].ID,
EndUnitID: req.Source.Units[1].ID,
StartUnitID: units[0].ID,
EndUnitID: units[len(units)-1].ID,
},
},
},
@@ -126,6 +225,31 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
}, nil
}
type compositionMerger struct{}
func (merger compositionMerger) Key() string {
return "generic-merger"
}
func (merger compositionMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
var candidates []artifacts.ArtifactCandidate
for _, chunkArtifacts := range req.ChunkArtifacts {
candidates = append(candidates, chunkArtifacts.Candidates...)
}
return contracts.MergeResult{Candidates: candidates}, nil
}
type compositionNormalizer struct{}
func (normalizer compositionNormalizer) Key() string {
return "generic-normalizer"
}
func (normalizer compositionNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
return contracts.NormalizeResult{Candidates: req.Candidates}, nil
}
type compositionValidator struct{}
func (validator compositionValidator) Name() string {
@@ -148,3 +272,28 @@ func (validator compositionValidator) Validate(ctx context.Context, req contract
Decisions: decisions,
}, nil
}
type compositionOutputEncoder struct{}
func (encoder compositionOutputEncoder) Key() string {
return "generic-output"
}
func (encoder compositionOutputEncoder) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
payload := struct {
RunID string `json:"run_id"`
ApprovedCount int `json:"approved_count"`
}{
RunID: req.Manifest.RunID,
ApprovedCount: len(req.Approved),
}
encoded, err := json.Marshal(payload)
if err != nil {
return contracts.OutputResult{}, err
}
return contracts.OutputResult{
Bytes: encoded,
ContentType: "application/json",
}, nil
}

89
internal/framework/contracts/contracts.go
View File

@@ -46,10 +46,35 @@ type InputAdapter interface {
Parse(ctx context.Context, req ParseRequest) (*source.SourceDocument, error)
}
type SourceChunk struct {
ID string `json:"id"`
SourceID string `json:"source_id"`
Index int `json:"index"`
Units []source.SourceUnit `json:"units"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ChunkRequest struct {
Source *source.SourceDocument `json:"-"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ChunkResult struct {
Chunks []SourceChunk `json:"chunks"`
Warnings []Warning `json:"warnings,omitempty"`
}
type Chunker interface {
Key() string
Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error)
}
type ExtractionRequest struct {
Source *source.SourceDocument `json:"-"`
LLMClient StructuredLLMClient `json:"-"`
Metadata map[string]any `json:"metadata,omitempty"`
Source *source.SourceDocument `json:"-"`
Chunk *SourceChunk `json:"chunk,omitempty"`
AmbientContext map[string]any `json:"ambient_context,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ExtractionResult struct {
@@ -65,6 +90,45 @@ type Extractor interface {
Extract(ctx context.Context, req ExtractionRequest) (ExtractionResult, error)
}
type ChunkArtifacts struct {
Chunk SourceChunk `json:"chunk"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
}
type MergeRequest struct {
Source *source.SourceDocument `json:"-"`
LaneID string `json:"lane_id"`
ChunkArtifacts []ChunkArtifacts `json:"chunk_artifacts"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type MergeResult struct {
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
Warnings []Warning `json:"warnings,omitempty"`
}
type Merger interface {
Key() string
Merge(ctx context.Context, req MergeRequest) (MergeResult, error)
}
type NormalizeRequest struct {
Source *source.SourceDocument `json:"-"`
LaneID string `json:"lane_id"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type NormalizeResult struct {
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
Warnings []Warning `json:"warnings,omitempty"`
}
type Normalizer interface {
Key() string
Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error)
}
type ValidationRequest struct {
Source *source.SourceDocument `json:"-"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
@@ -95,3 +159,22 @@ type Warning struct {
ReasonCode string `json:"reason_code"`
Message string `json:"message"`
}
type OutputRequest struct {
Manifest artifacts.RunManifest `json:"manifest"`
Approved []artifacts.Artifact `json:"approved,omitempty"`
Rejected []artifacts.RejectedArtifact `json:"rejected,omitempty"`
Warnings []Warning `json:"warnings,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type OutputResult struct {
Bytes []byte `json:"-"`
ContentType string `json:"content_type,omitempty"`
Warnings []Warning `json:"warnings,omitempty"`
}
type OutputEncoder interface {
Key() string
Encode(ctx context.Context, req OutputRequest) (OutputResult, error)
}

244
internal/framework/contracts/contracts_test.go
View File

@@ -10,9 +10,13 @@ import (
)
var _ InputAdapter = fakeAdapter{}
var _ Chunker = fakeChunker{}
var _ Extractor = fakeExtractor{}
var _ Merger = fakeMerger{}
var _ Normalizer = fakeNormalizer{}
var _ Validator = fakeValidator{}
var _ StructuredLLMClient = fakeLLMClient{}
var _ OutputEncoder = fakeOutputEncoder{}
func TestFakeExtractorReturnsCandidateAndValidator(t *testing.T) {
validator := fakeValidator{name: "generic-validator"}
@@ -74,6 +78,166 @@ func TestFakeExtractorReturnsCandidateAndValidator(t *testing.T) {
}
}
func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
doc := &source.SourceDocument{
ID: "source-1",
Kind: "document",
Format: "text/plain",
Digest: "sha256:abc123",
Units: []source.SourceUnit{
{ID: "u1", Kind: "section", Text: "Source text."},
},
}
chunker := fakeChunker{key: "generic-chunker"}
result, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc})
if err != nil {
t.Fatalf("Chunk() error = %v, want nil", err)
}
if chunker.Key() != "generic-chunker" {
t.Fatalf("Key() = %q, want generic-chunker", chunker.Key())
}
if len(result.Chunks) != 1 {
t.Fatalf("len(Chunks) = %d, want 1", len(result.Chunks))
}
chunk := result.Chunks[0]
if chunk.ID != "source-1:chunk:0" {
t.Fatalf("SourceChunk.ID = %q, want source-1:chunk:0", chunk.ID)
}
if chunk.SourceID != doc.ID {
t.Fatalf("SourceChunk.SourceID = %q, want %q", chunk.SourceID, doc.ID)
}
if chunk.Index != 0 {
t.Fatalf("SourceChunk.Index = %d, want 0", chunk.Index)
}
if len(chunk.Units) != 1 {
t.Fatalf("len(SourceChunk.Units) = %d, want 1", len(chunk.Units))
}
}
func TestFakeExtractorReceivesChunkAndAmbientContext(t *testing.T) {
extractor := fakeExtractor{
key: "generic-extractor",
artifactType: "generic-artifact",
schemaVersion: "v1",
}
doc := &source.SourceDocument{
ID: "source-1",
Kind: "document",
Format: "text/plain",
Digest: "sha256:abc123",
Units: []source.SourceUnit{
{ID: "u1", Kind: "section", Text: "First source text."},
{ID: "u2", Kind: "section", Text: "Second source text."},
},
}
chunk := SourceChunk{
ID: "source-1:chunk:1",
SourceID: doc.ID,
Index: 1,
Units: []source.SourceUnit{doc.Units[1]},
}
result, err := extractor.Extract(context.Background(), ExtractionRequest{
Source: doc,
Chunk: &chunk,
AmbientContext: map[string]any{"mode": "chunked"},
})
if err != nil {
t.Fatalf("Extract() error = %v, want nil", err)
}
if len(result.Candidates) != 1 {
t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates))
}
candidate := result.Candidates[0]
if string(candidate.Payload) != `{"value":"chunked"}` {
t.Fatalf("ArtifactCandidate.Payload = %s, want chunked payload", candidate.Payload)
}
if len(candidate.SourceRefs) != 1 {
t.Fatalf("len(SourceRefs) = %d, want 1", len(candidate.SourceRefs))
}
ref := candidate.SourceRefs[0]
if ref.StartUnitID != "u2" || ref.EndUnitID != "u2" {
t.Fatalf("SourceRef = %+v, want u2 range", ref)
}
}
func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
candidate := artifacts.ArtifactCandidate{
Index: 0,
ExtractorKey: "generic-extractor",
ArtifactType: "generic-artifact",
SchemaVersion: "v1",
Payload: json.RawMessage(`{"value":"example"}`),
}
chunk := SourceChunk{
ID: "source-1:chunk:0",
SourceID: "source-1",
Index: 0,
Units: []source.SourceUnit{
{ID: "u1", Kind: "section", Text: "Source text."},
},
}
merger := fakeMerger{key: "generic-merger"}
normalizer := fakeNormalizer{key: "generic-normalizer"}
encoder := fakeOutputEncoder{key: "generic-output"}
merged, err := merger.Merge(context.Background(), MergeRequest{
LaneID: "generic-artifact",
ChunkArtifacts: []ChunkArtifacts{
{
Chunk: chunk,
Candidates: []artifacts.ArtifactCandidate{candidate},
},
},
})
if err != nil {
t.Fatalf("Merge() error = %v, want nil", err)
}
if merger.Key() != "generic-merger" {
t.Fatalf("Merger.Key() = %q, want generic-merger", merger.Key())
}
if len(merged.Candidates) != 1 {
t.Fatalf("len(merged.Candidates) = %d, want 1", len(merged.Candidates))
}
normalized, err := normalizer.Normalize(context.Background(), NormalizeRequest{
LaneID: "generic-artifact",
Candidates: merged.Candidates,
})
if err != nil {
t.Fatalf("Normalize() error = %v, want nil", err)
}
if normalizer.Key() != "generic-normalizer" {
t.Fatalf("Normalizer.Key() = %q, want generic-normalizer", normalizer.Key())
}
if len(normalized.Candidates) != 1 {
t.Fatalf("len(normalized.Candidates) = %d, want 1", len(normalized.Candidates))
}
encoded, err := encoder.Encode(context.Background(), OutputRequest{
Manifest: artifacts.RunManifest{RunID: "run-1"},
Approved: []artifacts.Artifact{
artifacts.ArtifactFromCandidate(normalized.Candidates[0]),
},
})
if err != nil {
t.Fatalf("Encode() error = %v, want nil", err)
}
if encoder.Key() != "generic-output" {
t.Fatalf("OutputEncoder.Key() = %q, want generic-output", encoder.Key())
}
if encoded.ContentType != "application/json" {
t.Fatalf("ContentType = %q, want application/json", encoded.ContentType)
}
if string(encoded.Bytes) != `{"run_id":"run-1","approved_count":1}` {
t.Fatalf("Bytes = %s, want encoded output", encoded.Bytes)
}
}
type fakeAdapter struct {
key string
doc *source.SourceDocument
@@ -87,6 +251,27 @@ func (adapter fakeAdapter) Parse(ctx context.Context, req ParseRequest) (*source
return adapter.doc, nil
}
type fakeChunker struct {
key string
}
func (chunker fakeChunker) Key() string {
return chunker.key
}
func (chunker fakeChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
return ChunkResult{
Chunks: []SourceChunk{
{
ID: req.Source.ID + ":chunk:0",
SourceID: req.Source.ID,
Index: 0,
Units: append([]source.SourceUnit(nil), req.Source.Units...),
},
},
}, nil
}
type fakeExtractor struct {
key string
artifactType string
@@ -111,6 +296,15 @@ func (extractor fakeExtractor) Validators() []Validator {
}
func (extractor fakeExtractor) Extract(ctx context.Context, req ExtractionRequest) (ExtractionResult, error) {
units := req.Source.Units
if req.Chunk != nil {
units = req.Chunk.Units
}
payload := json.RawMessage(`{"value":"example"}`)
if req.AmbientContext["mode"] == "chunked" {
payload = json.RawMessage(`{"value":"chunked"}`)
}
return ExtractionResult{
Candidates: []artifacts.ArtifactCandidate{
{
@@ -118,12 +312,12 @@ func (extractor fakeExtractor) Extract(ctx context.Context, req ExtractionReques
ExtractorKey: extractor.key,
ArtifactType: extractor.artifactType,
SchemaVersion: extractor.schemaVersion,
Payload: json.RawMessage(`{"value":"example"}`),
Payload: payload,
SourceRefs: []source.SourceRef{
{
SourceID: req.Source.ID,
StartUnitID: req.Source.Units[0].ID,
EndUnitID: req.Source.Units[0].ID,
StartUnitID: units[0].ID,
EndUnitID: units[len(units)-1].ID,
},
},
},
@@ -131,6 +325,35 @@ func (extractor fakeExtractor) Extract(ctx context.Context, req ExtractionReques
}, nil
}
type fakeMerger struct {
key string
}
func (merger fakeMerger) Key() string {
return merger.key
}
func (merger fakeMerger) Merge(ctx context.Context, req MergeRequest) (MergeResult, error) {
var candidates []artifacts.ArtifactCandidate
for _, chunkArtifacts := range req.ChunkArtifacts {
candidates = append(candidates, chunkArtifacts.Candidates...)
}
return MergeResult{Candidates: candidates}, nil
}
type fakeNormalizer struct {
key string
}
func (normalizer fakeNormalizer) Key() string {
return normalizer.key
}
func (normalizer fakeNormalizer) Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error) {
return NormalizeResult{Candidates: req.Candidates}, nil
}
type fakeValidator struct {
name string
}
@@ -163,3 +386,18 @@ func (client fakeLLMClient) CompleteStructured(ctx context.Context, req Structur
Content: json.RawMessage(`{"value":"example"}`),
}, nil
}
type fakeOutputEncoder struct {
key string
}
func (encoder fakeOutputEncoder) Key() string {
return encoder.key
}
func (encoder fakeOutputEncoder) Encode(ctx context.Context, req OutputRequest) (OutputResult, error) {
return OutputResult{
Bytes: []byte(`{"run_id":"` + req.Manifest.RunID + `","approved_count":1}`),
ContentType: "application/json",
}, nil
}