Feat/prompt error

.changeset/loud-wasps-approve.md

@@ -0,0 +1,5 @@
+---
+'outpostkit': patch
+---
+
+added prompt response type and error handling

src/comet.ts

@@ -7,7 +7,9 @@ import type {
   ProvideMessageFeedbackPayload,
   IComet,
   ListSessionsPayload,
-  ICometSession,
+  GetSessionPayload,
+  TCometPromptResponse,
+  ICometSession
 } from './types';
 import { streamPromptWithAxios, streamPromptWithNativeFetch } from 'helpers';
 
@@ -38,11 +40,33 @@ export class Comet implements IComet {
     //TODO: better error handling
     if (payload.stream) {
       if (typeof window !== 'undefined') {
-        return await streamPromptWithNativeFetch(this.cometId, this.apiKey, payload, handleNewText);
-      } else return await streamPromptWithAxios(this.cometAPI, payload, handleNewText);
+        const resp = await streamPromptWithNativeFetch(
+          this.cometId,
+          this.apiKey,
+          payload,
+          handleNewText
+        );
+
+        // @ts-ignore
+        if (resp.error) {
+          // @ts-ignore
+          throw new Error(resp.error);
+        } else {
+          return resp as TCometPromptResponse;
+        }
+      } else {
+        const resp = await streamPromptWithAxios(this.cometAPI, payload, handleNewText);
+        // @ts-ignore
+        if (resp.error) {
+          // @ts-ignore
+          throw new Error(resp.error);
+        } else {
+          return resp as TCometPromptResponse;
+        }
+      }
     } else {
       const { data } = await this.cometAPI.post(`/prompt`, payload);
-      return data;
+      return data as TCometPromptResponse;
     }
   }
 

src/helpers.ts

@@ -1,7 +1,8 @@
 import { AxiosInstance } from 'axios';
 import { PROMPT_STREAM_RESPONSE_PREFIX } from './constants';
 import { API_V1_URL } from './constants';
-import { PromptPayload } from 'types';
+import { PromptPayload, TCometPromptResponse } from 'types';
+import { TCometPromptStreamResponseError } from './types';
 
 export function sanitizeFilename(filename: string): string {
   const splitFilename = filename.split('.');
@@ -18,7 +19,7 @@ export const streamPromptWithNativeFetch = (
   payload: PromptPayload,
   handleNewText?: (token: string) => void | Promise<void>
 ) => {
-  return new Promise<{ response: string; referencePaths: string[] }>((resolve, reject) => {
+  return new Promise<TCometPromptResponse | TCometPromptStreamResponseError>((resolve, reject) => {
     (async () => {
       const response = await fetch(`${API_V1_URL}/comets/${cometId}/prompt`, {
         method: 'POST',
@@ -56,12 +57,7 @@ export const streamPromptWithNativeFetch = (
           }
           if (done) {
             try {
-              return resolve(
-                JSON.parse(responseText) as {
-                  response: string;
-                  referencePaths: string[];
-                }
-              );
+              return resolve(JSON.parse(responseText));
             } catch (e) {
               return reject('Could not parse the response');
             }
@@ -79,7 +75,7 @@ export const streamPromptWithAxios = (
   payload: PromptPayload,
   handleNewText?: (token: string) => void | Promise<void>
 ) => {
-  return new Promise((resolve, reject) => {
+  return new Promise<TCometPromptResponse | TCometPromptStreamResponseError>((resolve, reject) => {
     (async () => {
       const { data: stream } = await cometAPI.post(`/prompt`, payload, {
         responseType: 'stream',
@@ -104,7 +100,7 @@ export const streamPromptWithAxios = (
       });
 
       stream.on('end', () => {
-        return resolve(responseText);
+        return resolve(JSON.parse(responseText));
       });
     })();
   });

src/types.ts

@@ -1,30 +1,30 @@
 // Memory Client
 export interface MemoryClient {
-  addMemory: (sessionId: string, payload: Memory) => Promise<Memory>
-  getMemory: (sessionId: string) => Promise<Memory>
-  deleteMemory: (sessionId: string) => Promise<void>
+  addMemory: (sessionId: string, payload: Memory) => Promise<Memory>;
+  getMemory: (sessionId: string) => Promise<Memory>;
+  deleteMemory: (sessionId: string) => Promise<void>;
 }
 
 export interface MemoryConfig {
-  apiKey?: string
-  clientId?: string
-  baseUrl?: string
+  apiKey?: string;
+  clientId?: string;
+  baseUrl?: string;
 }
 
 export interface Memory {
   messages: Array<{
-    content: string
-    role: 'Human' | 'AI'
-  }>
-  context?: string
+    content: string;
+    role: 'Human' | 'AI';
+  }>;
+  context?: string;
 }
 
 // Comet Client
 export interface IComet {
   prompt: (
     payload: PromptPayload,
     handleNewText?: (data: string) => void | Promise<void>
-  ) => Promise<any>;
+  ) => Promise<TCometPromptResponse>;
   updateConfig: (payload: UpdateConfigPayload) => Promise<void>;
   setGenerationModel: (payload: SetGenerationModelPayload) => Promise<void>;
   getMessage: (payload: GetMessagePayload) => Promise<ICometMessage>;
@@ -254,4 +254,23 @@ export interface ICometMessage {
   meta: Record<string, any> | null;
   conversationId: string;
   createdAt: string;
-}
+}
+
+export type TCometPromptResponse = {
+  text: string;
+  referencePaths?: string[];
+  referencesWithSources?: {
+    path: string;
+    source_id: string;
+  }[];
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+  };
+  conversationId?: string;
+  sessionId?: string;
+};
+
+export type TCometPromptStreamResponseError = {
+  error: string;
+};