feat(wing): LLM classifier with tolerant, fail-closed reply parsing (#205 S1)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,67 @@
|
||||
package wing_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/mathiasbq/hyperguild/ingestion/internal/wing"
|
||||
)
|
||||
|
||||
func completes(reply string, err error) wing.CompleteFunc {
|
||||
return func(_ context.Context, _, _ string) (string, error) { return reply, err }
|
||||
}
|
||||
|
||||
func TestLLMClassifier_ParsesTheReply(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
reply string
|
||||
want wing.Suggestion
|
||||
wantErr bool
|
||||
}{
|
||||
{"bare json", `{"wing":"tapir","confidence":0.9,"reason":"video summariser"}`, wing.Suggestion{Wing: "tapir", Confidence: 0.9, Reason: "video summariser"}, false},
|
||||
{"json inside prose and fences", "Sure!\n```json\n{\"wing\":\"homelab\",\"confidence\":0.7,\"reason\":\"k3s\"}\n```\nDone.", wing.Suggestion{Wing: "homelab", Confidence: 0.7, Reason: "k3s"}, false},
|
||||
{"none means abstain", `{"wing":"none","confidence":0.2,"reason":"generic"}`, wing.Suggestion{Wing: "", Confidence: 0.2, Reason: "generic"}, false},
|
||||
{"confidence above 1 is clamped", `{"wing":"tapir","confidence":7,"reason":"x"}`, wing.Suggestion{Wing: "tapir", Confidence: 1, Reason: "x"}, false},
|
||||
{"negative confidence is clamped", `{"wing":"tapir","confidence":-1,"reason":"x"}`, wing.Suggestion{Wing: "tapir", Confidence: 0, Reason: "x"}, false},
|
||||
{"empty reply (a reasoning model that ran out of tokens)", ``, wing.Suggestion{}, true},
|
||||
{"no json at all", `I think it is tapir`, wing.Suggestion{}, true},
|
||||
{"malformed json", `{"wing":"tapir","confidence":`, wing.Suggestion{}, true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := wing.NewLLMClassifier(completes(tt.reply, nil))
|
||||
got, err := c.Classify(context.Background(), "note", []string{"tapir", "homelab"})
|
||||
if tt.wantErr {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLLMClassifier_PropagatesTheModelError(t *testing.T) {
|
||||
boom := errors.New("model down")
|
||||
_, err := wing.NewLLMClassifier(completes("", boom)).Classify(context.Background(), "n", []string{"tapir"})
|
||||
assert.ErrorIs(t, err, boom)
|
||||
}
|
||||
|
||||
func TestLLMClassifier_PromptListsCandidatesAndBoundsTheNote(t *testing.T) {
|
||||
var gotUser string
|
||||
c := wing.NewLLMClassifier(func(_ context.Context, _, user string) (string, error) {
|
||||
gotUser = user
|
||||
return `{"wing":"tapir","confidence":0.9,"reason":"x"}`, nil
|
||||
})
|
||||
long := strings.Repeat("word ", 5000)
|
||||
_, err := c.Classify(context.Background(), long, []string{"tapir", "homelab"})
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, gotUser, "tapir")
|
||||
assert.Contains(t, gotUser, "homelab")
|
||||
assert.Less(t, len(gotUser), 12000, "a whole transcript must not be sent as the classification input")
|
||||
}
|
||||
Reference in New Issue
Block a user