463 lines
14 KiB
Go
463 lines
14 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
)
|
|
|
|
func TestWalkingSkeletonFixture(t *testing.T) {
|
|
inputBytes := readTestFixture(t, "testdata/walking_skeleton_input.json")
|
|
expectedBytes := readTestFixture(t, "testdata/walking_skeleton_output.json")
|
|
llmClient := &walkingSkeletonLLMClient{}
|
|
|
|
resolved, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, walkingSkeletonCatalog(t))
|
|
if err != nil {
|
|
t.Fatalf("ResolvePipeline() error = %v, want nil", err)
|
|
}
|
|
|
|
output, err := New(walkingSkeletonRegistries(t)).Run(context.Background(), RunInput{
|
|
Pipeline: resolved,
|
|
SourceID: "fixture-source",
|
|
Path: "walking_skeleton_input.json",
|
|
RawInput: inputBytes,
|
|
LLMClient: llmClient,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v, want nil", err)
|
|
}
|
|
if len(output.OutputFiles) != 1 {
|
|
t.Fatalf("len(OutputFiles) = %d, want 1", len(output.OutputFiles))
|
|
}
|
|
if output.OutputFiles[0].ContentType != "application/json" {
|
|
t.Fatalf("ContentType = %q, want application/json", output.OutputFiles[0].ContentType)
|
|
}
|
|
assertStructuralJSONEqual(t, output.OutputFiles[0].Bytes, expectedBytes)
|
|
if llmClient.calls != 3 {
|
|
t.Fatalf("LLM calls = %d, want extractor calls plus normalizer call", llmClient.calls)
|
|
}
|
|
}
|
|
|
|
func TestWalkingSkeletonResolutionRejectsMissingCapability(t *testing.T) {
|
|
catalog := walkingSkeletonCatalog(t)
|
|
catalog.Extractors = NewExtractorRegistry()
|
|
if err := catalog.Extractors.RegisterWithSpec(ModuleSpec{
|
|
Key: "fake/extract",
|
|
Stage: StageExtract,
|
|
Requires: []string{"missing"},
|
|
Provides: []string{"fake_artifacts"},
|
|
}, func() (contracts.Extractor, error) {
|
|
return walkingSkeletonExtractor{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("RegisterWithSpec() error = %v, want nil", err)
|
|
}
|
|
|
|
_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, catalog)
|
|
if err == nil {
|
|
t.Fatal("ResolvePipeline() error = nil, want error")
|
|
}
|
|
if !strings.Contains(err.Error(), "missing") {
|
|
t.Fatalf("ResolvePipeline() error = %q, want missing capability", err.Error())
|
|
}
|
|
}
|
|
|
|
func TestWalkingSkeletonResolutionRejectsUnknownOnlyLane(t *testing.T) {
|
|
_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{Only: []string{"missing"}}, walkingSkeletonCatalog(t))
|
|
if err == nil {
|
|
t.Fatal("ResolvePipeline() error = nil, want error")
|
|
}
|
|
if !strings.Contains(err.Error(), "missing") || !strings.Contains(err.Error(), "not declared") {
|
|
t.Fatalf("ResolvePipeline() error = %q, want unknown lane error", err.Error())
|
|
}
|
|
}
|
|
|
|
func walkingSkeletonProfile() PipelineProfile {
|
|
return PipelineProfile{
|
|
ID: "walking-skeleton",
|
|
Input: Binding("fake/input"),
|
|
Chunk: Binding("fake/chunk"),
|
|
Output: Binding("json"),
|
|
Artifacts: map[string]ArtifactLaneProfile{
|
|
"events": {Extract: Binding("fake/extract")},
|
|
},
|
|
}
|
|
}
|
|
|
|
func walkingSkeletonCatalog(t *testing.T) ModuleCatalog {
|
|
t.Helper()
|
|
|
|
catalog := ModuleCatalog{
|
|
Inputs: NewInputAdapterRegistry(),
|
|
Chunkers: NewChunkerRegistry(),
|
|
Extractors: NewExtractorRegistry(),
|
|
Mergers: NewMergerRegistry(),
|
|
Normalizers: NewNormalizerRegistry(),
|
|
ValidatorChains: NewValidatorChainRegistry(),
|
|
Outputs: NewOutputEncoderRegistry(),
|
|
}
|
|
if err := catalog.Inputs.RegisterWithSpec(ModuleSpec{
|
|
Key: "fake/input",
|
|
Stage: StageInput,
|
|
Provides: []string{"plain_text"},
|
|
}, func() (contracts.InputAdapter, error) {
|
|
return walkingSkeletonInput{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register fake input: %v", err)
|
|
}
|
|
if err := catalog.Chunkers.RegisterWithSpec(ModuleSpec{
|
|
Key: "fake/chunk",
|
|
Stage: StageChunk,
|
|
Requires: []string{"plain_text"},
|
|
Provides: []string{"chunks"},
|
|
}, func() (contracts.Chunker, error) {
|
|
return walkingSkeletonChunker{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register fake chunker: %v", err)
|
|
}
|
|
if err := catalog.Extractors.RegisterWithSpec(ModuleSpec{
|
|
Key: "fake/extract",
|
|
Stage: StageExtract,
|
|
Requires: []string{"chunks"},
|
|
Provides: []string{"fake_artifacts"},
|
|
}, func() (contracts.Extractor, error) {
|
|
return walkingSkeletonExtractor{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register fake extractor: %v", err)
|
|
}
|
|
if err := catalog.Mergers.RegisterWithSpec(ModuleSpec{
|
|
Key: DefaultMergeModule,
|
|
Stage: StageMerge,
|
|
Requires: []string{"fake_artifacts"},
|
|
}, func() (contracts.Merger, error) {
|
|
return walkingSkeletonMerger{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register append-order merger: %v", err)
|
|
}
|
|
if err := catalog.Normalizers.RegisterWithSpec(ModuleSpec{
|
|
Key: DefaultNormalizeModule,
|
|
Stage: StageNormalize,
|
|
}, func() (contracts.Normalizer, error) {
|
|
return walkingSkeletonNormalizer{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register no-op normalizer: %v", err)
|
|
}
|
|
if err := catalog.Outputs.RegisterWithSpec(ModuleSpec{
|
|
Key: "json",
|
|
Stage: StageOutput,
|
|
}, func() (contracts.OutputEncoder, error) {
|
|
return walkingSkeletonOutput{}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register fake output: %v", err)
|
|
}
|
|
return catalog
|
|
}
|
|
|
|
func walkingSkeletonRegistries(t *testing.T) Registries {
|
|
t.Helper()
|
|
|
|
catalog := walkingSkeletonCatalog(t)
|
|
return Registries{
|
|
Inputs: catalog.Inputs,
|
|
Chunkers: catalog.Chunkers,
|
|
Extractors: catalog.Extractors,
|
|
Mergers: catalog.Mergers,
|
|
Normalizers: catalog.Normalizers,
|
|
Outputs: catalog.Outputs,
|
|
}
|
|
}
|
|
|
|
type walkingSkeletonInput struct{}
|
|
|
|
func (input walkingSkeletonInput) Key() string {
|
|
return "fake/input"
|
|
}
|
|
|
|
func (input walkingSkeletonInput) Parse(ctx context.Context, req contracts.ParseRequest) (*source.SourceDocument, error) {
|
|
var fixture struct {
|
|
ID string `json:"id"`
|
|
Units []struct {
|
|
ID int `json:"id"`
|
|
Text string `json:"text"`
|
|
} `json:"units"`
|
|
}
|
|
if err := json.Unmarshal(req.Raw, &fixture); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
units := make([]source.SourceUnit, 0, len(fixture.Units))
|
|
for _, unit := range fixture.Units {
|
|
units = append(units, source.SourceUnit{
|
|
ID: unit.ID,
|
|
Kind: "unit",
|
|
Text: unit.Text,
|
|
})
|
|
}
|
|
return &source.SourceDocument{
|
|
ID: fixture.ID,
|
|
Kind: "fixture",
|
|
Format: "application/json",
|
|
Digest: rawDigest(req.Raw),
|
|
Units: units,
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonChunker struct{}
|
|
|
|
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")
|
|
}
|
|
return contracts.ChunkResult{
|
|
Chunks: []contracts.SourceChunk{
|
|
{
|
|
ID: req.Source.ID + ":chunk:0",
|
|
SourceID: req.Source.ID,
|
|
Index: 0,
|
|
StartUnitID: req.Source.Units[0].ID,
|
|
EndUnitID: req.Source.Units[1].ID,
|
|
Content: []byte(`{"units":[1,2]}`),
|
|
MediaType: "application/json",
|
|
Units: append([]source.SourceUnit(nil), req.Source.Units[:2]...),
|
|
},
|
|
{
|
|
ID: req.Source.ID + ":chunk:1",
|
|
SourceID: req.Source.ID,
|
|
Index: 1,
|
|
StartUnitID: req.Source.Units[2].ID,
|
|
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
|
Content: []byte(`{"units":[3]}`),
|
|
MediaType: "application/json",
|
|
Units: append([]source.SourceUnit(nil), req.Source.Units[2:]...),
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonExtractor struct{}
|
|
|
|
func (extractor walkingSkeletonExtractor) Key() string {
|
|
return "fake/extract"
|
|
}
|
|
|
|
func (extractor walkingSkeletonExtractor) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return nil
|
|
}
|
|
|
|
func (extractor walkingSkeletonExtractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
|
|
var response struct {
|
|
Call int `json:"call"`
|
|
}
|
|
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: "fake/extract",
|
|
PromptID: "fake.event",
|
|
PromptVersion: "v1",
|
|
}, &response); err != nil {
|
|
return contracts.ExtractionResult{}, err
|
|
}
|
|
|
|
payload, err := json.Marshal(map[string]any{
|
|
"chunk_id": req.Chunk.ID,
|
|
"llm_call": response.Call,
|
|
"text": chunkText(req.Chunk.Units),
|
|
})
|
|
if err != nil {
|
|
return contracts.ExtractionResult{}, err
|
|
}
|
|
|
|
return contracts.ExtractionResult{
|
|
Output: contracts.ExtractOutput{
|
|
Schema: contracts.ResponseSchema{ID: "fake_event", Name: "fake_event", Version: "v1"},
|
|
Payload: contracts.RawPayload{
|
|
Content: payload,
|
|
MediaType: "application/json",
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonLLMClient struct {
|
|
calls int
|
|
}
|
|
|
|
func (client *walkingSkeletonLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
client.calls++
|
|
if response, ok := out.(*struct {
|
|
Call int `json:"call"`
|
|
}); ok {
|
|
response.Call = client.calls
|
|
}
|
|
content, err := json.Marshal(map[string]any{"call": client.calls})
|
|
if err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
return contracts.StructuredCompletionResponse{
|
|
Content: content,
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonMerger struct{}
|
|
|
|
func (merger walkingSkeletonMerger) Key() string {
|
|
return DefaultMergeModule
|
|
}
|
|
|
|
func (merger walkingSkeletonMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
|
|
outputs := make([]json.RawMessage, 0, len(req.ExtractOutputs))
|
|
for _, output := range req.ExtractOutputs {
|
|
outputs = append(outputs, json.RawMessage(output.Payload.Content))
|
|
}
|
|
content, err := json.Marshal(map[string]any{"outputs": outputs})
|
|
if err != nil {
|
|
return contracts.MergeResult{}, err
|
|
}
|
|
return contracts.MergeResult{
|
|
Output: contracts.MergeOutput{
|
|
LaneID: req.LaneID,
|
|
SourceID: req.Source.ID,
|
|
Schema: contracts.ResponseSchema{ID: "fake_event", Name: "fake_event", Version: "v1"},
|
|
Payload: contracts.RawPayload{
|
|
Content: content,
|
|
MediaType: "application/json",
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonNormalizer struct{}
|
|
|
|
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) {
|
|
var response struct {
|
|
Call int `json:"call"`
|
|
}
|
|
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: "fake/normalize",
|
|
PromptID: "fake.normalize",
|
|
PromptVersion: "v1",
|
|
}, &response); err != nil {
|
|
return contracts.NormalizeResult{}, err
|
|
}
|
|
return contracts.NormalizeResult{
|
|
Output: contracts.NormalizeOutput{
|
|
LaneID: req.LaneID,
|
|
SourceID: req.MergeOutput.SourceID,
|
|
Schema: req.MergeOutput.Schema,
|
|
Payload: req.MergeOutput.Payload,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type walkingSkeletonOutput struct{}
|
|
|
|
func (output walkingSkeletonOutput) Key() string {
|
|
return "json"
|
|
}
|
|
|
|
func (output walkingSkeletonOutput) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
|
|
type rawOutput struct {
|
|
LaneID string `json:"lane_id"`
|
|
NormalizerKey string `json:"normalizer_key"`
|
|
SourceID string `json:"source_id"`
|
|
Schema contracts.ResponseSchema `json:"schema"`
|
|
MediaType string `json:"media_type"`
|
|
Content json.RawMessage `json:"content"`
|
|
}
|
|
rawOutputs := make([]rawOutput, 0, len(req.NormalizeOutputs))
|
|
for _, output := range req.NormalizeOutputs {
|
|
rawOutputs = append(rawOutputs, rawOutput{
|
|
LaneID: output.LaneID,
|
|
NormalizerKey: output.NormalizerKey,
|
|
SourceID: output.SourceID,
|
|
Schema: output.Schema,
|
|
MediaType: output.Payload.MediaType,
|
|
Content: json.RawMessage(output.Payload.Content),
|
|
})
|
|
}
|
|
encoded, err := json.Marshal(struct {
|
|
Manifest artifacts.RunManifest `json:"manifest"`
|
|
NormalizeOutputs []rawOutput `json:"normalize_outputs"`
|
|
}{
|
|
Manifest: artifacts.RunManifest{
|
|
PipelineID: req.Manifest.PipelineID,
|
|
PipelineDigest: req.Manifest.PipelineDigest,
|
|
ArtifactLanes: req.Manifest.ArtifactLanes,
|
|
ValidationStatus: req.Manifest.ValidationStatus,
|
|
},
|
|
NormalizeOutputs: rawOutputs,
|
|
})
|
|
if err != nil {
|
|
return contracts.OutputResult{}, err
|
|
}
|
|
return contracts.OutputResult{
|
|
Files: []contracts.OutputFile{
|
|
{Name: "output.json", ContentType: "application/json", Bytes: encoded},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
func readTestFixture(t *testing.T, path string) []byte {
|
|
t.Helper()
|
|
|
|
bytes, err := os.ReadFile(path)
|
|
if err != nil {
|
|
t.Fatalf("read fixture %q: %v", path, err)
|
|
}
|
|
return bytes
|
|
}
|
|
|
|
func assertStructuralJSONEqual(t *testing.T, gotBytes, wantBytes []byte) {
|
|
t.Helper()
|
|
|
|
var got any
|
|
if err := json.Unmarshal(gotBytes, &got); err != nil {
|
|
t.Fatalf("unmarshal actual JSON: %v\n%s", err, gotBytes)
|
|
}
|
|
var want any
|
|
if err := json.Unmarshal(wantBytes, &want); err != nil {
|
|
t.Fatalf("unmarshal expected JSON: %v\n%s", err, wantBytes)
|
|
}
|
|
if !reflect.DeepEqual(got, want) {
|
|
gotFormatted, _ := json.MarshalIndent(got, "", " ")
|
|
wantFormatted, _ := json.MarshalIndent(want, "", " ")
|
|
t.Fatalf("actual JSON:\n%s\nwant:\n%s", gotFormatted, wantFormatted)
|
|
}
|
|
}
|
|
|
|
func chunkText(units []source.SourceUnit) string {
|
|
parts := make([]string, 0, len(units))
|
|
for _, unit := range units {
|
|
parts = append(parts, unit.Text)
|
|
}
|
|
return strings.Join(parts, " ")
|
|
}
|
|
|
|
func rawDigest(raw []byte) string {
|
|
sum := sha256.Sum256(raw)
|
|
return "sha256:" + hex.EncodeToString(sum[:])
|
|
}
|