100 lines
2.8 KiB
Go
100 lines
2.8 KiB
Go
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
|
|
}
|