Added support for new messages API endpoints (#42)

Added Message API endpoints * Minimum server version: 0.21.0 * Bump version to 0.10.0 Signed-off-by: Peter Garbers <[email protected]> --------- Signed-off-by: Peter Garbers <[email protected]> Co-authored-by: Peter <[email protected]>
getzep · Dec 14, 2023 · f4d2faa · f4d2faa
1 parent 3dd5ae6
commit f4d2faa
Show file tree

Hide file tree

Showing 10 changed files with 381 additions and 62 deletions.
diff --git a/examples/memory/memory_example.ts b/examples/memory/memory_example.ts
@@ -100,6 +100,54 @@ async function main() {
       }
    }
 
+   // get session messages
+   let sessionMessages: any[] = [];
+   try {
+      sessionMessages = await client.message.getSessionMessages(sessionID, 10, 1);
+      console.debug("Session messages: ", JSON.stringify(sessionMessages));
+   } catch (error) {
+      if (error instanceof NotFoundError) {
+         console.error("Session not found:", error.message);
+      } else {
+         console.error("Got error:", error);
+      }
+   }
+
+   const firstSessionsMessageId = sessionMessages[0].uuid;
+
+   // Update session message metadata
+   try {
+      const metadata = { metadata: { foo: "bar" }};
+      const updatedMessage = await client.message.updateSessionMessageMetadata(
+         sessionID,
+         firstSessionsMessageId,
+         metadata
+      );
+      console.debug("Updated message: ", JSON.stringify(updatedMessage));
+   } catch (error) {
+      if (error instanceof NotFoundError) {
+         console.error("Session not found:", error.message);
+      } else {
+         console.error("Got error:", error);
+      }
+   }
+
+   // Get session message
+
+   try {
+      const message = await client.message.getSessionMessage(
+         sessionID,
+         firstSessionsMessageId
+      );
+      console.debug("Session message: ", JSON.stringify(message));
+   } catch (error) {
+      if (error instanceof NotFoundError) {
+         console.error("Session not found:", error.message);
+      } else {
+         console.error("Got error:", error);
+      }
+   }
+
    // Search messages in memory
    try {
       const searchText = "Name some books that are about dystopian futures.";

diff --git a/package.json b/package.json
@@ -1,7 +1,7 @@
 {
   "show-config": "tsc --showConfig",
   "name": "@getzep/zep-js",
-  "version": "0.9.0",
+  "version": "0.10.0",
   "description": "Zep: Fast, scalable building blocks for production LLM apps",
   "private": false,
   "publishConfig": {

diff --git a/src/index.ts b/src/index.ts
@@ -1,11 +1,13 @@
 import MemoryManager from "./memory_manager";
+import MessageManager from "./message_manager";
 import ZepClient from "./zep-client";
 import DocumentManager from "./document_manager";
 import DocumentCollection from "./document_collection";
 import UserManager from "./user_manager";
 
 export {
    MemoryManager,
+   MessageManager,
    ZepClient,
    DocumentManager,
    DocumentCollection,
@@ -27,8 +29,6 @@ export {
 } from "./document_models";
 export {
    Memory,
-   Message,
-   IMessage,
    Summary,
    ISummary,
    IMemory,
@@ -39,6 +39,7 @@ export {
    Session,
    ISession,
 } from "./memory_models";
+export { Message, IMessage } from "./message_models";
 export {
    IAddCollectionParams,
    IUpdateCollectionParams,

diff --git a/src/memory_manager.ts b/src/memory_manager.ts
@@ -2,9 +2,11 @@ import {
    Memory,
    MemorySearchPayload,
    MemorySearchResult,
-   Message,
    Session,
 } from "./memory_models";
+
+import { Message } from "./message_models";
+
 import { IZepClient } from "./interfaces";
 import { handleRequest } from "./utils";
 

diff --git a/src/memory_models.ts b/src/memory_models.ts
@@ -1,3 +1,5 @@
+import { IMessage, Message } from "./message_models";
+
 /* eslint-disable camelcase */
 /**
  * Interface for Session data.
@@ -61,63 +63,6 @@ export class Session {
    }
 }
 
-/**
- * IMessage interface for providing input to create a Message instance.
- */
-export interface IMessage {
-   uuid?: string;
-   created_at?: string;
-   role: string;
-   content: string;
-   token_count?: number;
-   metadata?: Record<string, any>;
-}
-
-/**
- * Represents a message in the memory.
- */
-export class Message {
-   uuid?: string;
-
-   created_at?: string;
-
-   role: string;
-
-   content: string;
-
-   token_count?: number;
-
-   metadata?: Record<string, any>;
-
-   /**
-    * Constructs a new Message instance.
-    * @param {IMessage} data - The data to create a message instance.
-    */
-   constructor(data: IMessage) {
-      this.uuid = data.uuid;
-      this.created_at = data.created_at;
-      this.role = data.role;
-      this.content = data.content;
-      this.token_count = data.token_count;
-      this.metadata = data.metadata;
-   }
-
-   /**
-    * Converts the Message instance to a dictionary.
-    * @returns {IMessage} A dictionary representation of Message instance.
-    */
-   toDict(): IMessage {
-      return {
-         uuid: this.uuid,
-         created_at: this.created_at,
-         role: this.role,
-         content: this.content,
-         token_count: this.token_count,
-         metadata: this.metadata,
-      };
-   }
-}
-
 /**
  * ISummary interface for providing input to create a Summary instance.
  */

diff --git a/src/message_manager.ts b/src/message_manager.ts
@@ -0,0 +1,107 @@
+import { Message } from "./message_models";
+
+import { IZepClient } from "./interfaces";
+import { handleRequest } from "./utils";
+
+export default class MessageManager {
+   client: IZepClient;
+
+   constructor(client: IZepClient) {
+      this.client = client;
+   }
+
+   async getSessionMessages(
+      sessionId: string,
+      limit?: number,
+      cursor?: number,
+   ): Promise<Message[]> {
+      if (!sessionId || sessionId.trim() === "") {
+         throw new Error("sessionId must be provided");
+      }
+
+      let url = this.client.getFullUrl(`/sessions/${sessionId}/messages`);
+
+      const params = new URLSearchParams();
+      if (limit) params.append("limit", limit.toString());
+      if (cursor) params.append("cursor", cursor.toString());
+
+      if (params.toString()) url += `?${params.toString()}`;
+
+      const response = await handleRequest(
+         fetch(url, {
+            method: "GET",
+            headers: {
+               ...this.client.headers,
+               "Content-Type": "application/json",
+            },
+         }),
+         `No session found for session ${sessionId}`,
+      );
+
+      const responseData = await response.json();
+      return responseData.messages.map((message: any) => new Message(message));
+   }
+
+   async getSessionMessage(
+      sessionId: string,
+      messageId: string,
+   ): Promise<Message> {
+      if (!sessionId || sessionId.trim() === "") {
+         throw new Error("sessionId must be provided");
+      }
+
+      if (!messageId || messageId.trim() === "") {
+         throw new Error("messageId must be provided");
+      }
+
+      const response = await handleRequest(
+         fetch(
+            this.client.getFullUrl(
+               `/sessions/${sessionId}/messages/${messageId}`,
+            ),
+            {
+               headers: this.client.headers,
+            },
+         ),
+         `No session found for session ${sessionId}, or message ${messageId}`,
+      );
+
+      const responseData = await response.json();
+
+      return new Message(responseData);
+   }
+
+   async updateSessionMessageMetadata(
+      sessionId: string,
+      messageId: string,
+      metadata: Record<string, any>,
+   ): Promise<Message> {
+      if (!sessionId || sessionId.trim() === "") {
+         throw new Error("sessionId must be provided");
+      }
+
+      if (!messageId || messageId.trim() === "") {
+         throw new Error("messageId must be provided");
+      }
+
+      const response = await handleRequest(
+         fetch(
+            this.client.getFullUrl(
+               `/sessions/${sessionId}/messages/${messageId}`,
+            ),
+            {
+               method: "PATCH",
+               headers: this.client.headers,
+               body: JSON.stringify({
+                  metadata,
+               }),
+            },
+         ),
+         `No session found for session ${sessionId}, or message ${messageId}`,
+      );
+
+      const responseData = await response.json();
+
+      return new Message(responseData);
+   }
+}
diff --git a/src/message_models.ts b/src/message_models.ts
@@ -0,0 +1,56 @@
+/**
+ * IMessage interface for providing input to create a Message instance.
+ */
+export interface IMessage {
+   uuid?: string;
+   created_at?: string;
+   role: string;
+   content: string;
+   token_count?: number;
+   metadata?: Record<string, any>;
+}
+
+/**
+ * Represents a message in the memory.
+ */
+export class Message {
+   uuid?: string;
+
+   created_at?: string;
+
+   role: string;
+
+   content: string;
+
+   token_count?: number;
+
+   metadata?: Record<string, any>;
+
+   /**
+    * Constructs a new Message instance.
+    * @param {IMessage} data - The data to create a message instance.
+    */
+   constructor(data: IMessage) {
+      this.uuid = data.uuid;
+      this.created_at = data.created_at;
+      this.role = data.role;
+      this.content = data.content;
+      this.token_count = data.token_count;
+      this.metadata = data.metadata;
+   }
+
+   /**
+    * Converts the Message instance to a dictionary.
+    * @returns {IMessage} A dictionary representation of Message instance.
+    */
+   toDict(): IMessage {
+      return {
+         uuid: this.uuid,
+         created_at: this.created_at,
+         role: this.role,
+         content: this.content,
+         token_count: this.token_count,
+         metadata: this.metadata,
+      };
+   }
+}