dify/sdks/nodejs-client/src/client/knowledge-base.test.js

import { beforeEach, describe, expect, it, vi } from "vitest";
import { KnowledgeBaseClient } from "./knowledge-base";
import { createHttpClientWithSpies } from "../../tests/test-utils";

describe("KnowledgeBaseClient", () => {
  beforeEach(() => {
    vi.restoreAllMocks();
  });

  it("handles dataset and tag operations", async () => {
    const { client, request } = createHttpClientWithSpies();
    const kb = new KnowledgeBaseClient(client);

    await kb.listDatasets({
      page: 1,
      limit: 2,
      keyword: "k",
      includeAll: true,
      tagIds: ["t1"],
    });
    await kb.createDataset({ name: "dataset" });
    await kb.getDataset("ds");
    await kb.updateDataset("ds", { name: "new" });
    await kb.deleteDataset("ds");
    await kb.updateDocumentStatus("ds", "enable", ["doc1"]);

    await kb.listTags();
    await kb.createTag({ name: "tag" });
    await kb.updateTag({ tag_id: "tag", name: "name" });
    await kb.deleteTag({ tag_id: "tag" });
    await kb.bindTags({ tag_ids: ["tag"], target_id: "doc" });
    await kb.unbindTags({ tag_id: "tag", target_id: "doc" });
    await kb.getDatasetTags("ds");

    expect(request).toHaveBeenCalledWith({
      method: "GET",
      path: "/datasets",
      query: {
        page: 1,
        limit: 2,
        keyword: "k",
        include_all: true,
        tag_ids: ["t1"],
      },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets",
      data: { name: "dataset" },
    });
    expect(request).toHaveBeenCalledWith({
      method: "PATCH",
      path: "/datasets/ds",
      data: { name: "new" },
    });
    expect(request).toHaveBeenCalledWith({
      method: "PATCH",
      path: "/datasets/ds/documents/status/enable",
      data: { document_ids: ["doc1"] },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/tags/binding",
      data: { tag_ids: ["tag"], target_id: "doc" },
    });
  });

  it("handles document operations", async () => {
    const { client, request } = createHttpClientWithSpies();
    const kb = new KnowledgeBaseClient(client);
    const form = { append: vi.fn(), getHeaders: () => ({}) };

    await kb.createDocumentByText("ds", { name: "doc", text: "text" });
    await kb.updateDocumentByText("ds", "doc", { name: "doc2" });
    await kb.createDocumentByFile("ds", form);
    await kb.updateDocumentByFile("ds", "doc", form);
    await kb.listDocuments("ds", { page: 1, limit: 20, keyword: "k" });
    await kb.getDocument("ds", "doc", { metadata: "all" });
    await kb.deleteDocument("ds", "doc");
    await kb.getDocumentIndexingStatus("ds", "batch");

    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/document/create_by_text",
      data: { name: "doc", text: "text" },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/documents/doc/update_by_text",
      data: { name: "doc2" },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/document/create_by_file",
      data: form,
    });
    expect(request).toHaveBeenCalledWith({
      method: "GET",
      path: "/datasets/ds/documents",
      query: { page: 1, limit: 20, keyword: "k", status: undefined },
    });
  });

  it("handles segments and child chunks", async () => {
    const { client, request } = createHttpClientWithSpies();
    const kb = new KnowledgeBaseClient(client);

    await kb.createSegments("ds", "doc", { segments: [{ content: "x" }] });
    await kb.listSegments("ds", "doc", { page: 1, limit: 10, keyword: "k" });
    await kb.getSegment("ds", "doc", "seg");
    await kb.updateSegment("ds", "doc", "seg", {
      segment: { content: "y" },
    });
    await kb.deleteSegment("ds", "doc", "seg");

    await kb.createChildChunk("ds", "doc", "seg", { content: "c" });
    await kb.listChildChunks("ds", "doc", "seg", { page: 1, limit: 10 });
    await kb.updateChildChunk("ds", "doc", "seg", "child", {
      content: "c2",
    });
    await kb.deleteChildChunk("ds", "doc", "seg", "child");

    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/documents/doc/segments",
      data: { segments: [{ content: "x" }] },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/documents/doc/segments/seg",
      data: { segment: { content: "y" } },
    });
    expect(request).toHaveBeenCalledWith({
      method: "PATCH",
      path: "/datasets/ds/documents/doc/segments/seg/child_chunks/child",
      data: { content: "c2" },
    });
  });

  it("handles metadata and retrieval", async () => {
    const { client, request } = createHttpClientWithSpies();
    const kb = new KnowledgeBaseClient(client);

    await kb.listMetadata("ds");
    await kb.createMetadata("ds", { name: "m", type: "string" });
    await kb.updateMetadata("ds", "mid", { name: "m2" });
    await kb.deleteMetadata("ds", "mid");
    await kb.listBuiltInMetadata("ds");
    await kb.updateBuiltInMetadata("ds", "enable");
    await kb.updateDocumentsMetadata("ds", {
      operation_data: [
        { document_id: "doc", metadata_list: [{ id: "m", name: "n" }] },
      ],
    });
    await kb.hitTesting("ds", { query: "q" });
    await kb.retrieve("ds", { query: "q" });

    expect(request).toHaveBeenCalledWith({
      method: "GET",
      path: "/datasets/ds/metadata",
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/metadata",
      data: { name: "m", type: "string" },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/hit-testing",
      data: { query: "q" },
    });
  });

  it("handles pipeline operations", async () => {
    const { client, request, requestStream } = createHttpClientWithSpies();
    const kb = new KnowledgeBaseClient(client);
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const form = { append: vi.fn(), getHeaders: () => ({}) };

    await kb.listDatasourcePlugins("ds", { isPublished: true });
    await kb.runDatasourceNode("ds", "node", {
      inputs: { input: "x" },
      datasource_type: "custom",
      is_published: true,
    });
    await kb.runPipeline("ds", {
      inputs: { input: "x" },
      datasource_type: "custom",
      datasource_info_list: [],
      start_node_id: "start",
      is_published: true,
      response_mode: "streaming",
    });
    await kb.runPipeline("ds", {
      inputs: { input: "x" },
      datasource_type: "custom",
      datasource_info_list: [],
      start_node_id: "start",
      is_published: true,
      response_mode: "blocking",
    });
    await kb.uploadPipelineFile(form);

    expect(warn).toHaveBeenCalled();
    expect(request).toHaveBeenCalledWith({
      method: "GET",
      path: "/datasets/ds/pipeline/datasource-plugins",
      query: { is_published: true },
    });
    expect(requestStream).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/pipeline/datasource/nodes/node/run",
      data: {
        inputs: { input: "x" },
        datasource_type: "custom",
        is_published: true,
      },
    });
    expect(requestStream).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/pipeline/run",
      data: {
        inputs: { input: "x" },
        datasource_type: "custom",
        datasource_info_list: [],
        start_node_id: "start",
        is_published: true,
        response_mode: "streaming",
      },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/ds/pipeline/run",
      data: {
        inputs: { input: "x" },
        datasource_type: "custom",
        datasource_info_list: [],
        start_node_id: "start",
        is_published: true,
        response_mode: "blocking",
      },
    });
    expect(request).toHaveBeenCalledWith({
      method: "POST",
      path: "/datasets/pipeline/file-upload",
      data: form,
    });
  });
});