petal/internal/suggestions/chat_test.go

package suggestions

import (
	"context"
	"encoding/json"
	"net/http"
	"strings"
	"testing"

	"gitea.parodia.dev/drwily/petal/internal/db"
	"gitea.parodia.dev/drwily/petal/internal/llm"
)

// streamClient serves a canned checkpoint (so a suggestion can be seeded through
// the public /check route) and a fixed token stream for chat. It records the
// last streamed request so a test can assert on the injected system context.
type streamClient struct {
	checkpoint string
	tokens     []string
	lastStream llm.CompletionRequest
}

func (c *streamClient) Complete(_ context.Context, _ llm.CompletionRequest) (string, error) {
	return c.checkpoint, nil
}

func (c *streamClient) Stream(_ context.Context, req llm.CompletionRequest) (<-chan string, error) {
	c.lastStream = req
	ch := make(chan string)
	go func() {
		defer close(ch)
		for _, t := range c.tokens {
			ch <- t
		}
	}()
	return ch, nil
}

func TestAskPetalChat(t *testing.T) {
	client := &streamClient{
		checkpoint: `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"}]}`,
		tokens:     []string{"Because ", "\"I\" ", "takes ", "\"have\"."},
	}
	srv, docID := newTestServer(t, client)

	// Seed one suggestion via the checkpoint route, then grab its id.
	do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	listed := do(t, srv, http.MethodGet, "/docs/"+docID+"/suggestions", "")
	var sugs []db.Suggestion
	if err := json.Unmarshal(listed.Body.Bytes(), &sugs); err != nil {
		t.Fatalf("decode suggestions: %v", err)
	}
	if len(sugs) != 1 {
		t.Fatalf("seed: want 1 suggestion, got %d", len(sugs))
	}
	sugID := sugs[0].ID

	rec := do(t, srv, http.MethodPost, "/suggestions/"+sugID+"/chat",
		`{"messages":[{"role":"user","content":"why is this wrong?"}]}`)
	if rec.Code != http.StatusOK {
		t.Fatalf("chat: code=%d body=%s", rec.Code, rec.Body)
	}
	if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" {
		t.Fatalf("chat: content-type=%q, want text/event-stream", ct)
	}

	// The streamed tokens should arrive concatenated across data: events, and the
	// stream should terminate with a done event.
	body := rec.Body.String()
	got := collectSSEText(body)
	if want := "Because \"I\" takes \"have\"."; got != want {
		t.Fatalf("streamed text = %q, want %q", got, want)
	}
	if !strings.Contains(body, "event: done") {
		t.Fatalf("chat stream missing done event:\n%s", body)
	}

	// Context injection: the system prompt (first message) must carry the
	// suggestion's original text and the surrounding paragraph — loaded
	// server-side, never from the client.
	msgs := client.lastStream.Messages
	if len(msgs) < 2 || msgs[0].Role != "system" {
		t.Fatalf("expected a leading system message, got %+v", msgs)
	}
	sys := msgs[0].Content
	if !strings.Contains(sys, "I has") {
		t.Fatalf("system prompt missing original text:\n%s", sys)
	}
	if !strings.Contains(sys, "I has two apple.") {
		t.Fatalf("system prompt missing surrounding paragraph:\n%s", sys)
	}
	// The client's user message rides after the system turn.
	if last := msgs[len(msgs)-1]; last.Role != "user" || last.Content != "why is this wrong?" {
		t.Fatalf("user message not forwarded: %+v", last)
	}

	// Chat sampling parameters from the spec.
	if client.lastStream.MaxTokens != 512 || client.lastStream.Temperature != 0.7 {
		t.Fatalf("unexpected chat params: %+v", client.lastStream)
	}
}

func TestAskPetalChatNotFound(t *testing.T) {
	client := &streamClient{checkpoint: `{"suggestions":[]}`}
	srv, _ := newTestServer(t, client)
	rec := do(t, srv, http.MethodPost, "/suggestions/does-not-exist/chat",
		`{"messages":[{"role":"user","content":"hi"}]}`)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("want 404 for unknown suggestion, got %d", rec.Code)
	}
}

func TestSurroundingParagraph(t *testing.T) {
	text := "First paragraph here.\n\nThe target sentence lives here.\n\nThird paragraph."
	from := strings.Index(text, "target")
	got := surroundingParagraph(text, from)
	if got != "The target sentence lives here." {
		t.Fatalf("paragraph = %q", got)
	}

	// Unknown offset falls back to the (trimmed) document.
	if got := surroundingParagraph(text, -1); !strings.Contains(got, "First paragraph") {
		t.Fatalf("fallback paragraph = %q", got)
	}
}

// collectSSEText concatenates the JSON text fields from every `event: token`
// frame in an SSE body.
func collectSSEText(body string) string {
	var b strings.Builder
	for _, frame := range strings.Split(body, "\n\n") {
		if !strings.Contains(frame, "event: token") {
			continue
		}
		for _, line := range strings.Split(frame, "\n") {
			data, ok := strings.CutPrefix(line, "data: ")
			if !ok {
				continue
			}
			var payload struct {
				Text string `json:"text"`
			}
			_ = json.Unmarshal([]byte(data), &payload)
			b.WriteString(payload.Text)
		}
	}
	return b.String()
}