save context

gregreindel · Jul 13, 2023 · f5f6654 · f5f6654
1 parent a790ae9
commit f5f6654
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/__tests__/unit/state/_base.ts b/__tests__/unit/state/_base.ts
@@ -161,9 +161,17 @@ describe("llm-exe:state/BaseState", () => {
     dialogue2.setUserMessage("This is the user message");
     dialogue2.setAssistantMessage("Hello, this is the assistant");
 
+    state.createContextItem(createStateItem("mock-item", {test: "val"}))
+
     expect(state.serialize()).toEqual({
       attributes: {},
-      context: {},
+      context: {
+        "mock-item": {
+          class: "BaseStateItem",
+          name: "mock-item",
+          value: {"mock-item": { test: "val" }},
+        }
+      },
       dialogues: {
         chat: {
           class: "Dialogue",

diff --git a/src/state/_base.ts b/src/state/_base.ts
@@ -63,7 +63,7 @@ export abstract class BaseState {
 
   serialize() {
     const dialogues: any = {};
-    const context: any = { ...this.context };
+    const context: any = {};
     const attributes: any = { ...this.attributes };
 
     const dialogueKeys = Object.keys(
@@ -73,6 +73,13 @@ export abstract class BaseState {
       dialogues[dialogueKey] = this.dialogues[dialogueKey].serialize();
     }
 
+    const contextKeys = Object.keys(
+      this.context
+    ) as (keyof typeof this.context)[];
+    for (const contextKey of contextKeys) {
+      context[contextKey] = this.context[contextKey].serialize();
+    }
+
     return {
       dialogues,
       context,