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

Implement raw module output contracts

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-07 18:58:23 +00:00
parent 9e3f8809b3
commit c05ecb58d8
33 changed files with 1166 additions and 1780 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,13 +19,9 @@ var _ Validator = fakeValidator{}
var _ StructuredLLMClient = fakeLLMClient{}
var _ OutputEncoder = fakeOutputEncoder{}
func TestFakeExtractorReturnsCandidateAndValidator(t *testing.T) {
validator := fakeValidator{name: "generic-validator"}
func TestFakeExtractorReturnsRawOutput(t *testing.T) {
extractor := fakeExtractor{
key: "generic-extractor",
artifactType: "generic-artifact",
schemaVersion: "v1",
validators: []Validator{validator},
key: "generic-extractor",
}
doc := &source.SourceDocument{
ID: "source-1",
@@ -45,37 +41,14 @@ func TestFakeExtractorReturnsCandidateAndValidator(t *testing.T) {
if extractor.Key() != "generic-extractor" {
t.Fatalf("Key() = %q, want generic-extractor", extractor.Key())
}
if extractor.ArtifactType() != "generic-artifact" {
t.Fatalf("ArtifactType() = %q, want generic-artifact", extractor.ArtifactType())
if result.Output.ExtractorKey != "" {
t.Fatalf("ExtractorKey = %q, want runner-owned empty value", result.Output.ExtractorKey)
}
if extractor.SchemaVersion() != "v1" {
t.Fatalf("SchemaVersion() = %q, want v1", extractor.SchemaVersion())
if result.Output.Schema.Version != "v1" {
t.Fatalf("Schema.Version = %q, want v1", result.Output.Schema.Version)
}
if len(extractor.Validators()) != 1 {
t.Fatalf("len(Validators()) = %d, want 1", len(extractor.Validators()))
}
if extractor.Validators()[0].Name() != "generic-validator" {
t.Fatalf("Validators()[0].Name() = %q, want generic-validator", extractor.Validators()[0].Name())
}
if len(result.Candidates) != 1 {
t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates))
}
candidate := result.Candidates[0]
if candidate.Index != 0 {
t.Fatalf("ArtifactCandidate.Index = %d, want 0", candidate.Index)
}
if candidate.ExtractorKey != extractor.Key() {
t.Fatalf("ArtifactCandidate.ExtractorKey = %q, want %q", candidate.ExtractorKey, extractor.Key())
}
if candidate.ArtifactType != extractor.ArtifactType() {
t.Fatalf("ArtifactCandidate.ArtifactType = %q, want %q", candidate.ArtifactType, extractor.ArtifactType())
}
if candidate.SchemaVersion != extractor.SchemaVersion() {
t.Fatalf("ArtifactCandidate.SchemaVersion = %q, want %q", candidate.SchemaVersion, extractor.SchemaVersion())
}
if string(candidate.Payload) != `{"value":"example"}` {
t.Fatalf("ArtifactCandidate.Payload = %s, want example payload", candidate.Payload)
if result.Output.Payload.MediaType != "application/json" || string(result.Output.Payload.Content) != `{"value":"example"}` {
t.Fatalf("payload = %q %s, want JSON raw output", result.Output.Payload.MediaType, result.Output.Payload.Content)
}
}
@@ -146,11 +119,7 @@ func TestFakeChunkerReceivesLLMClient(t *testing.T) {
}
func TestFakeExtractorReceivesChunkAndAmbientContext(t *testing.T) {
extractor := fakeExtractor{
key: "generic-extractor",
artifactType: "generic-artifact",
schemaVersion: "v1",
}
extractor := fakeExtractor{key: "generic-extractor"}
doc := &source.SourceDocument{
ID: "source-1",
Kind: "document",
@@ -180,20 +149,11 @@ func TestFakeExtractorReceivesChunkAndAmbientContext(t *testing.T) {
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))
if result.Output.ChunkID != "" || result.Output.ChunkIndex != 0 {
t.Fatalf("chunk provenance = %q/%d, want runner-owned zero values", result.Output.ChunkID, result.Output.ChunkIndex)
}
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 != 2 || ref.EndUnitID != 2 {
t.Fatalf("SourceRef = %+v, want unit 2 range", ref)
if string(result.Output.Payload.Content) != `{"value":"chunked"}` {
t.Fatalf("Payload.Content = %s, want chunked payload", result.Output.Payload.Content)
}
}
@@ -363,19 +323,17 @@ func TestLLMInputSetCloneCopiesContent(t *testing.T) {
}
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: 1, Kind: "section", Text: "Source text."},
extractOutput := ExtractOutput{
LaneID: "generic-lane",
ExtractorKey: "generic-extractor",
SourceID: "source-1",
ChunkID: "source-1:chunk:0",
ChunkIndex: 0,
Schema: ResponseSchema{ID: "schema-id", Name: "schema-name", Version: "v1"},
Payload: RawPayload{
Content: []byte(`{"value":"example"}`),
MediaType: "application/json",
Metadata: map[string]any{"confidence": 0.75},
},
}
merger := fakeMerger{key: "generic-merger"}
@@ -383,13 +341,8 @@ func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
encoder := fakeOutputEncoder{key: "generic-output"}
merged, err := merger.Merge(context.Background(), MergeRequest{
LaneID: "generic-artifact",
ChunkArtifacts: []ChunkArtifacts{
{
Chunk: chunk,
Candidates: []artifacts.ArtifactCandidate{candidate},
},
},
LaneID: "generic-lane",
ExtractOutputs: []ExtractOutput{extractOutput},
})
if err != nil {
t.Fatalf("Merge() error = %v, want nil", err)
@@ -397,13 +350,13 @@ func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
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))
if string(merged.Output.Payload.Content) != `{"value":"example"}` {
t.Fatalf("merged content = %s, want raw extract content", merged.Output.Payload.Content)
}
normalized, err := normalizer.Normalize(context.Background(), NormalizeRequest{
LaneID: "generic-artifact",
Candidates: merged.Candidates,
LaneID: "generic-lane",
MergeOutput: merged.Output,
})
if err != nil {
t.Fatalf("Normalize() error = %v, want nil", err)
@@ -411,15 +364,13 @@ func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
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))
if string(normalized.Output.Payload.Content) != `{"value":"example"}` {
t.Fatalf("normalized content = %s, want raw merge content", normalized.Output.Payload.Content)
}
encoded, err := encoder.Encode(context.Background(), OutputRequest{
Manifest: artifacts.RunManifest{RunID: "run-1"},
Approved: []artifacts.Artifact{
artifacts.ArtifactFromCandidate(normalized.Candidates[0]),
},
Manifest: artifacts.RunManifest{RunID: "run-1"},
NormalizeOutputs: []NormalizeOutput{normalized.Output},
})
if err != nil {
t.Fatalf("Encode() error = %v, want nil", err)
@@ -433,7 +384,7 @@ func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
if encoded.Files[0].ContentType != "application/json" {
t.Fatalf("ContentType = %q, want application/json", encoded.Files[0].ContentType)
}
if string(encoded.Files[0].Bytes) != `{"run_id":"run-1","approved_count":1}` {
if string(encoded.Files[0].Bytes) != `{"run_id":"run-1","output_count":1}` {
t.Fatalf("Bytes = %s, want encoded output", encoded.Files[0].Bytes)
}
}
@@ -529,57 +480,29 @@ func (chunker *recordingChunker) Chunk(ctx context.Context, req ChunkRequest) (C
}
type fakeExtractor struct {
key string
artifactType string
schemaVersion string
validators []Validator
key string
}
func (extractor fakeExtractor) Key() string {
return extractor.key
}
func (extractor fakeExtractor) ArtifactType() string {
return extractor.artifactType
}
func (extractor fakeExtractor) SchemaVersion() string {
return extractor.schemaVersion
}
func (extractor fakeExtractor) ReferenceSlots() []ReferenceSlot {
return nil
}
func (extractor fakeExtractor) Validators() []Validator {
return extractor.validators
}
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{
{
Index: 0,
ExtractorKey: extractor.key,
ArtifactType: extractor.artifactType,
SchemaVersion: extractor.schemaVersion,
Payload: payload,
SourceRefs: []source.SourceRef{
{
SourceID: req.Source.ID,
StartUnitID: units[0].ID,
EndUnitID: units[len(units)-1].ID,
},
},
Output: ExtractOutput{
Schema: ResponseSchema{ID: "schema-id", Name: "schema-name", Version: "v1"},
Payload: RawPayload{
Content: append([]byte(nil), payload...),
MediaType: "application/json",
},
},
}, nil
@@ -594,12 +517,14 @@ func (merger fakeMerger) Key() string {
}
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
output := req.ExtractOutputs[0]
return MergeResult{Output: MergeOutput{
LaneID: req.LaneID,
MergerKey: merger.key,
SourceID: output.SourceID,
Schema: output.Schema,
Payload: cloneTestRawPayload(output.Payload),
}}, nil
}
type fakeNormalizer struct {
@@ -615,7 +540,33 @@ func (normalizer fakeNormalizer) ReferenceSlots() []ReferenceSlot {
}
func (normalizer fakeNormalizer) Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error) {
return NormalizeResult{Candidates: req.Candidates}, nil
return NormalizeResult{Output: NormalizeOutput{
LaneID: req.LaneID,
NormalizerKey: normalizer.key,
SourceID: req.MergeOutput.SourceID,
Schema: req.MergeOutput.Schema,
Payload: cloneTestRawPayload(req.MergeOutput.Payload),
}}, nil
}
func cloneTestRawPayload(payload RawPayload) RawPayload {
return RawPayload{
Content: append([]byte(nil), payload.Content...),
MediaType: payload.MediaType,
Metadata: cloneTestMetadata(payload.Metadata),
Warnings: append([]Warning(nil), payload.Warnings...),
}
}
func cloneTestMetadata(metadata map[string]any) map[string]any {
if len(metadata) == 0 {
return nil
}
out := make(map[string]any, len(metadata))
for key, value := range metadata {
out[key] = value
}
return out
}
type fakeValidator struct {
@@ -665,7 +616,7 @@ func (encoder fakeOutputEncoder) Encode(ctx context.Context, req OutputRequest)
{
Name: "artifacts/generic.json",
ContentType: "application/json",
Bytes: []byte(`{"run_id":"` + req.Manifest.RunID + `","approved_count":1}`),
Bytes: []byte(`{"run_id":"` + req.Manifest.RunID + `","output_count":1}`),
},
},
}, nil