Skip to content

Commit 7cc4574

Browse files
authored
improvement(knowledge): search returns document name (#1167)
1 parent 3f90094 commit 7cc4574

File tree

5 files changed

+50
-30
lines changed

5 files changed

+50
-30
lines changed

apps/sim/app/api/knowledge/search/route.ts

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -77,14 +77,13 @@ export async function POST(request: NextRequest) {
7777
? validatedData.knowledgeBaseIds
7878
: [validatedData.knowledgeBaseIds]
7979

80-
// Check access permissions for each knowledge base using proper workspace-based permissions
81-
const accessibleKbIds: string[] = []
82-
for (const kbId of knowledgeBaseIds) {
83-
const accessCheck = await checkKnowledgeBaseAccess(kbId, userId)
84-
if (accessCheck.hasAccess) {
85-
accessibleKbIds.push(kbId)
86-
}
87-
}
80+
// Check access permissions in parallel for performance
81+
const accessChecks = await Promise.all(
82+
knowledgeBaseIds.map((kbId) => checkKnowledgeBaseAccess(kbId, userId))
83+
)
84+
const accessibleKbIds: string[] = knowledgeBaseIds.filter(
85+
(_, idx) => accessChecks[idx]?.hasAccess
86+
)
8887

8988
// Map display names to tag slots for filtering
9089
let mappedFilters: Record<string, string> = {}
@@ -137,7 +136,10 @@ export async function POST(request: NextRequest) {
137136

138137
// Generate query embedding only if query is provided
139138
const hasQuery = validatedData.query && validatedData.query.trim().length > 0
140-
const queryEmbedding = hasQuery ? await generateSearchEmbedding(validatedData.query!) : null
139+
// Start embedding generation early and await when needed
140+
const queryEmbeddingPromise = hasQuery
141+
? generateSearchEmbedding(validatedData.query!)
142+
: Promise.resolve(null)
141143

142144
// Check if any requested knowledge bases were not accessible
143145
const inaccessibleKbIds = knowledgeBaseIds.filter((id) => !accessibleKbIds.includes(id))
@@ -165,7 +167,7 @@ export async function POST(request: NextRequest) {
165167
// Tag + Vector search
166168
logger.debug(`[${requestId}] Executing tag + vector search with filters:`, mappedFilters)
167169
const strategy = getQueryStrategy(accessibleKbIds.length, validatedData.topK)
168-
const queryVector = JSON.stringify(queryEmbedding)
170+
const queryVector = JSON.stringify(await queryEmbeddingPromise)
169171

170172
results = await handleTagAndVectorSearch({
171173
knowledgeBaseIds: accessibleKbIds,
@@ -178,7 +180,7 @@ export async function POST(request: NextRequest) {
178180
// Vector-only search
179181
logger.debug(`[${requestId}] Executing vector-only search`)
180182
const strategy = getQueryStrategy(accessibleKbIds.length, validatedData.topK)
181-
const queryVector = JSON.stringify(queryEmbedding)
183+
const queryVector = JSON.stringify(await queryEmbeddingPromise)
182184

183185
results = await handleVectorOnlySearch({
184186
knowledgeBaseIds: accessibleKbIds,
@@ -213,24 +215,28 @@ export async function POST(request: NextRequest) {
213215
}
214216

215217
// Fetch tag definitions for display name mapping (reuse the same fetch from filtering)
218+
const tagDefsResults = await Promise.all(
219+
accessibleKbIds.map(async (kbId) => {
220+
try {
221+
const tagDefs = await getDocumentTagDefinitions(kbId)
222+
const map: Record<string, string> = {}
223+
tagDefs.forEach((def) => {
224+
map[def.tagSlot] = def.displayName
225+
})
226+
return { kbId, map }
227+
} catch (error) {
228+
logger.warn(
229+
`[${requestId}] Failed to fetch tag definitions for display mapping:`,
230+
error
231+
)
232+
return { kbId, map: {} as Record<string, string> }
233+
}
234+
})
235+
)
216236
const tagDefinitionsMap: Record<string, Record<string, string>> = {}
217-
for (const kbId of accessibleKbIds) {
218-
try {
219-
const tagDefs = await getDocumentTagDefinitions(kbId)
220-
221-
tagDefinitionsMap[kbId] = {}
222-
tagDefs.forEach((def) => {
223-
tagDefinitionsMap[kbId][def.tagSlot] = def.displayName
224-
})
225-
logger.debug(
226-
`[${requestId}] Display mapping - KB ${kbId} tag definitions:`,
227-
tagDefinitionsMap[kbId]
228-
)
229-
} catch (error) {
230-
logger.warn(`[${requestId}] Failed to fetch tag definitions for display mapping:`, error)
231-
tagDefinitionsMap[kbId] = {}
232-
}
233-
}
237+
tagDefsResults.forEach(({ kbId, map }) => {
238+
tagDefinitionsMap[kbId] = map
239+
})
234240

235241
return NextResponse.json({
236242
success: true,
@@ -260,6 +266,7 @@ export async function POST(request: NextRequest) {
260266
id: result.id,
261267
content: result.content,
262268
documentId: result.documentId,
269+
documentName: (result as any).documentName || undefined,
263270
chunkIndex: result.chunkIndex,
264271
tags, // Clean display name mapped tags
265272
similarity: hasQuery ? 1 - result.distance : 1, // Perfect similarity for tag-only searches

apps/sim/app/api/knowledge/search/utils.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
import { and, eq, inArray, sql } from 'drizzle-orm'
22
import { createLogger } from '@/lib/logs/console/logger'
33
import { db } from '@/db'
4-
import { embedding } from '@/db/schema'
4+
import { document, embedding } from '@/db/schema'
55

66
const logger = createLogger('KnowledgeSearchUtils')
77

88
export interface SearchResult {
99
id: string
1010
content: string
1111
documentId: string
12+
documentName: string | null
1213
chunkIndex: number
1314
tag1: string | null
1415
tag2: string | null
@@ -130,6 +131,7 @@ async function executeVectorSearchOnIds(
130131
id: embedding.id,
131132
content: embedding.content,
132133
documentId: embedding.documentId,
134+
documentName: document.filename,
133135
chunkIndex: embedding.chunkIndex,
134136
tag1: embedding.tag1,
135137
tag2: embedding.tag2,
@@ -142,6 +144,7 @@ async function executeVectorSearchOnIds(
142144
knowledgeBaseId: embedding.knowledgeBaseId,
143145
})
144146
.from(embedding)
147+
.innerJoin(document, eq(embedding.documentId, document.id))
145148
.where(
146149
and(
147150
inArray(embedding.id, embeddingIds),
@@ -173,6 +176,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
173176
id: embedding.id,
174177
content: embedding.content,
175178
documentId: embedding.documentId,
179+
documentName: document.filename,
176180
chunkIndex: embedding.chunkIndex,
177181
tag1: embedding.tag1,
178182
tag2: embedding.tag2,
@@ -185,6 +189,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
185189
knowledgeBaseId: embedding.knowledgeBaseId,
186190
})
187191
.from(embedding)
192+
.innerJoin(document, eq(embedding.documentId, document.id))
188193
.where(
189194
and(
190195
eq(embedding.knowledgeBaseId, kbId),
@@ -204,6 +209,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
204209
id: embedding.id,
205210
content: embedding.content,
206211
documentId: embedding.documentId,
212+
documentName: document.filename,
207213
chunkIndex: embedding.chunkIndex,
208214
tag1: embedding.tag1,
209215
tag2: embedding.tag2,
@@ -216,6 +222,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
216222
knowledgeBaseId: embedding.knowledgeBaseId,
217223
})
218224
.from(embedding)
225+
.innerJoin(document, eq(embedding.documentId, document.id))
219226
.where(
220227
and(
221228
inArray(embedding.knowledgeBaseId, knowledgeBaseIds),
@@ -247,6 +254,7 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise<Sear
247254
id: embedding.id,
248255
content: embedding.content,
249256
documentId: embedding.documentId,
257+
documentName: document.filename,
250258
chunkIndex: embedding.chunkIndex,
251259
tag1: embedding.tag1,
252260
tag2: embedding.tag2,
@@ -259,6 +267,7 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise<Sear
259267
knowledgeBaseId: embedding.knowledgeBaseId,
260268
})
261269
.from(embedding)
270+
.innerJoin(document, eq(embedding.documentId, document.id))
262271
.where(
263272
and(
264273
eq(embedding.knowledgeBaseId, kbId),
@@ -280,6 +289,7 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise<Sear
280289
id: embedding.id,
281290
content: embedding.content,
282291
documentId: embedding.documentId,
292+
documentName: document.filename,
283293
chunkIndex: embedding.chunkIndex,
284294
tag1: embedding.tag1,
285295
tag2: embedding.tag2,
@@ -292,6 +302,7 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise<Sear
292302
knowledgeBaseId: embedding.knowledgeBaseId,
293303
})
294304
.from(embedding)
305+
.innerJoin(document, eq(embedding.documentId, document.id))
295306
.where(
296307
and(
297308
inArray(embedding.knowledgeBaseId, knowledgeBaseIds),

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/create-menu/create-menu.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ const logger = createLogger('CreateMenu')
1818

1919
const TIMERS = {
2020
LONG_PRESS_DELAY: 500,
21-
CLOSE_DELAY: 150,
21+
CLOSE_DELAY: 300,
2222
} as const
2323

2424
interface CreateMenuProps {

apps/sim/tools/knowledge/search.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ export const knowledgeSearchTool: ToolConfig<any, KnowledgeSearchResponse> = {
113113
id: { type: 'string' },
114114
content: { type: 'string' },
115115
documentId: { type: 'string' },
116+
documentName: { type: 'string' },
116117
chunkIndex: { type: 'number' },
117118
similarity: { type: 'number' },
118119
metadata: { type: 'object' },

apps/sim/tools/knowledge/types.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ export interface KnowledgeSearchResult {
22
id: string
33
content: string
44
documentId: string
5+
documentName: string
56
chunkIndex: number
67
metadata: Record<string, any>
78
similarity: number

0 commit comments

Comments
 (0)