Merge pull request #53 from getzep/peter/zepai-299-add-roletype-role-…

…is-a-label-to-api

feat(Message model): introduce new RoleType and update files to handle it

.eslintrc.json

@@ -40,6 +40,7 @@
          }
       ],
       "max-classes-per-file": 0,
+      "camelcase": "warn",
       "multiline-ternary": 0,
       "no-unused-vars": "off",
       "no-underscore-dangle": "off",

examples/memory/history.ts

@@ -1,54 +1,72 @@
+import { RoleType } from "../../src/message_models";
+
 export const history = [
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Who was Octavia Butler?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Estelle Butler (June 22, 1947 – February 24, 2006) was an American" +
          " science fiction author.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Which books of hers were made into movies?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "The most well-known adaptation of Octavia Butler's work is the FX series" +
          " Kindred, based on her novel of the same name.",
    },
-   { role: "human", content: "Who were her contemporaries?" },
+   {
+      role: "human",
+      role_type: "user" as RoleType,
+      content: "Who were her contemporaries?" },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R." +
          " Delany, and Joanna Russ.",
    },
-   { role: "human", content: "What awards did she win?" },
+   {
+      role: "human",
+      role_type: "user" as RoleType,
+      content: "What awards did she win?" },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Butler won the Hugo Award, the Nebula Award, and the MacArthur" +
          " Fellowship.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Which other women sci-fi writers might I want to read?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content: "You might want to read Ursula K. Le Guin or Joanna Russ.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content:
          "Write a short synopsis of Butler's book, Parable of the Sower. What is it" +
          " about?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Parable of the Sower is a science fiction novel by Octavia Butler," +
          " published in 1993. It follows the story of Lauren Olamina, a young woman" +
@@ -57,31 +75,37 @@ export const history = [
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "What inspired Octavia Butler's writing style?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Butler was inspired by a wide range of influences, including her" +
          " passion for science fiction, her interest in sociology, and her desire to" +
          " explore issues of power, race, and gender in her works.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Did Octavia Butler write only science fiction?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "While Octavia Butler is best known for her contributions to science fiction," +
          " she also wrote in other genres, including fantasy and speculative fiction.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Tell me more about the Xenogenesis series.",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "The Xenogenesis series, also known as the Lilith's Brood trilogy, consists" +
          " of three novels: Dawn (1987), Adulthood Rites (1988), and Imago (1989)." +
@@ -90,21 +114,25 @@ export const history = [
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Are there any upcoming adaptations of Octavia Butler's work?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "As of my last update, there were discussions about potential adaptations," +
          " but no confirmed upcoming projects. It's advisable to check the latest" +
          " news sources for any recent developments.",
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "What is the Earthseed series about?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "The Earthseed series consists of Parable of the Sower (1993) and Parable of" +
          " the Talents (1998). It follows the protagonist, Lauren Olamina, as she" +
@@ -113,10 +141,12 @@ export const history = [
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "How did Octavia Butler impact the science fiction genre?",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Butler made significant contributions to the science fiction genre" +
          " by challenging conventions and addressing social issues. Her works" +
@@ -125,10 +155,12 @@ export const history = [
    },
    {
       role: "human",
+      role_type: "user" as RoleType,
       content: "Tell me about Octavia Butler's writing process.",
    },
    {
       role: "ai",
+      role_type: "assistant" as RoleType,
       content:
          "Octavia Butler was known for her disciplined and dedicated writing" +
          " process. She often wrote in the early morning hours and set specific" +

examples/memory/memory_example.ts

@@ -68,8 +68,8 @@ async function main() {
    // Add memory. We could do this in a batch, but we'll do it one by one rather to
    // ensure that summaries and other artifacts are generated correctly.
    try {
-      for (const { role, content } of history) {
-         const message = new Message({ role, content });
+      for (const { role, role_type, content } of history) {
+         const message = new Message({ role, role_type, content });
          const memory = new Memory({ messages: [message] });
 
          await client.memory.addMemory(sessionID, memory);

src/langchain/memory.ts

@@ -144,11 +144,11 @@ export class ZepMemory extends BaseChatMemory implements ZepMemoryInput {
       if (memory) {
          messages = messages.concat(
             memory.messages.map((message) => {
-               const { content, role } = message;
-               if (role === this.humanPrefix) {
+               const { content, role_type, role } = message;
+               if (role_type === "user") {
                   return new HumanMessage(content);
                }
-               if (role === this.aiPrefix) {
+               if (role_type === "assistant") {
                   return new AIMessage(content);
                }
                // default to generic ChatMessage
@@ -189,10 +189,12 @@ export class ZepMemory extends BaseChatMemory implements ZepMemoryInput {
          messages: [
             new Message({
                role: this.humanPrefix,
+               role_type: "user",
                content: `${input}`,
             }),
             new Message({
                role: this.aiPrefix,
+               role_type: "assistant",
                content: `${output}`,
             }),
          ],

src/langchain/message_history.ts

@@ -10,7 +10,7 @@ import { BaseChatMessageHistory } from "@langchain/core/chat_history";
 import ZepClient from "../zep-client";
 import { Memory } from "../memory_models";
 import { NotFoundError } from "../errors";
-import { Message } from "../message_models";
+import { getZepMessageRoleType, Message, RoleType } from "../message_models";
 import type { MemoryType } from "../interfaces";
 
 /**
@@ -22,6 +22,8 @@ interface ZepMemoryInput {
    sessionId: string;
    client: ZepClient;
    memoryType: MemoryType;
+   humanPrefix?: string;
+   aiPrefix?: string;
 }
 
 /**
@@ -43,11 +45,21 @@ export class ZepChatMessageHistory
 
    memoryType: MemoryType;
 
+   humanPrefix = "human";
+
+   aiPrefix = "ai";
+
    constructor(fields: ZepMemoryInput) {
       super();
       this.sessionId = fields.sessionId;
       this.memoryType = fields.memoryType;
       this.client = fields.client;
+      if (fields.humanPrefix) {
+         this.humanPrefix = fields.humanPrefix;
+      }
+      if (fields.aiPrefix) {
+         this.aiPrefix = fields.aiPrefix;
+      }
    }
 
    private async getMemory(): Promise<Memory | null> {
@@ -87,7 +99,8 @@ export class ZepChatMessageHistory
                   token_count: msg.token_count,
                   metadata: msg.metadata,
                };
-               if (msg.role === "ai") {
+
+               if (msg.role_type === ("assistant" as RoleType)) {
                   return new AIMessage(msg.content, metadata);
                }
                return new HumanMessage(msg.content, metadata);
@@ -98,13 +111,17 @@ export class ZepChatMessageHistory
       return messages;
    }
 
-   // @ts-ignore
    async addAIChatMessage(message: string, metadata?: any): Promise<void> {
       await this.addMessage(new AIMessage({ content: message }), metadata);
    }
 
-   // @ts-ignore
    async addMessage(message: BaseMessage, metadata?: any): Promise<void> {
+      const messageToSave = message;
+      if (message._getType() === "ai") {
+         messageToSave.name = this.aiPrefix;
+      } else if (message._getType() === "human") {
+         messageToSave.name = this.humanPrefix;
+      }
       if (message.content === null) {
          throw new Error("Message content cannot be null");
       }
@@ -115,8 +132,8 @@ export class ZepChatMessageHistory
 
       const zepMessage = new Message({
          content: message.content,
-         // eslint-disable-next-line no-underscore-dangle
-         role: message._getType(),
+         role: message.name ?? message._getType(),
+         role_type: getZepMessageRoleType(message._getType()),
          metadata,
       });
 

src/message_models.ts

@@ -1,10 +1,15 @@
+import { MessageType } from "@langchain/core/messages";
+
+export type RoleType = "user" | "assistant" | "system" | "function" | "tool";
+
 /**
  * IMessage interface for providing input to create a Message instance.
  */
 export interface IMessage {
    uuid?: string;
    created_at?: string;
    role: string;
+   role_type: RoleType;
    content: string;
    token_count?: number;
    metadata?: Record<string, any>;
@@ -20,6 +25,8 @@ export class Message {
 
    role: string;
 
+   role_type: RoleType;
+
    content: string;
 
    token_count?: number;
@@ -34,6 +41,7 @@ export class Message {
       this.uuid = data.uuid;
       this.created_at = data.created_at;
       this.role = data.role;
+      this.role_type = data.role_type;
       this.content = data.content;
       this.token_count = data.token_count;
       this.metadata = data.metadata;
@@ -48,9 +56,27 @@ export class Message {
          uuid: this.uuid,
          created_at: this.created_at,
          role: this.role,
+         role_type: this.role_type,
          content: this.content,
          token_count: this.token_count,
          metadata: this.metadata,
       };
    }
 }
+
+export const getZepMessageRoleType = (role: MessageType): RoleType => {
+   switch (role) {
+      case "human":
+         return "user";
+      case "ai":
+         return "assistant";
+      case "system":
+         return "system";
+      case "function":
+         return "function";
+      case "tool":
+         return "tool";
+      default:
+         return "system";
+   }
+};