177 lines
8.5 KiB
Go
177 lines
8.5 KiB
Go
package integration_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/config"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
frameworkllm "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
|
)
|
|
|
|
const (
|
|
concurrentChunkerKey = "test/concurrent-chunks"
|
|
concurrentExtractorKey = "test/concurrent-extractor"
|
|
concurrentValidatorKey = "test/concurrent-validator"
|
|
)
|
|
|
|
type integrationConcurrencyTracker struct {
|
|
extractActive atomic.Int32
|
|
extractMaximum atomic.Int32
|
|
providerActive atomic.Int32
|
|
providerMaximum atomic.Int32
|
|
providerCalls atomic.Int32
|
|
validatorCalls atomic.Int32
|
|
}
|
|
|
|
func updateMaximum(maximum *atomic.Int32, current int32) {
|
|
for {
|
|
seen := maximum.Load()
|
|
if current <= seen || maximum.CompareAndSwap(seen, current) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
type instrumentedProvider struct {
|
|
tracker *integrationConcurrencyTracker
|
|
}
|
|
|
|
func (p instrumentedProvider) CompleteStructured(ctx context.Context, _ contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
p.tracker.providerCalls.Add(1)
|
|
current := p.tracker.providerActive.Add(1)
|
|
defer p.tracker.providerActive.Add(-1)
|
|
updateMaximum(&p.tracker.providerMaximum, current)
|
|
select {
|
|
case <-time.After(5 * time.Millisecond):
|
|
case <-ctx.Done():
|
|
return contracts.StructuredCompletionResponse{}, ctx.Err()
|
|
}
|
|
content := []byte(`{"approved":true}`)
|
|
if out != nil {
|
|
if err := json.Unmarshal(content, out); err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
}
|
|
return contracts.StructuredCompletionResponse{Content: content}, nil
|
|
}
|
|
|
|
type concurrentChunker struct{}
|
|
|
|
func (concurrentChunker) Key() string { return concurrentChunkerKey }
|
|
func (concurrentChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
|
func (concurrentChunker) Plan(_ context.Context, request contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
|
ranges := make([]source.ChunkRange, len(request.Source.Units))
|
|
for i, unit := range request.Source.Units {
|
|
ranges[i] = source.ChunkRange{StartUnitID: unit.ID, EndUnitID: unit.ID}
|
|
}
|
|
return contracts.ChunkPlanResult{Plan: source.ChunkPlan{SourceDigest: request.Source.Digest, Ranges: ranges}}, nil
|
|
}
|
|
|
|
type concurrentExtractor struct {
|
|
client contracts.StructuredLLMClient
|
|
tracker *integrationConcurrencyTracker
|
|
failed atomic.Bool
|
|
}
|
|
|
|
func (*concurrentExtractor) Key() string { return concurrentExtractorKey }
|
|
func (*concurrentExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
|
func (e *concurrentExtractor) Extract(ctx context.Context, request contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
|
|
current := e.tracker.extractActive.Add(1)
|
|
defer e.tracker.extractActive.Add(-1)
|
|
updateMaximum(&e.tracker.extractMaximum, current)
|
|
var response map[string]any
|
|
if _, err := e.client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: concurrentExtractorKey}, &response); err != nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, err
|
|
}
|
|
if e.failed.CompareAndSwap(false, true) {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, errors.New("retry requested")
|
|
}
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{}}}, nil
|
|
}
|
|
|
|
type concurrentValidator struct {
|
|
client contracts.StructuredLLMClient
|
|
tracker *integrationConcurrencyTracker
|
|
}
|
|
|
|
func (*concurrentValidator) Name() string { return concurrentValidatorKey }
|
|
func (*concurrentValidator) ExecutionClass() contracts.ExecutionClass {
|
|
return contracts.ExecutionClassLLMBacked
|
|
}
|
|
func (v *concurrentValidator) Validate(ctx context.Context, _ contracts.TypedValidationRequest[dnd.SpellList]) (contracts.ValidationResult, error) {
|
|
v.tracker.validatorCalls.Add(1)
|
|
var response map[string]any
|
|
if _, err := v.client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: concurrentValidatorKey}, &response); err != nil {
|
|
return contracts.ValidationResult{}, err
|
|
}
|
|
return contracts.ValidationResult{Approved: true}, nil
|
|
}
|
|
|
|
func TestRunnerIndependentlyBoundsWorkersAndProviderCallsAcrossRegisteredModules(t *testing.T) {
|
|
tracker := &integrationConcurrencyTracker{}
|
|
scheduler, err := frameworkllm.NewScheduler(2)
|
|
if err != nil {
|
|
t.Fatalf("NewScheduler() error = %v", err)
|
|
}
|
|
client := frameworkllm.NewScheduledClient(instrumentedProvider{tracker: tracker}, scheduler)
|
|
catalog := dndSpellsTestCatalog(t, dndSpellsCatalogSpecs{})
|
|
if err := catalog.Chunkers.RegisterWithSpec(pipeline.ModuleSpec{Key: concurrentChunkerKey, Stage: pipeline.StageChunk, ExecutionClass: contracts.ExecutionClassDeterministic, Requires: []string{"source.transcript"}, Provides: []string{"chunks"}}, func() (contracts.Chunker, error) { return concurrentChunker{}, nil }); err != nil {
|
|
t.Fatalf("register chunker: %v", err)
|
|
}
|
|
validateOptions := func(options map[string]any) error { return pipeline.RejectUnknownOptions(options) }
|
|
if err := pipeline.RegisterExtractorBuilder(catalog.Extractors, pipeline.ModuleSpec{Key: concurrentExtractorKey, Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: dnd.SpellListKind, Requires: []string{"chunks", "source.transcript"}, Provides: []string{"dnd.spell_casts"}}, validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
|
|
return &concurrentExtractor{client: request.Dependencies.LLM, tracker: tracker}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register extractor: %v", err)
|
|
}
|
|
validators := pipeline.NewValidatorRegistry()
|
|
catalog.Validators = validators
|
|
if err := pipeline.RegisterTypedValidatorBuilder(validators, dnd.SpellListKind, pipeline.ValidatorSpec{Key: concurrentValidatorKey, ExecutionClass: contracts.ExecutionClassLLMBacked}, validateOptions, func(request pipeline.BuildRequest) (contracts.TypedValidator[dnd.SpellList], error) {
|
|
return &concurrentValidator{client: request.Dependencies.LLM, tracker: tracker}, nil
|
|
}); err != nil {
|
|
t.Fatalf("register validator: %v", err)
|
|
}
|
|
|
|
cfg := loadDNDSpellsPipelineConfig(t)
|
|
profile := cfg.Pipelines["dnd-spells-fixture"]
|
|
profile.Chunk = pipeline.Binding(concurrentChunkerKey)
|
|
validatorOverride := pipeline.ValidatorOverride{Set: true, Validators: []pipeline.ModuleBinding{pipeline.Binding(concurrentValidatorKey)}}
|
|
profile.Artifacts = map[string]pipeline.ArtifactLaneProfile{
|
|
"alpha": {Extract: pipeline.ModuleBinding{Module: concurrentExtractorKey, Retries: 1, Validators: validatorOverride}, Merge: pipeline.Binding(pipeline.DefaultMergeModule), Normalize: pipeline.Binding(pipeline.DefaultNormalizeModule)},
|
|
"beta": {Extract: pipeline.ModuleBinding{Module: concurrentExtractorKey, Retries: 1, Validators: validatorOverride}, Merge: pipeline.Binding(pipeline.DefaultMergeModule), Normalize: pipeline.Binding(pipeline.DefaultNormalizeModule)},
|
|
}
|
|
cfg.Pipelines["dnd-spells-fixture"] = profile
|
|
resolved, err := cfg.Resolve(config.ResolveInput{PipelineID: "dnd-spells-fixture", Catalog: catalog})
|
|
if err != nil {
|
|
t.Fatalf("Resolve() error = %v", err)
|
|
}
|
|
registries := pipeline.Registries{Inputs: catalog.Inputs, Chunkers: catalog.Chunkers, ArtifactCodecs: catalog.ArtifactCodecs, ArtifactEvidence: catalog.ArtifactEvidence, Extractors: catalog.Extractors, Mergers: catalog.Mergers, Normalizers: catalog.Normalizers, Validators: catalog.Validators, ValidatorChains: catalog.ValidatorChains, Outputs: catalog.Outputs}
|
|
output, err := runPreparedPipeline(t, registries, resolved.ResolvedPipeline, client, pipeline.RunInput{RawInput: readDNDSpellsFixture(t), ExtractWorkers: 3})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != 2 {
|
|
t.Fatalf("normalize outputs = %d, want 2", len(output.NormalizeOutputs))
|
|
}
|
|
if got := tracker.extractMaximum.Load(); got != 3 {
|
|
t.Fatalf("maximum extract jobs = %d, want 3", got)
|
|
}
|
|
if got := tracker.providerMaximum.Load(); got != 2 {
|
|
t.Fatalf("maximum provider calls = %d, want 2", got)
|
|
}
|
|
if got := tracker.validatorCalls.Load(); got == 0 {
|
|
t.Fatal("LLM-backed validator was not called")
|
|
}
|
|
if got := tracker.providerCalls.Load(); got <= tracker.validatorCalls.Load() {
|
|
t.Fatalf("provider calls = %d, validator calls = %d; want extractor, retry, and validator calls", got, tracker.validatorCalls.Load())
|
|
}
|
|
}
|