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

Connect processing to shared PromptKit backend capacity

Browse Source
This commit is contained in:
Eric Rakestraw
2026-09-24 02:18:31 +00:00
parent a4a72168d5
commit 4812032c9e
26 changed files with 406 additions and 1816 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
internal/framework/proposal_generation/generate.go
View File

@@ -62,7 +62,6 @@ type Request struct {
StageName string `json:"stage_name,omitempty"`
StartIndex int `json:"start_index"`
LLMClient contracts.StructuredLLMClient
Scheduler contracts.LLMScheduler
DiagnosticsDir string
DiagnosticsWriter InteractionDiagnosticsWriter
}
@@ -117,19 +116,11 @@ func GenerateCandidates(ctx context.Context, req Request) (Result, error) {
callErr error
artifacts InteractionArtifacts
)
call := func(callCtx context.Context) error {
completion, callErr = req.LLMClient.CompleteStructured(callCtx, contracts.StructuredCompletionRequest{
PromptID: req.PromptID, PromptVersion: req.PromptVersion,
Variables: req.Variables, Schema: responseschema.CorrectionSetKey,
Context: structuredoutput.DiagnosticContext{Role: "proposal", ModuleKey: req.ModuleKey, ModuleInstance: req.ModuleInstance, SectionIndex: sectionIndexPtr(req.Section)},
}, &response)
return callErr
}
if req.Scheduler != nil {
callErr = req.Scheduler.Run(ctx, call)
} else {
callErr = call(ctx)
}
completion, callErr = req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
PromptID: req.PromptID, PromptVersion: req.PromptVersion,
Variables: req.Variables, Schema: responseschema.CorrectionSetKey,
Context: structuredoutput.DiagnosticContext{Role: "proposal", ModuleKey: req.ModuleKey, ModuleInstance: req.ModuleInstance, SectionIndex: sectionIndexPtr(req.Section)},
}, &response)
requestMetadata := map[string]any{
"module_key": req.ModuleKey,

128
internal/framework/proposal_generation/generate_test.go
View File

@@ -7,17 +7,12 @@ import (
"os"
"path/filepath"
"reflect"
"runtime"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
ck "gitea.maximumdirect.net/eric/ConfigKit"
"gitea.maximumdirect.net/eric/audita/internal/core/config"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
@@ -51,15 +46,6 @@ func (f *fakeStructuredClient) CompleteStructured(ctx context.Context, req contr
return contracts.StructuredCompletionResponse{Prepared: &structuredoutput.PreparedSnapshot{Messages: []structuredoutput.Message{{Role: "user", Content: content.String()}}}}, nil
}
type countingScheduler struct {
runs int
}
func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
s.runs++
return fn(ctx)
}
type captureDiagnosticsWriter struct {
lastStage string
lastRequestMetadata any
@@ -75,45 +61,6 @@ func (w *captureDiagnosticsWriter) WriteInteraction(stage string, requestMetadat
return InteractionArtifacts{}, nil
}
type sleepingStructuredClient struct {
inFlight int32
maxInFlight int32
entered chan struct{}
release chan struct{}
}
func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
_ = req
current := atomic.AddInt32(&c.inFlight, 1)
for {
prior := atomic.LoadInt32(&c.maxInFlight)
if current <= prior || atomic.CompareAndSwapInt32(&c.maxInFlight, prior, current) {
break
}
}
if c.entered != nil {
c.entered <- struct{}{}
}
select {
case <-c.release:
case <-ctx.Done():
atomic.AddInt32(&c.inFlight, -1)
return contracts.StructuredCompletionResponse{}, ctx.Err()
}
atomic.AddInt32(&c.inFlight, -1)
target, ok := out.(*StructuredCorrectionSet)
if !ok {
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output type")
}
*target = StructuredCorrectionSet{
Corrections: []StructuredCorrectionProposal{
{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9},
},
}
return contracts.StructuredCompletionResponse{}, nil
}
func defaultRequest(t *testing.T) Request {
t.Helper()
cfg := config.Config{}
@@ -401,26 +348,6 @@ func TestGenerateCandidatesDiagnosticsWrittenAndRedacted(t *testing.T) {
}
}
func TestGenerateCandidatesSchedulerUsage(t *testing.T) {
client := &fakeStructuredClient{
responses: []StructuredCorrectionSet{
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9}}},
},
}
scheduler := &countingScheduler{}
req := defaultRequest(t)
req.LLMClient = client
req.Scheduler = scheduler
_, err := GenerateCandidates(context.Background(), req)
if err != nil {
t.Fatalf("GenerateCandidates error: %v", err)
}
if scheduler.runs != 1 {
t.Fatalf("expected scheduler to run once, got %d", scheduler.runs)
}
}
func TestGenerateCandidatesClientError(t *testing.T) {
client := &fakeStructuredClient{err: errors.New("boom")}
req := defaultRequest(t)
@@ -442,58 +369,3 @@ func TestGenerateCandidatesClientError(t *testing.T) {
t.Fatalf("expected error diagnostics artifact under %s", req.DiagnosticsDir)
}
}
func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
scheduler, err := llm.NewScheduler(2)
if err != nil {
t.Fatalf("NewScheduler: %v", err)
}
client := &sleepingStructuredClient{
entered: make(chan struct{}, 16),
release: make(chan struct{}),
}
baseReq := defaultRequest(t)
baseReq.LLMClient = client
baseReq.Scheduler = scheduler
var wg sync.WaitGroup
for i := 0; i < 8; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
req := baseReq
req.StartIndex = idx
if _, runErr := GenerateCandidates(context.Background(), req); runErr != nil {
t.Errorf("GenerateCandidates[%d] error: %v", idx, runErr)
}
}(i)
}
waitForEntries(t, client.entered, 2)
close(client.release)
wg.Wait()
if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
t.Fatalf("expected scheduler cap <= 2, got %d", got)
}
if got := atomic.LoadInt32(&client.maxInFlight); got < 2 {
t.Fatalf("expected observed concurrency of at least 2, got %d", got)
}
}
func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
t.Helper()
deadline := time.Now().Add(300 * time.Millisecond)
got := 0
for got < want && time.Now().Before(deadline) {
select {
case <-entered:
got++
default:
runtime.Gosched()
}
}
if got < want {
t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
}
}

1
internal/framework/proposal_generation/module_proposal.go
View File

@@ -68,7 +68,6 @@ func ExecuteModuleProposal(ctx context.Context, req ModuleProposalRequest) (cont
StageName: stagename.ModuleProposal(req.ProposalRequest.RunSpec.InstanceName, sectionIndexPtr(req.ProposalRequest.Section)),
StartIndex: 0,
LLMClient: req.ProposalRequest.LLMClient,
Scheduler: req.ProposalRequest.LLMScheduler,
DiagnosticsDir: req.ProposalRequest.DiagnosticsDir,
})
if err != nil {