mirror of
https://github.com/simstudioai/sim.git
synced 2026-10-24 11:36:37 +08:00
fix(knowledge): bound chunking separators so one config can't stall processing (#6735)
* fix(knowledge): bound chunking separators so one config can't stall processing `chunkingStrategyOptionsSchema.separators` accepted an arbitrary-length array of arbitrary-length strings, next to a `pattern` field already capped at 500 chars. `RecursiveChunker` splits the whole document once per separator and walks the list from the top for every oversized fragment, so a persisted config with thousands of non-matching separators cost seconds of synchronous CPU on every later document upload — work neither the processing `Promise.race` timeout nor the after-the-fact chunk-count cap can interrupt. Measured on a 21.3 MB document: 632 ms at 100 separators, 6.1 s at 1000, 36.8 s at 5000. - Bound `separators` to 32 entries of at most 100 characters on the write path. The largest built-in recipe (markdown) uses 16, so hand-tuned lists still fit. - Keep the stored/read shape tolerant, so a config written before the bound still lists instead of failing response validation. - Clamp in `RecursiveChunker` too, with a warning, so an already-persisted oversized list cannot reach the split loop. An over-long separator is dropped rather than truncated: a truncated separator matches where the configured one never did, silently re-cutting the document, while dropping it behaves like a separator that finds no match. A list left empty falls back to the recipe. - Walk non-matching separators iteratively instead of recursing, so stack depth no longer tracks the separator count. Verified behavior-preserving against the previous implementation over 4000 randomized configs — byte-identical output. - Validate in the create-base modal so the limit surfaces inline. After the fix the same 21.3 MB document costs ~300 ms at every separator count. * fix(knowledge): gate separator validation on the recursive strategy - The separator refines ran for every strategy, but the field only renders for `recursive` and only that strategy submits it, so a value left behind by a strategy switch could block submit with no visible field to clear. Gated the same way the regex-pattern refine already is. - Use absolute imports in the chunker test, per the repo convention.
This commit is contained in:
+38
-6
@@ -25,6 +25,7 @@ import { getErrorMessage } from '@sim/utils/errors'
|
||||
import { useParams } from 'next/navigation'
|
||||
import { type FieldErrors, useForm } from 'react-hook-form'
|
||||
import { z } from 'zod'
|
||||
import { MAX_CHUNKING_SEPARATOR_LENGTH, MAX_CHUNKING_SEPARATORS } from '@/lib/chunkers/constants'
|
||||
import type { StrategyOptions } from '@/lib/chunkers/types'
|
||||
import { KNOWLEDGE_BASE_DESCRIPTION_MAX_LENGTH } from '@/lib/knowledge/constants'
|
||||
import {
|
||||
@@ -57,6 +58,14 @@ const STRATEGY_OPTIONS = [
|
||||
{ value: 'regex', label: 'Regex (custom pattern)' },
|
||||
] as const
|
||||
|
||||
/** Splits the comma-separated separator field into the list the API receives. */
|
||||
function parseSeparators(value: string | undefined): string[] {
|
||||
if (!value?.trim()) return []
|
||||
return value
|
||||
.split(',')
|
||||
.map((separator) => separator.trim().replace(/\\n/g, '\n').replace(/\\t/g, '\t'))
|
||||
}
|
||||
|
||||
const STRATEGY_COMBOBOX_OPTIONS: ComboboxOption[] = STRATEGY_OPTIONS.map((o) => ({
|
||||
label: o.label,
|
||||
value: o.value,
|
||||
@@ -124,6 +133,31 @@ const FormSchema = z
|
||||
path: ['regexPattern'],
|
||||
}
|
||||
)
|
||||
/**
|
||||
* Gated on the strategy for the same reason the regex pattern is: the field only
|
||||
* renders for `recursive` and only that strategy submits it, so an out-of-bound
|
||||
* value left behind by a strategy switch must not block a submit that drops it.
|
||||
*/
|
||||
.refine(
|
||||
(data) =>
|
||||
data.strategy !== 'recursive' ||
|
||||
parseSeparators(data.customSeparators).length <= MAX_CHUNKING_SEPARATORS,
|
||||
{
|
||||
message: `At most ${MAX_CHUNKING_SEPARATORS} separators are allowed`,
|
||||
path: ['customSeparators'],
|
||||
}
|
||||
)
|
||||
.refine(
|
||||
(data) =>
|
||||
data.strategy !== 'recursive' ||
|
||||
parseSeparators(data.customSeparators).every(
|
||||
(separator) => separator.length <= MAX_CHUNKING_SEPARATOR_LENGTH
|
||||
),
|
||||
{
|
||||
message: `Each separator must be ${MAX_CHUNKING_SEPARATOR_LENGTH} characters or less`,
|
||||
path: ['customSeparators'],
|
||||
}
|
||||
)
|
||||
|
||||
type FormInputValues = z.input<typeof FormSchema>
|
||||
type FormValues = z.output<typeof FormSchema>
|
||||
@@ -265,11 +299,7 @@ export const CreateBaseModal = memo(function CreateBaseModal({
|
||||
...(data.regexStrictBoundaries && { strictBoundaries: true }),
|
||||
}
|
||||
: data.strategy === 'recursive' && data.customSeparators?.trim()
|
||||
? {
|
||||
separators: data.customSeparators
|
||||
.split(',')
|
||||
.map((s) => s.trim().replace(/\\n/g, '\n').replace(/\\t/g, '\t')),
|
||||
}
|
||||
? { separators: parseSeparators(data.customSeparators) }
|
||||
: undefined
|
||||
|
||||
const newKnowledgeBase = await createKnowledgeBaseMutation.mutateAsync({
|
||||
@@ -465,11 +495,13 @@ export const CreateBaseModal = memo(function CreateBaseModal({
|
||||
<ChipModalField
|
||||
type='custom'
|
||||
title='Custom Separators (optional)'
|
||||
hint='Comma-separated list of delimiters in priority order. Leave empty for default separators.'
|
||||
hint={`Comma-separated list of delimiters in priority order, up to ${MAX_CHUNKING_SEPARATORS}. Leave empty for default separators.`}
|
||||
error={errors.customSeparators?.message}
|
||||
>
|
||||
<ChipInput
|
||||
placeholder='e.g. \n\n, \n, . , '
|
||||
{...register('customSeparators')}
|
||||
error={Boolean(errors.customSeparators)}
|
||||
autoComplete='off'
|
||||
data-form-type='other'
|
||||
/>
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import {
|
||||
chunkingStrategyOptionsSchema,
|
||||
createKnowledgeBaseBodySchema,
|
||||
knowledgeBaseDataSchema,
|
||||
} from '@/lib/api/contracts/knowledge/base'
|
||||
import { MAX_CHUNKING_SEPARATOR_LENGTH, MAX_CHUNKING_SEPARATORS } from '@/lib/chunkers/constants'
|
||||
|
||||
const separators = (count: number) => Array.from({ length: count }, (_, i) => `@@sep${i}@@`)
|
||||
|
||||
describe('chunkingStrategyOptionsSchema.separators', () => {
|
||||
it('accepts a separator list at the bound', () => {
|
||||
const parsed = chunkingStrategyOptionsSchema.parse({
|
||||
separators: separators(MAX_CHUNKING_SEPARATORS),
|
||||
})
|
||||
expect(parsed.separators).toHaveLength(MAX_CHUNKING_SEPARATORS)
|
||||
})
|
||||
|
||||
it('rejects more separators than the bound', () => {
|
||||
const result = chunkingStrategyOptionsSchema.safeParse({
|
||||
separators: separators(MAX_CHUNKING_SEPARATORS + 1),
|
||||
})
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
|
||||
it('rejects a separator longer than the per-item bound', () => {
|
||||
const result = chunkingStrategyOptionsSchema.safeParse({
|
||||
separators: ['|'.repeat(MAX_CHUNKING_SEPARATOR_LENGTH + 1)],
|
||||
})
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
|
||||
it('rejects an oversized list on the knowledge base create body', () => {
|
||||
const result = createKnowledgeBaseBodySchema.safeParse({
|
||||
name: 'kb',
|
||||
workspaceId: 'ws',
|
||||
chunkingConfig: {
|
||||
maxSize: 1024,
|
||||
minSize: 100,
|
||||
overlap: 200,
|
||||
strategy: 'recursive',
|
||||
strategyOptions: { separators: separators(5000) },
|
||||
},
|
||||
})
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe('knowledgeBaseDataSchema.chunkingConfig', () => {
|
||||
it('still reads a stored config written before the separator bound', () => {
|
||||
const result = knowledgeBaseDataSchema.safeParse({
|
||||
id: 'kb-1',
|
||||
userId: 'u-1',
|
||||
name: 'kb',
|
||||
description: null,
|
||||
tokenCount: 0,
|
||||
embeddingModel: 'text-embedding-3-small',
|
||||
embeddingDimension: 1536,
|
||||
chunkingConfig: {
|
||||
maxSize: 1024,
|
||||
minSize: 100,
|
||||
overlap: 200,
|
||||
strategy: 'recursive',
|
||||
strategyOptions: { separators: separators(5000) },
|
||||
},
|
||||
createdAt: new Date().toISOString(),
|
||||
updatedAt: new Date().toISOString(),
|
||||
deletedAt: null,
|
||||
workspaceId: 'ws',
|
||||
folderId: null,
|
||||
})
|
||||
expect(result.success).toBe(true)
|
||||
})
|
||||
})
|
||||
@@ -11,6 +11,7 @@ import {
|
||||
workspaceIdSchema,
|
||||
} from '@/lib/api/contracts/primitives'
|
||||
import { defineRouteContract } from '@/lib/api/contracts/types'
|
||||
import { MAX_CHUNKING_SEPARATOR_LENGTH, MAX_CHUNKING_SEPARATORS } from '@/lib/chunkers/constants'
|
||||
import type { StrategyOptions } from '@/lib/chunkers/types'
|
||||
import {
|
||||
DEFAULT_CHUNKING_CONFIG,
|
||||
@@ -26,7 +27,13 @@ export const listKnowledgeBasesQuerySchema = z.object({
|
||||
scope: knowledgeScopeSchema.default('active'),
|
||||
})
|
||||
|
||||
export const chunkingStrategyOptionsSchema = z
|
||||
/**
|
||||
* Strategy options as they are stored. Reads stay tolerant of a `separators`
|
||||
* list written before {@link chunkingStrategyOptionsSchema} bounded it, so an
|
||||
* oversized legacy config lists instead of failing response validation. The
|
||||
* chunker clamps such a list at construction, so nothing reprocesses unbounded.
|
||||
*/
|
||||
export const storedChunkingStrategyOptionsSchema = z
|
||||
.object({
|
||||
pattern: z
|
||||
.string()
|
||||
@@ -48,6 +55,29 @@ export const chunkingStrategyOptionsSchema = z
|
||||
})
|
||||
.strict() satisfies z.ZodType<StrategyOptions>
|
||||
|
||||
/**
|
||||
* Strategy options accepted on writes. `separators` is bounded in both length
|
||||
* and item size: the recursive chunker rescans the whole document once per
|
||||
* separator, synchronously, so an unbounded list turns one persisted config
|
||||
* into seconds of uninterruptible CPU on every later document upload.
|
||||
*/
|
||||
export const chunkingStrategyOptionsSchema = storedChunkingStrategyOptionsSchema
|
||||
.extend({
|
||||
separators: z
|
||||
.array(
|
||||
z
|
||||
.string()
|
||||
.max(
|
||||
MAX_CHUNKING_SEPARATOR_LENGTH,
|
||||
`Each separator must be ${MAX_CHUNKING_SEPARATOR_LENGTH} characters or less`
|
||||
)
|
||||
)
|
||||
.max(MAX_CHUNKING_SEPARATORS, `At most ${MAX_CHUNKING_SEPARATORS} separators are allowed`)
|
||||
.optional()
|
||||
.describe('Ordered separators used to split content into chunks.'),
|
||||
})
|
||||
.strict() satisfies z.ZodType<StrategyOptions>
|
||||
|
||||
export const chunkingConfigSchema = z
|
||||
.object({
|
||||
maxSize: z.number().min(100).max(4000),
|
||||
@@ -116,7 +146,7 @@ const knowledgeChunkingConfigSchema = z
|
||||
minSize: z.number(),
|
||||
overlap: z.number(),
|
||||
strategy: z.enum(['auto', 'text', 'regex', 'recursive', 'sentence', 'token']).optional(),
|
||||
strategyOptions: chunkingStrategyOptionsSchema.optional(),
|
||||
strategyOptions: storedChunkingStrategyOptionsSchema.optional(),
|
||||
})
|
||||
.passthrough()
|
||||
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
/**
|
||||
* Bounds on the separator list a recursive chunking config may carry.
|
||||
*
|
||||
* `RecursiveChunker` scans the whole document once per separator and walks the
|
||||
* list from the top for every oversized fragment, so the separator count is a
|
||||
* direct multiplier on synchronous CPU per document. The work happens inside a
|
||||
* split loop, which neither the processing `Promise.race` timeout nor the
|
||||
* after-the-fact chunk-count cap can interrupt — the list has to be bounded on
|
||||
* the way in instead.
|
||||
*
|
||||
* The largest built-in recipe (`markdown`) uses 16 separators, so 32 leaves room
|
||||
* for a hand-tuned list without letting one config stall the processing tier.
|
||||
*/
|
||||
export const MAX_CHUNKING_SEPARATORS = 32
|
||||
|
||||
/** Max characters in a single chunking separator. Real delimiters are a few characters. */
|
||||
export const MAX_CHUNKING_SEPARATOR_LENGTH = 100
|
||||
@@ -3,7 +3,8 @@
|
||||
*/
|
||||
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { RecursiveChunker } from './recursive-chunker'
|
||||
import { MAX_CHUNKING_SEPARATOR_LENGTH, MAX_CHUNKING_SEPARATORS } from '@/lib/chunkers/constants'
|
||||
import { RecursiveChunker } from '@/lib/chunkers/recursive-chunker'
|
||||
|
||||
describe('RecursiveChunker', () => {
|
||||
describe('empty and whitespace input', () => {
|
||||
@@ -101,6 +102,63 @@ describe('RecursiveChunker', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('separator bounds', () => {
|
||||
it.concurrent('ignores separators past the list bound', async () => {
|
||||
const separators = [
|
||||
...Array.from({ length: MAX_CHUNKING_SEPARATORS }, (_, i) => `@@nomatch${i}@@`),
|
||||
'---',
|
||||
]
|
||||
const chunker = new RecursiveChunker({ chunkSize: 15, separators })
|
||||
const text =
|
||||
'Section one content here with words.---Section two content here with words.---Section three content here.'
|
||||
|
||||
const chunks = await chunker.chunk(text)
|
||||
|
||||
expect(chunks.length).toBeGreaterThan(1)
|
||||
expect(chunks.some((chunk) => chunk.text.includes('---'))).toBe(true)
|
||||
})
|
||||
|
||||
it.concurrent('splits on a separator that survives the clamp', async () => {
|
||||
const separators = [
|
||||
'---',
|
||||
...Array.from({ length: MAX_CHUNKING_SEPARATORS }, (_, i) => `@@nomatch${i}@@`),
|
||||
]
|
||||
const chunker = new RecursiveChunker({ chunkSize: 15, separators })
|
||||
const text =
|
||||
'Section one content here with words.---Section two content here with words.---Section three content here.'
|
||||
|
||||
const chunks = await chunker.chunk(text)
|
||||
|
||||
expect(chunks.length).toBeGreaterThan(1)
|
||||
expect(chunks.every((chunk) => !chunk.text.includes('---'))).toBe(true)
|
||||
})
|
||||
|
||||
it.concurrent('drops a separator longer than the per-item bound', async () => {
|
||||
const oversized = '|'.repeat(MAX_CHUNKING_SEPARATOR_LENGTH + 1)
|
||||
const chunker = new RecursiveChunker({ chunkSize: 15, separators: [oversized, '---'] })
|
||||
const text = `Section one content here.${oversized}Section two content.---Section three content.`
|
||||
|
||||
const chunks = await chunker.chunk(text)
|
||||
|
||||
expect(chunks.some((chunk) => chunk.text.includes('|'))).toBe(true)
|
||||
expect(chunks.every((chunk) => !chunk.text.includes('---'))).toBe(true)
|
||||
})
|
||||
|
||||
it.concurrent('falls back to the recipe when every separator is over the bound', async () => {
|
||||
const oversized = '|'.repeat(MAX_CHUNKING_SEPARATOR_LENGTH + 1)
|
||||
const text =
|
||||
'Section one content here with words.\n\nSection two content here with words.\n\nSection three content.'
|
||||
|
||||
const chunks = await new RecursiveChunker({ chunkSize: 15, separators: [oversized] }).chunk(
|
||||
text
|
||||
)
|
||||
const defaultChunks = await new RecursiveChunker({ chunkSize: 15 }).chunk(text)
|
||||
|
||||
expect(chunks.length).toBeGreaterThan(1)
|
||||
expect(chunks).toEqual(defaultChunks)
|
||||
})
|
||||
})
|
||||
|
||||
describe('recipe: plain', () => {
|
||||
it.concurrent('should use plain recipe by default', async () => {
|
||||
const chunker = new RecursiveChunker({ chunkSize: 20 })
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { MAX_CHUNKING_SEPARATOR_LENGTH, MAX_CHUNKING_SEPARATORS } from '@/lib/chunkers/constants'
|
||||
import type { Chunk, RecursiveChunkerOptions } from '@/lib/chunkers/types'
|
||||
import {
|
||||
addOverlap,
|
||||
@@ -62,8 +63,29 @@ export class RecursiveChunker {
|
||||
this.chunkSize = resolved.chunkSize
|
||||
this.chunkOverlap = resolved.chunkOverlap
|
||||
|
||||
if (options.separators && options.separators.length > 0) {
|
||||
this.separators = options.separators
|
||||
/**
|
||||
* Bounded here as well as at the API boundary: a config persisted before the
|
||||
* boundary bound existed would otherwise still cost one full document scan
|
||||
* per separator, synchronously, on every document it processes.
|
||||
*
|
||||
* An over-long separator is dropped rather than truncated — a truncated
|
||||
* separator matches where the configured one never did, silently re-cutting
|
||||
* the document, whereas dropping it behaves like a separator that finds no
|
||||
* match, which the split already handles.
|
||||
*/
|
||||
const requested = options.separators ?? []
|
||||
const usable = requested
|
||||
.filter((separator) => separator.length <= MAX_CHUNKING_SEPARATOR_LENGTH)
|
||||
.slice(0, MAX_CHUNKING_SEPARATORS)
|
||||
|
||||
if (usable.length < requested.length) {
|
||||
logger.warn(
|
||||
`Chunking config carries ${requested.length} separators; using ${usable.length} within the ${MAX_CHUNKING_SEPARATORS} × ${MAX_CHUNKING_SEPARATOR_LENGTH}-character bound`
|
||||
)
|
||||
}
|
||||
|
||||
if (usable.length > 0) {
|
||||
this.separators = usable
|
||||
} else {
|
||||
const recipe = options.recipe ?? 'plain'
|
||||
this.separators = [...RECIPES[recipe]]
|
||||
@@ -77,23 +99,36 @@ export class RecursiveChunker {
|
||||
return text.trim() ? [text] : []
|
||||
}
|
||||
|
||||
if (separatorIndex >= this.separators.length) {
|
||||
/**
|
||||
* Advance past separators that do not split this text. Iterating rather
|
||||
* than recursing keeps stack depth independent of the separator count.
|
||||
*/
|
||||
let index = separatorIndex
|
||||
let separator = ''
|
||||
let parts: string[] = []
|
||||
|
||||
while (index < this.separators.length) {
|
||||
separator = this.separators[index]
|
||||
|
||||
if (separator === '') {
|
||||
index = this.separators.length
|
||||
break
|
||||
}
|
||||
|
||||
parts = text.split(separator).filter((part) => part.trim())
|
||||
|
||||
if (parts.length > 1) {
|
||||
break
|
||||
}
|
||||
|
||||
index++
|
||||
}
|
||||
|
||||
if (index >= this.separators.length) {
|
||||
const chunkSizeChars = tokensToChars(this.chunkSize)
|
||||
return splitAtWordBoundaries(text, chunkSizeChars)
|
||||
}
|
||||
|
||||
const separator = this.separators[separatorIndex]
|
||||
|
||||
if (separator === '') {
|
||||
return this.splitRecursively(text, this.separators.length)
|
||||
}
|
||||
|
||||
const parts = text.split(separator).filter((part) => part.trim())
|
||||
|
||||
if (parts.length <= 1) {
|
||||
return this.splitRecursively(text, separatorIndex + 1)
|
||||
}
|
||||
|
||||
const chunks: string[] = []
|
||||
let currentChunk = ''
|
||||
|
||||
@@ -108,7 +143,7 @@ export class RecursiveChunker {
|
||||
}
|
||||
|
||||
if (estimateTokens(part) > this.chunkSize) {
|
||||
const subChunks = this.splitRecursively(part, separatorIndex + 1)
|
||||
const subChunks = this.splitRecursively(part, index + 1)
|
||||
for (const subChunk of subChunks) {
|
||||
chunks.push(subChunk)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user