feat: fix tests

AI21Labs · Dec 4, 2024 · 50a98fa · 50a98fa
1 parent 46ba724
commit 50a98fa
Show file tree

Hide file tree

Showing 17 changed files with 40 additions and 27 deletions.
diff --git a/examples/studio/conversational-rag/rag-engine.ts b/examples/studio/conversational-rag/rag-engine.ts
@@ -1,5 +1,5 @@
 import { AI21 } from 'ai21';
-import { FileResponse, UploadFileResponse } from '../../../src/types/rag';
+import { FileResponse, UploadFileResponse } from '../../../src/types';
 import path from 'path';
 
 function sleep(ms) {

diff --git a/src/APIClient.ts b/src/APIClient.ts
@@ -9,11 +9,11 @@ import {
   Headers,
   CrossPlatformResponse,
   UnifiedFormData,
+  FilePathOrFileObject,
 } from './types';
 import { AI21EnvConfig } from './EnvConfig';
 import { createFetchInstance, createFilesHandlerInstance } from './runtime';
 import { Fetch } from 'fetch';
-import { FilePathOrFileObject } from 'types/rag';
 import { BaseFilesHandler } from 'files/BaseFilesHandler';
 
 const validatePositiveInteger = (name: string, n: unknown): number => {

diff --git a/src/files/BaseFilesHandler.ts b/src/files/BaseFilesHandler.ts
@@ -1,5 +1,5 @@
+import { FilePathOrFileObject } from 'types';
 import { FormDataRequest } from 'types/API';
-import { FilePathOrFileObject } from 'types/rag';
 
 export abstract class BaseFilesHandler {
   abstract prepareFormDataRequest(file: FilePathOrFileObject): Promise<FormDataRequest>;

diff --git a/src/files/BrowserFilesHandler.ts b/src/files/BrowserFilesHandler.ts
@@ -1,4 +1,4 @@
-import { FilePathOrFileObject } from 'types/rag';
+import { FilePathOrFileObject } from 'types';
 import { BaseFilesHandler } from './BaseFilesHandler';
 import { FormDataRequest } from 'types/API';
 

diff --git a/src/files/NodeFilesHandler.ts b/src/files/NodeFilesHandler.ts
@@ -1,4 +1,4 @@
-import { FilePathOrFileObject } from 'types/rag';
+import { FilePathOrFileObject } from 'types';
 import { BaseFilesHandler } from './BaseFilesHandler';
 import { FormDataRequest } from 'types/API';
 

diff --git a/src/resources/rag/conversationalRag.ts b/src/resources/rag/conversationalRag.ts
@@ -1,7 +1,7 @@
 import * as Models from '../../types';
 import { APIResource } from '../../APIResource';
-import { ConversationalRagRequest } from '../../types/rag/ConversationalRagRequest';
-import { ConversationalRagResponse } from '../../types/rag/ConversationalRagResponse';
+import { ConversationalRagRequest } from '../../types';
+import { ConversationalRagResponse } from '../../types';
 
 export class ConversationalRag extends APIResource {
   create(body: ConversationalRagRequest, options?: Models.RequestOptions) {

diff --git a/src/resources/rag/files.ts b/src/resources/rag/files.ts
@@ -1,7 +1,7 @@
 import * as Models from '../../types';
 import { APIResource } from '../../APIResource';
-import { UploadFileResponse, UploadFileRequest, ListFilesFilters, UpdateFileRequest } from '../../types/rag';
-import { FileResponse } from 'types/rag/FileResponse';
+import { UploadFileResponse, UploadFileRequest, ListFilesFilters, UpdateFileRequest } from '../../types';
+import { FileResponse } from 'types/files/FileResponse';
 
 const FILES_PATH = '/library/files';
 

diff --git a/src/types/rag/FilePathOrFileObject.ts → src/types/files/FilePathOrFileObject.ts b/src/types/rag/FilePathOrFileObject.ts → src/types/files/FilePathOrFileObject.ts
diff --git a/src/types/rag/FileResponse.ts → src/types/files/FileResponse.ts b/src/types/rag/FileResponse.ts → src/types/files/FileResponse.ts
diff --git a/src/types/rag/ListFilesFilters.ts → src/types/files/ListFilesFilters.ts b/src/types/rag/ListFilesFilters.ts → src/types/files/ListFilesFilters.ts
diff --git a/src/types/rag/UpdateFileRequest.ts → src/types/files/UpdateFileRequest.ts b/src/types/rag/UpdateFileRequest.ts → src/types/files/UpdateFileRequest.ts
diff --git a/src/types/rag/UploadFileRequest.ts → src/types/files/UploadFileRequest.ts b/src/types/rag/UploadFileRequest.ts → src/types/files/UploadFileRequest.ts
diff --git a/src/types/rag/UploadFileResponse.ts → src/types/files/UploadFileResponse.ts b/src/types/rag/UploadFileResponse.ts → src/types/files/UploadFileResponse.ts
diff --git a/src/types/files/index.ts b/src/types/files/index.ts
@@ -0,0 +1,11 @@
+export { type UploadFileRequest } from './UploadFileRequest';
+
+export { type FileResponse } from './FileResponse';
+
+export { type UploadFileResponse } from './UploadFileResponse';
+
+export { type ListFilesFilters } from './ListFilesFilters';
+
+export { type UpdateFileRequest } from './UpdateFileRequest';
+
+export { type FilePathOrFileObject } from './FilePathOrFileObject';
diff --git a/src/types/index.ts b/src/types/index.ts
@@ -44,3 +44,12 @@ export {
   type ConversationalRagSource,
   type RetrievalStrategy,
 } from './rag';
+
+export {
+  type UploadFileRequest,
+  type UploadFileResponse,
+  type FileResponse,
+  type ListFilesFilters,
+  type UpdateFileRequest,
+  type FilePathOrFileObject,
+} from './files';
diff --git a/src/types/rag/index.ts b/src/types/rag/index.ts
@@ -5,15 +5,3 @@ 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 './UpdateFileRequest';
-
-export { type FilePathOrFileObject } from './FilePathOrFileObject';
diff --git a/tests/unittests/resources/rag-engine.test.ts b/tests/unittests/resources/rag-engine.test.ts
@@ -31,7 +31,7 @@ describe('RAGEngine', () => {
 
   it('should upload a file and return the fileId', async () => {
     const fileInput = 'path/to/file.txt';
-    const body = { file: fileInput, path: 'label' };
+    const body = { file: fileInput, path: 'path' };
     const expectedResponse = { fileId: '12345' };
 
     mockClient.upload.mockResolvedValue(expectedResponse);
@@ -41,7 +41,9 @@ describe('RAGEngine', () => {
     expect(mockClient.upload).toHaveBeenCalledWith(
       '/library/files',
       fileInput,
-      { body, headers: { 'Authorization': `Bearer ${dummyAPIKey}` } }
+      {
+        body: { path: 'path' },
+      }
     );
     expect(response).toEqual(expectedResponse);
   });
@@ -56,7 +58,7 @@ describe('RAGEngine', () => {
 
     expect(mockClient.get).toHaveBeenCalledWith(
       `/library/files/${fileId}`,
-      { headers: { 'Authorization': `Bearer ${dummyAPIKey}` } }
+      { ...options }
     );
     expect(response).toEqual(expectedResponse);
   });
@@ -70,7 +72,7 @@ describe('RAGEngine', () => {
 
     expect(mockClient.delete).toHaveBeenCalledWith(
       `/library/files/${fileId}`,
-      { headers: { 'Authorization': `Bearer ${dummyAPIKey}` } }
+      { ...options } 
     );
     expect(response).toBeNull();
   });
@@ -85,7 +87,7 @@ describe('RAGEngine', () => {
 
     expect(mockClient.put).toHaveBeenCalledWith(
       `/library/files/${fileId}`,
-      { body, headers: { 'Authorization': `Bearer ${dummyAPIKey}` } }
+      { body },
     );
     expect(response).toBeNull();
   });
@@ -100,7 +102,10 @@ describe('RAGEngine', () => {
 
     expect(mockClient.get).toHaveBeenCalledWith(
       '/library/files',
-      { query: filters, headers: { 'Authorization': `Bearer ${dummyAPIKey}` } }
+      {
+        query: filters,
+        headers: { 'Authorization': `Bearer ${dummyAPIKey}` }
+      }
     );
     expect(response).toEqual(expectedResponse);
   });