Update agent.module.ts

indexnetwork · Jul 23, 2024 · 0173929 · 0173929
1 parent 0e1a327
commit 0173929
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/indexer/src/app/modules/agent.module.ts b/indexer/src/app/modules/agent.module.ts
@@ -17,18 +17,19 @@ import { DynamicModule, Logger, Module } from '@nestjs/common';
 import { pull } from 'langchain/hub';
 import { HttpService } from '@nestjs/axios';
 
-const formatChatHistory = (chatHistory: string | string[]) => {
+const formatChatHistory = (chatHistory: any) => {
   if (Array.isArray(chatHistory)) {
     const updatedChat = chatHistory
-      .map((dialogTurn: any) => {
-        if (dialogTurn['role'] == 'user') {
-          return `Human: ${dialogTurn['content']}`;
+      .map((dialogTurn: HumanMessage | AIMessage | SystemMessage) => {
+        const mType = dialogTurn._getType();
+        if (mType == 'human') {
+          return `Human: ${dialogTurn.content}`;
         }
-        if (dialogTurn['role'] == 'assistant') {
-          return `AI: ${dialogTurn['content']}`;
+        if (mType === 'ai') {
+          return `AI: ${dialogTurn.content}`;
         }
-        if (dialogTurn['role'] == 'system') {
-          return `System: ${dialogTurn['content']}`;
+        if (mType == 'system') {
+          return `System: ${dialogTurn.content}`;
         }
       })
       .join('\n');