From a7b0ca1d9a5b0bd5482cfa32e698e79097ab786e Mon Sep 17 00:00:00 2001
From: Nicolas Brichet <nicolas.brichet@quantstack.net>
Date: Fri, 25 Oct 2024 10:00:53 +0200
Subject: [PATCH] Replace SystemMessage by AIMessage, and add the AIMessage to
 the history

---
 src/handler.ts | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/src/handler.ts b/src/handler.ts
index edf3097..5e9d8d3 100644
--- a/src/handler.ts
+++ b/src/handler.ts
@@ -12,9 +12,9 @@ import {
 import { UUID } from '@lumino/coreutils';
 import type { ChatMistralAI } from '@langchain/mistralai';
 import {
+  AIMessage,
   HumanMessage,
-  mergeMessageRuns,
-  SystemMessage
+  mergeMessageRuns
 } from '@langchain/core/messages';
 
 export type ConnectionMessage = {
@@ -39,21 +39,13 @@ export class CodestralHandler extends ChatModel {
     };
     this.messageAdded(msg);
     this._history.messages.push(msg);
-    // const response = await this._mistralClient.chat({
-    //   model: 'codestral-latest',
-    //   messages: this._history.messages.map(msg => {
-    //     return {
-    //       role: msg.sender === 'User' ? 'user' : 'assistant',
-    //       content: msg.body
-    //     };
-    //   })
-    // });
+
     const messages = mergeMessageRuns(
       this._history.messages.map(msg => {
         if (msg.sender.username === 'User') {
           return new HumanMessage(msg.body);
         }
-        return new SystemMessage(msg.body);
+        return new AIMessage(msg.body);
       })
     );
     const response = await this._mistralClient.invoke(messages);
@@ -67,6 +59,7 @@ export class CodestralHandler extends ChatModel {
       type: 'msg'
     };
     this.messageAdded(botMsg);
+    this._history.messages.push(botMsg);
     return true;
   }