Update index.ts

AbdelStark · Dec 19, 2024 · db43e13 · db43e13
1 parent 832a274
commit db43e13
Showing 1 changed file with 24 additions and 2 deletions.
diff --git a/packages/plugin-web-search/src/index.ts b/packages/plugin-web-search/src/index.ts
@@ -8,8 +8,30 @@ import {
     State,
 } from "@ai16z/eliza";
 import { generateWebSearch } from "@ai16z/eliza";
-
 import { SearchResult } from "@ai16z/eliza";
+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;
+  }
 
 const webSearch: Action = {
     name: "WEB_SEARCH",
@@ -68,7 +90,7 @@ const webSearch: Action = {
                 : "";
 
             callback({
-                text: responseList,
+                text: MaxTokens(responseList, DEFAULT_MAX_WEB_SEARCH_TOKENS),
             });
         } else {
             elizaLogger.error("search failed or returned no data.");