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

Author: amirai21

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

@@ -1,7 +1,11 @@
 import { AI21 } from 'ai21';
 import { FileResponse, UploadFileResponse } from '../../../src/types/rag';
 
-async function waitForFileProcessing(client: AI21, fileId: string, interval: number = 3000): Promise<FileResponse> {
+async function waitForFileProcessing(
+  client: AI21,
+  fileId: string,
+  interval: number = 3000,
+): Promise<FileResponse> {
   while (true) {
     const file: FileResponse = await client.ragEngine.get(fileId);
 
@@ -10,26 +14,31 @@ async function waitForFileProcessing(client: AI21, fileId: string, interval: num
     }
 
     console.log(`File status is '${file.status}'. Waiting for it to be 'PROCESSED'...`);
-    await new Promise(resolve => setTimeout(resolve, interval));
+    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
+      '/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"});
+    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");
+    console.log('Now deleting the file');
     await client.ragEngine.delete(uploadFileResponse.fileId);
   } catch (error) {
     console.error('Error:', error);
@@ -38,11 +47,10 @@ async function uploadQueryUpdateDelete() {
 
 async function listFiles() {
   const client = new AI21({ apiKey: process.env.AI21_API_KEY });
-  const files = await client.ragEngine.list({limit: 10});
+  const files = await client.ragEngine.list({ limit: 10 });
   console.log(files);
 }
 
 uploadQueryUpdateDelete().catch(console.error);
 
 listFiles().catch(console.error);
-

src/APIClient.ts

@@ -71,7 +71,11 @@ export abstract class APIClient {
     );
   }
 
-  protected makeFormDataRequest<Req>(path: string, filePath: string, opts?: RequestOptions<Req>): FinalRequestOptions {
+  protected makeFormDataRequest<Req>(
+    path: string,
+    filePath: string,
+    opts?: RequestOptions<Req>,
+  ): FinalRequestOptions {
     const formData = new FormData();
     const fileStream = createReadStream(filePath);
     const fileName = getBasename(filePath);
@@ -82,7 +86,7 @@ export abstract class APIClient {
       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));
+          value.forEach((item) => formData.append(key, item));
         } else {
           formData.append(key, value);
         }
@@ -91,12 +95,12 @@ export abstract class APIClient {
 
     const headers = {
       ...opts?.headers,
-      'Content-Type': `multipart/form-data; boundary=${formData.getBoundary()}`
+      'Content-Type': `multipart/form-data; boundary=${formData.getBoundary()}`,
     };
     console.log(headers);
-    console.log("-------------------------");
+    console.log('-------------------------');
     console.log(formData.getHeaders());
-    console.log("-------------------------");
+    console.log('-------------------------');
 
     const options: FinalRequestOptions = {
       method: 'post',

src/resources/rag/ragEngine.ts

@@ -3,10 +3,8 @@ 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, {
@@ -17,21 +15,29 @@ export class RAGEngine extends APIResource {
 
   get(fileId: string, options?: Models.RequestOptions) {
     return this.client.get<string, FileResponse>(
-      `${RAG_ENGINE_PATH}/${fileId}`, options as Models.RequestOptions<string>) as Promise<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>;
+      `${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[]>;
+    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>;
+    return this.client.put<UpdateFileRequest, null>(`${RAG_ENGINE_PATH}/${fileId}`, {
+      body,
+      ...options,
+    } as Models.RequestOptions<UpdateFileRequest>) as Promise<null>;
   }
 }

src/types/rag/UploadFileRequest.ts

@@ -1,8 +1,8 @@
 export interface UpdateFileRequest {
-    labels?: string[] | null;
-    publicUrl?: string| null;
-  }
+  labels?: string[] | null;
+  publicUrl?: string | null;
+}
 
-  export interface UploadFileRequest extends UpdateFileRequest {
-    path?: string | null;
-  }
+export interface UploadFileRequest extends UpdateFileRequest {
+  path?: string | null;
+}