textmachine/backend/internal/pipeline/fewshot_test.go

package pipeline

import (
	"fmt"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"textmachine/backend/internal/config"
	"textmachine/backend/internal/obs"
)

// TestFewShotSplitAndAssembly pins the D38.4 few-shot toggle at the render layer: a template
// splits into a CORE system + a ---FEWSHOT--- block, SystemFor(true) appends the examples (the
// P1a discourse default), SystemFor(false) is core-only (zero-shot, for a reasoning model whose
// CoT is disrupted by hand examples — deepseek-thinking swap-arm).
func TestFewShotSplitAndAssembly(t *testing.T) {
	tpl := writeTemplate(t, "CORE-система.\n---FEWSHOT---\nПРИМЕР-один\nПРИМЕР-два\n---USER---\n{{text}}")
	if tpl.System != "CORE-система." {
		t.Fatalf("System must be the core only, got %q", tpl.System)
	}
	if tpl.FewShot != "ПРИМЕР-один\nПРИМЕР-два" {
		t.Fatalf("FewShot must be the example block, got %q", tpl.FewShot)
	}
	if on := tpl.SystemFor(true); on != "CORE-система.\n\nПРИМЕР-один\nПРИМЕР-два" {
		t.Fatalf("few_shot ON must append the examples after a blank line, got %q", on)
	}
	if off := tpl.SystemFor(false); off != "CORE-система." || strings.Contains(off, "ПРИМЕР") {
		t.Fatalf("few_shot OFF must be core-only, got %q", off)
	}
}

// TestFewShotAbsentIsNoop: a template with no ---FEWSHOT--- block leaves FewShot empty and the
// toggle a no-op, so every existing prompt (translator, golden stub) is untouched.
func TestFewShotAbsentIsNoop(t *testing.T) {
	tpl := writeTemplate(t, "Только система.\n---USER---\n{{text}}")
	if tpl.FewShot != "" {
		t.Fatalf("no ---FEWSHOT--- must leave FewShot empty, got %q", tpl.FewShot)
	}
	if tpl.SystemFor(true) != tpl.System || tpl.SystemFor(false) != tpl.System {
		t.Fatal("the few_shot toggle must be a no-op without a ---FEWSHOT--- block")
	}
}

// TestFewShotEnabledDefault: an absent flag defaults to ON (preserve current behaviour); an
// explicit false/true resolves as written.
func TestFewShotEnabledDefault(t *testing.T) {
	if !fewShotEnabled(config.Stage{}) {
		t.Fatal("absent few_shot must default to ON")
	}
	off := false
	if fewShotEnabled(config.Stage{FewShot: &off}) {
		t.Fatal("few_shot:false must disable the examples")
	}
	on := true
	if !fewShotEnabled(config.Stage{FewShot: &on}) {
		t.Fatal("few_shot:true must enable the examples")
	}
}

// TestEditorFewShotToggleRealFile guards the boevoy editor.md (D38.4): it carries a
// ---FEWSHOT--- block, few_shot ON reproduces the ratified P1a discourse (core + examples +
// the 物是人非 chengyu-atom), and few_shot OFF keeps the discourse core (incl. the chengyu-atom)
// but drops the examples — never losing the meaning-preservation instructions.
func TestEditorFewShotToggleRealFile(t *testing.T) {
	tpl, err := LoadPromptTemplate(filepath.Join("..", "..", "prompts", "editor.md"))
	if err != nil {
		t.Fatalf("load editor.md: %v", err)
	}
	if tpl.FewShot == "" {
		t.Fatal("editor.md must carry a ---FEWSHOT--- block (D38.4 P1a examples)")
	}
	on, off := tpl.SystemFor(true), tpl.SystemFor(false)
	// Core discourse + chengyu-atom present in BOTH modes.
	for _, mode := range []struct {
		name string
		s    string
	}{{"few_shot ON", on}, {"few_shot OFF", off}} {
		if !strings.Contains(mode.s, "ДИСКУРС-ПЕРЕВЁРСТКА") {
			t.Errorf("%s must keep the discourse core", mode.name)
		}
		if !strings.Contains(mode.s, "物是人非") {
			t.Errorf("%s must keep the chengyu-atom (物是人非 — both semantic components)", mode.name)
		}
	}
	// Examples ONLY in the few-shot mode.
	if !strings.Contains(on, "Примеры требуемой") || !strings.Contains(on, "[focal-shift]") {
		t.Error("few_shot ON must include the P1a examples")
	}
	if strings.Contains(off, "Примеры требуемой") || strings.Contains(off, "[narrative") {
		t.Error("few_shot OFF (zero-shot) must drop the examples")
	}
}

// setupFewShotProject builds a minimal project whose EDITOR prompt has a ---FEWSHOT--- block
// and whose editor stage carries editStageExtra (e.g. ", few_shot: false"). It returns the
// resolved snapshot id and payload. The translator prompt has NO block (control).
func setupFewShotProject(t *testing.T, editStageExtra string) (string, string) {
	t.Helper()
	dir := t.TempDir()
	writeFile(t, filepath.Join(dir, "prompts", "translator.md"),
		"Переводи с {{source_lang}} на {{target_lang}}.\n---USER---\n{{text}}")
	writeFile(t, filepath.Join(dir, "prompts", "editor.md"),
		"CORE-редактор {{genre}}.\n---FEWSHOT---\n[ex] пример перевёрстки.\n---USER---\nИсходник: {{text}}\nЧерновик: {{draft}}")
	writeFile(t, filepath.Join(dir, "models.yaml"), fmt.Sprintf(`
prices_checked: %q
default_model: fake-model
providers:
  fake:
    kind: openai
    base_url: http://x
    timeouts: { attempt_s: 5, max_attempts: 2, backoff_cap_s: 1 }
models:
  fake-model:
    provider: fake
    price: { input_per_m: 1.0, cached_per_m: 0.1, cache_write_per_m: 0, output_per_m: 2.0 }
`, time.Now().UTC().Format("2006-01-02")))
	writeFile(t, filepath.Join(dir, "pipeline.yaml"), fmt.Sprintf(`
core: C1
version: 1
defaults: { max_output_ratio: 2.0, min_max_tokens: 512 }
retries: { regenerate_before_escalate: 0 }
context: { glossary_injection: selective, glossary_token_budget: 800 }
stages:
  - { name: draft, role: translator, model: fake-model, prompt: prompts/translator.md, prompt_version: v-test, temperature: 0.3, reasoning: "off" }
  - { name: edit, role: editor, model: fake-model, prompt: prompts/editor.md, prompt_version: v-test, temperature: 0.4, reasoning: "off"%s }
`, editStageExtra))
	writeFile(t, filepath.Join(dir, "source.txt"), "静かな朝。")
	writeFile(t, filepath.Join(dir, "book.yaml"), `
book_id: test-book
title: Тест
source_lang: ja
target_lang: ru
genre: ранобэ
audience: тест
venuti: 0.5
honorifics: keep
transcription: polivanov
footnotes: minimal
pipeline: pipeline.yaml
models: models.yaml
source_file: source.txt
ceilings: { book_usd: 1.0, day_usd: 2.0 }
`)
	r, err := NewRunner(filepath.Join(dir, "book.yaml"), obs.NewLogger())
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	id, payload, err := r.snapshotID()
	if err != nil {
		t.Fatal(err)
	}
	return id, payload
}

// TestFewShotSnapshotFold: flipping the editor stage's few_shot must be a LOUD --resnapshot (the
// snapshot id changes and the resolved flag appears in the payload), never a silent false-hit —
// the ---FEWSHOT--- block is inside PromptSHA256 but dropping it at render leaves that hash
// unchanged while the wire changes. The translator stage (no block) must NOT fold the field.
func TestFewShotSnapshotFold(t *testing.T) {
	idOn, payloadOn := setupFewShotProject(t, "")                    // default → ON
	idOff, payloadOff := setupFewShotProject(t, ", few_shot: false") // zero-shot

	if idOn == idOff {
		t.Fatal("flipping few_shot must change the snapshot id (loud --resnapshot)")
	}
	if !strings.Contains(payloadOn, `"few_shot":true`) {
		t.Fatalf("default few_shot must fold as true in the snapshot, got: %s", payloadOn)
	}
	if !strings.Contains(payloadOff, `"few_shot":false`) {
		t.Fatalf("few_shot:false must fold as false in the snapshot, got: %s", payloadOff)
	}
	// Only the editor stage carries a ---FEWSHOT--- block, so the field must appear exactly once
	// (the translator stage stays byte-identical in the snapshot).
	if n := strings.Count(payloadOn, `"few_shot"`); n != 1 {
		t.Fatalf("only the block-carrying stage must fold few_shot, want 1 occurrence, got %d: %s", n, payloadOn)
	}
}

// TestFewShotOffRendersCoreOnlyWire: a template with a block but few_shot resolved off must send
// a system message WITHOUT the examples — the wire-level guarantee behind the snapshot fold.
func TestFewShotOffRendersCoreOnlyWire(t *testing.T) {
	tpl := writeTemplate(t, "CORE.\n---FEWSHOT---\nПРИМЕР.\n---USER---\n{{text}}")
	tpl.System = tpl.SystemFor(false) // mirror loadTemplates' per-stage resolution
	msgs, err := Messages(tpl, RenderVars{Book: testBook(), Text: "x"})
	if err != nil {
		t.Fatal(err)
	}
	if got := msgs[0].Content; got != "CORE." {
		t.Fatalf("zero-shot system message must be core-only, got %q", got)
	}
}