Merge pull request #4 from cortexclick/jeremy/152

Implement listChats
cortexclick · Jun 20, 2024 · 416e4b6 · 416e4b6
2 parents 89522fc + 90139f2
commit 416e4b6
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 19 deletions.
diff --git a/chat.test.ts b/chat.test.ts
@@ -1,4 +1,4 @@
-import { expect, test } from 'vitest'
+import { expect, test } from 'vitest';
 import { CortexClient, TextDocument } from "./index";
 import { CatalogConfig } from "./catalog";
 import { Readable } from "stream";
@@ -61,18 +61,22 @@ test('e2e catalog, cortex, and sync chat', { timeout: 60000 }, async () => {
   const chat = await cortex.chat({ message: chatInput });
   expect(chat.messages[1].message.length).toBeGreaterThan(0);
 
-
   // get chat
   const getChatRes = await client.getChat(chat.id);
   expect(getChatRes.messages.length).toBe(2);
   expect(getChatRes.title).toBe(chatInput)
 
+  // respond to chat
   await chat.respond({ message: "what about customer verticals" });
   expect(chat.messages.length).toBe(4);
 
+  // list chats
+  const chatList = await client.listChats({ pageSize: 1 });
+  expect(chatList.chats.length).toBe(1);
 
-  // respond to chat
-
+  const nextPage = await chatList.nextPage();
+  expect(nextPage.chats.length).toBe(1);
+  expect(nextPage.chats[0].id).not.toBe(chat.id);
 
   // delete 
   await catalog.delete();
@@ -141,7 +145,7 @@ test('streaming chat', { timeout: 60000 }, async () => {
   statusStream.on('data', (data) => {
     const message = JSON.parse(data);
     expect(message.messageType).toBe("status");
-    switch(message.step) {
+    switch (message.step) {
       case "chat":
         sawChat = true;
         break;
@@ -152,22 +156,22 @@ test('streaming chat', { timeout: 60000 }, async () => {
 
   const chatResult = await chat;
 
-  expect(fullMessage).toBe(chatResult.messages[chatResult.messages.length -1 ].message);
+  expect(fullMessage).toBe(chatResult.messages[chatResult.messages.length - 1].message);
   expect(chatResult.messages.length).toBe(2);
   expect(sawChat).toBe(true);
 
   // respond to chat
   const respondResult = await chatResult.respond({ message: "what about customer verticals", stream: true, statusStream });
 
   const respondStream = respondResult.responseStream;
-  
+
   let fullResponse = ""
   respondStream.on('data', (data) => {
     fullResponse += data.toString();
   });
 
   const response = await respondResult.chat;
-  expect(fullResponse).toBe(response.messages[response.messages.length -1 ].message);
+  expect(fullResponse).toBe(response.messages[response.messages.length - 1].message);
   expect(response.messages.length).toBe(4);
 
   // delete 

diff --git a/chat.ts b/chat.ts
@@ -39,6 +39,21 @@ export interface RespondChatOptsSync extends RespondChatOptsBase {
     stream?: false
 }
 
+export interface ChatListItem {
+    title: string;
+    id: string;
+    messageCount: number;
+    Chat(): Promise<Chat>;
+}
+export interface ChatListResult {
+    chats: ChatListItem[];
+    nextPage: () => Promise<ChatListResult>;
+}
+export interface ChatListPaginationOpts {
+    cursor?: string;
+    pageSize?: number;
+}
+
 export type Message = {
     role: "user" | "cortex";
     message: string;
@@ -58,6 +73,33 @@ export class Chat {
         }
     }
 
+    static async list(client: CortexApiClient, paginationOpts?: ChatListPaginationOpts): Promise<ChatListResult> {
+        const chats: ChatListItem[] = [];
+
+        const query = new URLSearchParams();
+        if (paginationOpts?.cursor) {
+            query.set("cursor", paginationOpts.cursor);
+        }
+        query.set("pageSize", (paginationOpts?.pageSize || 50).toString());
+        const res = await client.GET(`/chats?${query.toString()}`);
+        if (res.status !== 200) {
+            throw new Error(`Failed to list chats: ${res.statusText}`);
+        }
+        const body = await res.json();
+        for (const chat of body.chats) {
+            chats.push({
+                title: chat.title,
+                id: chat.chatId,
+                messageCount: chat.messageCount,
+                Chat: () => { return Chat.get(client, chat.chatId) }
+            })
+        }
+
+        const cursor = body.cursor;
+        const pageSize = paginationOpts?.pageSize;
+        return { chats, nextPage: async () => { return Chat.list(client, { cursor, pageSize }) } };
+    }
+
     private static async createContentSync(opts: CreateChatOptsSync): Promise<Chat> {
         const { client, cortex, message } = opts;
         const res = await client.POST(`/chats`, { cortex: cortex.name, message });
@@ -87,7 +129,7 @@ export class Chat {
         const readableStream = new Readable({
             read() { }
         });
-    
+
         const chatPromise = processStream(reader, decoder, readableStream, opts.statusStream).then(content => {
             const messages: Message[] = [
                 {
@@ -102,8 +144,8 @@ export class Chat {
             return new Chat(client, id, title, messages);
         });
 
-        return { responseStream: readableStream, chat: chatPromise};
-     }
+        return { responseStream: readableStream, chat: chatPromise };
+    }
 
     static async get(client: CortexApiClient, id: string): Promise<Chat> {
         const res = await client.GET(`/chats/${id}`);
@@ -117,7 +159,7 @@ export class Chat {
     async respond(opts: RespondChatOptsSync): Promise<string>;
     async respond(opts: RespondChatOptsStreaming): Promise<StreamingChatResult>;
     async respond(opts: RespondChatOptsSync | RespondChatOptsStreaming): Promise<string | StreamingChatResult> {
-        if(isRespondChatOptsSync(opts)) {
+        if (isRespondChatOptsSync(opts)) {
             return this.respondChatSync(opts);
         } else {
             return this.respondChatStreaming(opts);
@@ -150,7 +192,7 @@ export class Chat {
         const readableStream = new Readable({
             read() { }
         });
-    
+
         const chatPromise = processStream(reader, decoder, readableStream, opts.statusStream).then(content => {
             this.messages.push(
                 {
@@ -164,7 +206,7 @@ export class Chat {
             return this;
         });
 
-        return { responseStream: readableStream, chat: chatPromise};
+        return { responseStream: readableStream, chat: chatPromise };
     }
 
 }

diff --git a/client.ts b/client.ts
@@ -45,6 +45,10 @@ export interface ClientListContentPaginationOpts {
     pageSize?: number;
     cursor?: string;
 }
+export interface ClientListChatPaginationOpts {
+    pageSize?: number;
+    cursor?: string;
+}
 
 const apiUrl = process.env.CORTEX_API_URL || "https://api.cortexclick.com";
 
@@ -114,7 +118,9 @@ export class CortexClient {
         return Content.list(this.apiClient, paginationOptions);
     }
 
-    async listChats() { }
+    async listChats(paginationOptions?: ClientListChatPaginationOpts) {
+        return Chat.list(this.apiClient, paginationOptions);
+    }
 
     async getCortex(name: string): Promise<Cortex> {
         return Cortex.get(this.apiClient, name)

diff --git a/content.ts b/content.ts
@@ -233,7 +233,7 @@ export class Content {
     }
 
     static async list(client: CortexApiClient, paginationOpts?: ContentListPaginationOptions): Promise<ContentListResult> {
-        const content: ContentListItem[] = [];
+        const contentList: ContentListItem[] = [];
 
         const query = new URLSearchParams();
         if (paginationOpts?.cursor) {
@@ -245,8 +245,8 @@ export class Content {
             throw new Error(`Failed to list content: ${res.statusText}`);
         }
         const body = await res.json();
-        for (let content of body.content) {
-            content.push({
+        for (const content of body.content) {
+            contentList.push({
                 title: content.title,
                 latestVersion: content.latestVersion,
                 id: content.contentId,
@@ -256,7 +256,7 @@ export class Content {
 
         const cursor = body.cursor;
         const pageSize = paginationOpts?.pageSize;
-        return { content, nextPage: async () => { return Content.list(client, { cursor, pageSize }) } };
+        return { content: contentList, nextPage: async () => { return Content.list(client, { cursor, pageSize }) } };
     }
 }