Skip to content

Commit 32c5ac1

Browse files
jonathanMLDevzhocursoragent
authored
split pinecone-client.ts with 3 sub files (cppalliance#86)
* split pinecone-client.ts with 3 sub files # Conflicts: # src/pinecone-client.ts * fixed typecheck errors * added some tests * fixed type check errors * fixed type check errors Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: zho <jornathanm910923@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 335bda1 commit 32c5ac1

8 files changed

Lines changed: 1050 additions & 546 deletions

File tree

src/pinecone-client.test.ts

Lines changed: 179 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
import { describe, it, expect, beforeEach } from 'vitest';
1+
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
22
import { PineconeClient } from './pinecone-client.js';
33
import type { SearchableIndex, PineconeHit } from './types.js';
4+
import * as rerankModule from './pinecone/rerank.js';
45

5-
/** Test double: client with stubbable ensureIndexes and searchIndex for hybrid tests */
6-
type PineconeClientTestDouble = PineconeClient & {
6+
/** Stubs for private methods (assigned at runtime; avoid intersecting private `PineconeClient` members). */
7+
type PineconeClientMethodStubs = {
78
ensureIndexes: () => Promise<{ denseIndex: SearchableIndex; sparseIndex: SearchableIndex }>;
89
searchIndex: (
910
index: SearchableIndex,
@@ -15,6 +16,10 @@ type PineconeClientTestDouble = PineconeClient & {
1516
) => Promise<PineconeHit[]>;
1617
};
1718

19+
function stubPineconeClient(client: PineconeClient): PineconeClientMethodStubs {
20+
return client as unknown as PineconeClientMethodStubs;
21+
}
22+
1823
describe('PineconeClient', () => {
1924
let client: PineconeClient;
2025

@@ -26,6 +31,12 @@ describe('PineconeClient', () => {
2631
});
2732
});
2833

34+
afterEach(() => {
35+
delete process.env['PINECONE_INDEX_NAME'];
36+
delete process.env['PINECONE_RERANK_MODEL'];
37+
delete process.env['PINECONE_TOP_K'];
38+
});
39+
2940
describe('constructor', () => {
3041
it('should initialize with provided config', () => {
3142
expect(client).toBeDefined();
@@ -64,7 +75,7 @@ describe('PineconeClient', () => {
6475
});
6576

6677
it('should continue hybrid search when one index fails', async () => {
67-
const testClient = client as PineconeClientTestDouble;
78+
const testClient = stubPineconeClient(client);
6879

6980
testClient.ensureIndexes = async () => ({
7081
denseIndex: {} as SearchableIndex,
@@ -99,7 +110,7 @@ describe('PineconeClient', () => {
99110
});
100111

101112
it('should throw when both dense and sparse searches fail', async () => {
102-
const testClient = client as PineconeClientTestDouble;
113+
const testClient = stubPineconeClient(client);
103114

104115
testClient.ensureIndexes = async () => ({
105116
denseIndex: {} as SearchableIndex,
@@ -122,7 +133,7 @@ describe('PineconeClient', () => {
122133

123134
describe('count', () => {
124135
it('should return unique document count using semantic search only with minimal fields', async () => {
125-
const testClient = client as PineconeClientTestDouble;
136+
const testClient = stubPineconeClient(client);
126137
testClient.ensureIndexes = async () => ({
127138
denseIndex: {} as SearchableIndex,
128139
sparseIndex: {} as SearchableIndex,
@@ -161,7 +172,7 @@ describe('PineconeClient', () => {
161172
});
162173

163174
it('should set truncated when hit limit is reached', async () => {
164-
const testClient = client as PineconeClientTestDouble;
175+
const testClient = stubPineconeClient(client);
165176
testClient.ensureIndexes = async () => ({
166177
denseIndex: {} as SearchableIndex,
167178
sparseIndex: {} as SearchableIndex,
@@ -178,5 +189,166 @@ describe('PineconeClient', () => {
178189
expect(result.count).toBe(10000);
179190
expect(result.truncated).toBe(true);
180191
});
192+
193+
it('falls back to chunk _id when no document identifier fields exist', async () => {
194+
const testClient = stubPineconeClient(client);
195+
testClient.ensureIndexes = async () => ({
196+
denseIndex: {} as SearchableIndex,
197+
sparseIndex: {} as SearchableIndex,
198+
});
199+
testClient.searchIndex = async () => [
200+
{ _id: 'chunk-only', _score: 1, fields: { chunk_text: 'x' } },
201+
];
202+
203+
const result = await client.count({ query: 'paper', namespace: 'ns' });
204+
205+
expect(result.count).toBe(1);
206+
expect(result.truncated).toBe(false);
207+
});
208+
});
209+
210+
describe('getSparseIndexName', () => {
211+
it('returns {indexName}-sparse derived from config indexName', () => {
212+
const c = new PineconeClient({ apiKey: 'k', indexName: 'my' });
213+
expect(c.getSparseIndexName()).toBe('my-sparse');
214+
});
215+
});
216+
217+
describe('query (rerank and fields)', () => {
218+
it('rejects non-finite topK', async () => {
219+
await expect(client.query({ query: 'q', namespace: 'n', topK: Number.NaN })).rejects.toThrow(
220+
'topK must be a finite number'
221+
);
222+
});
223+
224+
it('adds chunk_text to requested fields when reranking', async () => {
225+
const testClient = stubPineconeClient(client);
226+
const denseRef = {} as SearchableIndex;
227+
const sparseRef = {} as SearchableIndex;
228+
testClient.ensureIndexes = async () => ({
229+
denseIndex: denseRef,
230+
sparseIndex: sparseRef,
231+
});
232+
let fieldsPassed: string[] | undefined;
233+
testClient.searchIndex = async (_index, _q, _tk, _ns, _f, opts) => {
234+
fieldsPassed = opts?.fields;
235+
return [];
236+
};
237+
238+
await client.query({
239+
query: 'q',
240+
namespace: 'n',
241+
topK: 5,
242+
useReranking: true,
243+
fields: ['title', 'url'],
244+
});
245+
246+
expect(fieldsPassed).toBeDefined();
247+
expect(fieldsPassed).toContain('chunk_text');
248+
expect(fieldsPassed).toContain('title');
249+
});
250+
251+
it('uses rerankResults from pinecone/rerank when useReranking is true', async () => {
252+
const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue([
253+
{
254+
id: 'd1',
255+
content: 'from dense',
256+
score: 0.9,
257+
metadata: {},
258+
reranked: true,
259+
},
260+
]);
261+
try {
262+
const testClient = stubPineconeClient(client);
263+
const denseRef = {} as SearchableIndex;
264+
const sparseRef = {} as SearchableIndex;
265+
testClient.ensureIndexes = async () => ({
266+
denseIndex: denseRef,
267+
sparseIndex: sparseRef,
268+
});
269+
testClient.searchIndex = async (index) => {
270+
if (index === denseRef) {
271+
return [{ _id: 'd1', _score: 0.9, fields: { chunk_text: 'from dense' } }];
272+
}
273+
return [];
274+
};
275+
276+
const results = await client.query({
277+
query: 'q',
278+
namespace: 'n',
279+
topK: 5,
280+
useReranking: true,
281+
});
282+
283+
expect(results).toHaveLength(1);
284+
expect(results[0].reranked).toBe(true);
285+
expect(results[0].content).toBe('from dense');
286+
expect(spy).toHaveBeenCalled();
287+
} finally {
288+
spy.mockRestore();
289+
}
290+
});
291+
292+
it('dedupes hits with blank _id via synthetic keys', async () => {
293+
const testClient = stubPineconeClient(client);
294+
const denseRef = {} as SearchableIndex;
295+
const sparseRef = {} as SearchableIndex;
296+
testClient.ensureIndexes = async () => ({
297+
denseIndex: denseRef,
298+
sparseIndex: sparseRef,
299+
});
300+
testClient.searchIndex = async (index) => {
301+
if (index === denseRef) {
302+
return [
303+
{ _id: ' ', _score: 1, fields: { chunk_text: 'a' } },
304+
{ _id: '', _score: 0.5, fields: { chunk_text: 'b' } },
305+
];
306+
}
307+
return [];
308+
};
309+
310+
const results = await client.query({
311+
query: 'q',
312+
namespace: 'n',
313+
topK: 10,
314+
useReranking: false,
315+
});
316+
317+
expect(results.length).toBe(2);
318+
});
319+
});
320+
321+
describe('keywordSearch', () => {
322+
it('throws for empty query', async () => {
323+
await expect(client.keywordSearch({ query: ' ', namespace: 'n' })).rejects.toThrow(
324+
'Query cannot be empty'
325+
);
326+
});
327+
328+
it('searches sparse index only and maps hits', async () => {
329+
const testClient = stubPineconeClient(client);
330+
const denseRef = {} as SearchableIndex;
331+
const sparseRef = {} as SearchableIndex;
332+
testClient.ensureIndexes = async () => ({
333+
denseIndex: denseRef,
334+
sparseIndex: sparseRef,
335+
});
336+
testClient.searchIndex = async (index) => {
337+
if (index === sparseRef) {
338+
return [{ _id: 'k1', _score: 0.7, fields: { chunk_text: 'lexical', tag: 'x' } }];
339+
}
340+
return [];
341+
};
342+
343+
const results = await client.keywordSearch({
344+
query: 'find me',
345+
namespace: 'ns',
346+
topK: 3,
347+
});
348+
349+
expect(results).toHaveLength(1);
350+
expect(results[0].content).toBe('lexical');
351+
expect(results[0].metadata['tag']).toBe('x');
352+
});
181353
});
182354
});

0 commit comments

Comments
 (0)