From 09c3aa0a92e9672801665ce01e5bc89b1feb0eca Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Tue, 17 Jun 2025 15:01:02 +0800 Subject: [PATCH] chore(server): improve semantic search (#12838) ## Summary by CodeRabbit - **New Features** - Updated document-related tools in chat prompts, including new options for document keyword and semantic search. - **Refactor** - Renamed and reorganized document search tools for improved clarity in tool selection and results display. - **Bug Fixes** - Ensured tool identifiers and results are consistent across chat and search features. --- .../src/plugins/copilot/prompt/prompts.ts | 8 ++++---- .../src/plugins/copilot/providers/provider.ts | 17 +++++++++-------- .../src/plugins/copilot/providers/types.ts | 12 ++++++------ .../src/plugins/copilot/providers/utils.ts | 6 +++--- ...emantic-search.ts => doc-semantic-search.ts} | 13 +++++++++++-- .../server/src/plugins/copilot/tools/index.ts | 2 +- 6 files changed, 34 insertions(+), 24 deletions(-) rename packages/backend/server/src/plugins/copilot/tools/{semantic-search.ts => doc-semantic-search.ts} (77%) diff --git a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts index b9d9bd724..e18917e34 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts @@ -1792,11 +1792,11 @@ Below is the user's query. Please respond in the user's preferred language witho ], config: { tools: [ - 'readDoc', - 'editDoc', + 'docRead', + 'docEdit', + 'docKeywordSearch', + 'docSemanticSearch', 'webSearch', - 'keywordSearch', - 'semanticSearch', ], }, }; diff --git a/packages/backend/server/src/plugins/copilot/providers/provider.ts b/packages/backend/server/src/plugins/copilot/providers/provider.ts index e5d700d40..fc383fd9b 100644 --- a/packages/backend/server/src/plugins/copilot/providers/provider.ts +++ b/packages/backend/server/src/plugins/copilot/providers/provider.ts @@ -13,9 +13,9 @@ import { AccessController } from '../../../core/permission'; import { CopilotContextService } from '../context'; import { buildDocSearchGetter, + createDocSemanticSearchTool, createExaCrawlTool, createExaSearchTool, - createSemanticSearchTool, } from '../tools'; import { CopilotProviderFactory } from './factory'; import { @@ -131,20 +131,21 @@ export abstract class CopilotProvider { continue; } switch (tool) { - case 'webSearch': { - tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig); - tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig); - break; - } - case 'semanticSearch': { + case 'docSemanticSearch': { const ac = this.moduleRef.get(AccessController, { strict: false }); const context = this.moduleRef.get(CopilotContextService, { strict: false, }); const searchDocs = buildDocSearchGetter(ac, context); - tools.semantic_search = createSemanticSearchTool( + tools.doc_semantic_search = createDocSemanticSearchTool( searchDocs.bind(null, options) ); + break; + } + case 'webSearch': { + tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig); + tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig); + break; } } } diff --git a/packages/backend/server/src/plugins/copilot/providers/types.ts b/packages/backend/server/src/plugins/copilot/providers/types.ts index eb5bb20c6..bcec8d141 100644 --- a/packages/backend/server/src/plugins/copilot/providers/types.ts +++ b/packages/backend/server/src/plugins/copilot/providers/types.ts @@ -60,14 +60,14 @@ export const PromptConfigStrictSchema = z.object({ tools: z .enum([ // work with morph - 'editDoc', + 'docEdit', + // work with indexer + 'docRead', + 'docKeywordSearch', + // work with embeddings + 'docSemanticSearch', // work with exa/model internal tools 'webSearch', - // work with indexer - 'readDoc', - 'keywordSearch', - // work with embeddings - 'semanticSearch', ]) .array() .nullable() diff --git a/packages/backend/server/src/plugins/copilot/providers/utils.ts b/packages/backend/server/src/plugins/copilot/providers/utils.ts index 25b631261..068dab3c0 100644 --- a/packages/backend/server/src/plugins/copilot/providers/utils.ts +++ b/packages/backend/server/src/plugins/copilot/providers/utils.ts @@ -10,9 +10,9 @@ import { import { ZodType } from 'zod'; import { + createDocSemanticSearchTool, createExaCrawlTool, createExaSearchTool, - createSemanticSearchTool, } from '../tools'; import { PromptMessage } from './types'; @@ -380,9 +380,9 @@ export class CitationParser { } export interface CustomAITools extends ToolSet { + doc_semantic_search: ReturnType; web_search_exa: ReturnType; web_crawl_exa: ReturnType; - semantic_search: ReturnType; } type ChunkType = TextStreamPart['type']; @@ -429,7 +429,7 @@ export class TextStreamParser { case 'tool-result': { result = this.addPrefix(result); switch (chunk.toolName) { - case 'semantic_search': { + case 'doc_semantic_search': { if (Array.isArray(chunk.result)) { result += `\nFound ${chunk.result.length} document${chunk.result.length !== 1 ? 's' : ''} related to “${chunk.args.query}”.\n`; } diff --git a/packages/backend/server/src/plugins/copilot/tools/semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts similarity index 77% rename from packages/backend/server/src/plugins/copilot/tools/semantic-search.ts rename to packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 9e654c4a9..f14aa7fc8 100644 --- a/packages/backend/server/src/plugins/copilot/tools/semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -20,12 +20,21 @@ export const buildDocSearchGetter = ( .can('Workspace.Read'); if (!canAccess) return undefined; const chunks = await context.matchWorkspaceAll(options.workspace, query); - return chunks || undefined; + const docChunks = await ac + .user(options.user) + .workspace(options.workspace) + .docs( + chunks.filter(c => 'docId' in c), + 'Doc.Read' + ); + const fileChunks = chunks.filter(c => 'fileId' in c); + if (!docChunks.length && !fileChunks.length) return undefined; + return [...fileChunks, ...docChunks]; }; return searchDocs; }; -export const createSemanticSearchTool = ( +export const createDocSemanticSearchTool = ( searchDocs: (query: string) => Promise ) => { return tool({ diff --git a/packages/backend/server/src/plugins/copilot/tools/index.ts b/packages/backend/server/src/plugins/copilot/tools/index.ts index cde30ce1e..7a1e97c85 100644 --- a/packages/backend/server/src/plugins/copilot/tools/index.ts +++ b/packages/backend/server/src/plugins/copilot/tools/index.ts @@ -1,2 +1,2 @@ -export * from './semantic-search'; +export * from './doc-semantic-search'; export * from './web-search';