@future-agi/sdk
Version:
We help GenAI teams maintain high-accuracy for their Models in production.
625 lines • 26.8 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
const dataset_1 = require("../dataset");
const types_1 = require("../types");
const errors_1 = require("../../utils/errors");
const fs = __importStar(require("fs"));
const stream_1 = require("stream");
// Mock built-in modules
jest.mock('axios');
jest.mock('fs', () => (Object.assign(Object.assign({}, jest.requireActual('fs')), { existsSync: jest.fn(), createReadStream: jest.fn() })));
// Mock the APIKeyAuth base class
jest.mock('../../api/auth', () => ({
APIKeyAuth: class MockAPIKeyAuth {
constructor(config) {
this._baseUrl = 'https://api.futureagi.com';
// Mock constructor
}
request(config, responseHandler) {
return __awaiter(this, void 0, void 0, function* () {
// Mock request method - will be overridden in tests
return {};
});
}
},
ResponseHandler: class MockResponseHandler {
static parse(response, handlerClass) {
return handlerClass._parseSuccess(response);
}
static _parseSuccess(response) {
if (response.data) {
return response.data;
}
throw new Error("Mock _parseSuccess error: response.data is undefined");
}
static _handleError(response) {
throw new Error('Mock error');
}
},
}));
describe('DatasetResponseHandler', () => {
describe('_parseSuccess', () => {
test('should parse dataset names response', () => {
const mockResponse = {
data: {
result: {
datasets: [
{
datasetId: 'dataset-123',
name: 'test-dataset',
},
],
},
},
config: {
url: 'https://api.futureagi.com/model-hub/develops/get-datasets-names/',
headers: {},
},
};
const result = dataset_1.DatasetResponseHandler._parseSuccess(mockResponse);
expect(result).toEqual({
id: 'dataset-123',
name: 'test-dataset',
});
});
test('should throw error for empty datasets list', () => {
const mockResponse = {
data: {
result: {
datasets: [],
},
},
config: {
url: 'https://api.futureagi.com/model-hub/develops/get-datasets-names/',
headers: {},
},
};
expect(() => {
dataset_1.DatasetResponseHandler._parseSuccess(mockResponse);
}).toThrow(errors_1.DatasetNotFoundError);
});
test('should throw error for multiple datasets', () => {
const mockResponse = {
data: {
result: {
datasets: [
{ datasetId: '1', name: 'dataset1', },
{ datasetId: '2', name: 'dataset2', },
],
},
},
config: {
url: 'https://api.futureagi.com/model-hub/develops/get-datasets-names/',
headers: {},
},
};
expect(() => {
dataset_1.DatasetResponseHandler._parseSuccess(mockResponse);
}).toThrow('Multiple datasets found');
});
test('should parse dataset table response', () => {
const mockResponse = {
data: {
result: {
columnConfig: [
{
id: 'col-1',
name: 'name',
dataType: types_1.DataTypeChoices.TEXT,
originType: types_1.SourceChoices.OTHERS,
sourceId: null,
isFrozen: { isFrozen: false },
isVisible: true,
evalTag: [],
averageScore: null,
orderIndex: 0,
},
],
table: [
{
rowId: 'row-1',
order: 0,
'col-1': {
cellValue: 'John',
valueInfos: null,
metadata: null,
status: 'completed',
failureReason: null,
},
},
],
metadata: {},
},
},
config: {
url: 'https://api.futureagi.com/model-hub/develops/dataset-123/get-dataset-table/',
headers: {},
},
};
const result = dataset_1.DatasetResponseHandler._parseSuccess(mockResponse);
expect(result.id).toBe('dataset-123');
expect(result.columns).toHaveLength(1);
expect(result.rows).toHaveLength(1);
expect(result.columns[0].name).toBe('name');
expect(result.rows[0].cells).toHaveLength(1);
});
test('should parse dataset creation response', () => {
const mockResponse = {
data: {
result: {
datasetId: 'dataset-123',
datasetName: 'test-dataset',
},
},
config: {
url: 'https://api.futureagi.com/model-hub/develops/create-empty-dataset/',
headers: {},
},
};
const result = dataset_1.DatasetResponseHandler._parseSuccess(mockResponse);
expect(result).toEqual({
id: 'dataset-123',
name: 'test-dataset',
});
});
});
describe('_handleError', () => {
test('should handle different error status codes', () => {
const mockResponse = {
status: 404,
data: { message: 'Dataset not found' },
statusText: 'Not Found',
};
expect(() => {
dataset_1.DatasetResponseHandler._handleError(mockResponse);
}).toThrow(errors_1.DatasetNotFoundError);
});
test('should handle server errors', () => {
const mockResponse = {
status: 500,
data: { message: 'Internal server error' },
statusText: 'Internal Server Error',
};
expect(() => {
dataset_1.DatasetResponseHandler._handleError(mockResponse);
}).toThrow('Internal server error');
});
});
});
describe('Dataset', () => {
let dataset;
let mockRequest;
beforeEach(() => {
mockRequest = jest.fn();
dataset = new dataset_1.Dataset({
fiApiKey: 'test-key',
fiSecretKey: 'test-secret',
});
dataset.request = mockRequest;
});
describe('constructor', () => {
test('should create dataset instance', () => {
expect(dataset).toBeInstanceOf(dataset_1.Dataset);
});
test('should create dataset with config', () => {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const datasetWithConfig = new dataset_1.Dataset({
fiApiKey: 'test-key',
fiSecretKey: 'test-secret',
datasetConfig: config,
});
expect(datasetWithConfig.getConfig()).toEqual(config);
});
});
describe('getConfig', () => {
test('should return dataset config', () => {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
expect(dataset.getConfig()).toEqual(config);
});
test('should throw error when no config is set', () => {
expect(() => {
dataset.getConfig();
}).toThrow(errors_1.DatasetError);
});
});
describe('create', () => {
test('should create empty dataset', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
name: 'test-dataset',
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({
id: 'dataset-123',
name: 'test-dataset',
});
const result = yield dataset.create();
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
json: expect.objectContaining({
new_dataset_name: 'test-dataset',
}),
}), dataset_1.DatasetResponseHandler);
}));
test('should create dataset from file', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
name: 'test-dataset',
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({
id: 'dataset-123',
name: 'test-dataset',
});
// Mock fs.existsSync to avoid file system check
fs.existsSync.mockReturnValue(true);
fs.createReadStream.mockReturnValue(new stream_1.Readable());
const result = yield dataset.create('path/to/file.csv');
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
// Data is now a FormData object, which is harder to inspect directly
}), dataset_1.DatasetResponseHandler);
}));
test('should create dataset from Hugging Face', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
name: 'test-dataset',
};
const hfConfig = {
name: 'squad',
subset: 'plain_text',
split: 'train',
numRows: 1000,
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({
id: 'dataset-123',
name: 'test-dataset',
});
const result = yield dataset.create(hfConfig);
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
json: expect.objectContaining({
new_dataset_name: 'test-dataset',
huggingface_dataset_name: 'squad',
huggingface_dataset_config: 'plain_text',
huggingface_dataset_split: 'train',
num_rows: 1000,
}),
}), dataset_1.DatasetResponseHandler);
}));
test('should throw error when dataset config is not set', () => __awaiter(void 0, void 0, void 0, function* () {
yield expect(dataset.create()).rejects.toThrow(errors_1.DatasetError);
}));
test('should throw error when dataset already exists', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
yield expect(dataset.create()).rejects.toThrow(errors_1.DatasetError);
}));
});
describe('download', () => {
test('should download dataset', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({
id: 'dataset-123',
columns: [],
rows: [],
metadata: {},
});
const result = yield dataset.download();
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalled();
}));
test('should download dataset to memory', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const mockTable = {
id: 'dataset-123',
columns: [],
rows: [],
metadata: {},
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue(mockTable);
const result = yield dataset.download(undefined, true);
expect(result).toEqual(mockTable);
}));
test('should throw error when dataset name is not set', () => __awaiter(void 0, void 0, void 0, function* () {
yield expect(dataset.download()).rejects.toThrow(errors_1.DatasetError);
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
name: 'test-dataset',
};
dataset._datasetConfig = config;
yield expect(dataset.download()).rejects.toThrow(errors_1.DatasetError);
}));
});
describe('delete', () => {
test('should delete dataset', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({});
yield dataset.delete();
expect(dataset._datasetConfig).toBeNull();
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'DELETE',
json: { dataset_ids: ['dataset-123'] },
}), dataset_1.DatasetResponseHandler);
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
yield expect(dataset.delete()).rejects.toThrow(errors_1.DatasetError);
}));
});
describe('addColumns', () => {
test('should add columns to dataset', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const columns = [
(0, types_1.createColumn)({ name: 'col1', dataType: types_1.DataTypeChoices.TEXT }),
(0, types_1.createColumn)({ name: 'col2', dataType: types_1.DataTypeChoices.INTEGER }),
];
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({});
const result = yield dataset.addColumns(columns);
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
json: { new_columns_data: expect.any(Array) },
}), dataset_1.DatasetResponseHandler);
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
const columns = [
(0, types_1.createColumn)({ name: 'col1', dataType: types_1.DataTypeChoices.TEXT }),
];
yield expect(dataset.addColumns(columns)).rejects.toThrow(errors_1.DatasetError);
}));
test('should throw error for empty columns array', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
yield expect(dataset.addColumns([])).rejects.toThrow(errors_1.DatasetValidationError);
}));
});
describe('addRows', () => {
test('should add rows to dataset', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const rows = [
(0, types_1.createRow)({
cells: [
(0, types_1.createCell)({ columnId: 'col-1', rowId: 'row-1', value: 'test' }),
],
}),
];
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({});
const result = yield dataset.addRows(rows);
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
json: { rows: expect.any(Array) },
}), dataset_1.DatasetResponseHandler);
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
const rows = [
(0, types_1.createRow)({
cells: [
(0, types_1.createCell)({ columnId: 'col-1', rowId: 'row-1', value: 'test' }),
],
}),
];
yield expect(dataset.addRows(rows)).rejects.toThrow(errors_1.DatasetError);
}));
test('should throw error for empty rows array', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
yield expect(dataset.addRows([])).rejects.toThrow(errors_1.DatasetValidationError);
}));
});
describe('addRunPrompt', () => {
test('should add a run prompt with correct payload', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue({});
const options = {
name: 'my_prompt',
model: 'gpt-3.5-turbo',
messages: [{ role: 'user', content: 'Hello' }],
};
const result = yield dataset.addRunPrompt(options);
expect(result).toBe(dataset);
expect(mockRequest).toHaveBeenCalledWith(expect.objectContaining({
method: 'POST',
json: expect.objectContaining({
dataset_id: 'dataset-123',
name: 'my_prompt',
config: expect.objectContaining({
model: 'gpt-3.5-turbo',
messages: options.messages,
}),
}),
}), dataset_1.DatasetResponseHandler);
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
const options = {
name: 'my_prompt',
model: 'gpt-3.5-turbo',
messages: [{ role: 'user', content: 'Hello' }],
};
yield expect(dataset.addRunPrompt(options)).rejects.toThrow(errors_1.DatasetError);
}));
});
describe('getColumnId', () => {
test('should get column ID by name', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const mockTable = {
id: 'dataset-123',
columns: [
{
id: 'col-1',
name: 'test_column',
dataType: types_1.DataTypeChoices.TEXT,
source: types_1.SourceChoices.OTHERS,
metadata: {},
isFrozen: false,
isVisible: true,
evalTags: [],
orderIndex: 0,
},
],
rows: [],
metadata: {},
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue(mockTable);
const result = yield dataset.getColumnId('test_column');
expect(result).toBe('col-1');
}));
test('should return null for non-existent column', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
const mockTable = {
id: 'dataset-123',
columns: [],
rows: [],
metadata: {},
};
dataset._datasetConfig = config;
mockRequest.mockResolvedValue(mockTable);
const result = yield dataset.getColumnId('non_existent_column');
expect(result).toBeNull();
}));
test('should throw error when dataset ID is not set', () => __awaiter(void 0, void 0, void 0, function* () {
yield expect(dataset.getColumnId('test_column')).rejects.toThrow(errors_1.DatasetError);
}));
test('should throw error for empty column name', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
id: 'dataset-123',
name: 'test-dataset',
};
dataset._datasetConfig = config;
yield expect(dataset.getColumnId('')).rejects.toThrow(errors_1.DatasetValidationError);
}));
});
describe('static methods', () => {
test('should create dataset using static method', () => __awaiter(void 0, void 0, void 0, function* () {
const config = {
name: 'test-dataset',
};
const mockDatasetInstance = new dataset_1.Dataset();
const mockCreate = jest.fn().mockResolvedValue(mockDatasetInstance);
jest.spyOn(dataset_1.Dataset.prototype, 'create').mockImplementation(mockCreate);
const result = yield dataset_1.Dataset.createDataset(config, undefined, {
fiApiKey: 'test-key',
fiSecretKey: 'test-secret',
});
expect(result).toBeInstanceOf(dataset_1.Dataset);
// This assertion is tricky because create is called on a new instance
// A better test would be to mock the constructor
}));
test('should download dataset using static method', () => __awaiter(void 0, void 0, void 0, function* () {
const mockDownload = jest.fn().mockResolvedValue('csv content');
const mockFetchConfig = jest.fn().mockResolvedValue({
id: 'dataset-123',
name: 'test-dataset',
});
// Mock the internal methods on the prototype
jest.spyOn(dataset_1.Dataset.prototype, '_fetchDatasetConfig').mockImplementation(mockFetchConfig);
jest.spyOn(dataset_1.Dataset.prototype, 'download').mockImplementation(mockDownload);
const result = yield dataset_1.Dataset.downloadDataset('test-dataset');
expect(result).toBe('csv content');
expect(mockDownload).toHaveBeenCalled();
}));
test('should delete dataset using static method', () => __awaiter(void 0, void 0, void 0, function* () {
const mockDelete = jest.fn().mockResolvedValue(undefined);
const mockFetchConfig = jest.fn().mockResolvedValue({
id: 'dataset-123',
name: 'test-dataset',
});
jest.spyOn(dataset_1.Dataset.prototype, '_fetchDatasetConfig').mockImplementation(mockFetchConfig);
jest.spyOn(dataset_1.Dataset.prototype, 'delete').mockImplementation(mockDelete);
yield dataset_1.Dataset.deleteDataset('test-dataset');
expect(mockDelete).toHaveBeenCalled();
}));
});
});
//# sourceMappingURL=dataset.test.js.map