petal/internal/suggestions/chat.go

package suggestions

import (
	"database/sql"
	"encoding/json"
	"errors"
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"

	"gitea.parodia.dev/drwily/petal/internal/auth"
	"gitea.parodia.dev/drwily/petal/internal/httputil"
	"gitea.parodia.dev/drwily/petal/internal/llm"
)

// chatRequest is the body the AskPetal panel posts: the full conversation so
// far. The suggestion context is loaded server-side and never trusted from the
// client (spec Note #10).
type chatRequest struct {
	Messages []llm.Message `json:"messages"`
}

// chat streams an Ask Petal conversational reply over SSE. It loads the
// suggestion and its parent document's surrounding paragraph, injects them as
// the tutor system prompt, then relays the model's tokens to the browser as
// `data:` events. History persistence lives entirely in the client.
func (h *Handler) chat(w http.ResponseWriter, r *http.Request) {
	sugID := chi.URLParam(r, "id")

	var body chatRequest
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		httputil.ErrorJSON(w, http.StatusBadRequest, "invalid request body")
		return
	}

	// One query for the suggestion fields and the parent document's plain text,
	// scoped to the local user so a stray id can't read another user's doc.
	var (
		original, replacement, explanation, typ string
		fromPos                                 int
		contentText                             string
	)
	err := h.DB.QueryRow(
		`SELECT s.original, s.replacement, s.explanation, s.type, s.from_pos, d.content_text
		   FROM suggestions s
		   JOIN documents d ON d.id = s.doc_id
		  WHERE s.id = ? AND d.user_id = ?`,
		sugID, auth.UserID(r.Context()),
	).Scan(&original, &replacement, &explanation, &typ, &fromPos, &contentText)
	if errors.Is(err, sql.ErrNoRows) {
		httputil.ErrorJSON(w, http.StatusNotFound, "suggestion not found")
		return
	}
	if err != nil {
		httputil.ServerError(w, err)
		return
	}

	paragraph := surroundingParagraph(contentText, fromPos)
	systemPrompt := llm.AskPetalSystemPrompt(original, replacement, typ, explanation, paragraph)

	// SSE requires an unbuffered, flushable writer. chi's middleware writers pass
	// Flush through; bail with a plain error if somehow they don't.
	flusher, ok := w.(http.Flusher)
	if !ok {
		httputil.ServerError(w, errors.New("streaming unsupported"))
		return
	}

	ch, err := llm.StreamAskPetal(r.Context(), h.Client, systemPrompt, body.Messages)
	if err != nil {
		// The stream never opened (e.g. LLM unreachable) — a normal JSON error is
		// still appropriate since we haven't written SSE headers yet.
		httputil.ErrorJSON(w, http.StatusBadGateway, "chat failed: "+err.Error())
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no") // disable proxy buffering (e.g. nginx)
	w.WriteHeader(http.StatusOK)
	flusher.Flush()

	for chunk := range ch {
		writeSSE(w, "token", map[string]string{"text": chunk})
		flusher.Flush()
	}
	// Signal a clean end so the client can stop reading without waiting on EOF.
	writeSSE(w, "done", map[string]bool{"done": true})
	flusher.Flush()
}

// writeSSE emits one named SSE event with a JSON data payload. JSON-encoding the
// data keeps token text (which may contain newlines) from breaking SSE framing.
func writeSSE(w http.ResponseWriter, event string, data any) {
	payload, err := json.Marshal(data)
	if err != nil {
		return
	}
	_, _ = w.Write([]byte("event: " + event + "\ndata: "))
	_, _ = w.Write(payload)
	_, _ = w.Write([]byte("\n\n"))
}

// surroundingParagraph returns the paragraph of contentText containing the
// plain-text offset from. Tiptap flattens blocks with blank-line separators, so
// paragraphs are bounded by "\n\n". When the offset is unknown (original wasn't
// located, from == -1) it falls back to the latency-capped document so the tutor
// still has context to work with.
func surroundingParagraph(contentText string, from int) string {
	if from < 0 || from > len(contentText) {
		return strings.TrimSpace(llm.TruncateDoc(contentText))
	}
	start := strings.LastIndex(contentText[:from], "\n\n")
	if start < 0 {
		start = 0
	} else {
		start += 2
	}
	end := len(contentText)
	if rel := strings.Index(contentText[from:], "\n\n"); rel >= 0 {
		end = from + rel
	}
	return strings.TrimSpace(contentText[start:end])
}