notarius/internal/framework/contracts/composition_test.go

package contracts_test

import (
	"context"
	"encoding/json"
	"errors"
	"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"
)

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.StructuredLLMClient = compositionLLMClient{}
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 {
		t.Fatalf("Parse() error = %v, want nil", err)
	}
	if err := source.ValidateDocument(doc); err != nil {
		t.Fatalf("ValidateDocument() error = %v, want nil", err)
	}

	chunking, err := chunker.Chunk(ctx, contracts.ChunkRequest{
		Source:    doc,
		LLMClient: compositionLLMClient{},
		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)
	}
	if len(extraction.Candidates) != 1 {
		t.Fatalf("len(Candidates) = %d, want 1", len(extraction.Candidates))
	}

	candidate := extraction.Candidates[0]
	for _, ref := range candidate.SourceRefs {
		if err := source.ValidateRef(doc, ref); err != nil {
			t.Fatalf("ValidateRef() error = %v, want nil", err)
		}
	}

	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: normalize.Candidates,
	})
	if err != nil {
		t.Fatalf("Validate() error = %v, want nil", err)
	}
	if len(validation.Decisions) != 1 {
		t.Fatalf("len(Decisions) = %d, want 1", len(validation.Decisions))
	}

	decision := validation.Decisions[0]
	if !decision.Approved {
		t.Fatal("Approved = false, want true")
	}
	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 len(output.Files) != 1 {
		t.Fatalf("len(Files) = %d, want 1", len(output.Files))
	}
	if output.Files[0].ContentType != "application/json" {
		t.Fatalf("ContentType = %q, want application/json", output.Files[0].ContentType)
	}
	if len(output.Files[0].Bytes) == 0 {
		t.Fatal("len(Bytes) = 0, want encoded bytes")
	}
}

type compositionAdapter struct{}

func (adapter compositionAdapter) Key() string {
	return "generic-input"
}

func (adapter compositionAdapter) Parse(ctx context.Context, req contracts.ParseRequest) (*source.SourceDocument, error) {
	return &source.SourceDocument{
		ID:     req.SourceID,
		Kind:   "document",
		Format: "text/plain",
		Digest: "sha256:abc123",
		Units: []source.SourceUnit{
			{ID: "u1", Kind: "unit", Text: "First source unit."},
			{ID: "u2", Kind: "unit", Text: "Second source unit."},
		},
	}, nil
}

type compositionChunker struct{}

func (chunker compositionChunker) Key() string {
	return "generic-chunker"
}

func (chunker compositionChunker) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

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")
	}
	if req.LLMClient == nil {
		return contracts.ChunkResult{}, errors.New("structured llm client 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 compositionLLMClient struct{}

func (client compositionLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	return contracts.StructuredCompletionResponse{}, nil
}

type compositionExtractor struct{}

func (extractor compositionExtractor) Key() string {
	return "generic-extractor"
}

func (extractor compositionExtractor) ArtifactType() string {
	return "generic-artifact"
}

func (extractor compositionExtractor) SchemaVersion() string {
	return "v1"
}

func (extractor compositionExtractor) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (extractor compositionExtractor) Validators() []contracts.Validator {
	return []contracts.Validator{compositionValidator{}}
}

func (extractor compositionExtractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	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{
			{
				Index:         0,
				ExtractorKey:  extractor.Key(),
				ArtifactType:  extractor.ArtifactType(),
				SchemaVersion: extractor.SchemaVersion(),
				Payload:       json.RawMessage(`{"value":"example"}`),
				SourceRefs: []source.SourceRef{
					{
						SourceID:    req.Source.ID,
						StartUnitID: units[0].ID,
						EndUnitID:   units[len(units)-1].ID,
					},
				},
			},
		},
	}, 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) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

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 {
	return "generic-validator"
}

func (validator compositionValidator) Validate(ctx context.Context, req contracts.ValidationRequest) (contracts.ValidationResult, error) {
	decisions := make([]contracts.ValidationDecision, 0, len(req.Candidates))
	for _, candidate := range req.Candidates {
		decisions = append(decisions, contracts.ValidationDecision{
			CandidateIndex: candidate.Index,
			Approved:       true,
			ReasonCode:     "accepted",
			Message:        "candidate accepted",
		})
	}

	return contracts.ValidationResult{
		ValidatorName: validator.Name(),
		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{
		Files: []contracts.OutputFile{
			{
				Name:        "artifacts/generic.json",
				ContentType: "application/json",
				Bytes:       encoded,
			},
		},
	}, nil
}