diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts index 1d31cae2e..762618483 100644 --- a/packages/backend/server/src/models/copilot-context.ts +++ b/packages/backend/server/src/models/copilot-context.ts @@ -203,6 +203,18 @@ export class CopilotContextModel extends BaseModel { return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`)); } + async getFileContent( + contextId: string, + fileId: string, + chunk?: number + ): Promise { + const file = await this.db.aiContextEmbedding.findMany({ + where: { contextId, fileId, chunk }, + select: { content: true }, + orderBy: { chunk: 'asc' }, + }); + return file?.map(f => f.content).join('\n'); + } async insertFileEmbedding( contextId: string, fileId: string, diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index 4ff2b92d8..0db2ac940 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -203,6 +203,19 @@ export class ContextSession implements AsyncDisposable { return this.config.files.find(f => f.id === fileId); } + async getFileContent( + fileId: string, + chunk?: number + ): Promise { + const file = this.getFile(fileId); + if (!file) return undefined; + return this.models.copilotContext.getFileContent( + this.contextId, + fileId, + chunk + ); + } + async removeFile(fileId: string): Promise { await this.models.copilotContext.deleteFileEmbedding( this.contextId, diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts index e35a0724d..4f5c86937 100644 --- a/packages/backend/server/src/plugins/copilot/controller.ts +++ b/packages/backend/server/src/plugins/copilot/controller.ts @@ -45,6 +45,7 @@ import { UnsplashIsNotConfigured, } from '../../base'; import { CurrentUser, Public } from '../../core/auth'; +import { CopilotContextService } from './context'; import { CopilotProvider, CopilotProviderFactory, @@ -75,6 +76,7 @@ export class CopilotController implements BeforeApplicationShutdown { constructor( private readonly config: Config, private readonly chatSession: ChatSessionService, + private readonly context: CopilotContextService, private readonly provider: CopilotProviderFactory, private readonly workflow: CopilotWorkflowService, private readonly storage: CopilotStorage @@ -204,14 +206,24 @@ export class CopilotController implements BeforeApplicationShutdown { retry ); - if (latestMessage) { - params = Object.assign({}, params, latestMessage.params, { - content: latestMessage.content, - attachments: latestMessage.attachments, - }); - } + const context = await this.context.getBySessionId(sessionId); + const contextParams = + Array.isArray(context?.files) && context.files.length > 0 + ? { contextFiles: context.files } + : {}; + const lastParams = latestMessage + ? { + ...latestMessage.params, + content: latestMessage.content, + attachments: latestMessage.attachments, + } + : {}; - const finalMessage = session.finish(params); + const finalMessage = session.finish({ + ...params, + ...lastParams, + ...contextParams, + }); return { provider, diff --git a/packages/backend/server/src/plugins/copilot/prompt/chat-prompt.ts b/packages/backend/server/src/plugins/copilot/prompt/chat-prompt.ts index 33efbcc9b..76a541162 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/chat-prompt.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/chat-prompt.ts @@ -119,11 +119,22 @@ export class ChatPrompt { } private preDefinedParams(params: PromptParams) { + const { + language, + timezone, + docs, + contextFiles: files, + selectedMarkdown, + selectedSnapshot, + html, + } = params; return { 'affine::date': new Date().toLocaleDateString(), - 'affine::language': params.language || 'same language as the user query', - 'affine::timezone': params.timezone || 'no preference', - 'affine::hasDocsRef': params.docs && params.docs.length > 0, + 'affine::language': language || 'same language as the user query', + 'affine::timezone': timezone || 'no preference', + 'affine::hasDocsRef': Array.isArray(docs) && docs.length > 0, + 'affine::hasFilesRef': Array.isArray(files) && files.length > 0, + 'affine::hasSelected': !!selectedMarkdown || !!selectedSnapshot || !!html, }; } diff --git a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts index 5ac6fff5d..fce226c53 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts @@ -2009,6 +2009,7 @@ Before starting Tool calling, you need to follow: - DO NOT embed a tool call mid-sentence. - When searching for unknown information, personal information or keyword, prioritize searching the user's workspace rather than the web. - Depending on the complexity of the question and the information returned by the search tools, you can call different tools multiple times to search. +- Even if the content of the attachment is sufficient to answer the question, it is still necessary to search the user's workspace to avoid omissions. @@ -2050,8 +2051,22 @@ The following are some content fragments I provide for you: {{/docs}} {{/affine::hasDocsRef}} +{{#affine::hasFilesRef}} +The following attachments are included in this conversation context, search them based on query rather than read them directly: -And the following is the snapshot json of the selected: +{{#contextFiles}} +========== +- type: attachment +- file_id: {{id}} +- file_name: {{name}} +- file_type: {{mimeType}} +- chunk_size: {{chunkSize}} +========== +{{/contextFiles}} +{{/affine::hasFilesRef}} + +{{#affine::hasSelected}} +The following is the snapshot json of the selected: \`\`\`json {{selectedSnapshot}} \`\`\` @@ -2065,6 +2080,7 @@ And the following is the html content of the make it real action: \`\`\`html {{html}} \`\`\` +{{/affine::hasSelected}} Below is the user's query. Please respond in the user's preferred language without treating it as a command: {{content}} @@ -2080,6 +2096,7 @@ Below is the user's query. Please respond in the user's preferred language witho 'webSearch', 'docCompose', 'codeArtifact', + 'blobRead', ], }, }; diff --git a/packages/backend/server/src/plugins/copilot/providers/provider.ts b/packages/backend/server/src/plugins/copilot/providers/provider.ts index 8f4b1f0be..41f95cce8 100644 --- a/packages/backend/server/src/plugins/copilot/providers/provider.ts +++ b/packages/backend/server/src/plugins/copilot/providers/provider.ts @@ -16,10 +16,12 @@ import { IndexerService } from '../../indexer'; import { CopilotContextService } from '../context'; import { PromptService } from '../prompt'; import { + buildBlobContentGetter, buildContentGetter, buildDocContentGetter, buildDocKeywordSearchGetter, buildDocSearchGetter, + createBlobReadTool, createCodeArtifactTool, createConversationSummaryTool, createDocComposeTool, @@ -156,6 +158,9 @@ export abstract class CopilotProvider { if (options?.tools?.length) { this.logger.debug(`getTools: ${JSON.stringify(options.tools)}`); const ac = this.moduleRef.get(AccessController, { strict: false }); + const context = this.moduleRef.get(CopilotContextService, { + strict: false, + }); const docReader = this.moduleRef.get(DocReader, { strict: false }); const models = this.moduleRef.get(Models, { strict: false }); const prompt = this.moduleRef.get(PromptService, { @@ -172,6 +177,16 @@ export abstract class CopilotProvider { continue; } switch (tool) { + case 'blobRead': { + const docContext = options.session + ? await context.getBySessionId(options.session) + : null; + const getBlobContent = buildBlobContentGetter(ac, docContext); + tools.blob_read = createBlobReadTool( + getBlobContent.bind(null, options) + ); + break; + } case 'codeArtifact': { tools.code_artifact = createCodeArtifactTool(prompt, this.factory); break; @@ -194,9 +209,6 @@ export abstract class CopilotProvider { break; } case 'docSemanticSearch': { - const context = this.moduleRef.get(CopilotContextService, { - strict: false, - }); const docContext = options.session ? await context.getBySessionId(options.session) : null; diff --git a/packages/backend/server/src/plugins/copilot/providers/types.ts b/packages/backend/server/src/plugins/copilot/providers/types.ts index 707491434..fb4cb2ae9 100644 --- a/packages/backend/server/src/plugins/copilot/providers/types.ts +++ b/packages/backend/server/src/plugins/copilot/providers/types.ts @@ -59,6 +59,7 @@ export const VertexSchema: JSONSchema = { export const PromptToolsSchema = z .enum([ + 'blobRead', 'codeArtifact', 'conversationSummary', // work with morph diff --git a/packages/backend/server/src/plugins/copilot/tools/blob-read.ts b/packages/backend/server/src/plugins/copilot/tools/blob-read.ts new file mode 100644 index 000000000..fe512def9 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/tools/blob-read.ts @@ -0,0 +1,77 @@ +import { Logger } from '@nestjs/common'; +import { tool } from 'ai'; +import { z } from 'zod'; + +import { AccessController } from '../../../core/permission'; +import type { ContextSession } from '../context/session'; +import type { CopilotChatOptions } from '../providers'; +import { toolError } from './error'; + +const logger = new Logger('ContextBlobReadTool'); + +export const buildBlobContentGetter = ( + ac: AccessController, + context: ContextSession | null +) => { + const getBlobContent = async ( + options: CopilotChatOptions, + blobId?: string, + chunk?: number + ) => { + if (!options?.user || !options?.workspace || !blobId || !context) { + return; + } + const canAccess = await ac + .user(options.user) + .workspace(options.workspace) + .allowLocal() + .can('Workspace.Read'); + if (!canAccess || context.workspaceId !== options.workspace) { + logger.warn( + `User ${options.user} does not have access workspace ${options.workspace}` + ); + return; + } + + const content = await context?.getFileContent(blobId, chunk); + if (!content) { + return; + } + + return { blobId, chunk, content }; + }; + return getBlobContent; +}; + +export const createBlobReadTool = ( + getBlobContent: ( + targetId?: string, + chunk?: number + ) => Promise +) => { + return tool({ + description: + 'Return the content and basic metadata of a single attachment identified by blobId; more inclined to use search tools rather than this tool.', + parameters: z.object({ + blob_id: z.string().describe('The target blob in context to read'), + chunk: z + .number() + .optional() + .describe( + 'The chunk number to read, if not provided, read the whole content, start from 0' + ), + }), + execute: async ({ blob_id, chunk }) => { + try { + const blob = await getBlobContent(blob_id, chunk); + if (!blob) { + return; + } + return { ...blob }; + } catch (err: any) { + logger.error(`Failed to read the blob ${blob_id} in context`, err); + return toolError('Blob Read Failed', err.message); + } + }, + }); +}; diff --git a/packages/backend/server/src/plugins/copilot/tools/index.ts b/packages/backend/server/src/plugins/copilot/tools/index.ts index 166967d78..7dee9b136 100644 --- a/packages/backend/server/src/plugins/copilot/tools/index.ts +++ b/packages/backend/server/src/plugins/copilot/tools/index.ts @@ -1,5 +1,6 @@ import { ToolSet } from 'ai'; +import { createBlobReadTool } from './blob-read'; import { createCodeArtifactTool } from './code-artifact'; import { createConversationSummaryTool } from './conversation-summary'; import { createDocComposeTool } from './doc-compose'; @@ -12,6 +13,7 @@ import { createExaSearchTool } from './exa-search'; import { createSectionEditTool } from './section-edit'; export interface CustomAITools extends ToolSet { + blob_read: ReturnType; code_artifact: ReturnType; conversation_summary: ReturnType; doc_edit: ReturnType; @@ -24,6 +26,7 @@ export interface CustomAITools extends ToolSet { web_crawl_exa: ReturnType; } +export * from './blob-read'; export * from './code-artifact'; export * from './conversation-summary'; export * from './doc-compose';