feat: reorganized

Author: amirai21

examples/studio/conversational-rag/rag-engine.ts

@@ -1,40 +1,26 @@
 import { AI21 } from 'ai21';
 import { FileResponse, UploadFileResponse } from '../../../src/types/rag';
 
-async function waitForFileProcessing(
-  client: AI21,
-  fileId: string,
-  interval: number = 3000,
-): Promise<FileResponse> {
-  while (true) {
-    const file: FileResponse = await client.ragEngine.get(fileId);
-
-    if (file.status === 'PROCESSED') {
-      return file;
-    }
-
-    console.log(`File status is '${file.status}'. Waiting for it to be 'PROCESSED'...`);
-    await new Promise((resolve) => setTimeout(resolve, interval));
-  }
+function sleep(ms) {
+  return new Promise(resolve => setTimeout(resolve, ms));
 }
 
-async function uploadQueryUpdateDelete(fileInput, label) {
+async function uploadGetUpdateDelete(fileInput, label) {
   const client = new AI21({ apiKey: process.env.AI21_API_KEY });
   try {
     const uploadFileResponse: UploadFileResponse = await client.ragEngine.create(fileInput, {
       path: label,
     });
 
-    const fileId = uploadFileResponse.fileId;
-    let file: FileResponse = await waitForFileProcessing(client, fileId);
+    let file: FileResponse = await client.ragEngine.get(uploadFileResponse.fileId);
     console.log(file);
-
-    console.log('Now updating the file labels');
+    await sleep(1000); // Give it a sec to start process before updating
+    console.log('Now updating the file labels and publicUrl...');
     await client.ragEngine.update(uploadFileResponse.fileId, {
       labels: ['test99'],
       publicUrl: 'https://www.miri.com',
     });
-    file = await client.ragEngine.get(fileId);
+    file = await client.ragEngine.get(uploadFileResponse.fileId);
     console.log(file);
 
     console.log('Now deleting the file');
@@ -50,11 +36,13 @@ async function listFiles() {
   console.log(files);
 }
 
+/* Simulate a file upload passing file path */
 const filePath = '/Users/amirkoblyansky/Documents/ukraine.txt';
-const fileContent = Buffer.from('This is the content of the file.');
-const dummyFile = new File([fileContent], 'example.txt', { type: 'text/plain' });
+uploadGetUpdateDelete(filePath, Date.now().toString()).catch(console.error);
 
-uploadQueryUpdateDelete(filePath, "abc123").catch(console.error);
-uploadQueryUpdateDelete(dummyFile, "test2").catch(console.error);
+/* Simulate a file upload passing File instance */
+const fileContent = Buffer.from('Opossums are members of the marsupial order Didelphimorphia endemic to the Americas.');
+const dummyFile = new File([fileContent], 'example.txt', { type: 'text/plain' });
+uploadGetUpdateDelete(dummyFile, Date.now().toString()).catch(console.error);
 
 listFiles().catch(console.error);

src/APIClient.ts

@@ -8,12 +8,12 @@ import {
   HTTPMethod,
   Headers,
   CrossPlatformResponse,
-  UnifiedFormData,
 } from './types';
 import { AI21EnvConfig } from './EnvConfig';
-import { createFetchInstance } from './runtime';
+import { createFetchInstance, createFilesHandlerInstance } from './runtime';
 import { Fetch } from 'fetch';
-import { getBoundary, appendBodyToFormData } from './files/form-utils';
+import { FilePathOrFileObject } from 'types/rag';
+import { BaseFilesHandler } from 'files/BaseFilesHandler';
 
 const validatePositiveInteger = (name: string, n: unknown): number => {
   if (typeof n !== 'number' || !Number.isInteger(n)) {
@@ -30,22 +30,26 @@ export abstract class APIClient {
   protected maxRetries: number;
   protected timeout: number;
   protected fetch: Fetch;
+  protected filesHandler: BaseFilesHandler;
 
   constructor({
     baseURL,
     maxRetries = AI21EnvConfig.MAX_RETRIES,
     timeout = AI21EnvConfig.TIMEOUT_SECONDS,
     fetch = createFetchInstance(),
+    filesHandler = createFilesHandlerInstance(),
   }: {
     baseURL: string;
     maxRetries?: number | undefined;
     timeout: number | undefined;
     fetch?: Fetch;
+    filesHandler?: BaseFilesHandler;
   }) {
     this.baseURL = baseURL;
     this.maxRetries = validatePositiveInteger('maxRetries', maxRetries);
     this.timeout = validatePositiveInteger('timeout', timeout);
     this.fetch = fetch;
+    this.filesHandler = filesHandler;
   }
   get<Req, Rsp>(path: string, opts?: RequestOptions<Req>): Promise<Rsp> {
     return this.makeRequest('get', path, opts);
@@ -63,16 +67,17 @@ export abstract class APIClient {
     return this.makeRequest('delete', path, opts);
   }
 
-  async upload<Req, Rsp>(path: string, formData: UnifiedFormData, opts?: RequestOptions<Req>): Promise<Rsp> {
+  async upload<Req, Rsp>(path: string, file: FilePathOrFileObject, opts?: RequestOptions<Req>): Promise<Rsp> {
+    const formData = await this.filesHandler.createFormData(file);
+
     if (opts?.body) {
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      appendBodyToFormData(formData, opts.body as Record<string, any>);
+      this.filesHandler.appendBodyToFormData(formData, opts.body as Record<string, any>);
     }
 
-    const boundary = await getBoundary(formData);
     const headers = {
       ...opts?.headers,
-      'Content-Type': `multipart/form-data; boundary=${boundary}`,
+      ...this.filesHandler.getMultipartFormDataHeaders(formData),
     };
 
     const options: FinalRequestOptions = {
@@ -94,26 +99,16 @@ export abstract class APIClient {
   }
 
   protected defaultHeaders(opts: FinalRequestOptions): Headers {
-    return {
+    const defaultHeaders = {
       Accept: 'application/json',
-      'Content-Type': 'application/json',
       'User-Agent': this.getUserAgent(),
       ...this.authHeaders(opts),
     };
-    // if (opts?.body instanceof FormData) {
-    // return {
-    //   Accept: 'application/json',
-    //   'User-Agent': this.getUserAgent(),
-    //   ...this.authHeaders(opts),
-    //   };
-    // } else {
-    //   return {
-    //     Accept: 'application/json',
-    //     'Content-Type': 'application/json',
-    //     'User-Agent': this.getUserAgent(),
-    //     ...this.authHeaders(opts),
-    //   };
-    // }
+
+    if (opts?.body instanceof FormData) {
+      return defaultHeaders;
+    }
+    return { ...defaultHeaders, 'Content-Type': 'application/json' };
   }
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars

src/fetch/BaseFetch.ts

@@ -3,10 +3,7 @@ import { AI21Error } from '../errors';
 import { FinalRequestOptions, CrossPlatformResponse } from '../types';
 import { APIResponseProps } from '../types/API';
 
-export type APIResponse<T> = {
-  data?: T;
-  response: CrossPlatformResponse;
-};
+
 export abstract class BaseFetch {
   abstract call(url: string, options: FinalRequestOptions): Promise<CrossPlatformResponse>;
   async handleResponse<T>({ response, options }: APIResponseProps) {

src/files/BaseFilesHandler.ts

@@ -0,0 +1,21 @@
+import { UnifiedFormData } from "types";
+import { FilePathOrFileObject } from "types/rag";
+
+  export abstract class BaseFilesHandler {
+    abstract createFormData(file: FilePathOrFileObject): Promise<UnifiedFormData>;
+
+    abstract getMultipartFormDataHeaders(formData: UnifiedFormData): Record<string, string> | null;
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      appendBodyToFormData = (formData: UnifiedFormData, body: Record<string, any>): void => {
+        for (const [key, value] of Object.entries(body)) {
+          if (Array.isArray(value)) {
+            value.forEach((item) => formData.append(key, item));
+          } else {
+            formData.append(key, value);
+          }
+        }
+      };
+
+
+  }

src/files/BrowserFilesHandler.ts

@@ -0,0 +1,24 @@
+import { UnifiedFormData } from "types";
+import { FilePathOrFileObject } from "types/rag";
+import { BaseFilesHandler } from "./BaseFilesHandler";
+
+
+  export class BrowserFilesHandler extends BaseFilesHandler {
+
+    async createFormData(file: FilePathOrFileObject): Promise<UnifiedFormData> {
+      const formData = new FormData();
+  
+        if (file instanceof window.File) {
+          formData.append('file', file);
+        } else {
+          throw new Error('Unsupported file type in browser');
+        }
+
+        return formData;
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    getMultipartFormDataHeaders(formData: UnifiedFormData): Record<string, string> | null {
+      return {};
+    }
+  }

src/files/NodeFilesHandler.ts

@@ -0,0 +1,65 @@
+import { UnifiedFormData } from "types";
+import { FilePathOrFileObject } from "types/rag";
+import { BaseFilesHandler } from "./BaseFilesHandler";
+
+  export class NodeFilesHandler extends BaseFilesHandler  {
+
+    async convertReadableStream(whatwgStream: ReadableStream): Promise<NodeJS.ReadableStream> {
+      const { Readable } = await import('stream');
+        const reader = whatwgStream.getReader();
+  
+        return new Readable({
+          async read() {
+            const { done, value } = await reader.read();
+            if (done) {
+              this.push(null);
+            } else {
+              this.push(value);
+            }
+          },
+        });
+    }
+
+    async getBoundary(formData: UnifiedFormData): Promise<string | undefined> {
+      const { default: FormDataNode } = await import('form-data');
+        if (formData instanceof FormDataNode) {
+          return formData.getBoundary();
+        }
+        else {
+          throw new Error('getBoundary invoked with native browser FormData instance instead of NodeJS form-data');
+        }
+    }
+    
+
+    async createFormData(file: FilePathOrFileObject): Promise<UnifiedFormData> {
+      const { default: FormDataNode } = await import('form-data');
+        const formData = new FormDataNode();
+  
+        if (typeof file === 'string') {
+          const fs = (await import('fs')).default;
+          formData.append('file', fs.createReadStream(file), { filename: file.split('/').pop() });
+        } else if (Buffer.isBuffer(file)) {
+          formData.append('file', file, { filename: 'TODO - add filename to buffer flow' });
+        } else if (file instanceof File) {
+          const nodeStream = await this.convertReadableStream(file.stream());
+          formData.append('file', nodeStream, file.name);
+        } else {
+          throw new Error(`Unsupported file type for Node.js file upload flow: ${file}`);
+        }
+  
+        return formData;
+    }
+
+    getMultipartFormDataHeaders(formData: UnifiedFormData): Record<string, string> | null {
+      if (formData instanceof FormData) {
+        return null;
+      }
+
+      const boundary = formData.getBoundary();
+        return {
+          'Content-Type': `multipart/form-data; boundary=${boundary}`,
+        };
+      }
+  
+    }
+  

src/files/form-utils.ts

@@ -1 +1,2 @@
-export * from './form-utils';
+export { BrowserFilesHandler } from './BrowserFilesHandler';
+export { NodeFilesHandler } from './NodeFilesHandler';

src/files/index.ts

@@ -8,7 +8,6 @@ import {
   FilePathOrFileObject,
 } from '../../types/rag';
 import { FileResponse } from 'types/rag/FileResponse';
-import { CreateFormData } from '../../files';
 
 const RAG_ENGINE_PATH = '/library/files';
 
@@ -30,8 +29,7 @@ export class RAGEngine extends APIResource {
     body: UploadFileRequest,
     options?: Models.RequestOptions,
   ): Promise<UploadFileResponse> {
-    const formData = await new CreateFormData().createFormData(file);
-    return this.client.upload<Models.UnifiedFormData, UploadFileResponse>(RAG_ENGINE_PATH, formData, {
+    return this.client.upload<Models.UnifiedFormData, UploadFileResponse>(RAG_ENGINE_PATH, file, {
       body: body,
       ...options,
     } as Models.RequestOptions<Models.UnifiedFormData>) as Promise<UploadFileResponse>;