Added chunking logic to modules and added corresponding regression tests

This commit is contained in:
2026-05-12 12:25:14 -05:00
parent fb59cb21b9
commit cad172a758
14 changed files with 235 additions and 40 deletions

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

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{

View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -49,7 +50,12 @@ func (m *Module) Validators() []contracts.Validator {
}
func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
messages, err := BuildProposalMessages(req.WorkingTranscript, req.Glossary)
sectionTranscript := transcriptForSection(req.WorkingTranscript, req.Section)
sectionIndex := 0
if req.Section != nil {
sectionIndex = req.Section.Index
}
messages, err := BuildProposalMessages(sectionTranscript, req.Glossary, sectionIndex)
if err != nil {
return nil, err
}
@@ -74,3 +80,16 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
}
return generated.Corrections, nil
}
func transcriptForSection(transcript *schema.Transcript, section *contracts.SectionMetadata) *schema.Transcript {
if section == nil || transcript == nil {
return transcript
}
segments := make([]schema.Segment, 0, len(transcript.Segments))
for _, seg := range transcript.Segments {
if seg.ID >= section.StartSegmentID && seg.ID <= section.EndSegmentID {
segments = append(segments, seg)
}
}
return &schema.Transcript{Segments: segments}
}

View File

@@ -61,7 +61,7 @@ func tinyGlossary() *schema.Glossary {
}
func TestBuildProposalMessagesContainsGlossaryContextAndConstraints(t *testing.T) {
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary())
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary(), 0)
if err != nil {
t.Fatalf("BuildProposalMessages error: %v", err)
}
@@ -72,6 +72,7 @@ func TestBuildProposalMessagesContainsGlossaryContextAndConstraints(t *testing.T
for _, want := range []string{
"Glossary:",
"Transcript section:",
`"section_index": 0`,
`"Aliases":`,
`"Jester"`,
`"Category": "faction"`,

View File

@@ -22,14 +22,14 @@ type promptTranscriptSection struct {
Segments []promptSegment `json:"segments"`
}
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary) ([]contracts.LLMMessage, error) {
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary, sectionIndex int) ([]contracts.LLMMessage, error) {
glossaryJSON, err := json.MarshalIndent(glossary, "", " ")
if err != nil {
return nil, fmt.Errorf("marshal glossary prompt context: %w", err)
}
sectionPayload := promptTranscriptSection{
SectionIndex: 0,
SectionIndex: sectionIndex,
Segments: make([]promptSegment, 0),
}
if transcript != nil {

View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -49,7 +50,12 @@ func (m *Module) Validators() []contracts.Validator {
}
func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
messages, err := BuildProposalMessages(req.WorkingTranscript, req.Glossary)
sectionTranscript := transcriptForSection(req.WorkingTranscript, req.Section)
sectionIndex := 0
if req.Section != nil {
sectionIndex = req.Section.Index
}
messages, err := BuildProposalMessages(sectionTranscript, req.Glossary, sectionIndex)
if err != nil {
return nil, err
}
@@ -74,3 +80,16 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
}
return generated.Corrections, nil
}
func transcriptForSection(transcript *schema.Transcript, section *contracts.SectionMetadata) *schema.Transcript {
if section == nil || transcript == nil {
return transcript
}
segments := make([]schema.Segment, 0, len(transcript.Segments))
for _, seg := range transcript.Segments {
if seg.ID >= section.StartSegmentID && seg.ID <= section.EndSegmentID {
segments = append(segments, seg)
}
}
return &schema.Transcript{Segments: segments}
}

View File

@@ -71,7 +71,7 @@ func tinyGlossary() *schema.Glossary {
}
func TestBuildProposalMessagesConstraints(t *testing.T) {
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary())
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary(), 0)
if err != nil {
t.Fatalf("BuildProposalMessages error: %v", err)
}
@@ -82,6 +82,7 @@ func TestBuildProposalMessagesConstraints(t *testing.T) {
for _, want := range []string{
"Transcript section:",
"Protected glossary/context:",
`"section_index": 0`,
"punctuation, capitalization, spacing, and article cleanup only",
"Do not make word substitutions",
"Do not return speaker, start, or end fields",

View File

@@ -24,14 +24,14 @@ type promptTranscriptSection struct {
// BuildProposalMessages mirrors the Python grammar-module prompt intent:
// punctuation/capitalization/spacing cleanup only, with strict meaning guards.
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary) ([]contracts.LLMMessage, error) {
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary, sectionIndex int) ([]contracts.LLMMessage, error) {
glossaryJSON, err := json.MarshalIndent(glossary, "", " ")
if err != nil {
return nil, fmt.Errorf("marshal glossary prompt context: %w", err)
}
sectionPayload := promptTranscriptSection{
SectionIndex: 0,
SectionIndex: sectionIndex,
Segments: make([]promptSegment, 0),
}
if transcript != nil {

View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -49,7 +50,12 @@ func (m *Module) Validators() []contracts.Validator {
}
func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
messages, err := BuildProposalMessages(req.WorkingTranscript, req.Glossary)
sectionTranscript := transcriptForSection(req.WorkingTranscript, req.Section)
sectionIndex := 0
if req.Section != nil {
sectionIndex = req.Section.Index
}
messages, err := BuildProposalMessages(sectionTranscript, req.Glossary, sectionIndex)
if err != nil {
return nil, err
}
@@ -74,3 +80,16 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
}
return generated.Corrections, nil
}
func transcriptForSection(transcript *schema.Transcript, section *contracts.SectionMetadata) *schema.Transcript {
if section == nil || transcript == nil {
return transcript
}
segments := make([]schema.Segment, 0, len(transcript.Segments))
for _, seg := range transcript.Segments {
if seg.ID >= section.StartSegmentID && seg.ID <= section.EndSegmentID {
segments = append(segments, seg)
}
}
return &schema.Transcript{Segments: segments}
}

View File

@@ -62,7 +62,7 @@ func tinyGlossary() *schema.Glossary {
}
func TestBuildProposalMessagesContainsContextAndConservativeConstraints(t *testing.T) {
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary())
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary(), 0)
if err != nil {
t.Fatalf("BuildProposalMessages error: %v", err)
}
@@ -73,6 +73,7 @@ func TestBuildProposalMessagesContainsContextAndConservativeConstraints(t *testi
for _, want := range []string{
"Protected glossary/context:",
"Transcript section:",
`"section_index": 0`,
`"Aliases":`,
`"Category": "faction"`,
`"Summary": "Guild members"`,

View File

@@ -24,14 +24,14 @@ type promptTranscriptSection struct {
// BuildProposalMessages mirrors the Python homophones-module prompt intent:
// conservative homophone and mistranscription correction only.
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary) ([]contracts.LLMMessage, error) {
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary, sectionIndex int) ([]contracts.LLMMessage, error) {
glossaryJSON, err := json.MarshalIndent(glossary, "", " ")
if err != nil {
return nil, fmt.Errorf("marshal glossary prompt context: %w", err)
}
sectionPayload := promptTranscriptSection{
SectionIndex: 0,
SectionIndex: sectionIndex,
Segments: make([]promptSegment, 0),
}
if transcript != nil {

View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -49,7 +50,12 @@ func (m *Module) Validators() []contracts.Validator {
}
func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
messages, err := BuildProposalMessages(req.WorkingTranscript, req.Glossary)
sectionTranscript := transcriptForSection(req.WorkingTranscript, req.Section)
sectionIndex := 0
if req.Section != nil {
sectionIndex = req.Section.Index
}
messages, err := BuildProposalMessages(sectionTranscript, req.Glossary, sectionIndex)
if err != nil {
return nil, err
}
@@ -74,3 +80,16 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
}
return generated.Corrections, nil
}
func transcriptForSection(transcript *schema.Transcript, section *contracts.SectionMetadata) *schema.Transcript {
if section == nil || transcript == nil {
return transcript
}
segments := make([]schema.Segment, 0, len(transcript.Segments))
for _, seg := range transcript.Segments {
if seg.ID >= section.StartSegmentID && seg.ID <= section.EndSegmentID {
segments = append(segments, seg)
}
}
return &schema.Transcript{Segments: segments}
}

View File

@@ -62,7 +62,7 @@ func tinyGlossary() *schema.Glossary {
}
func TestBuildProposalMessagesContainsContextAndMeaningGuardrails(t *testing.T) {
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary())
msgs, err := BuildProposalMessages(tinyTranscript(), tinyGlossary(), 0)
if err != nil {
t.Fatalf("BuildProposalMessages error: %v", err)
}
@@ -73,6 +73,7 @@ func TestBuildProposalMessagesContainsContextAndMeaningGuardrails(t *testing.T)
for _, want := range []string{
"Protected glossary/context:",
"Transcript section:",
`"section_index": 0`,
`"Aliases":`,
`"Category": "faction"`,
`"Summary": "Guild members"`,

View File

@@ -24,14 +24,14 @@ type promptTranscriptSection struct {
// BuildProposalMessages mirrors the Python spoken_word-module prompt intent:
// conservative dysfluency cleanup with strict semantic preservation.
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary) ([]contracts.LLMMessage, error) {
func BuildProposalMessages(transcript *schema.Transcript, glossary *schema.Glossary, sectionIndex int) ([]contracts.LLMMessage, error) {
glossaryJSON, err := json.MarshalIndent(glossary, "", " ")
if err != nil {
return nil, fmt.Errorf("marshal glossary prompt context: %w", err)
}
sectionPayload := promptTranscriptSection{
SectionIndex: 0,
SectionIndex: sectionIndex,
Segments: make([]promptSegment, 0),
}
if transcript != nil {