diff --git a/packages/backend/server/src/plugins/copilot/providers/provider.ts b/packages/backend/server/src/plugins/copilot/providers/provider.ts index a2ee14f9dd..bfd390dfbd 100644 --- a/packages/backend/server/src/plugins/copilot/providers/provider.ts +++ b/packages/backend/server/src/plugins/copilot/providers/provider.ts @@ -141,6 +141,7 @@ export abstract class CopilotProvider { this.logger.debug(`getTools: ${JSON.stringify(options.tools)}`); const ac = this.moduleRef.get(AccessController, { strict: false }); const docReader = this.moduleRef.get(DocReader, { strict: false }); + const models = this.moduleRef.get(Models, { strict: false }); const prompt = this.moduleRef.get(PromptService, { strict: false, }); @@ -182,7 +183,12 @@ export abstract class CopilotProvider { const docContext = options.session ? await context.getBySessionId(options.session) : null; - const searchDocs = buildDocSearchGetter(ac, context, docContext); + const searchDocs = buildDocSearchGetter( + ac, + context, + docContext, + models + ); tools.doc_semantic_search = createDocSemanticSearchTool( searchDocs.bind(null, options) ); @@ -204,7 +210,6 @@ export abstract class CopilotProvider { break; } case 'docRead': { - const models = this.moduleRef.get(Models, { strict: false }); const getDoc = buildDocContentGetter(ac, docReader, models); tools.doc_read = createDocReadTool(getDoc.bind(null, options)); break; diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 20bcbd6c8e..7a0fe08a97 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -2,7 +2,7 @@ import { tool } from 'ai'; import { z } from 'zod'; import type { AccessController } from '../../../core/permission'; -import type { ChunkSimilarity } from '../../../models'; +import type { ChunkSimilarity, Models } from '../../../models'; import type { CopilotContextService } from '../context'; import type { ContextSession } from '../context/session'; import type { CopilotChatOptions } from '../providers'; @@ -11,7 +11,8 @@ import { toolError } from './error'; export const buildDocSearchGetter = ( ac: AccessController, context: CopilotContextService, - docContext: ContextSession | null + docContext: ContextSession | null, + models: Models ) => { const searchDocs = async ( options: CopilotChatOptions, @@ -45,7 +46,24 @@ export const buildDocSearchGetter = ( } if (!docChunks.length && !fileChunks.length) return `No results found for "${query}".`; - return [...fileChunks, ...docChunks]; + + const docMetas = await models.doc + .findAuthors( + docChunks.map(c => ({ + // oxlint-disable-next-line no-non-null-assertion + workspaceId: options.workspace!, + docId: c.docId, + })) + ) + .then(docs => new Map(docs.filter(d => !!d).map(doc => [doc.id, doc]))); + + return [ + ...fileChunks, + ...docChunks.map(c => ({ + ...c, + ...docMetas.get(c.docId), + })), + ] as ChunkSimilarity[]; }; return searchDocs; };