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

Added chunking logic to modules and added corresponding regression tests

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-12 12:25:14 -05:00
parent fb59cb21b9
commit cad172a758
14 changed files with 235 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
internal/framework/runner/runner.go
View File

@@ -6,6 +6,7 @@ import (
"path/filepath"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
"gitea.maximumdirect.net/eric/audita/internal/core/config"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
@@ -117,21 +118,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
}
policy := module.ReplacementPolicy()
proposed, err := module.Propose(ctx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: working,
Glossary: input.Glossary,
DiagnosticsDir: input.ProposalDiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: spec.ModuleKey,
InstanceName: spec.InstanceName,
ReplacementPolicy: policy,
},
LLMClient: input.ProposalLLMClient,
LLMScheduler: input.ProposalLLMScheduler,
})
sections, err := chunkWorkingTranscript(input.Config, working)
if err != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
@@ -143,19 +130,56 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, err)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q chunking failed: %w", spec.InstanceName, err)
}
enriched := make([]proposals.EnrichedCorrectionProposal, 0, len(proposed))
for i, p := range proposed {
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: p,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: i,
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
enriched := make([]proposals.EnrichedCorrectionProposal, 0)
nextProposalIndex := 0
for _, section := range sections {
sectionMeta := contracts.SectionMetadataFromSection(section)
sectionTranscript := transcriptFromSection(section)
proposed, proposeErr := module.Propose(ctx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: sectionTranscript,
Glossary: input.Glossary,
Section: &sectionMeta,
DiagnosticsDir: input.ProposalDiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: spec.ModuleKey,
InstanceName: spec.InstanceName,
ReplacementPolicy: policy,
},
LLMClient: input.ProposalLLMClient,
LLMScheduler: input.ProposalLLMScheduler,
})
if proposeErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, proposeErr)
}
for _, p := range proposed {
sectionIndex := sectionMeta.Index
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: p,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex,
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
nextProposalIndex++
}
}
validatorDecisions := make([]ValidatorDecisionRecord, 0)
@@ -272,6 +296,47 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return RunOutput{FinalTranscript: working, ModuleResults: results}, nil
}
func chunkWorkingTranscript(cfg *config.Config, transcript *schema.Transcript) ([]chunking.Section, error) {
if cfg == nil {
if transcript == nil || len(transcript.Segments) == 0 {
return []chunking.Section{}, nil
}
return []chunking.Section{
{
Index: 0,
EstimatedTokens: estimateTranscriptTokens(transcript),
Segments: append([]schema.Segment(nil), transcript.Segments...),
StartSegmentID: transcript.Segments[0].ID,
EndSegmentID: transcript.Segments[len(transcript.Segments)-1].ID,
},
}, nil
}
chunker := chunking.NewChunker(chunking.ChunkingConfig{
MaxSectionTokens: cfg.MaxSectionTokens,
MinSectionTokens: cfg.MinSectionTokens,
TargetSections: cfg.TargetSections,
})
return chunker.ChunkTranscript(transcript)
}
func estimateTranscriptTokens(transcript *schema.Transcript) int {
if transcript == nil || len(transcript.Segments) == 0 {
return 0
}
estimator := chunking.NewSimpleTokenEstimator()
total := 0
for _, segment := range transcript.Segments {
total += estimator.EstimateTokens(segment.Text)
}
return total
}
func transcriptFromSection(section chunking.Section) *schema.Transcript {
segments := make([]schema.Segment, len(section.Segments))
copy(segments, section.Segments)
return &schema.Transcript{Segments: segments}
}
func cloneTranscript(t *schema.Transcript) *schema.Transcript {
if t == nil {
return &schema.Transcript{}

50
internal/framework/runner/runner_test.go
View File

@@ -101,6 +101,56 @@ func TestRunnerModulesRunSequentially(t *testing.T) {
}
}
func TestRunnerProposalsExecutePerChunkSection(t *testing.T) {
cfg := config.Default()
cfg.MaxSectionTokens = 3
cfg.MinSectionTokens = 0
transcript := &schema.Transcript{Segments: []schema.Segment{
{ID: 1, Text: "one two"},
{ID: 2, Text: "three four"},
{ID: 3, Text: "five six"},
}}
seenSections := make([]contracts.SectionMetadata, 0)
seenSegmentCounts := make([]int, 0)
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
if req.Section == nil {
t.Fatalf("expected section metadata on proposal request")
}
seenSections = append(seenSections, *req.Section)
seenSegmentCounts = append(seenSegmentCounts, len(req.WorkingTranscript.Segments))
return nil, nil
}},
}})
out, err := r.Run(context.Background(), RunInput{
Config: &cfg,
Transcript: transcript,
ModuleSpecs: []contracts.ModuleRunSpec{
{ModuleKey: "m", InstanceName: "m"},
},
})
if err != nil {
t.Fatalf("Run error: %v", err)
}
if len(out.ModuleResults) != 1 {
t.Fatalf("expected one module result, got %+v", out.ModuleResults)
}
if len(seenSections) != 3 {
t.Fatalf("expected 3 chunked proposal calls, got %d", len(seenSections))
}
for i, section := range seenSections {
if section.Index != i {
t.Fatalf("expected section index %d, got %+v", i, section)
}
if seenSegmentCounts[i] != 1 {
t.Fatalf("expected one segment per section call, got %d at index %d", seenSegmentCounts[i], i)
}
}
}
func TestRunnerSkippedRecorded(t *testing.T) {
transcript := &schema.Transcript{Segments: []schema.Segment{{ID: 1, Speaker: "A", Start: 0, End: 1, Text: "word word"}}}
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{