feat(server): embedding search tool (#12810)

packages/backend/server/src/plugins/copilot/providers/anthropic/anthropic.ts

@@ -10,7 +10,6 @@ import {
   metrics,
   UserFriendlyError,
 } from '../../../../base';
-import { createExaCrawlTool, createExaSearchTool } from '../../tools';
 import { CopilotProvider } from '../provider';
 import type {
   CopilotChatOptions,
@@ -68,7 +67,7 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
         providerOptions: {
           anthropic: this.getAnthropicOptions(options, model.id),
         },
-        tools: this.getTools(),
+        tools: await this.getTools(options, model.id),
         maxSteps: this.MAX_STEPS,
         experimental_continueSteps: true,
       });
@@ -103,7 +102,7 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
         providerOptions: {
           anthropic: this.getAnthropicOptions(options, model.id),
         },
-        tools: this.getTools(),
+        tools: await this.getTools(options, model.id),
         maxSteps: this.MAX_STEPS,
         experimental_continueSteps: true,
       });
@@ -123,13 +122,6 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
     }
   }
 
-  private getTools() {
-    return {
-      web_search_exa: createExaSearchTool(this.AFFiNEConfig),
-      web_crawl_exa: createExaCrawlTool(this.AFFiNEConfig),
-    };
-  }
-
   private getAnthropicOptions(options: CopilotChatOptions, model: string) {
     const result: AnthropicProviderOptions = {};
     if (options?.reasoning && this.isReasoningModel(model)) {

packages/backend/server/src/plugins/copilot/providers/openai.ts

@@ -11,7 +11,7 @@ import {
   generateObject,
   generateText,
   streamText,
-  ToolSet,
+  Tool,
 } from 'ai';
 import { z } from 'zod';
 
@@ -21,10 +21,10 @@ import {
   metrics,
   UserFriendlyError,
 } from '../../../base';
-import { createExaCrawlTool, createExaSearchTool } from '../tools';
 import { CopilotProvider } from './provider';
 import type {
   CopilotChatOptions,
+  CopilotChatTools,
   CopilotEmbeddingOptions,
   CopilotImageOptions,
   CopilotStructuredOptions,
@@ -248,25 +248,14 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
     }
   }
 
-  private getTools(options: CopilotChatOptions, model: string): ToolSet {
-    const tools: ToolSet = {};
-    if (options?.tools?.length) {
-      for (const tool of options.tools) {
-        switch (tool) {
-          case 'webSearch': {
-            if (this.isReasoningModel(model)) {
-              tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig);
-              tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig);
-            } else {
-              tools.web_search_preview = openai.tools.webSearchPreview();
-            }
-            break;
-          }
-        }
-      }
-      return tools;
+  override getProviderSpecificTools(
+    toolName: CopilotChatTools,
+    model: string
+  ): [string, Tool] | undefined {
+    if (toolName === 'webSearch' && !this.isReasoningModel(model)) {
+      return ['web_search_preview', openai.tools.webSearchPreview()];
     }
-    return tools;
+    return;
   }
 
   async text(
@@ -297,7 +286,7 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
         providerOptions: {
           openai: this.getOpenAIOptions(options, model.id),
         },
-        tools: this.getTools(options, model.id),
+        tools: await this.getTools(options, model.id),
         maxSteps: this.MAX_STEPS,
         abortSignal: options.signal,
       });
@@ -338,7 +327,7 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
         providerOptions: {
           openai: this.getOpenAIOptions(options, model.id),
         },
-        tools: this.getTools(options, model.id),
+        tools: await this.getTools(options, model.id),
         maxSteps: this.MAX_STEPS,
         abortSignal: options.signal,
       });

packages/backend/server/src/plugins/copilot/providers/provider.ts

@@ -1,4 +1,6 @@
 import { Inject, Injectable, Logger } from '@nestjs/common';
+import { ModuleRef } from '@nestjs/core';
+import { Tool, ToolSet } from 'ai';
 import { z } from 'zod';
 
 import {
@@ -7,9 +9,18 @@ import {
   CopilotProviderNotSupported,
   OnEvent,
 } from '../../../base';
+import { AccessController } from '../../../core/permission';
+import { CopilotContextService } from '../context';
+import {
+  buildDocSearchGetter,
+  createExaCrawlTool,
+  createExaSearchTool,
+  createSemanticSearchTool,
+} from '../tools';
 import { CopilotProviderFactory } from './factory';
 import {
   type CopilotChatOptions,
+  CopilotChatTools,
   type CopilotEmbeddingOptions,
   type CopilotImageOptions,
   CopilotProviderModel,
@@ -33,6 +44,7 @@ export abstract class CopilotProvider<C = any> {
 
   @Inject() protected readonly AFFiNEConfig!: Config;
   @Inject() protected readonly factory!: CopilotProviderFactory;
+  @Inject() protected readonly moduleRef!: ModuleRef;
 
   get config(): C {
     return this.AFFiNEConfig.copilot.providers[this.type] as C;
@@ -98,6 +110,49 @@ export abstract class CopilotProvider<C = any> {
     );
   }
 
+  protected getProviderSpecificTools(
+    _toolName: CopilotChatTools,
+    _model: string
+  ): [string, Tool] | undefined {
+    return;
+  }
+
+  // use for tool use, shared between providers
+  protected async getTools(
+    options: CopilotChatOptions,
+    model: string
+  ): Promise<ToolSet> {
+    const tools: ToolSet = {};
+    if (options?.tools?.length) {
+      for (const tool of options.tools) {
+        const toolDef = this.getProviderSpecificTools(tool, model);
+        if (toolDef) {
+          tools[toolDef[0]] = toolDef[1];
+          continue;
+        }
+        switch (tool) {
+          case 'webSearch': {
+            tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig);
+            tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig);
+            break;
+          }
+          case 'semanticSearch': {
+            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(
+              searchDocs.bind(null, options)
+            );
+          }
+        }
+      }
+      return tools;
+    }
+    return tools;
+  }
+
   private handleZodError(ret: z.SafeParseReturnType<any, any>) {
     if (ret.success) return;
     const issues = ret.error.issues.map(i => {

packages/backend/server/src/plugins/copilot/providers/types.ts

@@ -57,7 +57,21 @@ export const VertexSchema: JSONSchema = {
 // ========== prompt ==========
 
 export const PromptConfigStrictSchema = z.object({
-  tools: z.enum(['webSearch']).array().nullable().optional(),
+  tools: z
+    .enum([
+      // work with morph
+      'editDoc',
+      // work with exa/model internal tools
+      'webSearch',
+      // work with indexer
+      'readDoc',
+      'keywordSearch',
+      // work with embeddings
+      'semanticSearch',
+    ])
+    .array()
+    .nullable()
+    .optional(),
   // params requirements
   requireContent: z.boolean().nullable().optional(),
   requireAttachment: z.boolean().nullable().optional(),
@@ -121,6 +135,7 @@ export type PromptParams = NonNullable<PromptMessage['params']>;
 const CopilotProviderOptionsSchema = z.object({
   signal: z.instanceof(AbortSignal).optional(),
   user: z.string().optional(),
+  workspace: z.string().optional(),
 });
 
 export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
@@ -133,6 +148,9 @@ export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
   .optional();
 
 export type CopilotChatOptions = z.infer<typeof CopilotChatOptionsSchema>;
+export type CopilotChatTools = NonNullable<
+  NonNullable<CopilotChatOptions>['tools']
+>[number];
 
 export const CopilotStructuredOptionsSchema =
   CopilotProviderOptionsSchema.merge(PromptConfigStrictSchema).optional();

packages/backend/server/src/plugins/copilot/providers/utils.ts

@@ -9,7 +9,11 @@ import {
 } from 'ai';
 import { ZodType } from 'zod';
 
-import { createExaCrawlTool, createExaSearchTool } from '../tools';
+import {
+  createExaCrawlTool,
+  createExaSearchTool,
+  createSemanticSearchTool,
+} from '../tools';
 import { PromptMessage } from './types';
 
 type ChatMessage = CoreUserMessage | CoreAssistantMessage;
@@ -378,6 +382,7 @@ export class CitationParser {
 export interface CustomAITools extends ToolSet {
   web_search_exa: ReturnType<typeof createExaSearchTool>;
   web_crawl_exa: ReturnType<typeof createExaCrawlTool>;
+  semantic_search: ReturnType<typeof createSemanticSearchTool>;
 }
 
 type ChunkType = TextStreamPart<CustomAITools>['type'];
@@ -424,6 +429,12 @@ export class TextStreamParser {
       case 'tool-result': {
         result = this.addPrefix(result);
         switch (chunk.toolName) {
+          case 'semantic_search': {
+            if (Array.isArray(chunk.result)) {
+              result += `\nFound ${chunk.result.length} document${chunk.result.length !== 1 ? 's' : ''} related to “${chunk.args.query}”.\n`;
+            }
+            break;
+          }
           case 'web_search_exa': {
             if (Array.isArray(chunk.result)) {
               result += `\n${this.getWebSearchLinks(chunk.result)}\n`;