feat(search-agent): add search agent flow

src/lib/agents/search/index.ts

@@ -3,6 +3,8 @@ import SessionManager from '@/lib/session';
 import Classifier from './classifier';
 import { WidgetRegistry } from './widgets';
 import Researcher from './researcher';
+import { getWriterPrompt } from '@/lib/prompts/search/writer';
+import fs from 'fs';
 
 class SearchAgent {
   async searchAsync(session: SessionManager, input: SearchAgentInput) {
@@ -15,15 +17,22 @@ class SearchAgent {
       llm: input.config.llm,
     });
 
-    session.emit('data', {
-      type: 'classification',
-      classification: classification,
-    });
-
     const widgetPromise = WidgetRegistry.executeAll(classification.widgets, {
       llm: input.config.llm,
       embedding: input.config.embedding,
       session: session,
+    }).then((widgetOutputs) => {
+      widgetOutputs.forEach((o) => {
+        session.emitBlock({
+          id: crypto.randomUUID(),
+          type: 'widget',
+          data: {
+            widgetType: o.type,
+            params: o.data,
+          },
+        });
+      });
+      return widgetOutputs;
     });
 
     let searchPromise: Promise<ResearcherOutput> | null = null;
@@ -42,6 +51,54 @@ class SearchAgent {
       widgetPromise,
       searchPromise,
     ]);
+
+    session.emit('data', {
+      type: 'researchComplete',
+    });
+
+    const finalContext =
+      searchResults?.findings
+        .filter((f) => f.type === 'search_results')
+        .flatMap((f) => f.results)
+        .map((f) => `${f.metadata.title}: ${f.content}`)
+        .join('\n') || '';
+
+    const widgetContext = widgetOutputs
+      .map((o) => {
+        return `${o.type}: ${JSON.stringify(o.data)}`;
+      })
+      .join('\n-------------\n');
+
+    const finalContextWithWidgets = `<search_results note="These are the search results and you can cite these">${finalContext}</search_results>\n<widgets_result noteForAssistant="Its output is already showed to the user, you can use this information to answer the query but do not CITE this as a souce">${widgetContext}</widgets_result>`;
+
+    const writerPrompt = getWriterPrompt(finalContextWithWidgets);
+
+    const answerStream = input.config.llm.streamText({
+      messages: [
+        {
+          role: 'system',
+          content: writerPrompt,
+        },
+        ...input.chatHistory,
+        {
+          role: 'user',
+          content: input.followUp,
+        },
+      ],
+    });
+
+    let accumulatedText = '';
+
+    for await (const chunk of answerStream) {
+      accumulatedText += chunk.contentChunk;
+
+      session.emit('data', {
+        type: 'response',
+        data: chunk.contentChunk,
+      });
+    }
+
+    session.emit('end', {});
   }
 }
 

src/lib/agents/search/researcher/actions/webSearch.ts

@@ -15,9 +15,9 @@ You have to use this action aggressively to find relevant information from the w
 When this action is present, you must use it to obtain current information from the web. 
 
 ### How to use:
-1. For fast search mode, you can use this action once. Make sure to cover all aspects of the user's query in that single search.
+1. For speed search mode, you can use this action once. Make sure to cover all aspects of the user's query in that single search.
 2. If you're on quality mode, you'll get to use this action up to two times. Use the first search to gather general information, and the second search to fill in any gaps or get more specific details based on the initial findings.
-3. If you're set on Deep research mode, then you will get to use this action multiple times to gather more information. Use your judgment to decide when additional searches are necessary to provide a thorough and accurate response.
+3. If you're set on quality mode, then you will get to use this action multiple times to gather more information. Use your judgment to decide when additional searches are necessary to provide a thorough and accurate response.
 
 Input: An array of search queries. Make sure the queries are relevant to the user's request and cover different aspects if necessary. You can include a maximum of 3 queries. Make sure the queries are SEO friendly and not sentences rather keywords which can be used to search a search engine like Google, Bing, etc.
 `;
@@ -32,6 +32,7 @@ const webSearchAction: ResearchAction<typeof actionSchema> = {
 
     const search = async (q: string) => {
       const res = await searchSearxng(q);
+
       res.results.forEach((r) => {
         results.push({
           content: r.content || r.title,

src/lib/agents/search/types.ts

@@ -10,7 +10,7 @@ export type SearchAgentConfig = {
   sources: SearchSources[];
   llm: BaseLLM<any>;
   embedding: BaseEmbedding<any>;
-  mode: 'fast' | 'balanced' | 'deep_research';
+  mode: 'speed' | 'balanced' | 'quality';
 };
 
 export type SearchAgentInput = {