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

import { beforeEach, describe, expect, it, vi } from 'vitest'
import { createHttpClientWithSpies } from '../../tests/test-utils'
import { FileUploadError, ValidationError } from '../errors/dify-error'
import { KnowledgeBaseClient } from './knowledge-base'

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 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(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,
    })
  })

  it('validates form-data and optional array filters', async () => {
    const { client } = createHttpClientWithSpies()
    const kb = new KnowledgeBaseClient(client)

    await expect(kb.createDocumentByFile('ds', {})).rejects.toBeInstanceOf(FileUploadError)
    await expect(
      kb.listSegments('ds', 'doc', { status: ['ok', 1] as unknown as string[] }),
    ).rejects.toBeInstanceOf(ValidationError)
    await expect(
      kb.hitTesting('ds', {
        query: 'q',
        attachment_ids: ['att-1', 2] as unknown as string[],
      }),
    ).rejects.toBeInstanceOf(ValidationError)
  })
})