dify/knowledge-fs/packages/api/src/retrieval-planner.test.ts

import { describe, expect, it } from "vitest";

import { createRetrievalPlanner, defaultRetrievalPlan } from "./retrieval-planner";
import { createInMemoryTraceRecorder } from "./tracing";

describe("retrieval planner", () => {
  it("plans resolved modes with bounded fanout and traced route provenance", () => {
    const traces = createInMemoryTraceRecorder();
    const planner = createRetrievalPlanner({ maxTopK: 100, traces });

    expect(
      planner.plan({
        mode: "fast",
        query: "contract renewal",
        topK: 5,
        traceId: "trace-fast",
      }),
    ).toEqual({
      denseTopK: 5,
      ftsTopK: 5,
      fusionLimit: 5,
      queryLanguage: "latin",
      requestedMode: "fast",
      rerankCandidateLimit: 5,
      resolvedMode: "fast",
      strategyVersion: "retrieval-planner-v1",
      topK: 5,
    });

    expect(
      planner.plan({
        mode: "auto",
        query: "比较 合同 ABC-123 的续约条款和 termination notice 风险",
        resolvedMode: "research",
        topK: 4,
        traceId: "trace-auto",
      }),
    ).toEqual(
      expect.objectContaining({
        denseTopK: 40,
        ftsTopK: 40,
        fusionLimit: 20,
        queryLanguage: "mixed-cjk-latin",
        requestedMode: "auto",
        rerankCandidateLimit: 20,
        resolvedMode: "research",
        strategyVersion: "retrieval-planner-v2",
        topK: 4,
      }),
    );

    expect(traces.spans).toEqual([
      {
        attributes: {
          denseTopK: 5,
          ftsTopK: 5,
          fusionLimit: 5,
          queryLanguage: "latin",
          requestedMode: "fast",
          resolvedMode: "fast",
          rerankCandidateLimit: 5,
          topK: 5,
          traceId: "trace-fast",
        },
        name: "retrieval.plan",
        status: "ok",
      },
      {
        attributes: {
          denseTopK: 40,
          ftsTopK: 40,
          fusionLimit: 20,
          queryLanguage: "mixed-cjk-latin",
          requestedMode: "auto",
          resolvedMode: "research",
          rerankCandidateLimit: 20,
          topK: 4,
          traceId: "trace-auto",
        },
        name: "retrieval.plan",
        status: "ok",
      },
    ]);
    expect(JSON.stringify(traces.spans)).not.toContain("termination notice");

    expect(() => planner.plan({ query: "too many", topK: 101 })).toThrow(
      "Retrieval planner topK exceeds maxTopK=100",
    );
    expect(() => createRetrievalPlanner({ maxTopK: 0 })).toThrow(
      "Retrieval planner maxTopK must be at least 1",
    );
  });

  it("requires Auto to be LLM-resolved before deterministic planning", () => {
    const traces = createInMemoryTraceRecorder();
    const planner = createRetrievalPlanner({ maxTopK: 100, traces });

    expect(
      planner.plan({
        mode: "deep",
        query: "contract renewal notice liability terms",
        topK: 7,
      }),
    ).toEqual(
      expect.objectContaining({
        denseTopK: 35,
        ftsTopK: 35,
        fusionLimit: 21,
        queryLanguage: "latin",
        requestedMode: "deep",
        rerankCandidateLimit: 21,
        resolvedMode: "deep",
        topK: 7,
      }),
    );
    expect(
      planner.plan({
        mode: "research",
        query: "research contract history",
        topK: 30,
      }),
    ).toEqual(
      expect.objectContaining({
        denseTopK: 100,
        ftsTopK: 100,
        fusionLimit: 100,
        requestedMode: "research",
        rerankCandidateLimit: 100,
        resolvedMode: "research",
        strategyVersion: "retrieval-planner-v2",
        topK: 30,
      }),
    );
    expect(
      planner.plan({
        mode: "auto",
        query: "合同续约条款",
        resolvedMode: "deep",
        topK: 3,
      }),
    ).toEqual(
      expect.objectContaining({
        denseTopK: 15,
        ftsTopK: 15,
        fusionLimit: 9,
        queryLanguage: "cjk",
        resolvedMode: "deep",
      }),
    );
    expect(
      planner.plan({
        mode: "auto",
        query: "合同编号是什么？",
        resolvedMode: "fast",
        topK: 2,
      }),
    ).toEqual(
      expect.objectContaining({
        queryLanguage: "cjk",
        requestedMode: "auto",
        resolvedMode: "fast",
      }),
    );
    expect(
      planner.plan({
        mode: "auto",
        query:
          "Analyze and explain the evidence in this deliberately long request while following the linked dependency chain",
        resolvedMode: "deep",
        topK: 2,
      }),
    ).toEqual(
      expect.objectContaining({
        queryLanguage: "latin",
        requestedMode: "auto",
        resolvedMode: "deep",
      }),
    );
    expect(planner.plan({ query: "σύμβαση", topK: 2 })).toEqual(
      expect.objectContaining({
        queryLanguage: "other",
        requestedMode: "fast",
        resolvedMode: "fast",
      }),
    );

    expect(() =>
      planner.plan({ mode: "auto", query: "must be classified upstream", topK: 1 }),
    ).toThrow("Retrieval planner auto mode requires an LLM-resolved mode");
    expect(() =>
      planner.plan({
        mode: "fast",
        query: "explicit mode cannot be overwritten",
        resolvedMode: "deep",
        topK: 1,
      }),
    ).toThrow("Retrieval planner resolved mode must match an explicit requested mode");
    expect(() => planner.plan({ query: "valid", topK: 0 })).toThrow(
      "Retrieval planner topK must be at least 1",
    );
    expect(() => planner.plan({ query: "   ", topK: 1 })).toThrow(
      "Retrieval planner query must not be empty",
    );
    expect(traces.spans.filter((span) => span.status === "error")).toHaveLength(4);
  });

  it("plans image-only retrieval while retaining explicit empty-input errors", () => {
    const planner = createRetrievalPlanner({ maxTopK: 10 });

    expect(planner.plan({ hasQueryImages: true, query: "", topK: 3 })).toEqual(
      expect.objectContaining({ queryLanguage: "other", topK: 3 }),
    );
    expect(defaultRetrievalPlan({ hasQueryImages: true, query: "", topK: 2 })).toEqual(
      expect.objectContaining({ queryLanguage: "other", topK: 2 }),
    );
    expect(() => defaultRetrievalPlan({ query: "", topK: 1 })).toThrow(
      "Retrieval planner query must not be empty",
    );
    expect(() => defaultRetrievalPlan({ query: "valid", topK: 0 })).toThrow(
      "Retrieval planner topK must be at least 1",
    );
    expect(() => planner.plan({ hasQueryImages: false, query: "", topK: 1 })).toThrow(
      "Retrieval planner query or query images must be provided",
    );
  });
});