Skip to content

Commit 7ca7640

Browse files
amirai21asafgardin
amirai21
authored and
asafgardin
committed
feat: add file path check before opening + rag examples improved
1 parent 382e69f commit 7ca7640

File tree

10 files changed

+66
-132
lines changed

10 files changed

+66
-132
lines changed

‎.prettierignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
*.txt

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
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

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,74 @@
11
import { AI21 } from 'ai21';
22
import { FileResponse, UploadFileResponse } from '../../../src/types/rag';
3+
import path from 'path';
34

45
function sleep(ms) {
56
return new Promise((resolve) => setTimeout(resolve, ms));
67
}
78

9+
async function waitForFileProcessing(
10+
client: AI21,
11+
fileId: string,
12+
timeout: number = 30000,
13+
interval: number = 1000,
14+
) {
15+
const startTime = Date.now();
16+
17+
while (Date.now() - startTime < timeout) {
18+
const file: FileResponse = await client.files.get(fileId);
19+
if (file.status !== 'PROCESSING') {
20+
return file;
21+
}
22+
await sleep(interval);
23+
}
24+
25+
throw new Error(`File processing timed out after ${timeout}ms`);
26+
}
27+
828
async function uploadGetUpdateDelete(fileInput, path) {
929
const client = new AI21({ apiKey: process.env.AI21_API_KEY });
1030
try {
11-
const uploadFileResponse: UploadFileResponse = await client.ragEngine.create({
31+
const uploadFileResponse: UploadFileResponse = await client.files.create({
1232
file: fileInput,
1333
path: path,
1434
});
1535
console.log(uploadFileResponse);
16-
let file: FileResponse = await client.ragEngine.get(uploadFileResponse.fileId);
17-
console.log(file);
18-
await sleep(1000); // Give it a sec to start process before updating
19-
console.log('Now updating the file labels and publicUrl...');
20-
await client.ragEngine.update({
21-
fileId: uploadFileResponse.fileId,
22-
labels: ['test99'],
23-
publicUrl: 'https://www.miri.com',
24-
});
25-
file = await client.ragEngine.get(uploadFileResponse.fileId);
36+
37+
let file: FileResponse = await waitForFileProcessing(client, uploadFileResponse.fileId);
2638
console.log(file);
2739

40+
if (file.status === 'PROCESSED') {
41+
console.log('Now updating the file labels and publicUrl...');
42+
await client.files.update({
43+
fileId: uploadFileResponse.fileId,
44+
labels: ['test99'],
45+
publicUrl: 'https://www.miri.com',
46+
});
47+
file = await client.files.get(uploadFileResponse.fileId);
48+
console.log(file);
49+
} else {
50+
console.log(`File did not processed well, ended with status ${file.status}`);
51+
}
52+
2853
console.log('Now deleting the file');
29-
await client.ragEngine.delete(uploadFileResponse.fileId);
54+
await client.files.delete(uploadFileResponse.fileId);
3055
} catch (error) {
3156
console.error('Error:', error);
3257
}
3358
}
3459

3560
async function listFiles() {
3661
const client = new AI21({ apiKey: process.env.AI21_API_KEY });
37-
const files = await client.ragEngine.list({ limit: 4 });
62+
const files = await client.files.list({ limit: 4 });
3863
console.log(files);
3964
}
4065

41-
/* Simulate a file upload passing file path */
42-
const filePath = '/Users/amirkoblyansky/Documents/ukraine.txt';
66+
/* Simulate file upload passing a path to file */
67+
const filePath = path.join(process.cwd(), 'examples/studio/conversational-rag/files', 'meerkat.txt'); // Use process.cwd() to get the current working directory
68+
4369
uploadGetUpdateDelete(filePath, Date.now().toString()).catch(console.error);
4470

45-
/* Simulate a file upload passing File instance */
71+
/* Simulate file upload passing File instance */
4672
const fileContent = Buffer.from(
4773
'Opossums are members of the marsupial order Didelphimorphia endemic to the Americas.',
4874
);

‎src/AI21.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { APIClient } from './APIClient';
66
import { Headers } from './types';
77
import * as Runtime from './runtime';
88
import { ConversationalRag } from './resources/rag/conversationalRag';
9-
import { RAGEngine } from './resources';
9+
import { Files } from './resources';
1010

1111
export interface ClientOptions {
1212
baseURL?: string | undefined;
@@ -68,7 +68,7 @@ export class AI21 extends APIClient {
6868
// Resources
6969
chat: Chat = new Chat(this);
7070
conversationalRag: ConversationalRag = new ConversationalRag(this);
71-
ragEngine: RAGEngine = new RAGEngine(this);
71+
files: Files = new Files(this);
7272

7373
// eslint-disable-next-line @typescript-eslint/no-unused-vars
7474
protected override authHeaders(_: Types.FinalRequestOptions): Types.Headers {

‎src/files/NodeFilesHandler.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ export class NodeFilesHandler extends BaseFilesHandler {
2525

2626
if (typeof file === 'string') {
2727
const fs = (await import('fs')).default;
28+
if (!fs.existsSync(file)) {
29+
throw new Error(`File not found: ${file}`);
30+
}
2831
formData.append('file', fs.createReadStream(file), { filename: file.split('/').pop() });
2932
} else if (file instanceof File) {
3033
const nodeStream = await this.convertReadableStream(file.stream());

‎src/resources/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
export { Chat, Completions } from './chat';
22
export { ConversationalRag } from './rag';
3-
export { RAGEngine } from './rag';
3+
export { Files } from './rag';

‎src/resources/rag/ragEngine.ts renamed to ‎src/resources/rag/files.ts

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,40 +3,40 @@ import { APIResource } from '../../APIResource';
33
import { UploadFileResponse, UploadFileRequest, ListFilesFilters, UpdateFileRequest } from '../../types/rag';
44
import { FileResponse } from 'types/rag/FileResponse';
55

6-
const RAG_ENGINE_PATH = '/library/files';
6+
const FILES_PATH = '/library/files';
77

8-
export class RAGEngine extends APIResource {
8+
export class Files extends APIResource {
99
async create(body: UploadFileRequest, options?: Models.RequestOptions): Promise<UploadFileResponse> {
1010
const { file, ...bodyWithoutFile } = body;
11-
return this.client.upload<Models.UnifiedFormData, UploadFileResponse>(RAG_ENGINE_PATH, file, {
11+
return this.client.upload<Models.UnifiedFormData, UploadFileResponse>(FILES_PATH, file, {
1212
body: bodyWithoutFile,
1313
...options,
1414
} as Models.RequestOptions<Models.UnifiedFormData>) as Promise<UploadFileResponse>;
1515
}
1616

1717
get(fileId: string, options?: Models.RequestOptions): Promise<FileResponse> {
1818
return this.client.get<string, FileResponse>(
19-
`${RAG_ENGINE_PATH}/${fileId}`,
19+
`${FILES_PATH}/${fileId}`,
2020
options as Models.RequestOptions<string>,
2121
) as Promise<FileResponse>;
2222
}
2323

2424
delete(fileId: string, options?: Models.RequestOptions): Promise<null> {
2525
return this.client.delete<string, null>(
26-
`${RAG_ENGINE_PATH}/${fileId}`,
26+
`${FILES_PATH}/${fileId}`,
2727
options as Models.RequestOptions<string>,
2828
) as Promise<null>;
2929
}
3030

3131
list(body?: ListFilesFilters, options?: Models.RequestOptions): Promise<FileResponse[]> {
32-
return this.client.get<ListFilesFilters | null, FileResponse[]>(RAG_ENGINE_PATH, {
32+
return this.client.get<ListFilesFilters | null, FileResponse[]>(FILES_PATH, {
3333
query: body,
3434
...options,
3535
} as Models.RequestOptions<ListFilesFilters | null>) as Promise<FileResponse[]>;
3636
}
3737

3838
update(body: UpdateFileRequest, options?: Models.RequestOptions): Promise<null> {
39-
return this.client.put<UpdateFileRequest, null>(`${RAG_ENGINE_PATH}/${body.fileId}`, {
39+
return this.client.put<UpdateFileRequest, null>(`${FILES_PATH}/${body.fileId}`, {
4040
body,
4141
...options,
4242
} as Models.RequestOptions<UpdateFileRequest>) as Promise<null>;

‎src/resources/rag/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
export { ConversationalRag } from './conversationalRag';
2-
export { RAGEngine } from './ragEngine';
2+
export { Files } from './files';

‎tests/unittests/resources/chat/conversational-rag.test.ts

Lines changed: 0 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import * as Models from '../../../../src/types';
22
import {ConversationalRag} from "../../../../src/resources/rag/conversationalRag";
33
import { APIClient } from '../../../../src/APIClient';
4-
import { RAGEngine } from '../../../../src/resources/rag/ragEngine';
54

65

76
class MockAPIClient extends APIClient {
@@ -101,99 +100,3 @@ describe('ConversationalRag', () => {
101100
expect(response).toEqual(expectedResponse);
102101
});
103102
});
104-
105-
describe('RAGEngine', () => {
106-
let ragEngine: RAGEngine;
107-
let mockClient: MockAPIClient;
108-
const dummyAPIKey = "test-api-key";
109-
110-
beforeEach(() => {
111-
mockClient = new MockAPIClient({
112-
baseURL: 'https://api.example.com',
113-
maxRetries: 3,
114-
timeout: 5000,
115-
});
116-
117-
ragEngine = new RAGEngine(mockClient);
118-
});
119-
120-
afterEach(() => {
121-
jest.clearAllMocks();
122-
});
123-
124-
it('should upload a file and return the response', async () => {
125-
const fileInput = 'path/to/file.txt';
126-
const body = { path: 'label' };
127-
const expectedResponse = { fileId: '12345' };
128-
129-
mockClient.upload.mockResolvedValue(expectedResponse);
130-
131-
const response = await ragEngine.create(fileInput, body);
132-
133-
expect(mockClient.upload).toHaveBeenCalledWith(
134-
'/library/files',
135-
fileInput,
136-
{ body, headers: undefined }
137-
);
138-
expect(response).toEqual(expectedResponse);
139-
});
140-
141-
it('should get a file by ID and return the response', async () => {
142-
const fileId = '12345';
143-
const expectedResponse = { id: fileId, name: 'file.txt' };
144-
145-
mockClient.get.mockResolvedValue(expectedResponse);
146-
147-
const response = await ragEngine.get(fileId);
148-
149-
expect(mockClient.get).toHaveBeenCalledWith(
150-
`/library/files/${fileId}`,
151-
undefined
152-
);
153-
expect(response).toEqual(expectedResponse);
154-
});
155-
156-
it('should delete a file by ID', async () => {
157-
const fileId = '12345';
158-
159-
mockClient.delete.mockResolvedValue(null);
160-
161-
const response = await ragEngine.delete(fileId);
162-
163-
expect(mockClient.delete).toHaveBeenCalledWith(
164-
`/library/files/${fileId}`,
165-
undefined
166-
);
167-
expect(response).toBeNull();
168-
});
169-
170-
it('should update a file by ID and return null', async () => {
171-
const fileId = '12345';
172-
const body = { labels: ['test'], publicUrl: 'https://example.com' };
173-
174-
mockClient.put.mockResolvedValue(null);
175-
176-
const response = await ragEngine.update(fileId, body);
177-
178-
expect(mockClient.put).toHaveBeenCalledWith(
179-
`/library/files/${fileId}`,
180-
{ body, headers: undefined }
181-
);
182-
expect(response).toBeNull();
183-
});
184-
185-
it('should list files and return the response', async () => {
186-
const filters = { limit: 4 };
187-
const expectedResponse = [{ id: '12345', name: 'file.txt' }];
188-
189-
mockClient.get.mockResolvedValue(expectedResponse);
190-
191-
const response = await ragEngine.list(filters);
192-
193-
expect(mockClient.get).toHaveBeenCalledWith(
194-
'/library/files',
195-
{ query: filters, headers: undefined }
196-
);
197-
expect(response).toEqual(expectedResponse);
198-
});
199-
});

‎tests/unittests/resources/rag-engine.test.ts

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import * as Models from '../../../src/types';
2-
import { RAGEngine } from '../../../src/resources/rag/ragEngine';
2+
import { Files } from '../../../src/resources/rag/files';
33
import { APIClient } from '../../../src/APIClient';
44

55
class MockAPIClient extends APIClient {
@@ -10,7 +10,7 @@ class MockAPIClient extends APIClient {
1010
}
1111

1212
describe('RAGEngine', () => {
13-
let ragEngine: RAGEngine;
13+
let files: Files;
1414
let mockClient: MockAPIClient;
1515
const dummyAPIKey = "test-api-key";
1616
const options: Models.RequestOptions = { headers: { 'Authorization': `Bearer ${dummyAPIKey}` } };
@@ -22,7 +22,7 @@ describe('RAGEngine', () => {
2222
timeout: 5000,
2323
});
2424

25-
ragEngine = new RAGEngine(mockClient);
25+
files = new Files(mockClient);
2626
});
2727

2828
afterEach(() => {
@@ -36,7 +36,7 @@ describe('RAGEngine', () => {
3636

3737
mockClient.upload.mockResolvedValue(expectedResponse);
3838

39-
const response = await ragEngine.create(body);
39+
const response = await files.create(body);
4040

4141
expect(mockClient.upload).toHaveBeenCalledWith(
4242
'/library/files',
@@ -52,7 +52,7 @@ describe('RAGEngine', () => {
5252

5353
mockClient.get.mockResolvedValue(expectedResponse);
5454

55-
const response = await ragEngine.get(fileId, options);
55+
const response = await files.get(fileId, options);
5656

5757
expect(mockClient.get).toHaveBeenCalledWith(
5858
`/library/files/${fileId}`,
@@ -66,7 +66,7 @@ describe('RAGEngine', () => {
6666

6767
mockClient.delete.mockResolvedValue(null);
6868

69-
const response = await ragEngine.delete(fileId, options);
69+
const response = await files.delete(fileId, options);
7070

7171
expect(mockClient.delete).toHaveBeenCalledWith(
7272
`/library/files/${fileId}`,
@@ -81,7 +81,7 @@ describe('RAGEngine', () => {
8181

8282
mockClient.put.mockResolvedValue(null);
8383

84-
const response = await ragEngine.update(body);
84+
const response = await files.update(body);
8585

8686
expect(mockClient.put).toHaveBeenCalledWith(
8787
`/library/files/${fileId}`,
@@ -96,7 +96,7 @@ describe('RAGEngine', () => {
9696

9797
mockClient.get.mockResolvedValue(expectedResponse);
9898

99-
const response = await ragEngine.list(filters, options);
99+
const response = await files.list(filters, options);
100100

101101
expect(mockClient.get).toHaveBeenCalledWith(
102102
'/library/files',

0 commit comments

Comments
 (0)