petal/internal/llm/checkpoint.go

package llm

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

// CheckpointInterval is the minimum time between grammar checkpoints for a
// single document. The frontend debounces at 4s; this is the server-side floor
// that protects the inference endpoint from rapid repeat checks.
const CheckpointInterval = 30 * time.Second

// checkpointMaxTokens caps the checkpoint's generated output. Each suggestion
// echoes its original sentence, a replacement, and a friendly explanation, so
// the full JSON for the prompt's "up to 5 issues" runs ~2,000 tokens on a long
// document (measured ~2,007 for a 380-word doc). The old 1,024 cap truncated
// the JSON mid-array — done_reason "length" — leaving an object that never
// closed, so ParseCheckpoint found no parseable JSON and the pass 502'd on
// every long doc. This is a ceiling, not a target: the model stops at its JSON
// close well before it, so the headroom costs nothing on shorter docs.
const checkpointMaxTokens = 4096

// RawSuggestion is one item as the model emits it. Positions are resolved
// server-side from Original; the frontend re-anchors by string at render time.
type RawSuggestion struct {
	Original    string `json:"original"`
	Replacement string `json:"replacement"`
	Explanation string `json:"explanation"`
	Type        string `json:"type"`
}

// checkpointResponse is the top-level JSON shape the checkpoint prompt requests.
type checkpointResponse struct {
	Suggestions []RawSuggestion `json:"suggestions"`
}

// RunCheckpoint sends the grammar checkpoint and parses the JSON result. It
// applies the latency-guard truncation and the checkpoint sampling parameters
// from the spec.
func RunCheckpoint(ctx context.Context, client LLMClient, contentText, tone string) ([]RawSuggestion, error) {
	raw, err := client.Complete(ctx, CompletionRequest{
		Messages:          CheckpointMessages(TruncateDoc(contentText), tone),
		MaxTokens:         checkpointMaxTokens,
		Temperature:       0.3,
		RepetitionPenalty: 1.15,
		TopP:              0.9,
		Stop:              []string{"```", "\n\n\n\n"},
	})
	if err != nil {
		return nil, err
	}
	return ParseCheckpoint(raw)
}

// ParseCheckpoint extracts the suggestions array from a model response. Smaller
// models sometimes wrap JSON in prose or markdown fences despite instructions,
// so we salvage the outermost {...} object before decoding. When the model
// truncates mid-array (it ran out of output budget) the outer object never
// closes and won't parse — we then recover the suggestion objects that *did*
// complete, so a long document still gets feedback instead of a hard failure.
func ParseCheckpoint(raw string) ([]RawSuggestion, error) {
	var suggestions []RawSuggestion
	parsed := false
	if jsonText := extractJSONObject(raw); jsonText != "" {
		var resp checkpointResponse
		if err := json.Unmarshal([]byte(jsonText), &resp); err == nil {
			suggestions = resp.Suggestions
			parsed = true
		}
	}
	if !parsed {
		// No closeable object, or it failed to decode (most often a truncated
		// array). Salvage the individual suggestion objects that completed.
		suggestions = salvageSuggestions(raw)
		if suggestions == nil {
			return nil, fmt.Errorf("checkpoint: no JSON object in model output: %q", truncateForError(raw))
		}
	}
	// Drop items the model returned with an empty original — they can't be
	// anchored — and normalize whitespace the model may have echoed.
	out := suggestions[:0]
	for _, s := range suggestions {
		s.Original = strings.TrimSpace(s.Original)
		s.Replacement = strings.TrimSpace(s.Replacement)
		if s.Original == "" {
			continue
		}
		out = append(out, s)
	}
	return out, nil
}

// salvageSuggestions recovers as many complete suggestion objects as it can from
// a response whose top-level JSON didn't parse — typically one truncated mid
// "suggestions" array. It scans the array region for brace-balanced {...}
// objects and keeps each that decodes into a RawSuggestion with a non-empty
// original; the final, cut-off object simply never balances and is dropped.
// Returns nil when there is no salvageable array (so the caller can error).
func salvageSuggestions(raw string) []RawSuggestion {
	key := strings.Index(raw, `"suggestions"`)
	if key < 0 {
		return nil
	}
	lb := strings.IndexByte(raw[key:], '[')
	if lb < 0 {
		return nil
	}
	rest := raw[key+lb+1:]
	out := []RawSuggestion{}
	for {
		obj, end := firstBalancedObject(rest)
		if end < 0 {
			break
		}
		var s RawSuggestion
		if err := json.Unmarshal([]byte(obj), &s); err == nil && strings.TrimSpace(s.Original) != "" {
			out = append(out, s)
		}
		rest = rest[end:]
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

// extractJSONObject returns the substring from the first '{' to its matching
// closing '}', or "" if none. Tolerates fences/preamble around the object.
func extractJSONObject(s string) string {
	obj, _ := firstBalancedObject(s)
	return obj
}

// firstBalancedObject returns the first brace-balanced {...} substring in s and
// the index in s just past its closing brace, or ("", -1) if no object closes
// (none present, or the only one is truncated). String contents — including
// braces inside quoted values — are skipped so they can't unbalance the count.
func firstBalancedObject(s string) (string, int) {
	start := strings.IndexByte(s, '{')
	if start < 0 {
		return "", -1
	}
	depth := 0
	inString := false
	escaped := false
	for i := start; i < len(s); i++ {
		c := s[i]
		switch {
		case escaped:
			escaped = false
		case c == '\\' && inString:
			escaped = true
		case c == '"':
			inString = !inString
		case inString:
			// ignore braces inside strings
		case c == '{':
			depth++
		case c == '}':
			depth--
			if depth == 0 {
				return s[start : i+1], i + 1
			}
		}
	}
	return "", -1
}

func truncateForError(s string) string {
	const max = 200
	if len(s) > max {
		return s[:max] + "…"
	}
	return s
}

// RateLimiter enforces a minimum interval between checkpoints per document.
// Concurrency-safe; one instance is shared across requests.
type RateLimiter struct {
	mu       sync.Mutex
	interval time.Duration
	last     map[string]time.Time
}

// NewRateLimiter constructs a limiter with the given per-document minimum gap.
func NewRateLimiter(interval time.Duration) *RateLimiter {
	return &RateLimiter{interval: interval, last: make(map[string]time.Time)}
}

// Allow reports whether a checkpoint may run for docID now. When allowed it
// records the time and returns (true, 0, recordedAt); when throttled it returns
// (false, retryAfter, zero) where retryAfter is the wait until the next allowed
// run. recordedAt lets a caller whose pass fails hand the exact timestamp to
// Release so it rolls back only its own slot.
func (rl *RateLimiter) Allow(docID string) (bool, time.Duration, time.Time) {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	now := time.Now()
	if last, ok := rl.last[docID]; ok {
		if elapsed := now.Sub(last); elapsed < rl.interval {
			return false, rl.interval - elapsed, time.Time{}
		}
	}
	rl.last[docID] = now
	return true, 0, now
}

// Release rolls back the slot a prior Allow recorded for docID, but only if no
// newer Allow has since claimed it. Called when an LLM pass fails: Allow runs
// before the model call, so without this a failed checkpoint would hold the
// per-document slot for the full interval and a retry (or the frontend's
// auto-retry) would hit the throttle path and get the stale/empty set back
// instead of re-running. `at` is the time the failed Allow returned.
func (rl *RateLimiter) Release(docID string, at time.Time) {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	if last, ok := rl.last[docID]; ok && last.Equal(at) {
		delete(rl.last, docID)
	}
}