compass_teams_nhs/internal/parse/importance.go

package parse

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"
)

// Importance labels stored in items.importance (FR-4).
const (
	importanceHigh   = "high"
	importanceNormal = "normal"
	importanceLow    = "low"
)

// ImportanceClassifiers names the classifier used for a run, so the choice is
// recorded rather than implied.
const (
	ClassifierNone          = "deterministic"
	ClassifierOmniRoute     = "omniroute"
	ClassifierJEV           = "jev"
	classifyBatchSize       = 40
	classifyMaxResponseChar = 16 * 1024
)

// Completer is one text-completion call. internal/llm supplies the OmniRoute
// client; tests supply a stub.
type Completer interface {
	Complete(ctx context.Context, system, prompt string) (string, error)
}

// ClassifyDeterministic is the offline importance rule. Overdue work and work
// due within the next week is high; everything else is normal. It is the
// fallback when no model is configured, and it is deliberately conservative: it
// never marks work "done".
func ClassifyDeterministic(items []Item, now time.Time) {
	week := now.AddDate(0, 0, 7)
	for i := range items {
		items[i].Importance = deterministicImportance(items[i], week)
	}
}

func deterministicImportance(it Item, week time.Time) string {
	// Only work carries a deadline that makes it important. A calendar entry's
	// due_at is its start time, and a lesson that starts this week is not
	// overdue work, so events and messages stay normal under the offline rule.
	if it.Type == TypeTask {
		if it.DueAt != nil && it.DueAt.Before(week) {
			return importanceHigh
		}
		if strings.Contains(strings.ToLower(it.Body), "overdue") {
			return importanceHigh
		}
	}
	return importanceNormal
}

// ClassifyModel asks a model to classify importance for every item (FR-4). It
// fails, loudly, on a model error or an unparseable answer; it never silently
// leaves a wrong label. Items are updated in place. now lets the model judge
// whether a due date is overdue.
func ClassifyModel(ctx context.Context, c Completer, items []Item, model string, now time.Time) error {
	if c == nil {
		return fmt.Errorf("classify: no model client configured")
	}
	const system = "You classify school items for a parent's morning brief. " +
		"Reply with JSON only. No prose, no code fences."

	for start := 0; start < len(items); start += classifyBatchSize {
		end := start + classifyBatchSize
		if end > len(items) {
			end = len(items)
		}
		batch := items[start:end]

		prompt := classifyPrompt(batch, model, now)
		raw, err := c.Complete(ctx, system, prompt)
		if err != nil {
			return fmt.Errorf("classify: items %d-%d: %w", start, end-1, err)
		}
		labels, err := parseLabels(raw)
		if err != nil {
			return fmt.Errorf("classify: items %d-%d: %w", start, end-1, err)
		}
		for i, label := range labels {
			if i >= len(batch) {
				return fmt.Errorf("classify: items %d-%d: model returned %d labels for %d items", start, end-1, len(labels), len(batch))
			}
			switch label {
			case importanceHigh, importanceNormal, importanceLow:
				batch[i].Importance = label
			default:
				return fmt.Errorf("classify: item %d: invalid importance %q", start+i, label)
			}
		}
		if len(labels) != len(batch) {
			return fmt.Errorf("classify: items %d-%d: model returned %d labels for %d items", start, end-1, len(labels), len(batch))
		}
	}
	return nil
}

// classifyPrompt builds the user message for one batch.
func classifyPrompt(batch []Item, model string, now time.Time) string {
	type row struct {
		I       int    `json:"i"`
		Child   string `json:"child"`
		Type    string `json:"type"`
		Subject string `json:"subject,omitempty"`
		Title   string `json:"title"`
		Due     string `json:"due,omitempty"`
		Detail  string `json:"detail,omitempty"`
	}
	rows := make([]row, 0, len(batch))
	for i, it := range batch {
		due := ""
		if it.DueAt != nil {
			due = it.DueAt.Format(time.RFC3339)
		}
		rows = append(rows, row{
			I:       i,
			Child:   it.Child,
			Type:    it.Type,
			Subject: it.Subject,
			Title:   it.Title,
			Due:     due,
			Detail:  detail(it.Body),
		})
	}
	payload, _ := json.Marshal(rows)
	var b strings.Builder
	fmt.Fprintf(&b, "Model: %s\n", model)
	fmt.Fprintf(&b, "Now: %s\n", now.Format(time.RFC3339))
	b.WriteString("Label every item with importance high, normal or low.\n")
	b.WriteString("high: overdue work, work due today or in the next 48 hours, or a school event that affects the family.\n")
	b.WriteString("low: routine recurring lessons, past submitted work, or notices with no action.\n")
	b.WriteString("normal: anything else.\n")
	b.WriteString("A task with a due date before Now and a submission status of Pending or Overdue is high.\n")
	b.WriteString("Return a JSON array with one object per item, in order, exactly: " +
		`{"i":<index>,"importance":"high|normal|low"}` + "\n\nItems:\n")
	b.Write(payload)
	return b.String()
}

// detail flattens an item body into one line for the model.
func detail(s string) string {
	return collapse(strings.ReplaceAll(s, "\n", "; "))
}

// parseLabels extracts the importance labels from a model reply. It tolerates
// code fences and surrounding prose, and returns the labels ordered by index.
func parseLabels(raw string) ([]string, error) {
	raw = strings.TrimSpace(raw)
	if len(raw) > classifyMaxResponseChar {
		return nil, fmt.Errorf("model reply too large (%d bytes)", len(raw))
	}
	raw = strings.TrimPrefix(raw, "```json")
	raw = strings.TrimPrefix(raw, "```")
	raw = strings.TrimSuffix(raw, "```")
	raw = strings.TrimSpace(raw)

	if i := strings.IndexByte(raw, '['); i >= 0 {
		if j := strings.LastIndexByte(raw, ']'); j > i {
			raw = raw[i : j+1]
		}
	}

	var entries []struct {
		I          *int   `json:"i"`
		Importance string `json:"importance"`
	}
	if err := json.Unmarshal([]byte(raw), &entries); err != nil {
		return nil, fmt.Errorf("model reply is not JSON: %w", err)
	}
	if len(entries) == 0 {
		return nil, fmt.Errorf("model reply has no items")
	}

	type indexed struct {
		i     int
		label string
	}
	got := make([]indexed, 0, len(entries))
	for pos, e := range entries {
		idx := pos
		if e.I != nil {
			idx = *e.I
		}
		got = append(got, indexed{i: idx, label: strings.ToLower(strings.TrimSpace(e.Importance))})
	}
	sort.Slice(got, func(a, b int) bool { return got[a].i < got[b].i })

	out := make([]string, len(got))
	for k, g := range got {
		if g.i < 0 || g.i >= len(got) {
			return nil, fmt.Errorf("model returned out-of-range index %d", g.i)
		}
		out[k] = g.label
	}
	return out, nil
}