Skip to content

Commit c8d09ca

Browse files
icecrasher321icecrasher321
committed
fix(search): bound retrieval and progressively read document evidence
1 parent 64bce6a commit c8d09ca

23 files changed

Lines changed: 1222 additions & 409 deletions

‎apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,10 @@ vi.mock('@/lib/knowledge/embeddings', () => ({
5959

6060
vi.mock('@/lib/knowledge/search/queries', () => ({
6161
generateSearchEmbedding: mocks.generateEmbedding,
62-
executeKnowledgeSearch: mocks.executeSearch,
62+
retrieveKnowledgeSearch: async (...args: unknown[]) => ({
63+
rows: await mocks.executeSearch(...args),
64+
retrieval: { status: 'complete', timedOutLegs: [] },
65+
}),
6366
getDocumentMetadataByIds: mocks.getDocumentMetadata,
6467
}))
6568

‎apps/sim/lib/api/contracts/knowledge/documents.ts‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -429,8 +429,23 @@ export const readSearchDocumentResultSchema = z.object({
429429
knowledgeBaseId: z.string().min(1),
430430
documentName: z.string().nullable(),
431431
sourceUrl: z.string().nullable(),
432-
chunks: z.array(z.object({ content: z.string(), chunkIndex: z.number().int().min(0) })).max(50),
432+
chunks: z
433+
.array(
434+
z.object({
435+
content: z.string().max(8000),
436+
chunkIndex: z.number().int().min(0),
437+
startOffset: z.number().int().min(0),
438+
endOffset: z.number().int().min(0),
439+
totalCharacters: z.number().int().min(0),
440+
})
441+
)
442+
.max(8),
433443
hasMore: z.boolean(),
434-
nextOffset: z.number().int().min(0).nullable(),
444+
next: z
445+
.object({
446+
startChunkIndex: z.number().int().min(0),
447+
startOffset: z.number().int().min(0),
448+
})
449+
.nullable(),
435450
})
436451
export type ReadSearchDocumentResult = z.output<typeof readSearchDocumentResultSchema>

‎apps/sim/lib/copilot/generated/tool-catalog-v1.ts‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4750,16 +4750,24 @@ export const ReadDocument: ToolCatalogEntry = {
47504750
type: 'string',
47514751
},
47524752
limit: {
4753-
default: 20,
4754-
description: 'Maximum number of chunks to read.',
4755-
maximum: 50,
4753+
default: 3,
4754+
description:
4755+
'Maximum chunks to read; the server may return fewer to fit its text budget. Follow next when more context is needed.',
4756+
maximum: 8,
47564757
minimum: 1,
47574758
type: 'integer',
47584759
},
4759-
offset: {
4760-
default: 0,
4761-
description: 'Number of chunks to skip.',
4762-
maximum: 5000,
4760+
startChunkIndex: {
4761+
description:
4762+
"Inclusive chunk index from search or a previous read's next object. Gaps from disabled chunks are skipped.",
4763+
maximum: 2147483647,
4764+
minimum: 0,
4765+
type: 'integer',
4766+
},
4767+
startOffset: {
4768+
description:
4769+
'UTF-16 character offset within startChunkIndex. Omit to read the chunk from its start, or copy next.startOffset to continue a partial chunk.',
4770+
maximum: 2147483647,
47634771
minimum: 0,
47644772
type: 'integer',
47654773
},
@@ -5616,7 +5624,8 @@ export const SearchWorkspace: ToolCatalogEntry = {
56165624
},
56175625
topK: {
56185626
default: 20,
5619-
description: 'Maximum number of matching chunks to return.',
5627+
description:
5628+
'Maximum number of matching passage previews to return. Retrieval ranking is independent of preview length.',
56205629
maximum: 50,
56215630
minimum: 1,
56225631
type: 'integer',

‎apps/sim/lib/copilot/generated/tool-schemas-v1.ts‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4694,16 +4694,24 @@ export const TOOL_RUNTIME_SCHEMAS: Record<string, ToolRuntimeSchemaEntry> = {
46944694
type: 'string',
46954695
},
46964696
limit: {
4697-
default: 20,
4698-
description: 'Maximum number of chunks to read.',
4699-
maximum: 50,
4697+
default: 3,
4698+
description:
4699+
'Maximum chunks to read; the server may return fewer to fit its text budget. Follow next when more context is needed.',
4700+
maximum: 8,
47004701
minimum: 1,
47014702
type: 'integer',
47024703
},
4703-
offset: {
4704-
default: 0,
4705-
description: 'Number of chunks to skip.',
4706-
maximum: 5000,
4704+
startChunkIndex: {
4705+
description:
4706+
"Inclusive chunk index from search or a previous read's next object. Gaps from disabled chunks are skipped.",
4707+
maximum: 2147483647,
4708+
minimum: 0,
4709+
type: 'integer',
4710+
},
4711+
startOffset: {
4712+
description:
4713+
'UTF-16 character offset within startChunkIndex. Omit to read the chunk from its start, or copy next.startOffset to continue a partial chunk.',
4714+
maximum: 2147483647,
47074715
minimum: 0,
47084716
type: 'integer',
47094717
},
@@ -5518,7 +5526,8 @@ export const TOOL_RUNTIME_SCHEMAS: Record<string, ToolRuntimeSchemaEntry> = {
55185526
},
55195527
topK: {
55205528
default: 20,
5521-
description: 'Maximum number of matching chunks to return.',
5529+
description:
5530+
'Maximum number of matching passage previews to return. Retrieval ranking is independent of preview length.',
55225531
maximum: 50,
55235532
minimum: 1,
55245533
type: 'integer',

‎apps/sim/lib/copilot/generated/trace-attributes-v1.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -478,6 +478,7 @@ export const TraceAttr = {
478478
LlmStreamChunks: 'llm.stream.chunks',
479479
LlmStreamFirstChunkBytes: 'llm.stream.first_chunk_bytes',
480480
LlmStreamFirstChunkMs: 'llm.stream.first_chunk_ms',
481+
LlmStreamFirstTokenMs: 'llm.stream.first_token_ms',
481482
LlmStreamOpenMs: 'llm.stream.open_ms',
482483
LlmStreamTotalMs: 'llm.stream.total_ms',
483484
LockAcquired: 'lock.acquired',
@@ -1159,6 +1160,7 @@ export const TraceAttrValues: readonly TraceAttrValue[] = [
11591160
'llm.stream.chunks',
11601161
'llm.stream.first_chunk_bytes',
11611162
'llm.stream.first_chunk_ms',
1163+
'llm.stream.first_token_ms',
11621164
'llm.stream.open_ms',
11631165
'llm.stream.total_ms',
11641166
'lock.acquired',

‎apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ describe('Assistant retrieval tools', () => {
6161
beforeEach(() => {
6262
vi.clearAllMocks()
6363
mocks.search.mockResolvedValue({
64+
retrieval: { status: 'complete', timedOutLegs: [] },
6465
knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
6566
results: [
6667
{
@@ -83,7 +84,7 @@ describe('Assistant retrieval tools', () => {
8384
sourceUrl: 'https://source.test/doc',
8485
chunks: [{ content: 'body', chunkIndex: 0 }],
8586
hasMore: false,
86-
nextOffset: null,
87+
next: null,
8788
})
8889
})
8990
it('pins organization and private chat while reusing the canonical search index and citations', async () => {
@@ -192,6 +193,7 @@ describe('Assistant retrieval tools', () => {
192193
async (count) => {
193194
const content = 'Confidential passage é🔎'.repeat(100)
194195
mocks.search.mockResolvedValueOnce({
196+
retrieval: { status: 'complete', timedOutLegs: [] },
195197
knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
196198
results: Array.from({ length: count }, (_, index) => ({
197199
knowledgeBaseId: 'index',
@@ -217,8 +219,9 @@ describe('Assistant retrieval tools', () => {
217219
expect.objectContaining({
218220
toolCallId: 'call',
219221
toolResultBytes: Buffer.byteLength(JSON.stringify(output)),
220-
passageBytes: count * Buffer.byteLength(content),
221-
maxPassageBytes: count ? Buffer.byteLength(content) : 0,
222+
passageBytes: count * Buffer.byteLength(content.slice(0, 1200)),
223+
originalPassageBytes: count * Buffer.byteLength(content),
224+
maxPassageBytes: count ? Buffer.byteLength(content.slice(0, 1200)) : 0,
222225
uniqueDocumentCount: Math.min(count, 4),
223226
})
224227
)
@@ -245,6 +248,7 @@ describe('Assistant retrieval tools', () => {
245248
})
246249
it('projects the provider name for connected-source citations instead of the index name', async () => {
247250
mocks.search.mockResolvedValueOnce({
251+
retrieval: { status: 'complete', timedOutLegs: [] },
248252
knowledgeBases: [{ id: 'index', name: 'Sim Search' }],
249253
results: [
250254
{
@@ -292,20 +296,20 @@ describe('Assistant retrieval tools', () => {
292296
})
293297
it('reads a selected document through the shared use case and rejects unbounded pages', async () => {
294298
expect(
295-
await readDocumentServerTool.execute({ documentId: 'doc', offset: 20 }, context)
299+
await readDocumentServerTool.execute({ documentId: 'doc', startChunkIndex: 20 }, context)
296300
).toMatchObject({ success: true })
297301
expect(mocks.read).toHaveBeenCalledWith(
298302
expect.objectContaining({
299303
input: expect.objectContaining({
300304
assertedWorkspaceId: 'workspace',
301305
filters: context.assistantSearch,
302-
offset: 20,
303-
limit: 20,
306+
startChunkIndex: 20,
307+
limit: 3,
304308
}),
305309
})
306310
)
307311
expect(
308-
await readDocumentServerTool.execute({ documentId: 'doc', limit: 10000 }, context)
312+
await readDocumentServerTool.execute({ documentId: 'doc', limit: 9 }, context)
309313
).toMatchObject({ success: false })
310314
expect(mocks.read).toHaveBeenCalledOnce()
311315
})

0 commit comments

Comments
 (0)