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{}