From 3348343d8d9c57daecab878e7545b34b4b27f28d Mon Sep 17 00:00:00 2001
From: Corina Gum <>
Date: Tue, 20 Aug 2024 14:59:56 -0700
Subject: [PATCH] Fix up ActionOutputMessage

---
 js/packages/teams-ai/src/AI.ts                |  2 +-
 .../src/prompts/ActionOutputMessage.ts        | 23 ++++++++++++-------
 .../teams-ai/src/prompts/PromptManager.ts     |  4 +++-
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/js/packages/teams-ai/src/AI.ts b/js/packages/teams-ai/src/AI.ts
index ae296488cb..401886b31b 100644
--- a/js/packages/teams-ai/src/AI.ts
+++ b/js/packages/teams-ai/src/AI.ts
@@ -454,7 +454,7 @@ export class AI<TState extends TurnState = TurnState> {
                 state.temp.inputFiles = [];
 
                 if (cmd.type === 'DO' && (cmd as PredictedDoCommand).actionId) {
-                    state.deleteValue(state.temp.input);
+                    state.deleteValue('temp.input');
                 } else {
                     state.temp.input = output;
                 }
diff --git a/js/packages/teams-ai/src/prompts/ActionOutputMessage.ts b/js/packages/teams-ai/src/prompts/ActionOutputMessage.ts
index ba85790e7f..0a4fecb7b9 100644
--- a/js/packages/teams-ai/src/prompts/ActionOutputMessage.ts
+++ b/js/packages/teams-ai/src/prompts/ActionOutputMessage.ts
@@ -14,21 +14,22 @@ import { Message } from './Message';
 import { PromptFunctions } from './PromptFunctions';
 import { RenderedPromptSection } from './PromptSection';
 import { PromptSectionBase } from './PromptSectionBase';
+import { ActionCall } from '../types';
 
 /**
  * A section capable of rendering user input text and images as a user message.
  */
 export class ActionOutputMessage extends PromptSectionBase {
-    private readonly _inputVariable: string;
+    private readonly _historyVariable: string;
 
     /**
      * Creates a new 'ActionOutputMessage' instance.
+     * @param {string} historyVariable Optional. Name of the variable containing the conversation history.
      * @param {number} tokens Optional. Sizing strategy for this section. Defaults to `auto`.
-     * @param {string} inputVariable Optional. Name of the variable containing the user input text. Defaults to `input`.
      */
-    public constructor(tokens: number = -1, inputVariable = 'actionOutputs') {
+    public constructor(historyVariable = 'temp.history', tokens: number = -1) {
         super(tokens, true, '\n', 'action: ');
-        this._inputVariable = inputVariable;
+        this._historyVariable = historyVariable;
     }
 
     /**
@@ -47,14 +48,20 @@ export class ActionOutputMessage extends PromptSectionBase {
         tokenizer: Tokenizer,
         maxTokens: number
     ): Promise<RenderedPromptSection<Message<any>[]>> {
-        const actionOutputs: Record<string, string> = memory.getValue(this._inputVariable) || {};
+        let actionOutputs: Record<string, string> = {};
+        let actionCalls: ActionCall[] = [];
+        const history: Message[] = memory.getValue(this._historyVariable) ?? [];
         const messages: Message<string>[] = [];
 
-        for (const action in actionOutputs) {
+        if (history.length > 1) {
+            actionOutputs = memory.getValue('temp.actionOutputs');
+            actionCalls = history[history.length - 1].action_calls ?? [];
+        }
+        for (const actionCall of actionCalls) {
             const message: Message<string> = {
                 role: 'tool',
-                content: actionOutputs[action],
-                action_call_id: action
+                content: actionOutputs[actionCall.id],
+                action_call_id: actionCall.id
             };
             messages.push(message);
         }
diff --git a/js/packages/teams-ai/src/prompts/PromptManager.ts b/js/packages/teams-ai/src/prompts/PromptManager.ts
index 967fae9557..1f6013381f 100644
--- a/js/packages/teams-ai/src/prompts/PromptManager.ts
+++ b/js/packages/teams-ai/src/prompts/PromptManager.ts
@@ -342,7 +342,9 @@ export class PromptManager implements PromptFunctions {
                 sections.push(new UserMessage('{{$temp.input}}', this.options.max_input_tokens));
             }
             if (template.config.augmentation && template.config.augmentation.augmentation_type === 'tools') {
-                sections.push(new ActionOutputMessage(this.options.max_input_tokens));
+                const includeHistory: boolean = template.config.completion.include_history;
+                const historyVariable = includeHistory ? `conversation.${name}_history` : 'temp.history';
+                sections.push(new ActionOutputMessage(historyVariable));
             }
 
             // Create prompt