Skip to content

Commit a718f3e

Browse files
icecrasher321icecrasher321
committed
fix(knowledge): optimize search and trace pipeline latency
1 parent e21b9a6 commit a718f3e

16 files changed

Lines changed: 1607 additions & 274 deletions

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

Lines changed: 88 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@ import {
1616
} from '@/lib/knowledge/application/workspace-search'
1717
import { sourceAuthor } from '@/lib/knowledge/search/author'
1818
import { createKnowledgeDocumentCitation } from '@/lib/knowledge/search/citation'
19+
import {
20+
measureSearchStage,
21+
recordSearchStageDuration,
22+
withSearchDiagnostics,
23+
} from '@/lib/knowledge/search/diagnostics'
1924
import { intersectWorkspaceSearchFilters } from '@/lib/knowledge/search/filters'
2025
import { connectorDisplayName } from '@/lib/sim-search/connectors'
2126
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
@@ -37,77 +42,91 @@ const CITATION_INSTRUCTION =
3742
export const searchWorkspaceServerTool: BaseServerTool = {
3843
name: 'search_workspace',
3944
async execute(raw, context?: ServerToolContext) {
40-
try {
41-
const scope = requireCopilotKnowledgeScope(context)
42-
const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw)
43-
const registry = context?.resolvedSecretTraceRegistry
44-
if (!registry) throw new Error('Knowledge result provenance is unavailable')
45-
const projected = projectResolvedSecretModelContent(query, registry)
46-
if (!projected.safe || typeof projected.value !== 'string') {
47-
return {
48-
success: false,
49-
message: 'Search query contains protected content. Rephrase the query.',
50-
}
51-
}
52-
const input = {
53-
query: projected.value,
54-
topK,
55-
filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch),
45+
return withSearchDiagnostics(
46+
{
5647
surface: context?.searchSurface ?? 'copilot',
57-
resultSecretRegistry: registry,
58-
signal: context?.abortSignal,
59-
} as const
60-
const result =
61-
scope.kind === 'organization'
62-
? await executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, {
63-
...input,
64-
organizationId: scope.organizationId,
65-
})
66-
: await executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, {
67-
...input,
68-
workspaceId: scope.workspaceId,
69-
})
70-
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
71-
return {
72-
success: true,
73-
message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`,
74-
data: {
75-
query,
76-
results: result.results.map((item) => ({
77-
documentId: item.documentId,
78-
knowledgeBaseId: item.knowledgeBaseId,
79-
knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '',
80-
siteName: item.connectorType
81-
? connectorDisplayName(item.connectorType)
82-
: names.get(item.knowledgeBaseId),
83-
documentName: item.documentName,
84-
sourceUrl: item.sourceUrl,
85-
connectorType: item.connectorType,
86-
sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null,
87-
author: sourceAuthor(item.metadata),
88-
content: item.content,
89-
chunkIndex: item.chunkIndex,
90-
similarity: item.similarity,
91-
...createKnowledgeDocumentCitation({
92-
scope,
93-
knowledgeBaseId: item.knowledgeBaseId,
94-
documentId: item.documentId,
95-
sourceUrl: item.sourceUrl,
96-
baseUrl: getBaseUrl(),
97-
}),
98-
})),
99-
},
100-
}
101-
} catch (error) {
102-
logger.error('Workspace search failed', { error })
103-
return {
104-
success: false,
105-
message:
106-
error instanceof z.ZodError
107-
? 'Invalid search arguments'
108-
: messageForCopilotKnowledgeError(error),
48+
toolCallId: context?.toolCallId,
49+
executionId: context?.executionId,
50+
},
51+
async () => {
52+
try {
53+
const inputStarted = performance.now()
54+
const scope = requireCopilotKnowledgeScope(context)
55+
const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw)
56+
const registry = context?.resolvedSecretTraceRegistry
57+
if (!registry) throw new Error('Knowledge result provenance is unavailable')
58+
const projected = projectResolvedSecretModelContent(query, registry)
59+
if (!projected.safe || typeof projected.value !== 'string') {
60+
return {
61+
success: false,
62+
message: 'Search query contains protected content. Rephrase the query.',
63+
}
64+
}
65+
const input = {
66+
query: projected.value,
67+
topK,
68+
filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch),
69+
surface: context?.searchSurface ?? 'copilot',
70+
resultSecretRegistry: registry,
71+
signal: context?.abortSignal,
72+
} as const
73+
recordSearchStageDuration('tool_input', performance.now() - inputStarted)
74+
const result = await measureSearchStage('tool_application', () =>
75+
scope.kind === 'organization'
76+
? executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, {
77+
...input,
78+
organizationId: scope.organizationId,
79+
})
80+
: executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, {
81+
...input,
82+
workspaceId: scope.workspaceId,
83+
})
84+
)
85+
return await measureSearchStage('tool_presentation', () => {
86+
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
87+
return {
88+
success: true,
89+
message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`,
90+
data: {
91+
query,
92+
results: result.results.map((item) => ({
93+
documentId: item.documentId,
94+
knowledgeBaseId: item.knowledgeBaseId,
95+
knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '',
96+
siteName: item.connectorType
97+
? connectorDisplayName(item.connectorType)
98+
: names.get(item.knowledgeBaseId),
99+
documentName: item.documentName,
100+
sourceUrl: item.sourceUrl,
101+
connectorType: item.connectorType,
102+
sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null,
103+
author: sourceAuthor(item.metadata),
104+
content: item.content,
105+
chunkIndex: item.chunkIndex,
106+
similarity: item.similarity,
107+
...createKnowledgeDocumentCitation({
108+
scope,
109+
knowledgeBaseId: item.knowledgeBaseId,
110+
documentId: item.documentId,
111+
sourceUrl: item.sourceUrl,
112+
baseUrl: getBaseUrl(),
113+
}),
114+
})),
115+
},
116+
}
117+
})
118+
} catch (error) {
119+
logger.error('Workspace search failed', { error })
120+
return {
121+
success: false,
122+
message:
123+
error instanceof z.ZodError
124+
? 'Invalid search arguments'
125+
: messageForCopilotKnowledgeError(error),
126+
}
127+
}
109128
}
110-
}
129+
)
111130
},
112131
}
113132

‎apps/sim/lib/knowledge/__integration__/application-acl.integration.ts‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,10 @@ import {
6262
seedKnowledgeMemberFixture,
6363
} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
6464
import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions'
65-
import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate'
65+
import {
66+
knowledgeAccessCondition,
67+
knowledgeMetadataCandidateAccessCondition,
68+
} from '@/lib/knowledge/access/predicate'
6669
import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope'
6770
import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks'
6871
import { readKnowledgeDocument } from '@/lib/knowledge/application/documents'
@@ -299,9 +302,16 @@ describe('indexed source content through real application access', () => {
299302
return result.results.map((row) => row.documentId)
300303
}
301304

302-
it.each(['workspace', 'admin', 'members'] as const)(
303-
'allows a remaining workspace ACL only in workspace mode, not during a %s transition',
304-
async (accessMode) => {
305+
it.each([
306+
['workspace', true],
307+
['workspace', false],
308+
['admin', true],
309+
['admin', false],
310+
['members', true],
311+
['members', false],
312+
] as const)(
313+
'requires settled workspace mode for a remaining workspace ACL (%s, rewrite pending=%s)',
314+
async (accessMode, accessRewritePending) => {
305315
const [savedConnector] = await db
306316
.select({
307317
accessMode: knowledgeConnector.accessMode,
@@ -324,18 +334,25 @@ describe('indexed source content through real application access', () => {
324334
.where(eq(document.id, documentId))
325335
await db
326336
.update(knowledgeConnector)
327-
.set({ accessMode, accessRewritePending: true })
337+
.set({ accessMode, accessRewritePending })
328338
.where(eq(knowledgeConnector.id, connectorId))
329-
const visible = await db
330-
.select({ id: document.id })
331-
.from(document)
332-
.where(
333-
and(
334-
eq(document.id, documentId),
335-
knowledgeAccessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] })
339+
for (const accessCondition of [
340+
knowledgeMetadataCandidateAccessCondition,
341+
knowledgeAccessCondition,
342+
]) {
343+
const visible = await db
344+
.select({ id: document.id })
345+
.from(document)
346+
.where(
347+
and(
348+
eq(document.id, documentId),
349+
accessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] })
350+
)
336351
)
352+
expect(visible.map((row) => row.id)).toEqual(
353+
accessMode === 'workspace' && !accessRewritePending ? [documentId] : []
337354
)
338-
expect(visible.map((row) => row.id)).toEqual(accessMode === 'workspace' ? [documentId] : [])
355+
}
339356
} finally {
340357
await db
341358
.update(knowledgeConnector)

0 commit comments

Comments
 (0)