Retain injected client diagnostic evidence
This commit is contained in:
@@ -91,8 +91,7 @@ func (c *Coordinator) generate(
|
||||
captureEvidence bool,
|
||||
) (*domain.GenerateResponse, error, GenerationEvidence) {
|
||||
if c.injected != nil {
|
||||
response, err := c.generateInjected(ctx, req)
|
||||
return response, err, GenerationEvidence{}
|
||||
return c.generateInjected(ctx, req, captureEvidence)
|
||||
}
|
||||
prepared, err := c.builtIn.PrepareGeneration(req)
|
||||
if err != nil {
|
||||
@@ -125,6 +124,7 @@ func (c *Coordinator) generate(
|
||||
if captureEvidence {
|
||||
attempt := AttemptEvidence{
|
||||
Dispatched: dispatched,
|
||||
ResponseReceived: outcome.Evidence.ResponseReceived,
|
||||
HTTP: cloneHTTPEvidence(outcome.Evidence),
|
||||
CapacityWait: capacityWait,
|
||||
RequestDuration: generationDuration,
|
||||
@@ -196,7 +196,19 @@ func (c *Coordinator) generate(
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
||||
func (c *Coordinator) generateInjected(
|
||||
ctx context.Context,
|
||||
req domain.GenerateRequest,
|
||||
captureEvidence bool,
|
||||
) (*domain.GenerateResponse, error, GenerationEvidence) {
|
||||
evidence := GenerationEvidence{}
|
||||
var usage observation.UsageAccumulator
|
||||
finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
|
||||
if captureEvidence {
|
||||
evidence.Usage = usage.Observed()
|
||||
}
|
||||
return response, err, evidence
|
||||
}
|
||||
recorder := observation.FromContext(ctx)
|
||||
generation := req.Generation
|
||||
if generation == 0 {
|
||||
@@ -206,18 +218,30 @@ func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateR
|
||||
release, err := c.manager.Acquire(ctx, req.Target.BackendID)
|
||||
capacityWait := c.capacityWait(req.Target.BackendID, start)
|
||||
if err != nil {
|
||||
if captureEvidence {
|
||||
evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
|
||||
CapacityWait: capacityWait,
|
||||
Failure: ClassifyFailure(err),
|
||||
})
|
||||
}
|
||||
if recorder != nil {
|
||||
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
|
||||
}
|
||||
return nil, err
|
||||
return finish(nil, err)
|
||||
}
|
||||
defer release()
|
||||
if err := ctx.Err(); err != nil {
|
||||
release()
|
||||
if captureEvidence {
|
||||
evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
|
||||
CapacityWait: capacityWait,
|
||||
Failure: ClassifyFailure(err),
|
||||
})
|
||||
}
|
||||
if recorder != nil {
|
||||
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
|
||||
}
|
||||
return nil, err
|
||||
return finish(nil, err)
|
||||
}
|
||||
dispatchStart := c.now()
|
||||
response, err := c.injected.Generate(ctx, req)
|
||||
@@ -225,6 +249,28 @@ func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateR
|
||||
if err == nil && response == nil {
|
||||
err = llm.ErrMalformedResponse
|
||||
}
|
||||
if captureEvidence {
|
||||
attempt := AttemptEvidence{
|
||||
Dispatched: true,
|
||||
ResponseReceived: response != nil,
|
||||
RequestDuration: duration,
|
||||
Failure: ClassifyFailure(err),
|
||||
GenerationCompleted: err == nil && response != nil,
|
||||
}
|
||||
requested := req.StructuredOutput != nil
|
||||
attempt.StructuredOutputRequested = &requested
|
||||
var reportedUsage domain.TokenUsage
|
||||
var reportedPresence *domain.TokenUsagePresence
|
||||
if err == nil && response != nil {
|
||||
reportedUsage = response.Usage
|
||||
reportedPresence = response.UsagePresence
|
||||
}
|
||||
usage.Add(reportedUsage, reportedPresence)
|
||||
var attemptUsage observation.UsageAccumulator
|
||||
attemptUsage.Add(reportedUsage, reportedPresence)
|
||||
attempt.Usage = attemptUsage.Observed()
|
||||
evidence.Attempts = append(evidence.Attempts, attempt)
|
||||
}
|
||||
release()
|
||||
if recorder != nil {
|
||||
data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
|
||||
@@ -236,12 +282,15 @@ func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateR
|
||||
recorder.Attempt(data)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return finish(nil, err)
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
if captureEvidence {
|
||||
return finish(response, err)
|
||||
}
|
||||
return finish(nil, err)
|
||||
}
|
||||
return response, nil
|
||||
return finish(response, nil)
|
||||
}
|
||||
|
||||
func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration, captureEvidence bool) (outcome llm.AttemptOutcome, dispatched bool, capacityWait, generationDuration time.Duration) {
|
||||
|
||||
@@ -19,6 +19,7 @@ type GenerationEvidence struct {
|
||||
// reached. Generated content remains on the separate response return value.
|
||||
type AttemptEvidence struct {
|
||||
Dispatched bool
|
||||
ResponseReceived bool
|
||||
HTTP llm.HTTPExecutionEvidence
|
||||
CapacityWait time.Duration
|
||||
RequestDuration time.Duration
|
||||
|
||||
@@ -9,7 +9,9 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
|
||||
)
|
||||
|
||||
func TestGenerateWithEvidenceRetainsAttemptsAndUsage(t *testing.T) {
|
||||
@@ -173,6 +175,137 @@ func TestGenerateWithEvidenceNoDispatchHasCompleteZeroUsage(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateWithEvidenceForInjectedClient(t *testing.T) {
|
||||
opaqueFailure := errors.New("injected client failed")
|
||||
presence := &domain.TokenUsagePresence{PromptTokens: true}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
response *domain.GenerateResponse
|
||||
err error
|
||||
wantResponse bool
|
||||
wantGeneration bool
|
||||
wantFailure domain.ExecutionFailureKind
|
||||
wantUsage bool
|
||||
wantStructured bool
|
||||
}{
|
||||
{
|
||||
name: "success",
|
||||
response: &domain.GenerateResponse{Content: "private output", Usage: domain.TokenUsage{PromptTokens: 4}, UsagePresence: presence},
|
||||
wantResponse: true,
|
||||
wantGeneration: true,
|
||||
wantUsage: true,
|
||||
wantStructured: true,
|
||||
},
|
||||
{
|
||||
name: "nil response",
|
||||
wantFailure: domain.ExecutionFailureMalformedResponse,
|
||||
},
|
||||
{
|
||||
name: "opaque failure",
|
||||
err: opaqueFailure,
|
||||
wantFailure: domain.ExecutionFailureGeneration,
|
||||
},
|
||||
{
|
||||
name: "response with error",
|
||||
response: &domain.GenerateResponse{Content: "must not be evaluated", Usage: domain.TokenUsage{PromptTokens: 9}, UsagePresence: presence},
|
||||
err: opaqueFailure,
|
||||
wantResponse: true,
|
||||
wantFailure: domain.ExecutionFailureGeneration,
|
||||
wantStructured: true,
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
client := &evidenceInjectedClient{response: tc.response, err: tc.err}
|
||||
coordinator := newEvidenceInjectedCoordinator(t, client)
|
||||
req := testRequest(7)
|
||||
if tc.wantStructured {
|
||||
req.StructuredOutput = &domain.StructuredOutputSpec{}
|
||||
}
|
||||
|
||||
result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), req)
|
||||
if client.calls != 1 || len(evidence.Attempts) != 1 || len(evidence.RetryDelays) != 0 {
|
||||
t.Fatalf("calls=%d attempts=%+v retry delays=%v", client.calls, evidence.Attempts, evidence.RetryDelays)
|
||||
}
|
||||
attempt := evidence.Attempts[0]
|
||||
if !attempt.Dispatched || attempt.ResponseReceived != tc.wantResponse || attempt.GenerationCompleted != tc.wantGeneration || attempt.Failure != tc.wantFailure {
|
||||
t.Fatalf("attempt evidence = %+v", attempt)
|
||||
}
|
||||
if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
|
||||
t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
|
||||
}
|
||||
if attempt.StructuredOutputRequested == nil || *attempt.StructuredOutputRequested != tc.wantStructured {
|
||||
t.Fatalf("structured-output request fact = %v", attempt.StructuredOutputRequested)
|
||||
}
|
||||
if tc.wantFailure == domain.ExecutionFailureNone {
|
||||
if err != nil || result == nil || result.Content != "private output" {
|
||||
t.Fatalf("result=%+v error=%v", result, err)
|
||||
}
|
||||
} else if result != nil || err == nil {
|
||||
t.Fatalf("result=%+v error=%v; expected failed call", result, err)
|
||||
}
|
||||
if attempt.Usage.PromptTokens.Known != tc.wantUsage || evidence.Usage.PromptTokens.Known != tc.wantUsage {
|
||||
t.Fatalf("attempt/aggregate usage = %+v / %+v", attempt.Usage, evidence.Usage)
|
||||
}
|
||||
if tc.wantUsage && (attempt.Usage.PromptTokens.Value != 4 || evidence.Usage.PromptTokens.Value != 4) {
|
||||
t.Fatalf("reported usage was not retained: %+v / %+v", attempt.Usage, evidence.Usage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateWithEvidenceRetainsInjectedSuccessAfterCancellation(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
client := &evidenceInjectedClient{
|
||||
response: &domain.GenerateResponse{Content: "private output"},
|
||||
afterGenerate: func() { cancel() },
|
||||
}
|
||||
coordinator := newEvidenceInjectedCoordinator(t, client)
|
||||
req := testRequest(5)
|
||||
req.StructuredOutput = &domain.StructuredOutputSpec{}
|
||||
|
||||
result, err, evidence := coordinator.GenerateWithEvidence(ctx, req)
|
||||
if !errors.Is(err, context.Canceled) || result == nil || result.Content != "private output" {
|
||||
t.Fatalf("result=%+v error=%v", result, err)
|
||||
}
|
||||
if client.calls != 1 || len(evidence.Attempts) != 1 {
|
||||
t.Fatalf("calls=%d attempts=%+v", client.calls, evidence.Attempts)
|
||||
}
|
||||
attempt := evidence.Attempts[0]
|
||||
if !attempt.ResponseReceived || !attempt.GenerationCompleted || attempt.Failure != domain.ExecutionFailureNone {
|
||||
t.Fatalf("successful generation evidence was lost: %+v", attempt)
|
||||
}
|
||||
if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
|
||||
t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
|
||||
}
|
||||
}
|
||||
|
||||
type evidenceInjectedClient struct {
|
||||
calls int
|
||||
response *domain.GenerateResponse
|
||||
err error
|
||||
afterGenerate func()
|
||||
}
|
||||
|
||||
func (c *evidenceInjectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
||||
c.calls++
|
||||
if c.afterGenerate != nil {
|
||||
c.afterGenerate()
|
||||
}
|
||||
return c.response, c.err
|
||||
}
|
||||
|
||||
func newEvidenceInjectedCoordinator(t *testing.T, client llm.Client) *Coordinator {
|
||||
t.Helper()
|
||||
manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
|
||||
"limited": {ConcurrencyLimit: 1},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return NewInjected(manager, client)
|
||||
}
|
||||
|
||||
type cancelOnDoneContext struct {
|
||||
context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
Reference in New Issue
Block a user