ingestion/internal/wing/llmclassifier.go

@@ -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
+}

ingestion/internal/wing/llmclassifier_test.go

@@ -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")
+}