1 files changed, 18 insertions, 17 deletions
diff --git a/src/client/views/nodes/ChatBox/Agent.ts b/src/client/views/nodes/ChatBox/Agent.ts
index 825cd831b..41c91b4c6 100644
--- a/src/client/views/nodes/ChatBox/Agent.ts
+++ b/src/client/views/nodes/ChatBox/Agent.ts
@@ -54,13 +54,13 @@ export class Agent {
         console.log(`System prompt: ${systemPrompt}`);
         this.interMessages = [{ role: 'system', content: systemPrompt }];
 
-        this.interMessages.push({ role: 'user', content: `<step1 role="user"><query>${question}</query></step>` });
+        this.interMessages.push({ role: 'user', content: `<stage number="1" role="user"><query>${question}</query></stage>` });
 
-        const parser = new XMLParser();
-        const builder = new XMLBuilder();
+        const parser = new XMLParser({ ignoreAttributes: false, attributeNamePrefix: '@_' });
+        const builder = new XMLBuilder({ ignoreAttributes: false, attributeNamePrefix: '@_' });
         let currentAction: string | undefined;
 
-        for (let i = 3; i < maxTurns; i += 2) {
+        for (let i = 2; i < maxTurns; i += 2) {
             console.log(`Turn ${i}/${maxTurns}`);
 
             const result = await this.execute();
@@ -75,19 +75,24 @@ export class Agent {
                 return '<error>Invalid response format.</error>';
             }
 
-            const step = parsedResult[Object.keys(parsedResult)[0]];
+            const stage = parsedResult.stage;
 
-            for (const key in step) {
+            if (!stage) {
+                console.log('Error: No stage found in response');
+                return '<error>Invalid response format: No stage found.</error>';
+            }
+
+            for (const key in stage) {
                 if (key === 'thought') {
-                    console.log(`Thought: ${step[key]}`);
+                    console.log(`Thought: ${stage[key]}`);
                 } else if (key === 'action') {
-                    currentAction = step[key] as string;
+                    currentAction = stage[key] as string;
                     console.log(`Action: ${currentAction}`);
                     if (this.tools[currentAction]) {
                         const nextPrompt = [
                             {
                                 type: 'text',
-                                text: `<step${i} role="user">` + builder.build({ action_rules: this.tools[currentAction].getActionRule() }) + `<\step>`,
+                                text: `<stage number="${i + 1}" role="user">` + builder.build({ action_rules: this.tools[currentAction].getActionRule() }) + `</stage>`,
                             },
                         ];
                         this.interMessages.push({ role: 'user', content: nextPrompt });
@@ -95,20 +100,16 @@ export class Agent {
                         break;
                     } else {
                         console.log('Error: No valid action');
-                        this.interMessages.push({ role: 'user', content: `<step${i}>No valid action, try again.</step>` });
+                        this.interMessages.push({ role: 'user', content: `<stage number="${i + 1}" role="system-error-reporter">No valid action, try again.</stage>` });
                         break;
                     }
                 } else if (key === 'action_input') {
-                    const actionInput = builder.build({ action_input: step[key] });
+                    const actionInput = builder.build({ action_input: stage[key] });
                     console.log(`Action input: ${actionInput}`);
                     if (currentAction) {
                         try {
-                            const observation = await this.processAction(currentAction, step[key]);
-                            // const stepElement = parsedResult.documentElement;
-                            // const rootTagName = stepElement.tagName;
-                            // const match = rootTagName.match(/step(\d+)/);
-                            // const currentStep = match ? parseInt(match[1]) + 1 : 1;
-                            const nextPrompt = [{ type: 'text', text: `<step${i}> <observation>` }, ...observation, { type: 'text', text: '</observation></step>' }];
+                            const observation = await this.processAction(currentAction, stage[key]);
+                            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 });
                             break;