28 lines
865 B
JavaScript
28 lines
865 B
JavaScript
const { VectorStore } = require('../src/vectorStore');
|
|
const { embed } = require('../src/utils');
|
|
|
|
describe('VectorStore', () => {
|
|
let store;
|
|
|
|
beforeAll(async () => {
|
|
store = new VectorStore();
|
|
await store.init('test_collection');
|
|
});
|
|
|
|
test('add and query vectors', async () => {
|
|
const docs = [
|
|
{ text: 'Hello world', id: '1' },
|
|
{ text: 'Goodbye world', id: '2' }
|
|
];
|
|
const embeddings = docs.map(d => embed(d.text));
|
|
const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
|
|
await store.add(embeddings, metadatas, docs.map(d => d.id));
|
|
|
|
const queryEmbedding = embed('Hello');
|
|
const results = await store.query(queryEmbedding, 2);
|
|
expect(results[0].metadatas.length).toBe(2);
|
|
const ids = results[0].metadatas.map(m => m.id);
|
|
expect(ids).toContain('1');
|
|
expect(ids).toContain('2');
|
|
});
|
|
}); |