From 022c05a727686a4b5260c6684759aafc01290631 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Tue, 8 Oct 2024 12:11:32 +0200
Subject: [PATCH] feat(memory): switch from LLM to  ChatLLM in SummarizeMemory

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 src/memory/summarizeMemory.ts | 53 +++++++++++++++--------------------
 1 file changed, 22 insertions(+), 31 deletions(-)

diff --git a/src/memory/summarizeMemory.ts b/src/memory/summarizeMemory.ts
index 4d630f5..fd7619b 100644
--- a/src/memory/summarizeMemory.ts
+++ b/src/memory/summarizeMemory.ts
@@ -14,45 +14,26 @@
  * limitations under the License.
  */
 
-import { BaseMessage } from "@/llms/primitives/message.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
 import { BaseMemory } from "@/memory/base.js";
 import { PromptTemplate } from "@/template.js";
-import { LLM } from "@/llms/llm.js";
-import { BaseLLMOutput } from "@/llms/base.js";
 import { shallowCopy } from "@/serializer/utils.js";
 import { z } from "zod";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
 
 export interface SummarizeMemoryInput {
-  llm: LLM<BaseLLMOutput>;
-  template: typeof SummarizeMemoryTemplate;
+  llm: ChatLLM<ChatLLMOutput>;
+  template?: typeof SummarizeMemoryTemplate;
 }
 
 export const SummarizeMemoryTemplate = new PromptTemplate({
   schema: z.object({
     summary: z.string(),
-    new_lines: z.string(),
   }),
   template: `Progressively summarize the lines of conversation provided, adding onto the previous summary returning a new summary.
 
-EXAMPLE
 Current summary:
-The human asks what the AI thinks of artificial intelligence. The AI thinks artificial intelligence is a force for good.
-
-New lines of conversation:
-Human: Why do you think artificial intelligence is a force for good?
-AI: Because artificial intelligence will help humans reach their full potential.
-
-New summary:
-The human asks what the AI thinks of artificial intelligence. The AI thinks artificial intelligence is a force for good because it will help humans reach their full potential.
-END OF EXAMPLE
-
-Current summary:
-{{summary}}
-
-New lines of conversation:
-{{new_lines}}
-
-New summary:`,
+{{summary}}`,
 });
 
 export class SummarizeMemory extends BaseMemory {
@@ -62,7 +43,7 @@ export class SummarizeMemory extends BaseMemory {
 
   constructor(config: SummarizeMemoryInput) {
     super();
-    this.template = config.template;
+    this.template = config.template ?? SummarizeMemoryTemplate;
     this.llm = config.llm;
   }
 
@@ -78,19 +59,29 @@ export class SummarizeMemory extends BaseMemory {
 
     return [
       BaseMessage.of({
-        role: "system",
+        role: Role.ASSISTANT,
         text: currentSummary,
       }),
     ];
   }
 
   async add(message: BaseMessage) {
-    const prompt = this.template.render({
-      summary: this.summary,
-      new_lines: `${message.role}: ${message.text}`,
-    });
+    const response = await this.llm.generate([
+      BaseMessage.of({
+        role: Role.SYSTEM,
+        text: this.template.render({
+          summary: this.summary,
+        }),
+      }),
+      BaseMessage.of({
+        role: Role.ASSISTANT,
+        text: `New lines of conversation:
+${message.role}: ${message.text}
 
-    const response = await this.llm.generate(prompt);
+New summary:
+`,
+      }),
+    ]);
     this.summary = response.getTextContent();
   }