src/pipeline/rerank.test.ts

import { describe, it, expect } from "vitest";
import { buildRerankerPrompt } from "./rerank";

// ---------------------------------------------------------------------------
// buildRerankerPrompt
// ---------------------------------------------------------------------------
describe("buildRerankerPrompt", () => {
  const SYSTEM =
    'Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".';

  const DEFAULT_INSTRUCTION =
    "Given a web search query, retrieve relevant passages that answer the query";

  it("uses default instruction when none is provided", () => {
    const prompt = buildRerankerPrompt("test query", "test document");
    expect(prompt).toContain(`<Instruct>: ${DEFAULT_INSTRUCTION}`);
    expect(prompt).toContain("<Query>: test query");
    expect(prompt).toContain("<Document>: test document");
  });

  it("uses custom instruction when provided", () => {
    const custom = "Find code examples that demonstrate this concept";
    const prompt = buildRerankerPrompt("test query", "test doc", custom);
    expect(prompt).toContain(`<Instruct>: ${custom}`);
    expect(prompt).not.toContain(DEFAULT_INSTRUCTION);
  });

  it("includes the system message", () => {
    const prompt = buildRerankerPrompt("q", "d");
    expect(prompt).toContain(`<|im_start|>system\n${SYSTEM}<|im_end|>`);
  });

  it("ends with assistant think block", () => {
    const prompt = buildRerankerPrompt("q", "d");
    expect(prompt.endsWith("<|im_start|>assistant\n<think>\n\n</think>\n")).toBe(
      true,
    );
  });

  it("wraps user content in im_start/im_end tags", () => {
    const prompt = buildRerankerPrompt("my query", "my doc");
    expect(prompt).toContain("<|im_start|>user\n");
    expect(prompt).toContain("<|im_end|>\n<|im_start|>assistant");
  });

  it("preserves query and document content exactly", () => {
    const query = "How does SQLite FTS5 work?";
    const doc = "FTS5 is an extension for full-text search.\nIt supports BM25.";
    const prompt = buildRerankerPrompt(query, doc);
    expect(prompt).toContain(`<Query>: ${query}`);
    expect(prompt).toContain(`<Document>: ${doc}`);
  });

  it("produces the exact expected format", () => {
    const prompt = buildRerankerPrompt("q", "d");
    const expected =
      `<|im_start|>system\n${SYSTEM}<|im_end|>\n` +
      `<|im_start|>user\n<Instruct>: ${DEFAULT_INSTRUCTION}\n\n<Query>: q\n\n<Document>: d<|im_end|>\n` +
      `<|im_start|>assistant\n<think>\n\n</think>\n`;
    expect(prompt).toBe(expected);
  });
});