Merge pull request #191 from miurla/improve-search

Improve web search skill
miurla · Nov 7, 2023 · e2ad638 · e2ad638 · vercel · Nov 7, 2023
2 parents 604f50b + 0ac36ba
commit e2ad638
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 4 deletions.
diff --git a/src/lib/agents/babyelfagi/tools/utils/textCompletion.ts b/src/lib/agents/babyelfagi/tools/utils/textCompletion.ts
@@ -1,4 +1,4 @@
-import { AgentMessage, AgentTask } from '@/types';
+import { AgentMessage, AgentTask, LLMParams } from '@/types';
 import { ChatOpenAI } from 'langchain/chat_models/openai';
 import { HumanChatMessage } from 'langchain/schema';
 
@@ -65,3 +65,48 @@ export const textCompletion = async (
     return 'Failed to generate text.';
   }
 };
+
+export const generateText = async (
+  prompt: string,
+  params: LLMParams = {},
+  userKey?: string,
+  signal?: AbortSignal,
+) => {
+  const modelName = 'gpt-3.5-turbo';
+  const temperature = 0.7;
+  const maxTokens = 1500;
+  const topP = 1;
+  const frequencyPenalty = 0;
+  const presencePenalty = 0;
+  const streaming = false;
+  const callbacks = undefined;
+
+  const llmParams = {
+    ...params,
+    modelName,
+    temperature,
+    maxTokens,
+    topP,
+    frequencyPenalty,
+    presencePenalty,
+    streaming,
+    callbacks,
+  };
+
+  const llm = new ChatOpenAI({
+    ...llmParams,
+    openAIApiKey: userKey,
+    verbose: false,
+  });
+
+  try {
+    const response = await llm.call([new HumanChatMessage(prompt)]);
+    return response.text;
+  } catch (error: any) {
+    if (error.name === 'AbortError') {
+      return null;
+    }
+    console.log('error: ', error);
+    return 'Failed to generate text.';
+  }
+};
diff --git a/src/lib/agents/babyelfagi/tools/webBrowsing.ts b/src/lib/agents/babyelfagi/tools/webBrowsing.ts
@@ -1,7 +1,7 @@
 import { simplifySearchResults } from '@/lib/agents/babyelfagi/tools/webSearch';
 import { AgentTask, AgentMessage } from '@/types';
 import { analystPrompt, searchQueryPrompt } from '../../../../utils/prompt';
-import { textCompletion } from './utils/textCompletion';
+import { textCompletion, generateText } from './utils/textCompletion';
 import { largeTextExtract } from './utils/largeTextExtract';
 import { v4 as uuidv4 } from 'uuid';
 import { webSearch } from './webSearch';
@@ -47,11 +47,13 @@ export const webBrowsing = async (
   let results = '';
   let index = 1;
   let completedCount = 0;
-  const MaxCompletedCount = 3;
+  const MaxRetryCount = 3;
+  let sufficientResultsObtained = false;
+
   // Loop through search results
   for (const searchResult of simplifiedSearchResults) {
     if (signal?.aborted) return '';
-    if (completedCount >= MaxCompletedCount) break;
+    if (completedCount >= MaxRetryCount || sufficientResultsObtained) break;
 
     // Extract the URL from the search result
     const url = searchResult.link;
@@ -107,6 +109,25 @@ export const webBrowsing = async (
     title = `${index}. Relevant info...`;
     callbackSearchStatus(id, message, task, messageCallback, verbose);
 
+    const reflect = `You are expart task manager.
+    Your Mission: You'll be provided with a goal for the team to accomplish and one task from the team's task list and its results.
+    You decide if the results of the task are satisfactory.
+    Your judgment will determine whether the task is complete or not.
+    RULE: If you judge the task to be complete, then OK. If it needs to be redone, then just reply with NG.
+    Do not include anything else in your response.
+    Objective: ${objective}
+    TASK: ${task.task}
+    Result: ${info}
+    `;
+
+    const judgment = await generateText(
+      reflect,
+      { temperature: 0.0 },
+      userApiKey,
+      signal,
+    );
+    sufficientResultsObtained = judgment === 'OK';
+
     results += `${info}. `;
     index += 1;
     completedCount += 1;