Merge pull request #24 from AI21Labs/library-upload

feat: RAG functionality

Author: Josephasafg

.github/workflows/integration-tests.yml

@@ -22,7 +22,7 @@ jobs:
     - name: Install dependencies
       run: |
         npm install
-        npm install ai21
+        npm run build
       
     - name: Run Integration Tests
       env:

.prettierignore

@@ -0,0 +1 @@
+*.txt

examples/studio/conversational-rag/files/meerkat.txt

@@ -0,0 +1 @@
+The meerkat (Suricata suricatta) or suricate is a small mongoose found in southern Africa. It is characterised by a broad head, large eyes, a pointed snout, long legs, a thin tapering tail, and a brindled coat pattern. The head-and-body length is around 24–35 cm (9.4–13.8 in), and the weight is typically between 0.62 and 0.97 kg (1.4 and 2.1 lb). The coat is light grey to yellowish-brown with alternate, poorly-defined light and dark bands on the back. Meerkats have foreclaws adapted for digging and have the ability to thermoregulate to survive in their harsh, dry habitat. Three subspecies are recognised.

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

@@ -0,0 +1,139 @@
+import { AI21, FileResponse, UploadFileResponse } from 'ai21';
+import path from 'path';
+import fs from 'fs';
+
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function waitForFileProcessing(
+  client: AI21,
+  fileId: string,
+  timeout: number = 30000,
+  interval: number = 1000,
+) {
+  const startTime = Date.now();
+
+  while (Date.now() - startTime < timeout) {
+    const file: FileResponse = await client.files.get(fileId);
+    if (file.status !== 'PROCESSING') {
+      return file;
+    }
+    await sleep(interval);
+  }
+
+  throw new Error(`File processing timed out after ${timeout}ms`);
+}
+
+async function uploadGetUpdateDelete(fileInput, path) {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+  try {
+    console.log(`Starting upload for file:`, typeof fileInput);
+    const uploadFileResponse: UploadFileResponse = await client.files.create({
+      file: fileInput,
+      path: path,
+    });
+    console.log(`✓ Upload completed. File ID: ${uploadFileResponse.fileId}`);
+
+    console.log('Waiting for file processing...');
+    let file: FileResponse = await waitForFileProcessing(client, uploadFileResponse.fileId);
+    console.log(`✓ File processing completed with status: ${file.status}`);
+
+    if (file.status === 'PROCESSED') {
+      console.log('Starting file update...');
+      await client.files.update({
+        fileId: uploadFileResponse.fileId,
+        labels: ['test99'],
+        publicUrl: 'https://www.miri.com',
+      });
+      file = await client.files.get(uploadFileResponse.fileId);
+      console.log('✓ File update completed');
+    } else {
+      console.log(`⚠ File processing failed with status ${file.status}`);
+      return; // Exit early if processing failed
+    }
+
+    console.log('Starting file deletion...');
+    await client.files.delete(uploadFileResponse.fileId);
+    console.log('✓ File deletion completed');
+
+    // Add buffer time between operations
+    await sleep(2000);
+  } catch (error) {
+    console.error('❌ Error in uploadGetUpdateDelete:', error);
+    throw error;
+  }
+}
+
+async function listFiles() {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+  const files = await client.files.list({ limit: 4 });
+  console.log(`Listed files: ${files}`);
+}
+
+const isBrowser = typeof window !== 'undefined' && typeof window.document !== 'undefined';
+
+const createNodeFile = (content: Buffer, filename: string, type: string) => {
+  if (process.platform === 'linux') {
+    console.log('Running on Linux (GitHub Actions)');
+    // Special handling for Linux (GitHub Actions)
+    return {
+      name: filename,
+      type: type,
+      buffer: content,
+      [Symbol.toStringTag]: 'File',
+    };
+  } else {
+    console.log('Running on other platforms');
+    // Regular handling for other platforms
+    return new File([content], filename, { type });
+  }
+};
+
+if (isBrowser) {
+  console.log('Cannot run upload examples in Browser environment');
+} else {
+  /* Log environment details */
+  console.log('=== Environment Information ===');
+  console.log(`Node.js Version: ${process.version}`);
+  console.log(`Platform: ${process.platform}`);
+  console.log(`Architecture: ${process.arch}`);
+  console.log(`Process ID: ${process.pid}`);
+  console.log(`Current Working Directory: ${process.cwd()}`);
+  console.log('===========================\n');
+
+  /* Run all operations sequentially */
+  (async () => {
+    try {
+      console.log('=== Starting first operation ===');
+      // First operation - upload file from path
+      const filePath = path.resolve(process.cwd(), 'examples/studio/conversational-rag/files', 'meerkat.txt');
+      if (!fs.existsSync(filePath)) {
+        throw new Error(`File not found: ${filePath}`);
+      } else {
+        console.log(`File found: ${filePath}`);
+      }
+
+      await uploadGetUpdateDelete(filePath, Date.now().toString());
+      console.log('=== First operation completed ===\n');
+      await sleep(2000);
+
+      console.log('=== Starting second operation ===');
+      // Second operation - upload file from File instance
+      const fileContent = Buffer.from(
+        'Opossums are members of the marsupial order Didelphimorphia endemic to the Americas.',
+      );
+      const dummyFile = createNodeFile(fileContent, 'example.txt', 'text/plain');
+      await uploadGetUpdateDelete(dummyFile, Date.now().toString());
+      console.log('=== Second operation completed ===\n');
+      await sleep(2000);
+
+      console.log('=== Starting file listing ===');
+      await listFiles();
+      console.log('=== File listing completed ===');
+    } catch (error) {
+      console.error('❌ Main execution error:', error);
+      process.exit(1); // Exit with error code if something fails
+    }
+  })();
+}

src/AI21.ts

@@ -5,7 +5,8 @@ import { Chat } from './resources/chat';
 import { APIClient } from './APIClient';
 import { Headers } from './types';
 import * as Runtime from './runtime';
-import { ConversationalRag } from './resources/rag/conversationalRag';
+import { ConversationalRag } from './resources/rag/conversational-rag';
+import { Files } from './resources';
 
 export interface ClientOptions {
   baseURL?: string | undefined;
@@ -67,6 +68,7 @@ export class AI21 extends APIClient {
   // Resources
   chat: Chat = new Chat(this);
   conversationalRag: ConversationalRag = new ConversationalRag(this);
+  files: Files = new Files(this);
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   protected override authHeaders(_: Types.FinalRequestOptions): Types.Headers {

src/APIClient.ts

@@ -8,10 +8,14 @@ import {
   HTTPMethod,
   Headers,
   CrossPlatformResponse,
+  UnifiedFormData,
+  FilePathOrFileObject,
 } from './types';
 import { AI21EnvConfig } from './EnvConfig';
-import { createFetchInstance } from './runtime';
+import { createFetchInstance, createFilesHandlerInstance } from './factory';
 import { Fetch } from 'fetch';
+import { BaseFilesHandler } from 'files/BaseFilesHandler';
+import { FormDataRequest } from 'types/API';
 
 const validatePositiveInteger = (name: string, n: unknown): number => {
   if (typeof n !== 'number' || !Number.isInteger(n)) {
@@ -23,42 +27,80 @@ const validatePositiveInteger = (name: string, n: unknown): number => {
   return n;
 };
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const 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);
+    }
+  }
+};
+
 export abstract class APIClient {
   protected baseURL: string;
   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);
+    return this.prepareAndExecuteRequest('get', path, opts);
   }
 
   post<Req, Rsp>(path: string, opts?: RequestOptions<Req>): Promise<Rsp> {
-    return this.makeRequest('post', path, opts);
+    return this.prepareAndExecuteRequest('post', path, opts);
   }
 
   put<Req, Rsp>(path: string, opts?: RequestOptions<Req>): Promise<Rsp> {
-    return this.makeRequest('put', path, opts);
+    return this.prepareAndExecuteRequest('put', path, opts);
   }
 
   delete<Req, Rsp>(path: string, opts?: RequestOptions<Req>): Promise<Rsp> {
-    return this.makeRequest('delete', path, opts);
+    return this.prepareAndExecuteRequest('delete', path, opts);
+  }
+
+  upload<Req, Rsp>(path: string, file: FilePathOrFileObject, opts?: RequestOptions<Req>): Promise<Rsp> {
+    return this.filesHandler.prepareFormDataRequest(file).then((formDataRequest: FormDataRequest) => {
+      if (opts?.body) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        appendBodyToFormData(formDataRequest.formData, opts.body as Record<string, any>);
+      }
+
+      const headers = {
+        ...opts?.headers,
+        ...formDataRequest.headers,
+      };
+
+      const options: FinalRequestOptions = {
+        method: 'post',
+        path: path,
+        body: formDataRequest.formData,
+        headers,
+      };
+
+      return this.performRequest(options).then((response) => this.fetch.handleResponse<Rsp>(response) as Rsp);
+    });
   }
 
   protected getUserAgent(): string {
@@ -70,38 +112,56 @@ 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),
     };
+
+    return { ...defaultHeaders, ...opts.headers };
   }
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   protected authHeaders(opts: FinalRequestOptions): Headers {
     return {};
   }
 
-  private makeRequest<Req, Rsp>(method: HTTPMethod, path: string, opts?: RequestOptions<Req>): Promise<Rsp> {
+  private buildFullUrl(path: string, query?: Record<string, unknown>): string {
+    let url = `${this.baseURL}${path}`;
+    if (query) {
+      const queryString = new URLSearchParams(query as Record<string, string>).toString();
+      url += `?${queryString}`;
+    }
+    return url;
+  }
+
+  private prepareAndExecuteRequest<Req, Rsp>(
+    method: HTTPMethod,
+    path: string,
+    opts?: RequestOptions<Req>,
+  ): Promise<Rsp> {
     const options = {
       method,
       path,
       ...opts,
-    };
+    } as FinalRequestOptions;
 
-    return this.performRequest(options as FinalRequestOptions).then(
-      (response) => this.fetch.handleResponse<Rsp>(response) as Rsp,
-    );
+    if (options?.body) {
+      options.body = JSON.stringify(options.body);
+      options.headers = { ...options.headers, 'Content-Type': 'application/json' };
+    }
+
+    return this.performRequest(options).then((response) => this.fetch.handleResponse<Rsp>(response) as Rsp);
   }
 
   private async performRequest(options: FinalRequestOptions): Promise<APIResponseProps> {
-    const url = `${this.baseURL}${options.path}`;
+    const url = this.buildFullUrl(options.path, options.query as Record<string, unknown>);
 
     const headers = {
       ...this.defaultHeaders(options),
       ...options.headers,
     };
+
     const response = await this.fetch.call(url, { ...options, headers });
 
     if (!response.ok) {

src/factory.ts

@@ -0,0 +1,21 @@
+import { BrowserFilesHandler } from './files/BrowserFilesHandler';
+import { BrowserFetch, Fetch, NodeFetch } from './fetch';
+import { NodeFilesHandler } from './files/NodeFilesHandler';
+import { BaseFilesHandler } from './files/BaseFilesHandler';
+import { isBrowser, isWebWorker } from './runtime';
+
+export function createFetchInstance(): Fetch {
+  if (isBrowser || isWebWorker) {
+    return new BrowserFetch();
+  }
+
+  return new NodeFetch();
+}
+
+export function createFilesHandlerInstance(): BaseFilesHandler {
+  if (isBrowser || isWebWorker) {
+    return new BrowserFilesHandler();
+  }
+
+  return new NodeFilesHandler();
+}

src/fetch/BaseFetch.ts

@@ -7,6 +7,7 @@ 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/fetch/BrowserFetch.ts

@@ -9,7 +9,7 @@ export class BrowserFetch extends BaseFetch {
     return fetch(url, {
       method: options.method,
       headers: options?.headers ? (options.headers as HeadersInit) : undefined,
-      body: options?.body ? JSON.stringify(options.body) : undefined,
+      body: options?.body ? (options.body as BodyInit) : undefined,
       signal: controller.signal,
     });
   }

src/fetch/NodeFetch.ts

@@ -1,6 +1,7 @@
 import { FinalRequestOptions, CrossPlatformResponse } from 'types';
 import { BaseFetch } from './BaseFetch';
 import { Stream, NodeSSEDecoder } from '../streaming';
+import { NodeHTTPBody } from 'types/API';
 
 export class NodeFetch extends BaseFetch {
   async call(url: string, options: FinalRequestOptions): Promise<CrossPlatformResponse> {
@@ -10,7 +11,7 @@ export class NodeFetch extends BaseFetch {
     return nodeFetch(url, {
       method: options.method,
       headers: options?.headers ? (options.headers as Record<string, string>) : undefined,
-      body: options?.body ? JSON.stringify(options.body) : undefined,
+      body: options?.body ? (options.body as NodeHTTPBody) : undefined,
     });
   }