Streaming wiht thoughts and actions working much better but still get error for web search

1 files changed, 34 insertions, 78 deletions

diff --git a/src/client/views/nodes/ChatBox/Agent.ts b/src/client/views/nodes/ChatBox/Agent.ts
index 413ecbd41..ae08271ee 100644
--- a/src/client/views/nodes/ChatBox/Agent.ts
+++ b/src/client/views/nodes/ChatBox/Agent.ts
@@ -1,5 +1,5 @@
 import OpenAI from 'openai';
-import { Tool, AgentMessage, AssistantMessage, TEXT_TYPE, CHUNK_TYPE, ASSISTANT_ROLE } from './types';
+import { Tool, AgentMessage, AssistantMessage, TEXT_TYPE, CHUNK_TYPE, ASSISTANT_ROLE, ProcessingInfo, PROCESSING_TYPE } from './types';
 import { getReactPrompt } from './prompts';
 import { XMLParser, XMLBuilder } from 'fast-xml-parser';
 import { Vectorstore } from './vectorstore/Vectorstore';
@@ -29,6 +29,8 @@ export class Agent {
     private _csvData: () => { filename: string; id: string; text: string }[];
     private actionNumber: number = 0;
     private thoughtNumber: number = 0;
+    private processingNumber: number = 0;
+    private processingInfo: ProcessingInfo[] = [];
 
     constructor(_vectorstore: Vectorstore, summaries: () => string, history: () => string, csvData: () => { filename: string; id: string; text: string }[], addLinkedUrlDoc: (url: string, id: string) => void) {
         this.client = new OpenAI({ apiKey: process.env.OPENAI_KEY, dangerouslyAllowBrowser: true });
@@ -46,7 +48,7 @@ export class Agent {
         };
     }
 
-    async askAgent(question: string, maxTurns: number = 30, onUpdate: (update: AssistantMessage) => void): Promise<AssistantMessage> {
+    async askAgent(question: string, onUpdate: (update: ProcessingInfo[]) => void, maxTurns: number = 30): Promise<AssistantMessage> {
         console.log(`Starting query: ${question}`);
         this.messages.push({ role: 'user', content: question });
         const chatHistory = this._history();
@@ -57,48 +59,34 @@ export class Agent {
         const builder = new XMLBuilder({ ignoreAttributes: false, attributeNamePrefix: '@_' });
         let currentAction: string | undefined;
 
-        let assistantMessage: AssistantMessage = {
-            role: ASSISTANT_ROLE.ASSISTANT,
-            content: [],
-            thoughts: [],
-            actions: [],
-            citations: [],
-        };
+        this.processingInfo = [];
 
         for (let i = 2; i < maxTurns; i += 2) {
             console.log(`Turn ${i}/${maxTurns}`);
 
-            const result = await this.execute(assistantMessage, onUpdate);
+            const result = await this.execute(onUpdate);
             this.interMessages.push({ role: 'assistant', content: result });
 
             let parsedResult;
             try {
                 parsedResult = parser.parse(result);
             } catch (error) {
-                console.log('Error: Invalid XML response from bot');
-                assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
-                return assistantMessage;
+                throw new Error(`Error parsing response: ${error}`);
             }
 
             const stage = parsedResult.stage;
 
             if (!stage) {
-                console.log('Error: No stage found in response');
-                assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
-                return assistantMessage;
+                throw new Error(`Error: No stage found in response`);
             }
 
             for (const key in stage) {
-                if (!assistantMessage.actions) {
-                    assistantMessage.actions = [];
-                }
                 if (key === 'thought') {
                     console.log(`Thought: ${stage[key]}`);
-                    this.thoughtNumber++;
+                    this.processingNumber++;
                 } else if (key === 'action') {
                     currentAction = stage[key] as string;
                     console.log(`Action: ${currentAction}`);
-                    onUpdate({ ...assistantMessage });
                     if (this.tools[currentAction]) {
                         const nextPrompt = [
                             {
@@ -118,35 +106,29 @@ export class Agent {
                     console.log(`Action input: ${actionInput}`);
                     if (currentAction) {
                         try {
-                            const observation = await this.processAction(currentAction, stage[key]);
+                            const observation = await this.processAction(currentAction, stage[key].inputs);
                             const nextPrompt = [{ type: 'text', text: `<stage number="${i + 1}" role="user"> <observation>` }, ...observation, { type: 'text', text: '</observation></stage>' }];
                             console.log(observation);
                             this.interMessages.push({ role: 'user', content: nextPrompt });
-                            this.actionNumber++; //might not work with no tool
+                            this.processingNumber++;
                             break;
                         } catch (error) {
-                            console.log(`Error processing action: ${error}`);
-                            assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
-                            return assistantMessage;
+                            throw new Error(`Error processing action: ${error}`);
                         }
                     } else {
-                        console.log('Error: Action input without a valid action');
-                        assistantMessage.content.push({ index: assistantMessage.content.length, type: TEXT_TYPE.ERROR, text: 'Invalid response from bot', citation_ids: null });
-                        return assistantMessage;
+                        throw new Error('Error: Action input without a valid action');
                     }
                 } else if (key === 'answer') {
                     console.log('Answer found. Ending query.');
-                    const parsedAnswer = AnswerParser.parse(result, assistantMessage);
-                    onUpdate({ ...parsedAnswer });
+                    const parsedAnswer = AnswerParser.parse(result, this.processingInfo);
                     return parsedAnswer;
                 }
             }
         }
-        console.log('Reached maximum turns. Ending query.');
-        return assistantMessage;
+        throw new Error('Reached maximum turns. Ending query.');
     }
 
-    private async execute(assistantMessage: AssistantMessage, onUpdate: (update: AssistantMessage) => void): Promise<string> {
+    private async execute(onUpdate: (update: ProcessingInfo[]) => void): Promise<string> {
         const stream = await this.client.chat.completions.create({
             model: 'gpt-4o',
             messages: this.interMessages as ChatCompletionMessageParam[],
@@ -158,12 +140,6 @@ export class Agent {
         let currentTag: string = '';
         let currentContent: string = '';
         let isInsideTag: boolean = false;
-        let isInsideActionInput: boolean = false;
-        let actionInputContent: string = '';
-
-        if (!assistantMessage.actions) {
-            assistantMessage.actions = [];
-        }
 
         for await (const chunk of stream) {
             const content = chunk.choices[0]?.delta?.content || '';
@@ -172,40 +148,20 @@ export class Agent {
             for (const char of content) {
                 if (char === '<') {
                     isInsideTag = true;
-                    if (currentTag && currentContent) {
-                        if (currentTag === 'action_input') {
-                            assistantMessage.actions[assistantMessage.actions.length - 1].action_input = actionInputContent;
-                            actionInputContent = '';
-                        } else {
-                            this.processStreamedContent(currentTag, currentContent, assistantMessage);
-                        }
-                        onUpdate({ ...assistantMessage });
-                    }
                     currentTag = '';
                     currentContent = '';
                 } else if (char === '>') {
                     isInsideTag = false;
-                    if (currentTag === 'action_input') {
-                        isInsideActionInput = true;
-                    } else if (currentTag === '/action_input') {
-                        isInsideActionInput = false;
-                        console.log('Action input:', actionInputContent);
-                        assistantMessage.actions[assistantMessage.actions.length - 1].action_input = actionInputContent;
-                        actionInputContent = '';
-                        onUpdate({ ...assistantMessage });
-                    }
                     if (currentTag.startsWith('/')) {
                         currentTag = '';
                     }
                 } else if (isInsideTag) {
                     currentTag += char;
-                } else if (isInsideActionInput) {
-                    actionInputContent += char;
                 } else {
                     currentContent += char;
-                    if (currentTag === 'thought' || currentTag === 'action') {
-                        this.processStreamedContent(currentTag, currentContent, assistantMessage);
-                        onUpdate({ ...assistantMessage });
+                    if (currentTag === 'thought' || currentTag === 'action_input_description') {
+                        this.processStreamedContent(currentTag, currentContent);
+                        onUpdate(this.processingInfo);
                     }
                 }
             }
@@ -214,24 +170,24 @@ export class Agent {
         return fullResponse;
     }
 
-    private processStreamedContent(tag: string, content: string, assistantMessage: AssistantMessage) {
-        if (!assistantMessage.thoughts) {
-            assistantMessage.thoughts = [];
-        }
-        if (!assistantMessage.actions) {
-            assistantMessage.actions = [];
-        }
+    private processStreamedContent(tag: string, streamed_content: string) {
+        const current_info = this.processingInfo.find(info => info.index === this.processingNumber);
         switch (tag) {
             case 'thought':
-                assistantMessage.thoughts[this.thoughtNumber] = content;
-                break;
-            case 'action':
-                assistantMessage.actions[this.actionNumber].action = content;
-
-                break;
-            case 'action_input':
-                assistantMessage.actions[this.actionNumber].action_input = content;
+                if (current_info) {
+                    current_info.content = streamed_content;
+                } else {
+                    console.log(`Adding thought: ${streamed_content}`);
+                    this.processingInfo.push({ index: this.processingNumber, type: PROCESSING_TYPE.THOUGHT, content: streamed_content.trim() });
+                }
                 break;
+            case 'action_input_description':
+                if (current_info) {
+                    current_info.content = streamed_content;
+                } else {
+                    console.log(`Adding thought: ${streamed_content}`);
+                    this.processingInfo.push({ index: this.processingNumber, type: PROCESSING_TYPE.ACTION, content: streamed_content.trim() });
+                }
         }
     }