From bd71bdefd9e64686183a425d8549a28f7584b9c4 Mon Sep 17 00:00:00 2001
From: Luis Otavio Martins <luisotaviomgr1@gmail.com>
Date: Tue, 19 Mar 2024 20:25:10 -0300
Subject: [PATCH] feat: Wikipedia tool (#287)

* feat: wikipedia tool, bug fix in not logged messages

* Update version to 2.6.0 in package.json
---
 package.json                                  |  2 +-
 src/clients/tools-openrouter.ts               |  8 ++++++
 .../llm-models/completion-open-router.ts      |  2 +-
 src/handlers/message/index.ts                 | 28 ++++++++++---------
 4 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/package.json b/package.json
index 113a958..9297527 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "whatsapp-ai-assistant",
-  "version": "2.5.0",
+  "version": "2.6.0",
   "description": "WhatsApp chatbot",
   "module": "src/index.ts",
   "type": "module",
diff --git a/src/clients/tools-openrouter.ts b/src/clients/tools-openrouter.ts
index bb7b6d3..698a999 100644
--- a/src/clients/tools-openrouter.ts
+++ b/src/clients/tools-openrouter.ts
@@ -3,6 +3,7 @@ import {
   GoogleCalendarViewTool,
 } from "@langchain/community/tools/google_calendar";
 import { SearchApi } from "@langchain/community/tools/searchapi";
+import { WikipediaQueryRun } from "@langchain/community/tools/wikipedia_query_run";
 import { ChatOpenAI, DallEAPIWrapper, OpenAI, OpenAIEmbeddings } from "@langchain/openai";
 import { Calculator } from "langchain/tools/calculator";
 import { WebBrowser } from "langchain/tools/webbrowser";
@@ -81,12 +82,19 @@ if (SEARCH_API !== '') {
 
 const calculatorTool = new Calculator()
 
+const wikipediaTool = new WikipediaQueryRun({
+  topKResults: 3,
+  maxDocContentLength: 4000,
+});
+
+
 export const tools = [
   ...(searchTool ? [searchTool] : []),
   ...(webBrowserTool ? [webBrowserTool] : []),
   ...(googleCalendarCreateTool ? [googleCalendarCreateTool] : []),
   ...(googleCalendarViewTool ? [googleCalendarViewTool] : []),
   ...(dalleTool ? [dalleTool] : []),
+  wikipediaTool,
   calculatorTool
 ];
 export const toolNames = tools.map((tool) => tool.name);
\ No newline at end of file
diff --git a/src/handlers/llm-models/completion-open-router.ts b/src/handlers/llm-models/completion-open-router.ts
index c5de2d3..42a234e 100644
--- a/src/handlers/llm-models/completion-open-router.ts
+++ b/src/handlers/llm-models/completion-open-router.ts
@@ -1,4 +1,5 @@
 import { Message } from "whatsapp-web.js";
+import { createExecutorForOpenRouter } from "../../clients/open-router";
 import {
   BOT_PREFIX,
   DEBUG_SUMMARY,
@@ -13,7 +14,6 @@ import {
   updateOpenRouterConversation,
 } from "../../crud/conversation";
 import { handleAudioMessage } from "../audio-message";
-import { createExecutorForOpenRouter } from "../../clients/open-router";
 
 export async function getCompletionWithOpenRouter(
   message: Message,
diff --git a/src/handlers/message/index.ts b/src/handlers/message/index.ts
index 2f04a49..0cf366e 100644
--- a/src/handlers/message/index.ts
+++ b/src/handlers/message/index.ts
@@ -1,23 +1,22 @@
 import { Message } from "whatsapp-web.js";
-import { setStatusFor } from "../../helpers/message";
-import { createContextFromMessage } from "../context";
-import {
-  getCompletionWithBing,
-  getSources,
-  getSuggestions,
-} from "../llm-models/completion-bing.ts";
-import { log } from "../../helpers/utils";
 import {
-  BING_COOKIES,
   BOT_PREFIX,
   ENABLE_REMINDERS,
   ENABLE_SOURCES,
   ENABLE_SUGGESTIONS,
-  OPENROUTER_API_KEY,
+  OPENROUTER_API_KEY
 } from "../../constants";
-import { handleReminderFor } from "../reminder/reminder.ts";
 import { getLLMModel, updateWaMessageId } from "../../crud/conversation";
+import { setStatusFor } from "../../helpers/message";
+import { log } from "../../helpers/utils";
+import { createContextFromMessage } from "../context";
+import {
+  getCompletionWithBing,
+  getSources,
+  getSuggestions,
+} from "../llm-models/completion-bing.ts";
 import { getCompletionWithOpenRouter } from "../llm-models/completion-open-router.ts";
+import { handleReminderFor } from "../reminder/reminder.ts";
 
 export async function handleMessage(message: Message) {
   await log(message);
@@ -61,8 +60,11 @@ export async function handleMessage(message: Message) {
 
     if (!response) throw new Error("No response from LLM");
 
-    // @ts-ignore
-    const finalReply = await streamingReply.edit(response);
+    let finalReply = null;
+
+    while (finalReply === null) {
+      finalReply = await streamingReply.edit(response);
+    }
 
     await log(finalReply, true);
     await setStatusFor(message, "done");