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;
   }