✨ feat: add builtin search with SearXNG

Author: arvinxx

src/config/tools.ts

@@ -0,0 +1,16 @@
+import { createEnv } from '@t3-oss/env-nextjs';
+import { z } from 'zod';
+
+export const getToolsConfig = () => {
+  return createEnv({
+    runtimeEnv: {
+      SEARXNG_URL: process.env.SEARXNG_URL,
+    },
+
+    server: {
+      SEARXNG_URL: z.string().url().optional(),
+    },
+  });
+};
+
+export const toolsEnv = getToolsConfig();

src/libs/trpc/client/index.ts

@@ -1,3 +1,4 @@
 export { asyncClient } from './async';
 export { edgeClient } from './edge';
 export * from './lambda';
+export * from './tools';

src/libs/trpc/client/tools.ts

@@ -0,0 +1,20 @@
+import { createTRPCClient, httpBatchLink } from '@trpc/client';
+import superjson from 'superjson';
+
+import type { ToolsRouter } from '@/server/routers/tools';
+
+export const toolsClient = createTRPCClient<ToolsRouter>({
+  links: [
+    httpBatchLink({
+      headers: async () => {
+        // dynamic import to avoid circular dependency
+        const { createHeaderWithAuth } = await import('@/services/_auth');
+
+        return createHeaderWithAuth();
+      },
+      maxURLLength: 2083,
+      transformer: superjson,
+      url: '/trpc/tools',
+    }),
+  ],
+});

src/locales/default/tool.ts

@@ -7,4 +7,19 @@ export default {
     images: '图片:',
     prompt: '提示词',
   },
+  search: {
+    createNewSearch: '创建新的搜索记录',
+    emptyResult: '没有搜索到结果，请修改关键词后重试',
+    includedTooltip: '当前搜索结果会进入会话的上下文中',
+    keywords: '关键词：',
+    scoreTooltip: '相关性分数，该分数越高说明与查询关键词越相关',
+    searchBar: {
+      button: '搜索',
+      placeholder: '关键词',
+      tooltip: '将会重新获取搜索结果，并创建一条新的总结消息',
+    },
+    searchEngine: '搜索引擎：',
+    searchResult: '搜索数量：',
+    viewMoreResults: '查看更多 {{results}} 个结果',
+  },
 };

src/server/modules/SearXNG.ts

@@ -0,0 +1,29 @@
+import qs from 'query-string';
+import urlJoin from 'url-join';
+
+import { SearchResponse } from '@/types/tool/search';
+
+export class SearXNGClient {
+  private baseUrl: string;
+
+  constructor(baseUrl: string) {
+    this.baseUrl = baseUrl;
+  }
+
+  async search(query: string, engines?: string[]): Promise<SearchResponse> {
+    try {
+      const searchParams = qs.stringify({
+        engines: engines?.join(','),
+        format: 'json',
+        q: query,
+      });
+
+      const response = await fetch(urlJoin(this.baseUrl, `/search?${searchParams}`));
+
+      return await response.json();
+    } catch (error) {
+      console.error('Error searching:', error);
+      throw error;
+    }
+  }
+}

src/server/routers/tools/__tests__/fixtures/searXNG.ts

@@ -0,0 +1,41 @@
+// @vitest-environment node
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { toolsEnv } from '@/config/tools';
+
+/**
+ * This file contains the root router of your tRPC-backend
+ */
+import { createCallerFactory } from '@/libs/trpc';
+import { AuthContext, createContextInner } from '@/server/context';
+import { SearXNGClient } from '@/server/modules/SearXNG';
+
+import { searchRouter } from '../search';
+import { hetongxue } from './fixtures/searXNG';
+
+vi.mock('@/config/tools', () => ({
+  toolsEnv: vi.fn(),
+}));
+
+const createCaller = createCallerFactory(searchRouter);
+let ctx: AuthContext;
+let router: ReturnType<typeof createCaller>;
+
+beforeEach(async () => {
+  vi.resetAllMocks();
+  ctx = await createContextInner({ userId: 'mock' });
+  router = createCaller(ctx);
+});
+
+describe('searchRouter', () => {
+  describe('search', () => {
+    it('搜索结果超过10个', async () => {
+      vi.spyOn(SearXNGClient.prototype, 'search').mockResolvedValueOnce(hetongxue);
+
+      const results = await router.query({ query: '何同学' });
+
+      // Assert
+      expect(results.results.length).toEqual(10);
+    });
+  });
+});

src/server/routers/tools/__tests__/search.test.ts

@@ -1,7 +1,10 @@
 import { publicProcedure, router } from '@/libs/trpc';
 
+import { searchRouter } from './search';
+
 export const toolsRouter = router({
   healthcheck: publicProcedure.query(() => "i'm live!"),
+  search: searchRouter,
 });
 
 export type ToolsRouter = typeof toolsRouter;

src/server/routers/tools/index.ts

@@ -0,0 +1,28 @@
+import { TRPCError } from '@trpc/server';
+import { z } from 'zod';
+
+import { toolsEnv } from '@/config/tools';
+import { isServerMode } from '@/const/version';
+import { authedProcedure, passwordProcedure, router } from '@/libs/trpc';
+import { SearXNGClient } from '@/server/modules/SearXNG';
+
+const searchProcedure = isServerMode ? authedProcedure : passwordProcedure;
+
+export const searchRouter = router({
+  query: searchProcedure
+    .input(
+      z.object({
+        query: z.string(),
+        searchEngine: z.array(z.string()).optional(),
+      }),
+    )
+    .query(async ({ input }) => {
+      if (!toolsEnv.SEARXNG_URL) {
+        throw new TRPCError({ code: 'SERVICE_UNAVAILABLE', message: 'SearXNG is not configured' });
+      }
+
+      const client = new SearXNGClient(toolsEnv.SEARXNG_URL);
+
+      return await client.search(input.query, input.searchEngine);
+    }),
+});

src/server/routers/tools/search.ts

@@ -0,0 +1,9 @@
+import { toolsClient } from '@/libs/trpc/client';
+
+class SearchService {
+  search(query: string, searchEngine?: string[]) {
+    return toolsClient.search.query.query({ query, searchEngine });
+  }
+}
+
+export const searchService = new SearchService();

src/services/search.ts

@@ -5,13 +5,17 @@ import { StateCreator } from 'zustand/vanilla';
 
 import { useClientDataSWR } from '@/libs/swr';
 import { fileService } from '@/services/file';
+import { searchService } from '@/services/search';
 import { imageGenerationService } from '@/services/textToImage';
 import { uploadService } from '@/services/upload';
 import { chatSelectors } from '@/store/chat/selectors';
 import { ChatStore } from '@/store/chat/store';
 import { useFileStore } from '@/store/file';
+import { CreateMessageParams } from '@/types/message';
 import { DallEImageItem } from '@/types/tool/dalle';
+import { SearchContent, SearchQuery } from '@/types/tool/search';
 import { setNamespace } from '@/utils/storeDebug';
+import { nanoid } from '@/utils/uuid';
 
 const n = setNamespace('tool');
 
@@ -21,8 +25,25 @@ const SWR_FETCH_KEY = 'FetchImageItem';
  */
 export interface ChatBuiltinToolAction {
   generateImageFromPrompts: (items: DallEImageItem[], id: string) => Promise<void>;
+  /**
+   * 重新发起搜索
+   * @description 会更新插件的 arguments 参数，然后再次搜索
+   */
+  reSearchWithSearXNG: (
+    id: string,
+    data: SearchQuery,
+    options?: { aiSummary: boolean },
+  ) => Promise<void>;
+  saveSearXNGSearchResult: (id: string) => Promise<void>;
+  searchWithSearXNG: (
+    id: string,
+    data: SearchQuery,
+    aiSummary?: boolean,
+  ) => Promise<void | boolean>;
   text2image: (id: string, data: DallEImageItem[]) => Promise<void>;
+
   toggleDallEImageLoading: (key: string, value: boolean) => void;
+  toggleSearchLoading: (id: string, loading: boolean) => void;
   updateImageItem: (id: string, updater: (data: DallEImageItem[]) => void) => Promise<void>;
   useFetchDalleImageItem: (id: string) => SWRResponse;
 }
@@ -82,19 +103,92 @@ export const chatToolSlice: StateCreator<
       });
     });
   },
+  reSearchWithSearXNG: async (id, data, options) => {
+    get().toggleSearchLoading(id, true);
+    await get().updatePluginArguments(id, data);
+
+    await get().searchWithSearXNG(id, data, options?.aiSummary);
+  },
+  saveSearXNGSearchResult: async (id) => {
+    const message = chatSelectors.getMessageById(id)(get());
+    if (!message || !message.plugin) return;
+
+    const { internal_addToolToAssistantMessage, internal_createMessage, openToolUI } = get();
+    // 1. 创建一个新的 tool call message
+    const newToolCallId = `tool_call_${nanoid()}`;
+
+    const toolMessage: CreateMessageParams = {
+      content: message.content,
+      id: undefined,
+      parentId: message.parentId,
+      plugin: message.plugin,
+      pluginState: message.pluginState,
+      role: 'tool',
+      sessionId: get().activeId,
+      tool_call_id: newToolCallId,
+      topicId: get().activeTopicId,
+    };
+
+    const addToolItem = async () => {
+      if (!message.parentId || !message.plugin) return;
+
+      await internal_addToolToAssistantMessage(message.parentId, {
+        id: newToolCallId,
+        ...message.plugin,
+      });
+    };
+
+    const [newMessageId] = await Promise.all([
+      // 1. 添加 tool message
+      internal_createMessage(toolMessage),
+      // 2. 将这条 tool call message 插入到 ai 消息的 tools 中
+      addToolItem(),
+    ]);
+
+    // 将新创建的 tool message 激活
+    openToolUI(newMessageId, message.plugin.identifier);
+  },
+  searchWithSearXNG: async (id, params, aiSummary = true) => {
+    get().toggleSearchLoading(id, true);
+    const data = await searchService.search(params.query, params.searchEngines);
+    await get().updatePluginState(id, data);
+
+    get().toggleSearchLoading(id, false);
+
+    // 只取前 5 个结果作为上下文
+    const searchContent: SearchContent[] = data.results.slice(0, 5).map((item) => ({
+      content: item.content,
+      title: item.title,
+      url: item.url,
+    }));
+
+    await get().internal_updateMessageContent(id, JSON.stringify(searchContent));
+
+    // 如果没搜索到结果，那么不触发 ai 总结
+    if (searchContent.length === 0) return;
+
+    // 如果 aiSummary 为 true，则会自动触发总结
+    return aiSummary;
+  },
   text2image: async (id, data) => {
     // const isAutoGen = settingsSelectors.isDalleAutoGenerating(useGlobalStore.getState());
     // if (!isAutoGen) return;
 
     await get().generateImageFromPrompts(data, id);
   },
+
   toggleDallEImageLoading: (key, value) => {
     set(
       { dalleImageLoading: { ...get().dalleImageLoading, [key]: value } },
       false,
       n('toggleDallEImageLoading'),
     );
   },
+
+  toggleSearchLoading: (id, loading) => {
+    set({ searchLoading: { ...get().searchLoading, [id]: loading } }, false, 'toggleSearchLoading');
+  },
+
   updateImageItem: async (id, updater) => {
     const message = chatSelectors.getMessageById(id)(get());
     if (!message) return;
@@ -104,6 +198,7 @@ export const chatToolSlice: StateCreator<
     const nextContent = produce(data, updater);
     await get().internal_updateMessageContent(id, JSON.stringify(nextContent));
   },
+
   useFetchDalleImageItem: (id) =>
     useClientDataSWR([SWR_FETCH_KEY, id], async () => {
       const item = await fileService.getFile(id);

src/store/chat/slices/builtinTool/action.ts

@@ -3,9 +3,11 @@ import { FileItem } from '@/types/files';
 export interface ChatToolState {
   dalleImageLoading: Record<string, boolean>;
   dalleImageMap: Record<string, FileItem>;
+  searchLoading: Record<string, boolean>;
 }
 
 export const initialToolState: ChatToolState = {
   dalleImageLoading: {},
   dalleImageMap: {},
+  searchLoading: {},
 };

src/store/chat/slices/builtinTool/initialState.ts

@@ -5,7 +5,10 @@ const isDallEImageGenerating = (id: string) => (s: ChatStoreState) => s.dalleIma
 const isGeneratingDallEImage = (s: ChatStoreState) =>
   Object.values(s.dalleImageLoading).some(Boolean);
 
+const isSearXNGSearching = (id: string) => (s: ChatStoreState) => s.searchLoading[id];
+
 export const chatToolSelectors = {
   isDallEImageGenerating,
   isGeneratingDallEImage,
+  isSearXNGSearching,
 };

src/store/chat/slices/builtinTool/selectors.ts

@@ -2,6 +2,7 @@ import { LobeBuiltinTool } from '@/types/tool';
 
 import { ArtifactsManifest } from './artifacts';
 import { DalleManifest } from './dalle';
+import { WebBrowsingManifest } from './web-browsing';
 
 export const builtinTools: LobeBuiltinTool[] = [
   {
@@ -14,4 +15,9 @@ export const builtinTools: LobeBuiltinTool[] = [
     manifest: DalleManifest,
     type: 'builtin',
   },
+  {
+    identifier: WebBrowsingManifest.identifier,
+    manifest: WebBrowsingManifest,
+    type: 'builtin',
+  },
 ];

src/tools/index.ts

@@ -1,3 +1,8 @@
 import { BuiltinPortal } from '@/types/tool';
 
-export const BuiltinToolsPortals: Record<string, BuiltinPortal> = {};
+import { WebBrowsingManifest } from './web-browsing';
+import WebBrowsing from './web-browsing/Portal';
+
+export const BuiltinToolsPortals: Record<string, BuiltinPortal> = {
+  [WebBrowsingManifest.identifier]: WebBrowsing as BuiltinPortal,
+};

src/tools/portals.ts

@@ -2,7 +2,10 @@ import { BuiltinRender } from '@/types/tool';
 
 import { DalleManifest } from './dalle';
 import DalleRender from './dalle/Render';
+import { WebBrowsingManifest } from './web-browsing';
+import WebBrowsing from './web-browsing/Render';
 
 export const BuiltinToolsRenders: Record<string, BuiltinRender> = {
   [DalleManifest.identifier]: DalleRender as BuiltinRender,
+  [WebBrowsingManifest.identifier]: WebBrowsing as BuiltinRender,
 };