✨ feat: add builtin search with SearXNG

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 || 'https://searx.tiekoetter.com/',
+ },
+
+ server: {
+ SEARXNG_URL: z.string().url(),
+ },
+ });
+};
+
+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,13 @@
+import { createTRPCClient, httpBatchLink } from '@trpc/client';
+import superjson from 'superjson';
+
+import type { ToolsRouter } from '@/server/routers/tools';
+
+export const toolsClient = createTRPCClient<ToolsRouter>({
+ links: [
+ httpBatchLink({
+ 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,28 @@
+import qs from 'query-string';
+
+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(`${this.baseUrl}/search?${searchParams}`);
+
+ return await response.json();
+ } catch (error) {
+ console.error('Error searching:', error);
+ throw error;
+ }
+ }
+}

src/server/routers/tools/__tests__/search.test.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/index.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/search.ts

@@ -0,0 +1,20 @@
+import { z } from 'zod';
+
+import { toolsEnv } from '@/config/tools';
+import { authedProcedure, router } from '@/libs/trpc';
+import { SearXNGClient } from '@/server/modules/SearXNG';
+
+export const searchRouter = router({
+ query: authedProcedure
+ .input(
+ z.object({
+ query: z.string(),
+ searchEngine: z.array(z.string()).optional(),
+ }),
+ )
+ .query(async ({ input }) => {
+ const client = new SearXNGClient(toolsEnv.SEARXNG_URL);
+
+ return await client.search(input.query, input.searchEngine);
+ }),
+});

src/services/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/store/chat/slices/builtinTool/action.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;
 }
@@ -83,19 +104,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;
@@ -105,6 +199,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/initialState.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/selectors.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/tools/index.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/portals.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/renders.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,
 };