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,99 @@
|
||||
package wing
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// CompleteFunc is a chat completion, matching llm.Client.Complete. It must be
|
||||
// backed by a local model: note content is sent to it.
|
||||
type CompleteFunc func(ctx context.Context, system, user string) (string, error)
|
||||
|
||||
// maxNoteChars bounds what is sent for classification; the start of a note
|
||||
// carries its subject, and a whole transcript is cost with no benefit.
|
||||
const maxNoteChars = 6000
|
||||
|
||||
const classifySystem = `You file notes into a personal knowledge base. Pick the single wing the note belongs to, from the candidate list only.
|
||||
Answer with ONLY a JSON object: {"wing": "<name or none>", "confidence": <0..1>, "reason": "<one short sentence>"}.
|
||||
Use "none" when no candidate clearly fits. Do not invent wings.`
|
||||
|
||||
// LLMClassifier asks a model to pick a wing.
|
||||
type LLMClassifier struct{ complete CompleteFunc }
|
||||
|
||||
// NewLLMClassifier wraps a chat completion as a Classifier.
|
||||
func NewLLMClassifier(c CompleteFunc) *LLMClassifier { return &LLMClassifier{complete: c} }
|
||||
|
||||
type classifyReply struct {
|
||||
Wing string `json:"wing"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// Classify implements Classifier. Any reply it cannot parse is an error, which
|
||||
// the Resolver turns into an Exception: an unreadable answer is not a guess.
|
||||
func (l *LLMClassifier) Classify(ctx context.Context, text string, candidates []string) (Suggestion, error) {
|
||||
if len(text) > maxNoteChars {
|
||||
text = text[:maxNoteChars]
|
||||
}
|
||||
user := "Candidate wings: " + strings.Join(candidates, ", ") + "\n\nNote:\n" + text
|
||||
out, err := l.complete(ctx, classifySystem, user)
|
||||
if err != nil {
|
||||
return Suggestion{}, fmt.Errorf("classify: %w", err)
|
||||
}
|
||||
raw, ok := firstJSONObject(out)
|
||||
if !ok {
|
||||
return Suggestion{}, errors.New("classify: reply contains no JSON object")
|
||||
}
|
||||
var r classifyReply
|
||||
if err := json.Unmarshal([]byte(raw), &r); err != nil {
|
||||
return Suggestion{}, fmt.Errorf("classify: parse reply: %w", err)
|
||||
}
|
||||
w := strings.TrimSpace(r.Wing)
|
||||
if strings.EqualFold(w, "none") {
|
||||
w = ""
|
||||
}
|
||||
return Suggestion{Wing: w, Confidence: clamp01(r.Confidence), Reason: r.Reason}, nil
|
||||
}
|
||||
|
||||
// firstJSONObject returns the first balanced {...} in s, ignoring braces
|
||||
// inside JSON strings, so prose or code fences around the reply are tolerated.
|
||||
func firstJSONObject(s string) (string, bool) {
|
||||
start := strings.IndexByte(s, '{')
|
||||
if start < 0 {
|
||||
return "", false
|
||||
}
|
||||
depth, inStr, esc := 0, false, false
|
||||
for i := start; i < len(s); i++ {
|
||||
c := s[i]
|
||||
switch {
|
||||
case esc:
|
||||
esc = false
|
||||
case inStr && c == '\\':
|
||||
esc = true
|
||||
case c == '"':
|
||||
inStr = !inStr
|
||||
case inStr:
|
||||
case c == '{':
|
||||
depth++
|
||||
case c == '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return s[start : i+1], true
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func clamp01(f float64) float64 {
|
||||
switch {
|
||||
case f < 0:
|
||||
return 0
|
||||
case f > 1:
|
||||
return 1
|
||||
}
|
||||
return f
|
||||
}
|
||||
@@ -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