diff --git a/internal/framework/contracts/composition_test.go b/internal/framework/contracts/composition_test.go index 19349a6..85a5e15 100644 --- a/internal/framework/contracts/composition_test.go +++ b/internal/framework/contracts/composition_test.go @@ -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 +} diff --git a/internal/framework/contracts/contracts.go b/internal/framework/contracts/contracts.go index e19dbee..421c263 100644 --- a/internal/framework/contracts/contracts.go +++ b/internal/framework/contracts/contracts.go @@ -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) +} diff --git a/internal/framework/contracts/contracts_test.go b/internal/framework/contracts/contracts_test.go index e091686..8c59dcc 100644 --- a/internal/framework/contracts/contracts_test.go +++ b/internal/framework/contracts/contracts_test.go @@ -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 +}