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

// 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:         1024,
		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.
func ParseCheckpoint(raw string) ([]RawSuggestion, error) {
	jsonText := extractJSONObject(raw)
	if jsonText == "" {
		return nil, fmt.Errorf("checkpoint: no JSON object in model output: %q", truncateForError(raw))
	}
	var parsed checkpointResponse
	if err := json.Unmarshal([]byte(jsonText), &parsed); err != nil {
		return nil, fmt.Errorf("checkpoint: parse JSON: %w (got %q)", err, truncateForError(jsonText))
	}
	// Drop items the model returned with an empty original — they can't be
	// anchored — and normalize whitespace the model may have echoed.
	out := parsed.Suggestions[:0]
	for _, s := range parsed.Suggestions {
		s.Original = strings.TrimSpace(s.Original)
		s.Replacement = strings.TrimSpace(s.Replacement)
		if s.Original == "" {
			continue
		}
		out = append(out, s)
	}
	return out, nil
}

// 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 {
	start := strings.IndexByte(s, '{')
	if start < 0 {
		return ""
	}
	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]
			}
		}
	}
	return ""
}

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)
	}
}