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

Preserve structured repair settings across pipeline boundaries

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-25 23:46:52 +00:00
parent 3d3f16db4a
commit 916d9210fd
7 changed files with 166 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
internal/framework/llm/promptkit_client.go
View File

@@ -144,6 +144,9 @@ func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts.
if req.StructuredOutputRepairAttempts != nil {
inspection, err := c.engine.InspectPrompt(ctx, promptID, strings.TrimSpace(req.PromptVersion))
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return contracts.StructuredCompletionResponse{}, ctxErr
}
return contracts.StructuredCompletionResponse{}, fmt.Errorf("inspect PromptKit prompt %q: %v", promptID, redactPromptKitError(err))
}
contract := inspection.OutputContract

36
internal/framework/llm/promptkit_client_test.go
View File

@@ -1147,19 +1147,31 @@ func TestPromptKitClientTranslatesBackendCapacityExhaustion(t *testing.T) {
}
func TestPromptKitClientContextCancellationIsRespected(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
attempts := 1
for _, test := range []struct {
name string
attempts *int
}{
{name: "during preparation"},
{name: "during override inspection", attempts: &attempts},
} {
t.Run(test.name, func(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
var out map[string]any
_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) {
t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
var out map[string]any
_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
StructuredOutputRepairAttempts: test.attempts,
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) {
t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
}
})
}
}

2
internal/framework/pipeline/merger_registry.go
View File

@@ -85,7 +85,7 @@ func RegisterMergerBuilder[T any](registry *MergerRegistry, spec ModuleSpec, val
}
outputs[i] = contracts.ExtractArtifact[T]{LaneID: output.LaneID, ExtractorKey: output.ExtractorKey, SourceID: output.SourceID, ChunkID: output.ChunkID, ChunkIndex: output.ChunkIndex, ChunkRef: output.ChunkRef, Value: value}
}
result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, Metadata: request.Metadata})
result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
if err != nil {
return erasedTypedResult{}, err
}

2
internal/framework/pipeline/normalizer_registry.go
View File

@@ -76,7 +76,7 @@ func RegisterNormalizerBuilder[T any](registry *NormalizerRegistry, spec ModuleS
if err != nil {
return erasedTypedResult{}, err
}
result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, Metadata: request.Metadata})
result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
if err != nil {
return erasedTypedResult{}, err
}

4
internal/framework/pipeline/runner.go
View File

@@ -493,9 +493,9 @@ func (r *Runner) validateChunks(ctx context.Context, doc *source.SourceDocument,
}
switch item.resolved.Target {
case ValidatorTargetChunk:
result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks})
result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks})
case ValidatorTargetSerialized:
result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)})
result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)})
default:
return nil, nil, fmt.Errorf("validator %q is incompatible with chunk validation", binding.Module)
}

40
internal/framework/pipeline/runner_terminal_debug_test.go
View File

@@ -41,6 +41,10 @@ type observingChunkValidator struct {
request contracts.ChunkValidationRequest
}
type observingSerializedValidator struct {
request contracts.SerializedValidationRequest
}
func (*observingChunkValidator) Name() string { return "observing/chunk-validator" }
func (*observingChunkValidator) ExecutionClass() contracts.ExecutionClass {
return contracts.ExecutionClassDeterministic
@@ -50,6 +54,15 @@ func (v *observingChunkValidator) Validate(_ context.Context, request contracts.
return contracts.ValidationResult{Approved: true}, nil
}
func (*observingSerializedValidator) Name() string { return "observing/serialized-validator" }
func (*observingSerializedValidator) ExecutionClass() contracts.ExecutionClass {
return contracts.ExecutionClassDeterministic
}
func (v *observingSerializedValidator) Validate(_ context.Context, request contracts.SerializedValidationRequest) (contracts.ValidationResult, error) {
v.request = request
return contracts.ValidationResult{Approved: true}, nil
}
func (terminalChunkValidator) Name() string { return "terminal/chunk-validator" }
func (terminalChunkValidator) ExecutionClass() contracts.ExecutionClass {
@@ -168,6 +181,33 @@ func TestRunnerMaterializesAnnotatedPlanBeforeChunkValidation(t *testing.T) {
}
}
func TestRunnerForwardsDetachedRepairAttemptsToChunkValidators(t *testing.T) {
prepared, plan := preparedTerminalDebugPipeline(t)
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
chunkValidator := &observingChunkValidator{}
serializedValidator := &observingSerializedValidator{}
chunkAttempts := 2
serializedAttempts := 3
chunkBinding := Binding(chunkValidator.Name())
chunkBinding.StructuredOutputRepairAttempts = &chunkAttempts
serializedBinding := Binding(serializedValidator.Name())
serializedBinding.StructuredOutputRepairAttempts = &serializedAttempts
prepared.chunkValidators.validators = []preparedValidator{
{resolved: ResolvedValidator{Binding: chunkBinding, Target: ValidatorTargetChunk}, chunk: chunkValidator},
{resolved: ResolvedValidator{Binding: serializedBinding, Target: ValidatorTargetSerialized}, serialized: serializedValidator},
}
if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if got := chunkValidator.request.StructuredOutputRepairAttempts; got == nil || *got != chunkAttempts || got == chunkBinding.StructuredOutputRepairAttempts {
t.Fatalf("chunk validator repair attempts = %v, want detached value %d", got, chunkAttempts)
}
if got := serializedValidator.request.StructuredOutputRepairAttempts; got == nil || *got != serializedAttempts || got == serializedBinding.StructuredOutputRepairAttempts {
t.Fatalf("serialized validator repair attempts = %v, want detached value %d", got, serializedAttempts)
}
}
func TestRunnerRetriesMalformedPlanWithDebugEnabled(t *testing.T) {
prepared, plan := preparedTerminalDebugPipeline(t)
plan.Ranges[0].Annotations = source.ChunkAnnotations{"broken": []byte(`{"value":`)}

95
internal/framework/pipeline/typed_registry_request_test.go Normal file
View File

@@ -0,0 +1,95 @@
package pipeline
import (
"context"
"testing"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)
type repairObservingNotesMerger struct {
attempts *int
}
func (*repairObservingNotesMerger) Key() string { return "test/repair-observing-merge" }
func (m *repairObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) {
m.attempts = request.StructuredOutputRepairAttempts
return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merged"}}}, nil
}
type repairObservingNotesNormalizer struct {
attempts *int
}
func (*repairObservingNotesNormalizer) Key() string { return "test/repair-observing-normalize" }
func (*repairObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (n *repairObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) {
n.attempts = request.StructuredOutputRepairAttempts
return contracts.TypedNormalizeResult[codecNotes]{Value: request.MergeOutput.Value}, nil
}
func TestTypedRegistryErasurePreservesStructuredOutputRepairAttempts(t *testing.T) {
const artifactKind contracts.ArtifactKind = "test/notes"
t.Run("merge", func(t *testing.T) {
implementation := &repairObservingNotesMerger{}
registry := NewMergerRegistry()
if err := RegisterMerger(registry, ModuleSpec{Key: implementation.Key(), Stage: StageMerge, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Merger[codecNotes], error) {
return implementation, nil
}); err != nil {
t.Fatalf("RegisterMerger() error = %v", err)
}
entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
if !ok {
t.Fatal("typed merger entry missing")
}
built, err := entry.builder(BuildRequest{})
if err != nil {
t.Fatalf("builder() error = %v", err)
}
attempts := 2
_, err = entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{
ExtractOutputs: []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extracted"}}}},
StructuredOutputRepairAttempts: &attempts,
})
if err != nil {
t.Fatalf("merge() error = %v", err)
}
if implementation.attempts == nil || *implementation.attempts != attempts {
t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
}
})
t.Run("normalize", func(t *testing.T) {
implementation := &repairObservingNotesNormalizer{}
registry := NewNormalizerRegistry()
if err := RegisterNormalizer(registry, ModuleSpec{Key: implementation.Key(), Stage: StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Normalizer[codecNotes], error) {
return implementation, nil
}); err != nil {
t.Fatalf("RegisterNormalizer() error = %v", err)
}
entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
if !ok {
t.Fatal("typed normalizer entry missing")
}
built, err := entry.builder(BuildRequest{})
if err != nil {
t.Fatalf("builder() error = %v", err)
}
attempts := 3
_, err = entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{
MergeOutput: contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merged"}}},
StructuredOutputRepairAttempts: &attempts,
})
if err != nil {
t.Fatalf("normalize() error = %v", err)
}
if implementation.attempts == nil || *implementation.attempts != attempts {
t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
}
})
}