diff --git a/internal/ai/harness/harness.go b/internal/ai/harness/harness.go index c7015b1..bc0dcb6 100644 --- a/internal/ai/harness/harness.go +++ b/internal/ai/harness/harness.go @@ -40,8 +40,9 @@ type Skill[In, Out any] interface { // The problems are sent to the model in the repair turn, so each must // say what to fix. Validate(out Out) []string - // PostProcess finishes a valid out, e.g. normalising names. Its error - // fails the run without a repair turn. + // PostProcess finishes a valid out, e.g. normalising names. Its result + // is validated again; problems found then are repaired like any other. + // Its error fails the run without a repair turn. PostProcess(out Out) (Out, error) } @@ -112,22 +113,20 @@ func Run[In, Out any](ctx context.Context, cfg Config, s Skill[In, Out], in In) return zero, &BadOutputError{Skill: s.Name(), Version: s.Version(), First: first, Repair: second} } } - - out, err = s.PostProcess(out) - if err != nil { - return zero, fmt.Errorf("harness: skill %s v%d: post-process: %w", s.Name(), s.Version(), err) - } return out, nil } -// attempt sends req once and decodes and validates the reply. It returns -// the reply text for the repair turn and the problems found; err is set -// only for a provider failure, which is not repairable. +// attempt sends req once, then decodes, validates, post-processes and +// validates the reply again. It returns the reply for the repair turn and +// the problems found; err is set only for a provider or post-process +// failure, neither of which is repairable. func attempt[In, Out any](ctx context.Context, p llm.Provider, s Skill[In, Out], req llm.Request) (out Out, reply string, problems []string, err error) { resp, err := p.Generate(ctx, req) - reply = resp.Text + // The repair turn must quote what failed: the structured output when + // there is any, otherwise the text. + reply = string(resp.JSON) if reply == "" { - reply = string(resp.JSON) + reply = resp.Text } if err != nil { if errors.Is(err, llm.ErrBadOutput) { @@ -138,6 +137,13 @@ func attempt[In, Out any](ctx context.Context, p llm.Provider, s Skill[In, Out], if err := decode(resp.JSON, &out); err != nil { return out, reply, []string{err.Error()}, nil } + if problems := s.Validate(out); len(problems) > 0 { + return out, reply, problems, nil + } + out, err = s.PostProcess(out) + if err != nil { + return out, "", nil, fmt.Errorf("harness: skill %s v%d: post-process: %w", s.Name(), s.Version(), err) + } return out, reply, s.Validate(out), nil } diff --git a/internal/ai/harness/harness_test.go b/internal/ai/harness/harness_test.go index 40b1c4d..7c68000 100644 --- a/internal/ai/harness/harness_test.go +++ b/internal/ai/harness/harness_test.go @@ -79,6 +79,53 @@ func TestRunInvalidThenValid(t *testing.T) { } } +func TestRunPostProcessedOutputIsValidated(t *testing.T) { + var f llm.Fake + f.QueueGenerate(reply(`{"echo":" "}`), reply(`{"echo":"taco"}`)) + out, err := run(t, &f) + if err != nil { + t.Fatal(err) + } + if out.Echo != "taco" { + t.Fatalf("Echo = %q", out.Echo) + } + msgs := assertCalls(t, &f, 2)[1].Messages + if len(msgs) != 3 || msgs[1] != (llm.Message{Role: llm.RoleAssistant, Content: `{"echo":" "}`}) || + !strings.Contains(msgs[2].Content, "echo must not be empty") { + t.Fatalf("repair messages = %+v", msgs) + } +} + +func TestRunWhitespaceTwiceIsBadOutput(t *testing.T) { + var f llm.Fake + f.QueueGenerate(reply(`{"echo":" "}`), reply(`{"echo":" "}`)) + _, err := run(t, &f) + var bad *BadOutputError + if !errors.Is(err, llm.ErrBadOutput) || !errors.As(err, &bad) { + t.Fatalf("err = %v, want *BadOutputError", err) + } + if !slices.Equal(bad.First, []string{"echo must not be empty"}) || + !slices.Equal(bad.Repair, []string{"echo must not be empty"}) { + t.Fatalf("BadOutputError = %+v", bad) + } + assertCalls(t, &f, 2) +} + +func TestRunRepairQuotesJSONOverText(t *testing.T) { + var f llm.Fake + f.QueueGenerate( + llm.GenerateResult{Response: llm.Response{Text: "Here you go!", JSON: json.RawMessage(`{"echo":"TACO"}`)}}, + reply(`{"echo":"taco"}`), + ) + if out, err := run(t, &f); err != nil || out.Echo != "taco" { + t.Fatalf("Run = %+v, %v", out, err) + } + msgs := assertCalls(t, &f, 2)[1].Messages + if len(msgs) != 3 || msgs[1] != (llm.Message{Role: llm.RoleAssistant, Content: `{"echo":"TACO"}`}) { + t.Fatalf("repair messages = %+v, want the failed JSON as the assistant reply", msgs) + } +} + func TestRunDecodeFailureIsRepaired(t *testing.T) { for name, first := range map[string]llm.GenerateResult{ "unknown field": reply(`{"echo":"taco","extra":1}`),