drwily/petal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Phase 3: LLM grammar checkpoint

Browse Source 
Backend (internal/llm): backend-agnostic LLMClient interface + factory
with vLLM (OpenAI-compat) and Ollama (native) clients, each Complete +
Stream. prompts.go holds the checkpoint and Ask Petal templates;
checkpoint.go salvages JSON from model output (brace-matched), enforces a
per-doc 30s RateLimiter, and truncates the doc to a latency cap.

internal/suggestions: POST /api/docs/:id/check runs a checkpoint and
replaces the doc's pending suggestions in one tx (accepted/rejected kept
as history); GET /api/docs/:id/suggestions lists pending;
POST /api/suggestions/:id/{accept,dismiss} resolves one. Throttled checks
return the current set rather than erroring.

Frontend: useCheckpoint (4s debounce, loads existing on open, stale-guard
tokens); SuggestionHighlight renders ProseMirror decorations re-anchored
by the `original` string on every doc change (not stored marks), with
precise textblock-offset→PM-position mapping; SuggestionCard shows the
type tag + diff + explanation and applies the replacement in-editor on
accept; breathing rose checkpoint dot in the StatusBar; fade-float +
breathe animations.

Tests: llm parse/rate-limit/truncate; suggestions full flow + rate-limit
over httptest with a stub client. Smoke-tested end-to-end against a fake
vLLM endpoint (anchoring verified) and the LLM-unreachable 502 path.

Claude-Session: https://claude.ai/code/session_016Yr6jELuRc7hyzYLccQKZd
This commit is contained in:
prosolis
2026-06-25 20:45:30 -07:00
parent 5e00cdce88
commit a4069d5755
18 changed files with 1640 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
internal/llm/checkpoint.go Normal file
View File

@@ -0,0 +1,143 @@
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 string) ([]RawSuggestion, error) {
raw, err := client.Complete(ctx, CompletionRequest{
Messages: CheckpointMessages(TruncateDoc(contentText)),
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); when throttled it returns
// (false, retryAfter) where retryAfter is the wait until the next allowed run.
func (rl *RateLimiter) Allow(docID string) (bool, time.Duration) {
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
}
}
rl.last[docID] = now
return true, 0
}

93
internal/llm/checkpoint_test.go Normal file
View File

@@ -0,0 +1,93 @@
package llm
import (
"testing"
"time"
)
func TestParseCheckpoint(t *testing.T) {
tests := []struct {
name string
raw string
want int
wantErr bool
}{
{
name: "clean json",
raw: `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"}]}`,
want: 1,
},
{
name: "wrapped in markdown fence",
raw: "Here you go:\n```json\n{\"suggestions\":[{\"original\":\"a apple\",\"replacement\":\"an apple\",\"explanation\":\"use an before a vowel\",\"type\":\"grammar\"}]}\n```",
want: 1,
},
{
name: "empty suggestions",
raw: `{"suggestions":[]}`,
want: 0,
},
{
name: "drops items with empty original",
raw: `{"suggestions":[{"original":"","replacement":"x","explanation":"e","type":"grammar"},{"original":"teh","replacement":"the","explanation":"typo","type":"grammar"}]}`,
want: 1,
},
{
name: "no json at all",
raw: "I could not find any issues!",
wantErr: true,
},
{
name: "braces inside string values",
raw: `{"suggestions":[{"original":"use {x}","replacement":"use x","explanation":"drop the braces","type":"clarity"}]}`,
want: 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := ParseCheckpoint(tt.raw)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got none")
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(got) != tt.want {
t.Fatalf("got %d suggestions, want %d", len(got), tt.want)
}
})
}
}
func TestRateLimiter(t *testing.T) {
rl := NewRateLimiter(30 * time.Second)
if ok, _ := rl.Allow("doc1"); !ok {
t.Fatal("first call should be allowed")
}
if ok, retry := rl.Allow("doc1"); ok || retry <= 0 {
t.Fatalf("immediate second call should be throttled, got ok=%v retry=%v", ok, retry)
}
// A different document is independent.
if ok, _ := rl.Allow("doc2"); !ok {
t.Fatal("different doc should be allowed")
}
}
func TestTruncateDoc(t *testing.T) {
short := "hello"
if got := TruncateDoc(short); got != short {
t.Fatalf("short doc should be unchanged")
}
long := make([]byte, maxDocChars+500)
for i := range long {
long[i] = 'a'
}
got := TruncateDoc(string(long))
if len(got) != maxDocChars {
t.Fatalf("truncated length = %d, want %d", len(got), maxDocChars)
}
}

110
internal/llm/client.go Normal file
View File

@@ -0,0 +1,110 @@
// Package llm wraps the local inference endpoint behind a backend-agnostic
// interface. Two concrete backends — vLLM (OpenAI-compatible) and Ollama
// (native) — implement LLMClient; the rest of the app (the grammar checkpoint,
// Ask Petal) calls the interface only and never knows which one is active.
package llm
import (
"context"
"time"
"gitea.parodia.dev/drwily/petal/internal/config"
)
// Message is one chat turn sent to the model.
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
}
// CompletionRequest is the backend-neutral request shape. Each backend maps
// these fields onto its own wire format (see the parameter cheatsheet in the
// spec). Zero-valued sampling fields fall back to backend defaults.
type CompletionRequest struct {
Messages []Message
MaxTokens int
Temperature float64
RepetitionPenalty float64
TopP float64
Stop []string
Stream bool
}
// LLMClient is the single surface handler code depends on.
type LLMClient interface {
// Complete returns the full response in one shot (used for the checkpoint
// JSON, which we want whole before parsing).
Complete(ctx context.Context, req CompletionRequest) (string, error)
// Stream returns a channel of text chunks (used for Ask Petal SSE). The
// channel closes when generation ends; on a mid-stream error it closes
// early and the error is reported via the returned error of a future call.
Stream(ctx context.Context, req CompletionRequest) (<-chan string, error)
}
// Truncation caps. These guard checkpoint latency (prefill time scales with
// input length), not the model window — Qwen 3.5 ships 256K, far larger.
const (
// maxDocChars ~ 10,000 tokens at ~4 chars/token. Grammar checkpoint only;
// the voice pass sends the whole document uncut.
maxDocChars = 40000
// maxHistoryMsgs caps the rolling Ask Petal history (5 turns); oldest pairs
// drop first.
maxHistoryMsgs = 10
)
// TruncateDoc keeps the recent end of a document — the user is actively writing
// there — when it exceeds the grammar-checkpoint latency cap. Most documents
// fit uncut. The voice pass deliberately does NOT call this.
func TruncateDoc(contentText string) string {
if len(contentText) > maxDocChars {
return contentText[len(contentText)-maxDocChars:]
}
return contentText
}
// TrimHistory keeps the most recent maxHistoryMsgs messages, dropping the
// oldest first so a long Ask Petal chat stays within budget.
func TrimHistory(msgs []Message) []Message {
if len(msgs) > maxHistoryMsgs {
return msgs[len(msgs)-maxHistoryMsgs:]
}
return msgs
}
// NewLLMClient selects a backend from config. LLM_CHAT_MODEL falls back to
// LLM_MODEL here, in the factory, so handlers never deal with the fallback.
func NewLLMClient(cfg *config.Config) LLMClient {
chatModel := cfg.LLMChatModel
if chatModel == "" {
chatModel = cfg.LLMModel
}
base := backend{
endpoint: cfg.LLMEndpoint,
checkpointModel: cfg.LLMModel,
chatModel: chatModel,
timeout: cfg.LLMTimeout,
}
switch cfg.LLMBackend {
case "ollama":
return &OllamaClient{base}
default: // "vllm"
return &VLLMClient{base}
}
}
// backend holds the fields shared by both concrete clients.
type backend struct {
endpoint string
checkpointModel string // LLM_MODEL — small/fast checkpoint model
chatModel string // LLM_CHAT_MODEL (or LLM_MODEL) — Ask Petal model
timeout time.Duration
}
// model picks the right model id for a request. Streaming requests are Ask
// Petal (chat); one-shot Complete calls are the grammar checkpoint.
func (b backend) model(req CompletionRequest) string {
if req.Stream {
return b.chatModel
}
return b.checkpointModel
}

124
internal/llm/ollama.go Normal file
View File

@@ -0,0 +1,124 @@
package llm
import (
"bufio"
"context"
"encoding/json"
"fmt"
"net/http"
)
// OllamaClient talks to Ollama's native /api/chat endpoint. Selected when
// LLM_BACKEND=ollama.
type OllamaClient struct {
backend
}
// ollamaRequest mirrors Ollama's /api/chat body. Sampling parameters live under
// "options" (see the cheatsheet in the spec).
type ollamaRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Stream bool `json:"stream"`
Options ollamaOptions `json:"options"`
}
type ollamaOptions struct {
NumPredict int `json:"num_predict"`
Temperature float64 `json:"temperature"`
RepeatPenalty float64 `json:"repeat_penalty"`
TopP float64 `json:"top_p"`
Stop []string `json:"stop,omitempty"`
}
func (c *OllamaClient) body(req CompletionRequest) ollamaRequest {
return ollamaRequest{
Model: c.model(req),
Messages: req.Messages,
Stream: req.Stream,
Options: ollamaOptions{
NumPredict: req.MaxTokens,
Temperature: req.Temperature,
RepeatPenalty: req.RepetitionPenalty,
TopP: req.TopP,
Stop: req.Stop,
},
}
}
// Complete sends a non-streaming request and returns the full message content.
func (c *OllamaClient) Complete(ctx context.Context, req CompletionRequest) (string, error) {
req.Stream = false
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
resp, err := c.post(ctx, c.body(req))
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", httpError("ollama", resp)
}
var out struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
}
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
return "", fmt.Errorf("ollama: decode response: %w", err)
}
return out.Message.Content, nil
}
// Stream sends a streaming request. Ollama emits newline-delimited JSON objects;
// we forward each message.content and stop when done:true.
func (c *OllamaClient) Stream(ctx context.Context, req CompletionRequest) (<-chan string, error) {
req.Stream = true
resp, err := c.post(ctx, c.body(req))
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
defer resp.Body.Close()
return nil, httpError("ollama", resp)
}
ch := make(chan string)
go func() {
defer close(ch)
defer resp.Body.Close()
sc := bufio.NewScanner(resp.Body)
sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for sc.Scan() {
line := sc.Bytes()
if len(line) == 0 {
continue
}
var chunk struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
Done bool `json:"done"`
}
if err := json.Unmarshal(line, &chunk); err != nil {
continue
}
if chunk.Message.Content != "" {
select {
case ch <- chunk.Message.Content:
case <-ctx.Done():
return
}
}
if chunk.Done {
return
}
}
}()
return ch, nil
}
func (c *OllamaClient) post(ctx context.Context, body any) (*http.Response, error) {
return postJSON(ctx, c.endpoint+"/api/chat", body)
}

63
internal/llm/prompts.go Normal file
View File

@@ -0,0 +1,63 @@
package llm
import "fmt"
// checkpointSystemPrompt is the grammar-checkpoint instruction. It asks for
// strict JSON (no fences, no preamble) so Complete's output parses directly.
const checkpointSystemPrompt = `You are a warm, encouraging writing assistant helping someone who speaks English as a second language. ` +
`Analyze the text below and identify up to 5 issues: grammar errors, unnatural phrasing, ` +
`incorrect idiom usage, or unclear sentences that are common ESL patterns.
Be specific, friendly, and explain WHY each suggestion improves the writing.
Respond ONLY with valid JSON. No preamble, no markdown fences. Format:
{
"suggestions": [
{
"original": "exact text from the document that needs fixing",
"replacement": "corrected version",
"explanation": "friendly one-sentence explanation",
"type": "grammar|phrasing|idiom|clarity"
}
]
}
If the writing looks good, return: {"suggestions": []}`
// CheckpointMessages builds the message array for a grammar checkpoint over the
// given (already-truncated) document text.
func CheckpointMessages(contentText string) []Message {
return []Message{
{Role: "system", Content: checkpointSystemPrompt},
{Role: "user", Content: contentText},
}
}
// askPetalSystemTemplate is the Ask Petal tutor prompt. The suggestion context
// is interpolated in; the user's own messages are appended after this system
// turn by the caller.
const askPetalSystemTemplate = `You are Petal, a warm and patient English writing tutor helping someone who is learning English ` +
`as a second language. You are currently discussing a specific writing suggestion.
Suggestion context:
- Original text: "%s"
- Suggested replacement: "%s"
- Issue type: %s
- Initial explanation: "%s"
- Surrounding paragraph: "%s"
The user wants to understand this suggestion better. Detect the language of the user's message ` +
`and respond in that same language. If they write in Mandarin Chinese, respond entirely in ` +
`Mandarin. If they write in English, respond in English. Never mix languages in a single response.
Explain clearly and kindly. Use simple language appropriate to the user's message. Give examples ` +
`when helpful. If they ask "why" (or "为什么"), explain the grammar rule or idiom behind it. ` +
`If they suggest an alternative phrasing, evaluate it honestly.
Keep responses concise (2-4 sentences). This is a chat, not an essay. Be encouraging — ` +
`learning a language is hard and they're doing great.`
// AskPetalSystemPrompt fills the tutor prompt with one suggestion's context.
func AskPetalSystemPrompt(original, replacement, suggestionType, explanation, paragraph string) string {
return fmt.Sprintf(askPetalSystemTemplate, original, replacement, suggestionType, explanation, paragraph)
}

156
internal/llm/vllm.go Normal file
View File

@@ -0,0 +1,156 @@
package llm
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
)
// VLLMClient talks to an OpenAI-compatible /v1/chat/completions endpoint
// (vLLM's API server). It is the default backend.
type VLLMClient struct {
backend
}
// vllmRequest is the OpenAI-compatible request body. repetition_penalty is a
// vLLM extension to the OpenAI schema, which vLLM accepts.
type vllmRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
MaxTokens int `json:"max_tokens"`
Temperature float64 `json:"temperature"`
RepetitionPenalty float64 `json:"repetition_penalty"`
TopP float64 `json:"top_p"`
Stop []string `json:"stop,omitempty"`
Stream bool `json:"stream"`
}
func (c *VLLMClient) body(req CompletionRequest) vllmRequest {
return vllmRequest{
Model: c.model(req),
Messages: req.Messages,
MaxTokens: req.MaxTokens,
Temperature: req.Temperature,
RepetitionPenalty: req.RepetitionPenalty,
TopP: req.TopP,
Stop: req.Stop,
Stream: req.Stream,
}
}
// Complete sends a non-streaming request and returns the full message content.
func (c *VLLMClient) Complete(ctx context.Context, req CompletionRequest) (string, error) {
req.Stream = false
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
resp, err := c.post(ctx, c.body(req))
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", httpError("vllm", resp)
}
var out struct {
Choices []struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
} `json:"choices"`
}
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
return "", fmt.Errorf("vllm: decode response: %w", err)
}
if len(out.Choices) == 0 {
return "", fmt.Errorf("vllm: empty choices in response")
}
return out.Choices[0].Message.Content, nil
}
// Stream sends a streaming request and emits delta content chunks on the
// returned channel, which closes when the stream ends ([DONE]) or ctx is done.
func (c *VLLMClient) Stream(ctx context.Context, req CompletionRequest) (<-chan string, error) {
req.Stream = true
resp, err := c.post(ctx, c.body(req))
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
defer resp.Body.Close()
return nil, httpError("vllm", resp)
}
ch := make(chan string)
go func() {
defer close(ch)
defer resp.Body.Close()
sc := bufio.NewScanner(resp.Body)
sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for sc.Scan() {
line := strings.TrimSpace(sc.Text())
if !strings.HasPrefix(line, "data:") {
continue
}
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "[DONE]" {
return
}
var chunk struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
} `json:"choices"`
}
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue // skip keep-alives / malformed frames
}
if len(chunk.Choices) == 0 {
continue
}
if text := chunk.Choices[0].Delta.Content; text != "" {
select {
case ch <- text:
case <-ctx.Done():
return
}
}
}
}()
return ch, nil
}
func (c *VLLMClient) post(ctx context.Context, body any) (*http.Response, error) {
return postJSON(ctx, c.endpoint+"/v1/chat/completions", body)
}
// --- shared HTTP helpers (used by both backends) ---------------------------
// llmHTTP has no client-level timeout: Complete bounds itself with a context
// deadline (LLM_TIMEOUT), and streaming requests must stay open as long as the
// model is generating. Cancellation flows through the request context.
var llmHTTP = &http.Client{}
func postJSON(ctx context.Context, url string, body any) (*http.Response, error) {
buf, err := json.Marshal(body)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
return llmHTTP.Do(req)
}
func httpError(backend string, resp *http.Response) error {
b, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
return fmt.Errorf("%s: %s: %s", backend, resp.Status, strings.TrimSpace(string(b)))
}