Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion benchmarks/latency.ts
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,11 @@ function createBenchPineconeMock(): PineconeClient {

return {
async query() {
return mockQueryResults;
return {
results: mockQueryResults,
degraded: false,
hybrid_leg_failed: null,
};
},
async count() {
return { count: 42, truncated: false };
Expand Down
10 changes: 6 additions & 4 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,12 @@
* `process.env` directly anymore — they receive their slice of the config.
*/

import { DEFAULT_INDEX_NAME, DEFAULT_RERANK_MODEL, FLOW_CACHE_TTL_MS } from './constants.js';
import {
DEFAULT_INDEX_NAME,
DEFAULT_RERANK_MODEL,
DEFAULT_TOP_K,
FLOW_CACHE_TTL_MS,
} from './constants.js';

/** Allowed log levels, in ascending severity. */
export type LogLevel = 'DEBUG' | 'INFO' | 'WARN' | 'ERROR';
Expand Down Expand Up @@ -49,9 +54,6 @@ export interface ServerConfig {
/** Default per-call timeout for Pinecone requests, in milliseconds. */
export const DEFAULT_REQUEST_TIMEOUT_MS = 15_000;

/** Default top-k mirrors constants.DEFAULT_TOP_K but is duplicated here to avoid a cycle. */
const DEFAULT_TOP_K = 10;

function asLogLevel(value: string | undefined, fallback: LogLevel): LogLevel {
const allowed: LogLevel[] = ['DEBUG', 'INFO', 'WARN', 'ERROR'];
return allowed.includes(value as LogLevel) ? (value as LogLevel) : fallback;
Expand Down
143 changes: 104 additions & 39 deletions src/pinecone-client.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { PineconeClient } from './pinecone-client.js';
import { resolveConfig } from './config.js';
import type { SearchableIndex, PineconeHit } from './types.js';
import * as rerankModule from './pinecone/rerank.js';

Expand Down Expand Up @@ -31,26 +32,40 @@ describe('PineconeClient', () => {
});
});

afterEach(() => {
delete process.env['PINECONE_INDEX_NAME'];
delete process.env['PINECONE_RERANK_MODEL'];
delete process.env['PINECONE_TOP_K'];
});

describe('constructor', () => {
it('should initialize with provided config', () => {
expect(client).toBeDefined();
});

it('should use environment variables as fallbacks', () => {
process.env['PINECONE_INDEX_NAME'] = 'env-index';
process.env['PINECONE_RERANK_MODEL'] = 'env-model';

const envClient = new PineconeClient({
apiKey: 'test-api-key',
});

expect(envClient).toBeDefined();
it('honors resolveConfig overrides without PINECONE_* env on the client path', () => {
const prevIndex = process.env['PINECONE_INDEX_NAME'];
const prevModel = process.env['PINECONE_RERANK_MODEL'];
const prevTopK = process.env['PINECONE_TOP_K'];
delete process.env['PINECONE_INDEX_NAME'];
delete process.env['PINECONE_RERANK_MODEL'];
delete process.env['PINECONE_TOP_K'];
try {
const resolved = resolveConfig({
apiKey: 'test-api-key',
indexName: 'resolved-index',
rerankModel: 'resolved-model',
defaultTopK: 42,
});
const c = new PineconeClient({
apiKey: resolved.apiKey,
indexName: resolved.indexName,
rerankModel: resolved.rerankModel,
defaultTopK: resolved.defaultTopK,
});
expect(c.getSparseIndexName()).toBe('resolved-index-sparse');
} finally {
if (prevIndex !== undefined) process.env['PINECONE_INDEX_NAME'] = prevIndex;
else delete process.env['PINECONE_INDEX_NAME'];
if (prevModel !== undefined) process.env['PINECONE_RERANK_MODEL'] = prevModel;
else delete process.env['PINECONE_RERANK_MODEL'];
if (prevTopK !== undefined) process.env['PINECONE_TOP_K'] = prevTopK;
else delete process.env['PINECONE_TOP_K'];
}
});
});

Expand All @@ -76,16 +91,16 @@ describe('PineconeClient', () => {

it('should continue hybrid search when one index fails', async () => {
const testClient = stubPineconeClient(client);
const denseRef = {} as SearchableIndex;
const sparseRef = {} as SearchableIndex;

testClient.ensureIndexes = async () => ({
denseIndex: {} as SearchableIndex,
sparseIndex: {} as SearchableIndex,
denseIndex: denseRef,
sparseIndex: sparseRef,
});

let searchCall = 0;
testClient.searchIndex = async () => {
searchCall += 1;
if (searchCall === 1) {
testClient.searchIndex = async (index) => {
if (index === denseRef) {
throw new Error('dense failure');
}
return [
Expand All @@ -97,16 +112,18 @@ describe('PineconeClient', () => {
];
};

const results = await client.query({
const out = await client.query({
query: 'hybrid search',
namespace: 'test',
topK: 5,
useReranking: false,
});

expect(results).toHaveLength(1);
expect(results[0].content).toBe('hybrid content');
expect(results[0].metadata.author).toBe('tester');
expect(out.results).toHaveLength(1);
expect(out.results[0]?.content).toBe('hybrid content');
expect(out.results[0]?.metadata.author).toBe('tester');
expect(out.hybrid_leg_failed).toBe('dense');
expect(out.degraded).toBe(false);
});

it('should throw when both dense and sparse searches fail', async () => {
Expand Down Expand Up @@ -249,15 +266,18 @@ describe('PineconeClient', () => {
});

it('uses rerankResults from pinecone/rerank when useReranking is true', async () => {
const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue([
{
id: 'd1',
content: 'from dense',
score: 0.9,
metadata: {},
reranked: true,
},
]);
const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue({
results: [
{
id: 'd1',
content: 'from dense',
score: 0.9,
metadata: {},
reranked: true,
},
],
degraded: false,
});
try {
const testClient = stubPineconeClient(client);
const denseRef = {} as SearchableIndex;
Expand All @@ -280,9 +300,54 @@ describe('PineconeClient', () => {
useReranking: true,
});

expect(results).toHaveLength(1);
expect(results[0].reranked).toBe(true);
expect(results[0].content).toBe('from dense');
expect(results.results).toHaveLength(1);
expect(results.results[0]?.reranked).toBe(true);
expect(results.results[0]?.content).toBe('from dense');
expect(spy).toHaveBeenCalled();
} finally {
spy.mockRestore();
}
});

it('propagates rerank degradation to hybrid query outcome', async () => {
const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue({
results: [
{
id: 'd1',
content: 'from dense',
score: 0.9,
metadata: {},
reranked: false,
},
],
degraded: true,
degradation_reason: 'rerank_failed: timeout',
});
try {
const testClient = stubPineconeClient(client);
const denseRef = {} as SearchableIndex;
const sparseRef = {} as SearchableIndex;
testClient.ensureIndexes = async () => ({
denseIndex: denseRef,
sparseIndex: sparseRef,
});
testClient.searchIndex = async (index) => {
if (index === denseRef) {
return [{ _id: 'd1', _score: 0.9, fields: { chunk_text: 'from dense' } }];
}
return [];
};

const out = await client.query({
query: 'q',
namespace: 'n',
topK: 5,
useReranking: true,
});

expect(out.degraded).toBe(true);
expect(out.degradation_reason).toBe('rerank_failed: timeout');
expect(out.results[0]?.reranked).toBe(false);
expect(spy).toHaveBeenCalled();
} finally {
spy.mockRestore();
Expand Down Expand Up @@ -314,7 +379,7 @@ describe('PineconeClient', () => {
useReranking: false,
});

expect(results.length).toBe(2);
expect(results.results.length).toBe(2);
});
});

Expand Down
49 changes: 38 additions & 11 deletions src/pinecone-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import type {
KeywordSearchParams,
KeywordIndexNamespacesResult,
SearchableIndex,
HybridQueryResult,
HybridLegFailed,
} from './types.js';
import {
DEFAULT_INDEX_NAME,
Expand All @@ -35,16 +37,16 @@ export class PineconeClient {
private defaultTopK: number;
private readonly indexSession: PineconeIndexSession;

/** Create a client with the given config; env vars override index name, rerank model, and top-k. */
/**
* Create a client from a resolved {@link PineconeClientConfig}.
* Index name, rerank model, and default top-k come only from this object (typically
* built via {@link resolveConfig} / CLI); this class does not read `process.env`.
*/
constructor(config: PineconeClientConfig) {
const indexName = config.indexName || process.env['PINECONE_INDEX_NAME'] || DEFAULT_INDEX_NAME;
const indexName = config.indexName ?? DEFAULT_INDEX_NAME;
this.indexSession = new PineconeIndexSession(config.apiKey, indexName);
this.rerankModel =
config.rerankModel || process.env['PINECONE_RERANK_MODEL'] || DEFAULT_RERANK_MODEL;
const envTopK = process.env['PINECONE_TOP_K'];
const parsedEnvTopK = envTopK !== undefined ? parseInt(envTopK, 10) : NaN;
this.defaultTopK =
config.defaultTopK ?? (Number.isFinite(parsedEnvTopK) ? parsedEnvTopK : DEFAULT_TOP_K);
this.rerankModel = config.rerankModel ?? DEFAULT_RERANK_MODEL;
this.defaultTopK = config.defaultTopK ?? DEFAULT_TOP_K;
}

/** Sparse index name `{indexName}-sparse` (keyword / hybrid sparse). */
Expand Down Expand Up @@ -105,7 +107,7 @@ export class PineconeClient {
return searchIndexImpl(index, query, topK, namespace, metadataFilter, options);
}

async query(params: QueryParams): Promise<SearchResult[]> {
async query(params: QueryParams): Promise<HybridQueryResult> {
const {
query,
topK: requestedTopK,
Expand Down Expand Up @@ -148,17 +150,37 @@ export class PineconeClient {
throw new Error('Hybrid search failed: both dense and sparse index searches failed.');
}

let hybridLegFailed: HybridLegFailed = null;
if (
denseResult.status === 'rejected' &&
sparseResult.status === 'fulfilled' &&
sparseHits.length > 0
) {
hybridLegFailed = 'dense';
} else if (
sparseResult.status === 'rejected' &&
denseResult.status === 'fulfilled' &&
denseHits.length > 0
) {
hybridLegFailed = 'sparse';
}
Comment thread
jonathanMLDev marked this conversation as resolved.

const mergedResults = mergeResults(denseHits, sparseHits);

let degraded = false;
let degradation_reason: string | undefined;
let documents: SearchResult[];
if (useReranking) {
documents = await rerankResultsImpl(
const rerankOut = await rerankResultsImpl(
this.indexSession.ensureClient(),
this.rerankModel,
query,
mergedResults,
topK
);
documents = rerankOut.results;
degraded = rerankOut.degraded;
degradation_reason = rerankOut.degradation_reason;
} else {
documents = sliceMergedHitsToSearchResults(mergedResults, topK);
}
Expand All @@ -167,7 +189,12 @@ export class PineconeClient {
`Retrieved ${documents.length} documents from hybrid search (dense: ${denseHits.length}, sparse: ${sparseHits.length})`
);

return documents;
return {
results: documents,
degraded,
...(degradation_reason !== undefined ? { degradation_reason } : {}),
hybrid_leg_failed: hybridLegFailed,
};
}

async keywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
Expand Down
26 changes: 15 additions & 11 deletions src/pinecone/rerank.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ const sampleMerged: MergedHit[] = [
];

describe('rerankResults', () => {
it('returns empty array when there are no merged hits', async () => {
it('returns empty outcome when there are no merged hits', async () => {
const pc = {} as Parameters<typeof rerankResults>[0];
const out = await rerankResults(pc, 'any-model', 'q', [], 5);
expect(out).toEqual([]);
expect(out.results).toEqual([]);
expect(out.degraded).toBe(false);
});

it('maps successful inference.rerank response', async () => {
Expand All @@ -26,21 +27,24 @@ describe('rerankResults', () => {

const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5);

expect(out).toHaveLength(1);
expect(out[0]?.reranked).toBe(true);
expect(out[0]?.id).toBe('1');
expect(out[0]?.content).toBe('hello');
expect(out[0]?.score).toBeCloseTo(0.99);
expect(out.results).toHaveLength(1);
expect(out.degraded).toBe(false);
expect(out.results[0]?.reranked).toBe(true);
expect(out.results[0]?.id).toBe('1');
expect(out.results[0]?.content).toBe('hello');
expect(out.results[0]?.score).toBeCloseTo(0.99);
});

it('returns unreranked slice when rerank throws', async () => {
it('returns unreranked slice with degraded when rerank throws', async () => {
const rerank = vi.fn().mockRejectedValue(new Error('rerank unavailable'));
const pc = { inference: { rerank } } as Parameters<typeof rerankResults>[0];

const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5);

expect(out).toHaveLength(1);
expect(out[0]?.reranked).toBe(false);
expect(out[0]?.content).toBe('hello');
expect(out.results).toHaveLength(1);
expect(out.degraded).toBe(true);
expect(out.degradation_reason).toMatch(/^rerank_failed:/);
expect(out.results[0]?.reranked).toBe(false);
expect(out.results[0]?.content).toBe('hello');
});
});
Loading
Loading