changed from step tags to stage tags

Changed from step1, step2, etc. tags to more straightforward (and not similar to other step tags usage in XML) stage tags with stage number as an attribute.

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;