refactor websearch into a service

chuasonglin1995 · chuasonglin1995 · commit 3c8c3bfa0a28 · 2025-01-12T18:21:41.000+08:00
diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
@@ -49,9 +49,10 @@ import {
     //VerifiableInferenceProvider,
     TelemetrySettings,
     TokenizerType,
+    IWebSearchService,
+    SearchOptions,
 } from "./types.ts";
 import { fal } from "@fal-ai/client";
-import { tavily } from "@tavily/core";
 
 type Tool = CoreTool<any, any>;
 type StepResult = AIStepResult<any>;
@@ -1725,23 +1726,26 @@ export const generateCaption = async (
 };
 
 export const generateWebSearch = async (
-    query: string,
+    data: {
+        query: string
+        options?: SearchOptions
+    },
     runtime: IAgentRuntime
 ): Promise<SearchResponse> => {
     try {
-        const apiKey = runtime.getSetting("TAVILY_API_KEY") as string;
-        if (!apiKey) {
-            throw new Error("TAVILY_API_KEY is not set");
+        const { query, options } = data;
+        const webSearchService =
+        runtime.getService<IWebSearchService>(
+            ServiceType.WEB_SEARCH
+        );
+
+        if (!webSearchService) {
+            throw new Error("Web search service not found");
         }
-        const tvly = tavily({ apiKey });
-        const response = await tvly.search(query, {
-            includeAnswer: true,
-            maxResults: 3, // 5 (default)
-            topic: "general", // "general"(default) "news"
-            searchDepth: "basic", // "basic"(default) "advanced"
-            includeImages: false, // false (default) true
-        });
+
+        const response = await webSearchService.search(query, runtime, options);
         return response;
+
     } catch (error) {
         elizaLogger.error("Error:", error);
     }
diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
@@ -1328,6 +1328,23 @@ export interface IPdfService extends Service {
     convertPdfToText(pdfBuffer: Buffer): Promise<string>;
 }
 
+export interface IWebSearchService extends Service {
+    search(
+        query: string,
+        runtime: IAgentRuntime,
+        options?: SearchOptions,
+    ): Promise<SearchResponse>;
+}
+
+export interface SearchOptions {
+    limit?: number;
+    type?: "news" | "general";
+    includeAnswer?: boolean;
+    searchDepth?: "basic" | "advanced";
+    includeImages?: boolean;
+    days?: number; // 1 means current day, 2 means last 2 days
+}
+
 export interface IAwsS3Service extends Service {
     uploadFile(
         imagePath: string,
@@ -1431,6 +1448,7 @@ export enum ServiceType {
     IRYS = "irys",
     TEE_LOG = "tee_log",
     GOPLUS_SECURITY = "goplus_security",
+    WEB_SEARCH = "web_search",
 }
 
 export enum LoggingLevel {
diff --git a/packages/plugin-node/src/services/index.ts b/packages/plugin-node/src/services/index.ts
@@ -6,6 +6,7 @@ import { PdfService } from "./pdf.ts";
 import { SpeechService } from "./speech.ts";
 import { TranscriptionService } from "./transcription.ts";
 import { VideoService } from "./video.ts";
+import { WebSearchService } from "./webSearch.ts";
 
 export {
     AwsS3Service,
@@ -16,4 +17,5 @@ export {
     SpeechService,
     TranscriptionService,
     VideoService,
+    WebSearchService
 };
diff --git a/packages/plugin-node/src/services/webSearch.ts b/packages/plugin-node/src/services/webSearch.ts
@@ -0,0 +1,47 @@
+import {
+    Service,
+    IWebSearchService,
+    ServiceType,
+    IAgentRuntime,
+    SearchResponse,
+    SearchOptions,
+} from "@elizaos/core";
+import { tavily } from "@tavily/core";
+
+export class WebSearchService extends Service implements IWebSearchService {
+    static serviceType: ServiceType = ServiceType.WEB_SEARCH;
+
+    async initialize(_runtime: IAgentRuntime): Promise<void> {}
+
+    getInstance(): IWebSearchService {
+        return WebSearchService.getInstance();
+    }
+
+    async search(
+        query: string,
+        runtime: IAgentRuntime,
+        options?: SearchOptions,
+    ): Promise<SearchResponse> {
+        try {
+            const apiKey = runtime.getSetting("TAVILY_API_KEY") as string;
+            if (!apiKey) {
+                throw new Error("TAVILY_API_KEY is not set");
+            }
+
+            const tvly = tavily({ apiKey });
+            const response = await tvly.search(query, {
+                includeAnswer: options?.includeAnswer || true,
+                maxResults: options?.limit || 3,
+                topic: options?.type || "general",
+                searchDepth: options?.searchDepth || "basic",
+                includeImages: options?.includeImages || false,
+                days: options?.days || 3,
+            });
+
+            return response;
+        } catch (error) {
+            console.error("Web search error:", error);
+            throw error;
+        }
+    }
+}
diff --git a/packages/plugin-web-search/src/actions.ts b/packages/plugin-web-search/src/actions.ts
@@ -0,0 +1,201 @@
+import { elizaLogger } from "@elizaos/core";
+import {
+    Action,
+    HandlerCallback,
+    IAgentRuntime,
+    Memory,
+    State,
+} from "@elizaos/core";
+import { generateWebSearch } from "@elizaos/core";
+import { SearchResult } from "@elizaos/core";
+import { encodingForModel, TiktokenModel } from "js-tiktoken";
+
+const DEFAULT_MAX_WEB_SEARCH_TOKENS = 4000;
+const DEFAULT_MODEL_ENCODING = "gpt-3.5-turbo";
+
+function getTotalTokensFromString(
+    str: string,
+    encodingName: TiktokenModel = DEFAULT_MODEL_ENCODING
+) {
+    const encoding = encodingForModel(encodingName);
+    return encoding.encode(str).length;
+}
+
+function MaxTokens(
+    data: string,
+    maxTokens: number = DEFAULT_MAX_WEB_SEARCH_TOKENS
+): string {
+    if (getTotalTokensFromString(data) >= maxTokens) {
+        return data.slice(0, maxTokens);
+    }
+    return data;
+}
+
+export const webSearch: Action = {
+    name: "WEB_SEARCH",
+    similes: [
+        "SEARCH_WEB",
+        "INTERNET_SEARCH",
+        "LOOKUP",
+        "QUERY_WEB",
+        "FIND_ONLINE",
+        "SEARCH_ENGINE",
+        "WEB_LOOKUP",
+        "ONLINE_SEARCH",
+        "FIND_INFORMATION",
+    ],
+    suppressInitialMessage: true,
+    description:
+        "Perform a web search to find information related to the message.",
+    validate: async (runtime: IAgentRuntime, message: Memory) => {
+        const tavilyApiKeyOk = !!runtime.getSetting("TAVILY_API_KEY");
+
+        return tavilyApiKeyOk;
+    },
+    handler: async (
+        runtime: IAgentRuntime,
+        message: Memory,
+        state: State,
+        options: any,
+        callback: HandlerCallback
+    ) => {
+        elizaLogger.log("Composing state for message:", message);
+        state = (await runtime.composeState(message)) as State;
+        const userId = runtime.agentId;
+        elizaLogger.log("User ID:", userId);
+
+        const webSearchPrompt = message.content.text;
+        elizaLogger.log("web search prompt received:", webSearchPrompt);
+
+        elizaLogger.log("Generating image with prompt:", webSearchPrompt);
+        const searchResponse = await generateWebSearch(
+            webSearchPrompt,
+            runtime
+        );
+
+        if (searchResponse && searchResponse.results.length) {
+            const responseList = searchResponse.answer
+                ? `${searchResponse.answer}${
+                      Array.isArray(searchResponse.results) &&
+                      searchResponse.results.length > 0
+                          ? `\n\nFor more details, you can check out these resources:\n${searchResponse.results
+                                .map(
+                                    (result: SearchResult, index: number) =>
+                                        `${index + 1}. [${result.title}](${result.url})`
+                                )
+                                .join("\n")}`
+                          : ""
+                  }`
+                : "";
+
+            callback({
+                text: MaxTokens(responseList, DEFAULT_MAX_WEB_SEARCH_TOKENS),
+            });
+        } else {
+            elizaLogger.error("search failed or returned no data.");
+        }
+    },
+    examples: [
+        [
+            {
+                user: "{{user1}}",
+                content: {
+                    text: "Find the latest news about SpaceX launches.",
+                },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here is the latest news about SpaceX launches:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: {
+                    text: "Can you find details about the iPhone 16 release?",
+                },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here are the details I found about the iPhone 16 release:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: {
+                    text: "What is the schedule for the next FIFA World Cup?",
+                },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here is the schedule for the next FIFA World Cup:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: { text: "Check the latest stock price of Tesla." },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here is the latest stock price of Tesla I found:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: {
+                    text: "What are the current trending movies in the US?",
+                },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here are the current trending movies in the US:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: {
+                    text: "What is the latest score in the NBA finals?",
+                },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here is the latest score from the NBA finals:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+        [
+            {
+                user: "{{user1}}",
+                content: { text: "When is the next Apple keynote event?" },
+            },
+            {
+                user: "{{agentName}}",
+                content: {
+                    text: "Here is the information about the next Apple keynote event:",
+                    action: "WEB_SEARCH",
+                },
+            },
+        ],
+    ],
+} as Action;
diff --git a/packages/plugin-web-search/src/index.ts b/packages/plugin-web-search/src/index.ts
