Skip to content

Commit 64f9c44

Browse files
alaa-eddinealaa-eddine
committed
add vectorDB integration tests
1 parent 4ea533d commit 64f9c44

File tree

3 files changed

+484
-0
lines changed

3 files changed

+484
-0
lines changed

‎packages/core/tests/integration/connectors/vectordb/milvus.test.ts‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest';
2+
import { setupSRE } from '../../../utils/sre';
3+
import { ConnectorService } from '@sre/Core/ConnectorsService';
4+
import { AccessCandidate } from '@sre/Security/AccessControl/AccessCandidate.class';
5+
6+
// Deterministic, offline embedding mock
7+
vi.mock('@sre/IO/VectorDB.service/embed', async () => {
8+
const base = await vi.importActual<any>('@sre/IO/VectorDB.service/embed/BaseEmbedding');
9+
10+
function deterministicVector(text: string, dimensions: number): number[] {
11+
const dims = dimensions || 8;
12+
const vec = Array(dims).fill(0);
13+
for (let i = 0; i < (text || '').length; i++) {
14+
const code = text.charCodeAt(i);
15+
vec[code % dims] += (code % 13) + 1;
16+
}
17+
return vec;
18+
}
19+
20+
class TestEmbeds extends base.BaseEmbedding {
21+
constructor(cfg?: any) {
22+
super(cfg);
23+
if (!this.dimensions) this.dimensions = 8;
24+
}
25+
async embedText(text: string): Promise<number[]> {
26+
return deterministicVector(text, this.dimensions as number);
27+
}
28+
async embedTexts(texts: string[]): Promise<number[][]> {
29+
return texts.map((t) => deterministicVector(t, this.dimensions as number));
30+
}
31+
}
32+
33+
return {
34+
EmbeddingsFactory: {
35+
create: (_provider: any, config: any) => new TestEmbeds(config),
36+
},
37+
};
38+
});
39+
40+
function makeVector(text: string, dimensions = 8): number[] {
41+
const vec = Array(dimensions).fill(0);
42+
for (let i = 0; i < (text || '').length; i++) {
43+
const code = text.charCodeAt(i);
44+
vec[code % dimensions] += (code % 13) + 1;
45+
}
46+
return vec;
47+
}
48+
49+
const MILVUS_ADDRESS = process.env.MILVUS_ADDRESS as string; // e.g. localhost:19530
50+
const MILVUS_TOKEN = process.env.MILVUS_TOKEN as string | undefined;
51+
const MILVUS_USER = process.env.MILVUS_USER as string | undefined;
52+
const MILVUS_PASSWORD = process.env.MILVUS_PASSWORD as string | undefined;
53+
const MILVUS_DIMENSIONS = Number(process.env.MILVUS_DIMENSIONS || 1024);
54+
55+
beforeAll(() => {
56+
const credentials = MILVUS_TOKEN
57+
? { address: MILVUS_ADDRESS, token: MILVUS_TOKEN }
58+
: { address: MILVUS_ADDRESS, user: MILVUS_USER, password: MILVUS_PASSWORD };
59+
60+
setupSRE({
61+
VectorDB: {
62+
Connector: 'Milvus',
63+
Settings: {
64+
credentials,
65+
embeddings: {
66+
provider: 'OpenAI',
67+
model: 'text-embedding-3-large',
68+
params: { dimensions: MILVUS_DIMENSIONS },
69+
},
70+
},
71+
},
72+
Log: { Connector: 'ConsoleLog' },
73+
});
74+
});
75+
76+
afterEach(() => {
77+
vi.clearAllMocks();
78+
});
79+
80+
describe('Milvus - VectorDB connector', () => {
81+
it('should create namespace, add/list/get/delete datasource, search by string/vector', async () => {
82+
const vdb = ConnectorService.getVectorDBConnector('Milvus');
83+
const user = AccessCandidate.user('test-user');
84+
const client = vdb.requester(user);
85+
86+
// Create namespace and verify
87+
await client.createNamespace('docs', { env: 'test' });
88+
await expect(client.namespaceExists('docs')).resolves.toBe(true);
89+
90+
// Create datasource with chunking
91+
const ds = await client.createDatasource('docs', {
92+
id: 'mv-ds1',
93+
label: 'MV DS1',
94+
text: 'ABCDEFGHIJKLMNOPQRSTUVWXYZ',
95+
chunkSize: 10,
96+
chunkOverlap: 2,
97+
metadata: { provider: 'milvus' },
98+
});
99+
expect(ds.id).toBe('mv-ds1');
100+
expect(ds.vectorIds.length).toBeGreaterThan(0);
101+
102+
// get/list datasource metadata
103+
const got = await client.getDatasource('docs', 'mv-ds1');
104+
expect(got?.id).toBe('mv-ds1');
105+
const list = await client.listDatasources('docs');
106+
expect(list.map((d) => d.id)).toContain('mv-ds1');
107+
108+
// Search by string
109+
const resText = await client.search('docs', 'KLM', { topK: 3, includeMetadata: true });
110+
expect(resText.length).toBeGreaterThan(0);
111+
112+
// Search by vector
113+
const qv = makeVector('KLM', MILVUS_DIMENSIONS);
114+
const resVec = await client.search('docs', qv, { topK: 1 });
115+
expect(resVec.length).toBe(1);
116+
117+
// topK behavior and sorting
118+
const top1 = await client.search('docs', 'ALPHA', { topK: 1 });
119+
expect(top1.length).toBe(1);
120+
const top3 = await client.search('docs', 'ALPHA', { topK: 3 });
121+
for (let i = 1; i < top3.length; i++) {
122+
expect((top3[i - 1].score || 0) >= (top3[i].score || 0)).toBe(true);
123+
}
124+
125+
// Delete datasource and verify
126+
await client.deleteDatasource('docs', 'mv-ds1');
127+
const maybeDeleted = await client.getDatasource('docs', 'mv-ds1');
128+
expect(maybeDeleted).toBeUndefined();
129+
130+
// Delete namespace
131+
await client.deleteNamespace('docs');
132+
await expect(client.namespaceExists('docs')).resolves.toBe(false);
133+
}, 60000);
134+
});

‎packages/core/tests/integration/connectors/vectordb/pinecone.test.ts‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest';
2+
import { setupSRE } from '../../../utils/sre';
3+
import { ConnectorService } from '@sre/Core/ConnectorsService';
4+
import { AccessCandidate } from '@sre/Security/AccessControl/AccessCandidate.class';
5+
6+
// Deterministic, offline embedding mock
7+
vi.mock('@sre/IO/VectorDB.service/embed', async () => {
8+
const base = await vi.importActual<any>('@sre/IO/VectorDB.service/embed/BaseEmbedding');
9+
10+
function deterministicVector(text: string, dimensions: number): number[] {
11+
const dims = dimensions || 8;
12+
const vec = Array(dims).fill(0);
13+
for (let i = 0; i < (text || '').length; i++) {
14+
const code = text.charCodeAt(i);
15+
vec[code % dims] += (code % 13) + 1;
16+
}
17+
return vec;
18+
}
19+
20+
class TestEmbeds extends base.BaseEmbedding {
21+
constructor(cfg?: any) {
22+
super(cfg);
23+
if (!this.dimensions) this.dimensions = 8;
24+
}
25+
async embedText(text: string): Promise<number[]> {
26+
return deterministicVector(text, this.dimensions as number);
27+
}
28+
async embedTexts(texts: string[]): Promise<number[][]> {
29+
return texts.map((t) => deterministicVector(t, this.dimensions as number));
30+
}
31+
}
32+
33+
return {
34+
EmbeddingsFactory: {
35+
create: (_provider: any, config: any) => new TestEmbeds(config),
36+
},
37+
};
38+
});
39+
40+
function makeVector(text: string, dimensions = 8): number[] {
41+
const vec = Array(dimensions).fill(0);
42+
for (let i = 0; i < (text || '').length; i++) {
43+
const code = text.charCodeAt(i);
44+
vec[code % dimensions] += (code % 13) + 1;
45+
}
46+
return vec;
47+
}
48+
49+
const PINECONE_API_KEY = process.env.PINECONE_API_KEY as string;
50+
const PINECONE_INDEX_NAME = process.env.PINECONE_INDEX_NAME as string;
51+
const PINECONE_DIMENSIONS = Number(process.env.PINECONE_DIMENSIONS || 1024);
52+
53+
beforeAll(() => {
54+
setupSRE({
55+
VectorDB: {
56+
Connector: 'Pinecone',
57+
Settings: {
58+
apiKey: PINECONE_API_KEY,
59+
indexName: PINECONE_INDEX_NAME,
60+
embeddings: {
61+
provider: 'OpenAI',
62+
model: 'text-embedding-3-large',
63+
params: { dimensions: PINECONE_DIMENSIONS },
64+
},
65+
},
66+
},
67+
Log: { Connector: 'ConsoleLog' },
68+
});
69+
});
70+
71+
afterEach(() => {
72+
vi.clearAllMocks();
73+
});
74+
75+
describe('Pinecone - VectorDB connector', () => {
76+
it('should create namespace, add/list/get/delete datasource, search by string/vector', async () => {
77+
const vdb = ConnectorService.getVectorDBConnector('Pinecone');
78+
const user = AccessCandidate.user('test-user');
79+
const client = vdb.requester(user);
80+
81+
// Create namespace and verify
82+
await client.createNamespace('docs', { env: 'test' });
83+
await expect(client.namespaceExists('docs')).resolves.toBe(true);
84+
85+
// Create datasource with chunking
86+
const ds = await client.createDatasource('docs', {
87+
id: 'pc-ds1',
88+
label: 'PC DS1',
89+
text: 'ABCDEFGHIJKLMNOPQRSTUVWXYZ',
90+
chunkSize: 10,
91+
chunkOverlap: 2,
92+
metadata: { provider: 'pinecone' },
93+
});
94+
expect(ds.id).toBe('pc-ds1');
95+
expect(ds.vectorIds.length).toBeGreaterThan(0);
96+
97+
// get/list datasource metadata (stored via NKV)
98+
const got = await client.getDatasource('docs', 'pc-ds1');
99+
expect(got.id).toBe('pc-ds1');
100+
const list = await client.listDatasources('docs');
101+
expect(list.map((d) => d.id)).toContain('pc-ds1');
102+
103+
// Search by string
104+
const resText = await client.search('docs', 'KLM', { topK: 3, includeMetadata: true });
105+
expect(resText.length).toBeGreaterThan(0);
106+
107+
// Search by vector
108+
const qv = makeVector('KLM', PINECONE_DIMENSIONS);
109+
const resVec = await client.search('docs', qv, { topK: 1 });
110+
expect(resVec.length).toBe(1);
111+
112+
// topK behavior and sorting
113+
const top1 = await client.search('docs', 'ALPHA', { topK: 1 });
114+
expect(top1.length).toBe(1);
115+
const top3 = await client.search('docs', 'ALPHA', { topK: 3 });
116+
for (let i = 1; i < top3.length; i++) {
117+
expect((top3[i - 1].score || 0) >= (top3[i].score || 0)).toBe(true);
118+
}
119+
120+
// Delete datasource and verify
121+
await client.deleteDatasource('docs', 'pc-ds1');
122+
await expect(client.getDatasource('docs', 'pc-ds1')).rejects.toThrow('Data source not found');
123+
124+
// Delete namespace
125+
await client.deleteNamespace('docs');
126+
await expect(client.namespaceExists('docs')).resolves.toBe(false);
127+
}, 60000);
128+
});

0 commit comments

Comments
 (0)