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

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
}