feat: rag-engine impl. - with examples, no tests yet

amirai21 · asafgardin · commit 4713d0229528 · 2024-12-04T14:56:22.000+02:00
diff --git a/examples/studio/conversational-rag/rag-engine.ts b/examples/studio/conversational-rag/rag-engine.ts
@@ -0,0 +1,48 @@
+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));
+  }
+}
+
+async function uploadQueryUpdateDelete() {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+  try {
+    const uploadFileResponse: UploadFileResponse = await client.ragEngine.create(
+      '/Users/amirkoblyansky/Documents/ukraine.txt', {path: "test10"});
+    
+    const fileId = uploadFileResponse.fileId
+    let file: FileResponse = await waitForFileProcessing(client, fileId);
+    console.log(file);
+
+    console.log("Now updating the file labels");
+    await client.ragEngine.update(uploadFileResponse.fileId, {labels: ["test99"], publicUrl: "https://www.miri.com"});
+    file = await client.ragEngine.get(fileId);
+    console.log(file);
+
+    console.log("Now deleting the file");
+    await client.ragEngine.delete(uploadFileResponse.fileId);
+  } catch (error) {
+    console.error('Error:', error);
+  }
+}
+
+async function listFiles() {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+  const files = await client.ragEngine.list({limit: 10});
+  console.log(files);
+}
+
+uploadQueryUpdateDelete().catch(console.error);
+
+listFiles().catch(console.error);
+
diff --git a/src/AI21.ts b/src/AI21.ts
@@ -6,6 +6,7 @@ import { APIClient } from './APIClient';
 import { Headers } from './types';
 import * as Runtime from './runtime';
 import { ConversationalRag } from './resources/rag/conversationalRag';
+import { RAGEngine } 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);
+  ragEngine: RAGEngine = new RAGEngine(this);
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   protected override authHeaders(_: Types.FinalRequestOptions): Types.Headers {
diff --git a/src/APIClient.ts b/src/APIClient.ts
@@ -12,6 +12,9 @@ import {
 import { AI21EnvConfig } from './EnvConfig';
 import { createFetchInstance } from './runtime';
 import { Fetch } from 'fetch';
+import { createReadStream } from 'fs';
+import { basename as getBasename } from 'path';
+import FormData from 'form-data';
 
 const validatePositiveInteger = (name: string, n: unknown): number => {
   if (typeof n !== 'number' || !Number.isInteger(n)) {
@@ -61,6 +64,49 @@ export abstract class APIClient {
     return this.makeRequest('delete', path, opts);
   }
 
+  upload<Req, Rsp>(path: string, filePath: string, opts?: RequestOptions<Req>): Promise<Rsp> {
+    const formDataRequest = this.makeFormDataRequest(path, filePath, opts);
+    return this.performRequest(formDataRequest).then(
+      (response) => this.fetch.handleResponse<Rsp>(response) as Rsp,
+    );
+  }
+
+  protected makeFormDataRequest<Req>(path: string, filePath: string, opts?: RequestOptions<Req>): FinalRequestOptions {
+    const formData = new FormData();
+    const fileStream = createReadStream(filePath);
+    const fileName = getBasename(filePath);
+
+    formData.append('file', fileStream, fileName);
+
+    if (opts?.body) {
+      const body = opts.body as Record<string, string>;
+      for (const [key, value] of Object.entries(body)) {
+        if (Array.isArray(value)) {
+          value.forEach(item => formData.append(key, item));
+        } else {
+          formData.append(key, value);
+        }
+      }
+    }
+
+    const headers = {
+      ...opts?.headers,
+      'Content-Type': `multipart/form-data; boundary=${formData.getBoundary()}`
+    };
+    console.log(headers);
+    console.log("-------------------------");
+    console.log(formData.getHeaders());
+    console.log("-------------------------");
+
+    const options: FinalRequestOptions = {
+      method: 'post',
+      path: path,
+      body: formData,
+      headers,
+    };
+    return options;
+  }
+
   protected getUserAgent(): string {
     const platform =
       this.isRunningInBrowser() ?
@@ -96,12 +142,18 @@ export abstract class APIClient {
   }
 
   private async performRequest(options: FinalRequestOptions): Promise<APIResponseProps> {
-    const url = `${this.baseURL}${options.path}`;
+    let url = `${this.baseURL}${options.path}`;
+
+    if (options.query) {
+      const queryString = new URLSearchParams(options.query as Record<string, string>).toString();
+      url += `?${queryString}`;
+    }
 
     const headers = {
       ...this.defaultHeaders(options),
       ...options.headers,
     };
+
     const response = await this.fetch.call(url, { ...options, headers });
 
     if (!response.ok) {
diff --git a/src/fetch/NodeFetch.ts b/src/fetch/NodeFetch.ts
@@ -1,16 +1,19 @@
 import { FinalRequestOptions, CrossPlatformResponse } from 'types';
 import { BaseFetch } from './BaseFetch';
 import { Stream, NodeSSEDecoder } from '../streaming';
+import FormData from 'form-data';
 
 export class NodeFetch extends BaseFetch {
   async call(url: string, options: FinalRequestOptions): Promise<CrossPlatformResponse> {
     const nodeFetchModule = await import('node-fetch');
     const nodeFetch = nodeFetchModule.default;
 
+    const body = options.body instanceof FormData ? options.body : JSON.stringify(options.body);
+
     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,
     });
   }
 
diff --git a/src/resources/index.ts b/src/resources/index.ts
@@ -1,2 +1,3 @@
 export { Chat, Completions } from './chat';
 export { ConversationalRag } from './rag';
+export { RAGEngine } from './rag';
diff --git a/src/resources/rag/index.ts b/src/resources/rag/index.ts
@@ -1 +1,2 @@
 export { ConversationalRag } from './conversationalRag';
+export { RAGEngine } from './ragEngine';
diff --git a/src/resources/rag/ragEngine.ts b/src/resources/rag/ragEngine.ts
@@ -0,0 +1,37 @@
+import * as Models from '../../types';
+import { APIResource } from '../../APIResource';
+import { UploadFileResponse, UploadFileRequest, ListFilesFilters, UpdateFileRequest } from '../../types/rag';
+import { FileResponse } from 'types/rag/FileResponse';
+
+
+const RAG_ENGINE_PATH = '/library/files';
+
+
+export class RAGEngine extends APIResource {
+  create(filePath: string, body: UploadFileRequest, options?: Models.RequestOptions) {
+    return this.client.upload<UploadFileRequest, UploadFileResponse>(RAG_ENGINE_PATH, filePath, {
+      body: body,
+      ...options,
+    } as Models.RequestOptions<UploadFileRequest>) as Promise<UploadFileResponse>;
+  }
+
+  get(fileId: string, options?: Models.RequestOptions) {
+    return this.client.get<string, FileResponse>(
+      `${RAG_ENGINE_PATH}/${fileId}`, options as Models.RequestOptions<string>) as Promise<FileResponse>;
+  }
+
+  delete(fileId: string, options?: Models.RequestOptions) {
+    return this.client.delete<string, null>(
+      `${RAG_ENGINE_PATH}/${fileId}`, options as Models.RequestOptions<string>) as Promise<null>;
+  }
+
+  list(body: ListFilesFilters | null, options?: Models.RequestOptions) {
+    return this.client.get<ListFilesFilters | null, FileResponse[]>(
+      RAG_ENGINE_PATH, {query: body, ...options} as Models.RequestOptions<ListFilesFilters | null>) as Promise<FileResponse[]>;
+  }
+
+  update(fileId: string, body: UpdateFileRequest, options?: Models.RequestOptions) {
+    return this.client.put<UpdateFileRequest, null>(
+      `${RAG_ENGINE_PATH}/${fileId}`, {body, ...options} as Models.RequestOptions<UpdateFileRequest>) as Promise<null>;
+  }
+}
diff --git a/src/types/API.ts b/src/types/API.ts
@@ -10,7 +10,7 @@ export type RequestOptions<Req = unknown | Record<string, unknown> | ArrayBuffer
   method?: HTTPMethod;
   path?: string;
   query?: Req | undefined;
-  body?: Req | null | undefined;
+  body?: Req | FormData | null | undefined;
   headers?: Headers | undefined;
 
   maxRetries?: number;
diff --git a/src/types/rag/FileResponse.ts b/src/types/rag/FileResponse.ts
@@ -0,0 +1,15 @@
+export interface FileResponse {
+  fileId: string;
+  name: string;
+  fileType: string;
+  sizeBytes: number;
+  createdBy: string;
+  creationDate: Date;
+  lastUpdated: Date;
+  status: string;
+  path?: string | null;
+  labels?: string[] | null;
+  publicUrl?: string | null;
+  errorCode?: number | null;
+  errorMessage?: string | null;
+}
diff --git a/src/types/rag/ListFilesFilters.ts b/src/types/rag/ListFilesFilters.ts
@@ -0,0 +1,4 @@
+export interface ListFilesFilters {
+  offset?: number | null;
+  limit?: number | null;
+}
diff --git a/src/types/rag/UploadFileRequest.ts b/src/types/rag/UploadFileRequest.ts
@@ -0,0 +1,8 @@
+export interface UpdateFileRequest {
+    labels?: string[] | null;
+    publicUrl?: string| null;
+  }
+
+  export interface UploadFileRequest extends UpdateFileRequest {
+    path?: string | null;
+  }
diff --git a/src/types/rag/UploadFileResponse.ts b/src/types/rag/UploadFileResponse.ts
@@ -0,0 +1,4 @@
+export interface UploadFileResponse {
+  id: string;
+  fileId: string;
+}
diff --git a/src/types/rag/index.ts b/src/types/rag/index.ts
@@ -5,3 +5,13 @@ export { type ConversationalRagSource } from './ConversationalRagSource';
 export { type ConversationalRagResponse } from './ConversationalRagResponse';
 
 export { type RetrievalStrategy } from './RetrievalStrategy';
+
+export { type UploadFileRequest } from './UploadFileRequest';
+
+export { type FileResponse } from './FileResponse';
+
+export { type UploadFileResponse } from './UploadFileResponse';
+
+export { type ListFilesFilters } from './ListFilesFilters';
+
+export { type UpdateFileRequest } from './UploadFileRequest';

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`export { Chat, Completions } from './chat';`
`2`	`2`	`export { ConversationalRag } from './rag';`
	`3`	`+export { RAGEngine } from './rag';`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`	`1`	`export { ConversationalRag } from './conversationalRag';`
	`2`	`+export { RAGEngine } from './ragEngine';`