diff --git a/src/pinecone-client.test.ts b/src/pinecone-client.test.ts index 559e43b..98f1523 100644 --- a/src/pinecone-client.test.ts +++ b/src/pinecone-client.test.ts @@ -1,9 +1,10 @@ -import { describe, it, expect, beforeEach } from 'vitest'; +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; import { PineconeClient } from './pinecone-client.js'; import type { SearchableIndex, PineconeHit } from './types.js'; +import * as rerankModule from './pinecone/rerank.js'; -/** Test double: client with stubbable ensureIndexes and searchIndex for hybrid tests */ -type PineconeClientTestDouble = PineconeClient & { +/** Stubs for private methods (assigned at runtime; avoid intersecting private `PineconeClient` members). */ +type PineconeClientMethodStubs = { ensureIndexes: () => Promise<{ denseIndex: SearchableIndex; sparseIndex: SearchableIndex }>; searchIndex: ( index: SearchableIndex, @@ -15,6 +16,10 @@ type PineconeClientTestDouble = PineconeClient & { ) => Promise; }; +function stubPineconeClient(client: PineconeClient): PineconeClientMethodStubs { + return client as unknown as PineconeClientMethodStubs; +} + describe('PineconeClient', () => { let client: PineconeClient; @@ -26,6 +31,12 @@ 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(); @@ -64,7 +75,7 @@ describe('PineconeClient', () => { }); it('should continue hybrid search when one index fails', async () => { - const testClient = client as PineconeClientTestDouble; + const testClient = stubPineconeClient(client); testClient.ensureIndexes = async () => ({ denseIndex: {} as SearchableIndex, @@ -99,7 +110,7 @@ describe('PineconeClient', () => { }); it('should throw when both dense and sparse searches fail', async () => { - const testClient = client as PineconeClientTestDouble; + const testClient = stubPineconeClient(client); testClient.ensureIndexes = async () => ({ denseIndex: {} as SearchableIndex, @@ -122,7 +133,7 @@ describe('PineconeClient', () => { describe('count', () => { it('should return unique document count using semantic search only with minimal fields', async () => { - const testClient = client as PineconeClientTestDouble; + const testClient = stubPineconeClient(client); testClient.ensureIndexes = async () => ({ denseIndex: {} as SearchableIndex, sparseIndex: {} as SearchableIndex, @@ -161,7 +172,7 @@ describe('PineconeClient', () => { }); it('should set truncated when hit limit is reached', async () => { - const testClient = client as PineconeClientTestDouble; + const testClient = stubPineconeClient(client); testClient.ensureIndexes = async () => ({ denseIndex: {} as SearchableIndex, sparseIndex: {} as SearchableIndex, @@ -178,5 +189,166 @@ describe('PineconeClient', () => { expect(result.count).toBe(10000); expect(result.truncated).toBe(true); }); + + it('falls back to chunk _id when no document identifier fields exist', async () => { + const testClient = stubPineconeClient(client); + testClient.ensureIndexes = async () => ({ + denseIndex: {} as SearchableIndex, + sparseIndex: {} as SearchableIndex, + }); + testClient.searchIndex = async () => [ + { _id: 'chunk-only', _score: 1, fields: { chunk_text: 'x' } }, + ]; + + const result = await client.count({ query: 'paper', namespace: 'ns' }); + + expect(result.count).toBe(1); + expect(result.truncated).toBe(false); + }); + }); + + describe('getSparseIndexName', () => { + it('returns {indexName}-sparse derived from config indexName', () => { + const c = new PineconeClient({ apiKey: 'k', indexName: 'my' }); + expect(c.getSparseIndexName()).toBe('my-sparse'); + }); + }); + + describe('query (rerank and fields)', () => { + it('rejects non-finite topK', async () => { + await expect(client.query({ query: 'q', namespace: 'n', topK: Number.NaN })).rejects.toThrow( + 'topK must be a finite number' + ); + }); + + it('adds chunk_text to requested fields when reranking', async () => { + const testClient = stubPineconeClient(client); + const denseRef = {} as SearchableIndex; + const sparseRef = {} as SearchableIndex; + testClient.ensureIndexes = async () => ({ + denseIndex: denseRef, + sparseIndex: sparseRef, + }); + let fieldsPassed: string[] | undefined; + testClient.searchIndex = async (_index, _q, _tk, _ns, _f, opts) => { + fieldsPassed = opts?.fields; + return []; + }; + + await client.query({ + query: 'q', + namespace: 'n', + topK: 5, + useReranking: true, + fields: ['title', 'url'], + }); + + expect(fieldsPassed).toBeDefined(); + expect(fieldsPassed).toContain('chunk_text'); + expect(fieldsPassed).toContain('title'); + }); + + 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, + }, + ]); + 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 results = await client.query({ + query: 'q', + namespace: 'n', + topK: 5, + useReranking: true, + }); + + expect(results).toHaveLength(1); + expect(results[0].reranked).toBe(true); + expect(results[0].content).toBe('from dense'); + expect(spy).toHaveBeenCalled(); + } finally { + spy.mockRestore(); + } + }); + + it('dedupes hits with blank _id via synthetic keys', async () => { + 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: ' ', _score: 1, fields: { chunk_text: 'a' } }, + { _id: '', _score: 0.5, fields: { chunk_text: 'b' } }, + ]; + } + return []; + }; + + const results = await client.query({ + query: 'q', + namespace: 'n', + topK: 10, + useReranking: false, + }); + + expect(results.length).toBe(2); + }); + }); + + describe('keywordSearch', () => { + it('throws for empty query', async () => { + await expect(client.keywordSearch({ query: ' ', namespace: 'n' })).rejects.toThrow( + 'Query cannot be empty' + ); + }); + + it('searches sparse index only and maps hits', async () => { + 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 === sparseRef) { + return [{ _id: 'k1', _score: 0.7, fields: { chunk_text: 'lexical', tag: 'x' } }]; + } + return []; + }; + + const results = await client.keywordSearch({ + query: 'find me', + namespace: 'ns', + topK: 3, + }); + + expect(results).toHaveLength(1); + expect(results[0].content).toBe('lexical'); + expect(results[0].metadata['tag']).toBe('x'); + }); }); }); diff --git a/src/pinecone-client.ts b/src/pinecone-client.ts index 14ff59e..37c5f0b 100644 --- a/src/pinecone-client.ts +++ b/src/pinecone-client.ts @@ -1,20 +1,6 @@ -/** - * Pinecone client for hybrid search retrieval. - * - * Performs hybrid search (dense + sparse) with optional reranking. Designed - * for high performance with connection pooling, lazy initialization, and - * bounded retry + timeout around every Pinecone call so a flaky 5xx does - * not bubble up immediately to the caller. - */ - -import { Pinecone } from '@pinecone-database/pinecone'; -import { - debug as logDebug, - error as logError, - info as logInfo, - warn as logWarn, -} from './logger.js'; -import { withRetry, withTimeout } from './server/retry.js'; +/** Hybrid dense+sparse query client with optional reranking (facade over `src/pinecone/*`). */ + +import { error as logError, info as logInfo } from './logger.js'; import type { PineconeClientConfig, SearchResult, @@ -23,11 +9,8 @@ import type { CountParams, CountResult, KeywordSearchParams, - MergedHit, - NamespaceHandle, - SearchableIndex, - PineconeMetadataValue, KeywordIndexNamespacesResult, + SearchableIndex, } from './types.js'; import { DEFAULT_INDEX_NAME, @@ -37,146 +20,38 @@ import { COUNT_TOP_K, COUNT_FIELDS, } from './constants.js'; +import { PineconeIndexSession } from './pinecone/indexes.js'; +import { + countUniqueDocumentsFromHits, + mapSparseHitsToSearchResults, + mergeResults, + searchIndex as searchIndexImpl, + sliceMergedHitsToSearchResults, +} from './pinecone/search.js'; +import { rerankResults as rerankResultsImpl } from './pinecone/rerank.js'; -/** - * Infers a human-readable metadata field type for namespace discovery. - * Distinguishes Pinecone-supported list type (string[]) from other arrays. - */ -function inferMetadataFieldType(value: unknown): string { - if (value === null || value === undefined) { - return 'unknown'; - } - if (Array.isArray(value)) { - if (value.length === 0) return 'array'; - if (value.every((item) => typeof item === 'string')) return 'string[]'; - return 'array'; - } - const t = typeof value; - if (t === 'string' || t === 'number' || t === 'boolean') return t; - return 'object'; -} - -/** Extract chunk_text and metadata from a Pinecone hit (single mapping path for merge + keyword + rerank). */ -function extractHitContent(hit: PineconeHit): { - content: string; - metadata: Record; -} { - const fields = hit.fields || {}; - let content = ''; - const metadata: Record = {}; - for (const [key, value] of Object.entries(fields)) { - if (key === 'chunk_text') { - content = typeof value === 'string' ? value : ''; - } else { - metadata[key] = value as PineconeMetadataValue; - } - } - return { content, metadata }; -} - -function pineconeHitToSearchResult(hit: PineconeHit, reranked: boolean): SearchResult { - const { content, metadata } = extractHitContent(hit); - return { - id: hit._id || '', - content, - score: hit._score || 0, - metadata, - reranked, - }; -} - -function mergedHitToSearchResult(result: MergedHit, reranked: boolean): SearchResult { - return { - id: result._id || '', - content: result.chunk_text || '', - score: result._score || 0, - metadata: result.metadata || {}, - reranked, - }; -} - -/** - * Hybrid Pinecone retrieval client. - * - * The constructor accepts the same `PineconeClientConfig` it always has, - * plus optional `sparseIndexName` / `requestTimeoutMs` knobs that flow in - * from the unified `ServerConfig`. - */ export class PineconeClient { - private apiKey: string; - private indexName: string; - private sparseIndexName: string; private rerankModel: string; private defaultTopK: number; - private requestTimeoutMs: number; - - // Lazy initialization - private pc: Pinecone | null = null; - private denseIndex: SearchableIndex | null = null; - private sparseIndex: SearchableIndex | null = null; - private initialized = false; + private readonly indexSession: PineconeIndexSession; - /** Create a client. `sparseIndexName` defaults to `${indexName}-sparse` for backwards compatibility. */ + /** Create a client with the given config; env vars override index name, rerank model, and top-k. */ constructor(config: PineconeClientConfig) { - this.apiKey = config.apiKey; - this.indexName = config.indexName || DEFAULT_INDEX_NAME; - this.sparseIndexName = config.sparseIndexName || `${this.indexName}-sparse`; - this.rerankModel = config.rerankModel || DEFAULT_RERANK_MODEL; - this.defaultTopK = config.defaultTopK ?? DEFAULT_TOP_K; - this.requestTimeoutMs = config.requestTimeoutMs ?? 15_000; + const indexName = config.indexName || process.env['PINECONE_INDEX_NAME'] || 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); } - /** Returns the configured sparse index name (used for hybrid sparse and keyword_search). */ + /** Sparse index name `{indexName}-sparse` (keyword / hybrid sparse). */ getSparseIndexName(): string { - return this.sparseIndexName; - } - - /** Returns the configured dense (hybrid) index name. */ - getIndexName(): string { - return this.indexName; - } - - /** - * Verify both dense and sparse indexes exist by issuing a `describeIndexStats` - * call against each. Used by the `--check-indexes` startup probe. - * - * @returns ok=true on success; ok=false with `errors[]` listing the indexes - * that failed (and why), so operators can fix configuration without - * waiting for the first user query. - */ - async checkIndexes(): Promise<{ ok: boolean; errors: string[] }> { - const errors: string[] = []; - try { - const { denseIndex, sparseIndex } = await this.ensureIndexes(); - try { - if (typeof denseIndex.describeIndexStats === 'function') { - await this.runWithRetryTimeout((signal) => { - void signal; - return denseIndex.describeIndexStats!(); - }, `describeIndexStats(dense:${this.indexName})`); - } - } catch (e) { - errors.push(`dense index "${this.indexName}": ${(e as Error).message}`); - } - try { - if (typeof sparseIndex.describeIndexStats === 'function') { - await this.runWithRetryTimeout((signal) => { - void signal; - return sparseIndex.describeIndexStats!(); - }, `describeIndexStats(sparse:${this.sparseIndexName})`); - } - } catch (e) { - errors.push(`sparse index "${this.sparseIndexName}": ${(e as Error).message}`); - } - } catch (e) { - errors.push(`Failed to initialize Pinecone client: ${(e as Error).message}`); - } - return { ok: errors.length === 0, errors }; + return this.indexSession.getSparseIndexName(); } - /** - * Normalize and clamp topK from request (validates >= 1, caps at MAX_TOP_K). - */ private clampTopK(requested: number | undefined): number { if (requested !== undefined && !Number.isFinite(requested)) { throw new Error('topK must be a finite number >= 1'); @@ -191,103 +66,19 @@ export class PineconeClient { return topK; } - /** - * Ensure Pinecone client is initialized - */ - private ensureClient(): Pinecone { - if (!this.pc) { - if (!this.apiKey) { - throw new Error( - 'Pinecone API key is required. Set PINECONE_API_KEY environment variable or pass apiKey parameter.' - ); - } - this.pc = new Pinecone({ apiKey: this.apiKey }); - logInfo('Pinecone client initialized'); - } - return this.pc; - } - - /** - * Ensure Pinecone indexes are initialized and return them - */ private async ensureIndexes(): Promise<{ denseIndex: SearchableIndex; sparseIndex: SearchableIndex; }> { - if (this.initialized && this.denseIndex !== null && this.sparseIndex !== null) { - return { denseIndex: this.denseIndex, sparseIndex: this.sparseIndex }; - } - - const pc = this.ensureClient(); - const denseName = this.indexName; - const sparseName = this.sparseIndexName; - - const dense = pc.index(denseName) as unknown as SearchableIndex; - const sparse = pc.index(sparseName) as unknown as SearchableIndex; - this.denseIndex = dense; - this.sparseIndex = sparse; - this.initialized = true; - - logInfo(`Connected to indexes: ${denseName} and ${sparseName}`); - return { denseIndex: dense, sparseIndex: sparse }; - } - - /** - * Wrap a Pinecone call with the configured request timeout and a small - * bounded retry. All outbound network calls go through this helper. - * The callback receives an {@link AbortSignal} that is aborted when the - * per-call timeout fires (for cooperative cancellation when stacks support it). - */ - private async runWithRetryTimeout( - operation: (signal: AbortSignal) => Promise, - label: string - ): Promise { - return withRetry(() => withTimeout(operation, { timeoutMs: this.requestTimeoutMs, label }), { - retries: 2, - backoffMs: 250, - onRetry: (attempt: number, err: unknown) => - logWarn( - `${label}: retry ${attempt} after error: ${ - err instanceof Error ? err.message : String(err) - }` - ), - }); + return this.indexSession.ensureIndexes(); } - /** - * List namespaces present on the sparse index (same index used for hybrid sparse and keyword_search). - * Use this to choose a namespace for sparse-only queries instead of the dense index list. - */ + /** Namespaces on the sparse (keyword) index with record counts. */ async listNamespacesFromKeywordIndex(): Promise { - try { - const { sparseIndex } = await this.ensureIndexes(); - const stats = sparseIndex.describeIndexStats - ? await this.runWithRetryTimeout((signal) => { - void signal; - return sparseIndex.describeIndexStats!(); - }, 'describeIndexStats(sparse)') - : undefined; - const namespaces = stats?.namespaces ?? {}; - return { - ok: true, - namespaces: Object.entries(namespaces).map(([namespace, info]) => ({ - namespace, - recordCount: info?.recordCount ?? 0, - })), - }; - } catch (error) { - const msg = error instanceof Error ? error.message : String(error); - logError('Error listing namespaces from keyword index', error); - return { ok: false, error: msg }; - } + return this.indexSession.listNamespacesFromKeywordIndex(); } - /** - * List all available namespaces with their metadata information - * - * Fetches namespaces from the index stats and samples records to discover - * available metadata fields and their types. - */ + /** Dense index namespaces with sampled metadata field types. */ async listNamespacesWithMetadata(): Promise< Array<{ namespace: string; @@ -295,98 +86,14 @@ export class PineconeClient { metadata: Record; }> > { - try { - const { denseIndex } = await this.ensureIndexes(); - - // Get index stats to find namespaces - const stats = denseIndex.describeIndexStats - ? await this.runWithRetryTimeout((signal) => { - void signal; - return denseIndex.describeIndexStats!(); - }, 'describeIndexStats(dense)') - : undefined; - const namespaces = stats?.namespaces ? Object.keys(stats.namespaces) : []; - - logInfo(`Found ${namespaces.length} namespace(s)`); - - // Get metadata info for each namespace by sampling records - const namespacesInfo = await Promise.all( - namespaces.map(async (ns: string) => { - try { - const recordCount = stats?.namespaces?.[ns]?.recordCount || 0; - const metadataFields: Record = {}; - - // Sample a few records to discover metadata fields - if (recordCount > 0 && denseIndex.namespace) { - try { - const nsObj: NamespaceHandle = denseIndex.namespace(ns); - const dim = stats?.dimension; - if (dim === undefined || !Number.isFinite(dim) || dim <= 0) { - logWarn( - `Skipping metadata sampling for namespace "${ns}": index stats did not report a finite vector dimension` - ); - } else if (typeof nsObj.query === 'function') { - const sampleQuery = await this.runWithRetryTimeout((signal) => { - void signal; - return nsObj.query!({ - topK: 5, - vector: Array(dim).fill(0), - includeMetadata: true, - }); - }, `namespace.query(${ns})`); - - // Collect unique metadata fields and infer types (including string[]) - if (sampleQuery?.matches) { - sampleQuery.matches.forEach((match: { metadata?: Record }) => { - if (match.metadata) { - Object.entries(match.metadata).forEach(([key, value]) => { - const inferredType = inferMetadataFieldType(value); - if (!(key in metadataFields)) { - metadataFields[key] = inferredType; - } else if ( - (metadataFields[key] === 'object' || metadataFields[key] === 'array') && - inferredType === 'string[]' - ) { - // Prefer array type over generic object when we see it in another sample - metadataFields[key] = inferredType; - } - }); - } - }); - } - } - } catch (queryError) { - logError(`Error sampling records for namespace ${ns}`, queryError); - } - } - - return { - namespace: ns, - recordCount, - metadata: metadataFields, - }; - } catch (error) { - logError(`Error processing namespace ${ns}`, error); - return { - namespace: ns, - recordCount: 0, - metadata: {}, - }; - } - }) - ); + return this.indexSession.listNamespacesWithMetadata(); + } - return namespacesInfo; - } catch (error) { - logError('Error listing namespaces', error); - return []; - } + /** Probe dense + sparse indexes (describeIndexStats) for startup checks. */ + async checkIndexes(): Promise<{ ok: boolean; errors: string[] }> { + return this.indexSession.checkIndexes(); } - /** - * Search a Pinecone index using text query with optional metadata filtering. - * When options.fields is set, only those fields are requested (e.g. for count: no chunk_text). - */ private async searchIndex( index: SearchableIndex, query: string, @@ -395,167 +102,9 @@ export class PineconeClient { metadataFilter?: Record, options?: { fields?: string[] } ): Promise { - // Build query payload in the same shape as Python implementation. - const queryPayload: Record = { - top_k: topK, - inputs: { text: query }, - }; - - // Include filter when explicitly provided (matches Python behavior). - if (metadataFilter !== undefined) { - queryPayload['filter'] = metadataFilter; - logDebug('Applying metadata filter', metadataFilter); - } - - try { - // Preferred path: Pinecone search API. - if (typeof index.search === 'function') { - const searchOpts: { - namespace?: string; - query: Record; - fields?: string[]; - } = { - namespace, - query: queryPayload, - }; - if (options?.fields?.length) { - searchOpts.fields = options.fields; - } - const result = await this.runWithRetryTimeout( - (signal) => { - void signal; - return index.search!(searchOpts); - }, - `index.search(ns=${namespace ?? 'default'})` - ); - return result?.result?.hits || []; - } - - // Backward-compatible fallback for older API shapes. - const target = namespace && index.namespace ? index.namespace(namespace) : index; - const queryParams: { query: Record; fields?: string[] } = { - query: { - topK, - inputs: { text: query }, - }, - }; - if (metadataFilter !== undefined) { - queryParams.query['filter'] = metadataFilter; - } - if (options?.fields?.length) { - queryParams.fields = options.fields; - } - const result = target.searchRecords - ? await this.runWithRetryTimeout( - (signal) => { - void signal; - return target.searchRecords!(queryParams); - }, - `index.searchRecords(ns=${namespace ?? 'default'})` - ) - : { result: { hits: [] as PineconeHit[] } }; - return result?.result?.hits || []; - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error); - throw new Error( - `Pinecone search failed for namespace "${namespace ?? 'default'}": ${errorMessage}` - ); - } + return searchIndexImpl(index, query, topK, namespace, metadataFilter, options); } - /** - * Merge and deduplicate results from dense and sparse searches - * - * Uses the higher score when duplicates are found. - */ - private mergeResults(denseHits: PineconeHit[], sparseHits: PineconeHit[]): MergedHit[] { - const deduped: Record = {}; - let syntheticKeySeq = 0; - - for (const hit of [...denseHits, ...sparseHits]) { - const rawId = hit._id?.trim(); - const hitId = rawId && rawId.length > 0 ? rawId : `__missing_${syntheticKeySeq++}`; - if (!rawId || rawId.length === 0) { - logWarn('mergeResults: hit missing or empty _id; using synthetic dedup key'); - } - const hitScore = hit._score || 0; - - const existing = deduped[hitId]; - if (existing !== undefined && (existing._score || 0) >= hitScore) { - continue; - } - - const { content, metadata: hitMetadata } = extractHitContent(hit); - - deduped[hitId] = { - _id: hitId, - _score: hitScore, - chunk_text: content, - metadata: hitMetadata, - }; - } - - return Object.values(deduped).sort((a, b) => (b._score || 0) - (a._score || 0)); - } - - /** - * Rerank results using Pinecone's reranking model - */ - private async rerankResults( - query: string, - results: MergedHit[], - topN: number - ): Promise { - if (!results || results.length === 0) { - return []; - } - - const pc = this.ensureClient(); - - try { - // The Pinecone SDK types constrain document values to `Record`, - // but the underlying HTTP API accepts any JSON value. We pass MergedHit objects - // (metadata may contain number/boolean/string[]) and only `chunk_text` — which is - // always a string — is accessed via rankFields. The double cast via `as unknown` - // is intentional: it bypasses the SDK's over-narrow type without stringifying - // metadata values that we need to read back from the returned documents. - const rerankDocs = results as unknown as Array>; - const rerankResult = await this.runWithRetryTimeout((signal) => { - void signal; - return pc.inference.rerank({ - model: this.rerankModel, - query, - documents: rerankDocs, - topN, - rankFields: ['chunk_text'], - returnDocuments: true, - parameters: { truncate: 'END' }, - }); - }, `inference.rerank(${this.rerankModel})`); - - const reranked: SearchResult[] = []; - for (const item of rerankResult.data || []) { - const document = (item.document || {}) as MergedHit; - const row = mergedHitToSearchResult(document, true); - reranked.push({ - ...row, - score: parseFloat(String(item.score ?? row.score)), - }); - } - return reranked; - } catch (error) { - logError('Error reranking results', error); - // Fall back to returning unreranked results - return results.slice(0, topN).map((result) => mergedHitToSearchResult(result, false)); - } - } - - /** - * Query Pinecone indexes using hybrid search with optional reranking - * - * Performs parallel searches on dense and sparse indexes, merges results, - * and optionally reranks using the configured reranking model. - */ async query(params: QueryParams): Promise { const { query, @@ -566,25 +115,21 @@ export class PineconeClient { fields: requestedFields, } = params; - // Validate inputs if (!query || !query.trim()) { throw new Error('Query cannot be empty'); } const topK = this.clampTopK(requestedTopK); - // When reranking, Pinecone requires chunk_text in returned fields; add it if user specified fields without it const searchFields = requestedFields?.length && useReranking && !requestedFields.includes('chunk_text') ? [...requestedFields, 'chunk_text'] : requestedFields; - // Ensure indexes are ready const { denseIndex, sparseIndex } = await this.ensureIndexes(); const searchOptions = searchFields?.length ? { fields: searchFields } : undefined; - // Perform hybrid search const [denseResult, sparseResult] = await Promise.allSettled([ this.searchIndex(denseIndex, query, topK, namespace, metadataFilter, searchOptions), this.searchIndex(sparseIndex, query, topK, namespace, metadataFilter, searchOptions), @@ -603,17 +148,19 @@ export class PineconeClient { throw new Error('Hybrid search failed: both dense and sparse index searches failed.'); } - // Merge results - const mergedResults = this.mergeResults(denseHits, sparseHits); + const mergedResults = mergeResults(denseHits, sparseHits); - // Optionally rerank let documents: SearchResult[]; if (useReranking) { - documents = await this.rerankResults(query, mergedResults, topK); + documents = await rerankResultsImpl( + this.indexSession.ensureClient(), + this.rerankModel, + query, + mergedResults, + topK + ); } else { - documents = mergedResults - .slice(0, topK) - .map((result) => mergedHitToSearchResult(result, false)); + documents = sliceMergedHitsToSearchResults(mergedResults, topK); } logInfo( @@ -623,11 +170,6 @@ export class PineconeClient { return documents; } - /** - * Keyword (sparse-only) search against the dedicated sparse index. - * Performs lexical/keyword retrieval only—no dense index, no reranking. - * Use for exact or keyword-style queries on the configured sparse index. - */ async keywordSearch(params: KeywordSearchParams): Promise { const { query, @@ -655,7 +197,7 @@ export class PineconeClient { searchOptions ); - const documents: SearchResult[] = hits.map((hit) => pineconeHitToSearchResult(hit, false)); + const documents = mapSparseHitsToSearchResults(hits); logInfo( `Keyword search returned ${documents.length} results from ${this.getSparseIndexName()}` @@ -663,11 +205,6 @@ export class PineconeClient { return documents; } - /** - * Return the number of unique documents matching the query and optional metadata filter. - * Uses semantic search only (dense index), requests minimal fields (document_number, url, doc_id) - * to avoid transferring chunk content, and deduplicates by document for a document-level count. - */ async count(params: CountParams): Promise { if (!params.query || !params.query.trim()) { throw new Error('Query cannot be empty'); @@ -683,36 +220,6 @@ export class PineconeClient { { fields: [...COUNT_FIELDS] } ); - const docKeys = new Set(); - let idFallbackCount = 0; - for (const hit of hits) { - const fields = hit.fields || {}; - const docNumber = fields['document_number']; - const url = fields['url']; - const docId = fields['doc_id']; - const docKey = - (typeof docNumber === 'string' ? docNumber : undefined) ?? - (typeof url === 'string' ? url : undefined) ?? - (typeof docId === 'string' ? docId : undefined); - if (docKey !== undefined) { - docKeys.add(docKey); - } else { - // Fall back to chunk ID — this yields a chunk count, not a document count - idFallbackCount++; - docKeys.add(hit._id ?? ''); - } - } - if (idFallbackCount > 0) { - logWarn( - `count(): ${idFallbackCount} hit(s) in namespace "${params.namespace}" had none of the ` + - `identifier fields (${COUNT_FIELDS.join(', ')}); fell back to chunk ID — result may overcount documents` - ); - } - - const count = docKeys.size; - return { - count, - truncated: hits.length >= COUNT_TOP_K, - }; + return countUniqueDocumentsFromHits(hits, params.namespace); } } diff --git a/src/pinecone/indexes.test.ts b/src/pinecone/indexes.test.ts new file mode 100644 index 0000000..a70d93b --- /dev/null +++ b/src/pinecone/indexes.test.ts @@ -0,0 +1,174 @@ +import { describe, it, expect, vi } from 'vitest'; +import { PineconeIndexSession } from './indexes.js'; +import type { SearchableIndex } from '../types.js'; + +/** Subclass so tests inject index handles without calling the real Pinecone SDK. */ +class PineconeIndexSessionTestDouble extends PineconeIndexSession { + constructor(private readonly pair: { dense: SearchableIndex; sparse: SearchableIndex }) { + super('test-api-key', 'test-index'); + } + + override async ensureIndexes(): Promise<{ + denseIndex: SearchableIndex; + sparseIndex: SearchableIndex; + }> { + return { denseIndex: this.pair.dense, sparseIndex: this.pair.sparse }; + } +} + +class ThrowingEnsureSession extends PineconeIndexSession { + constructor() { + super('test-api-key', 'test-index'); + } + + override async ensureIndexes(): Promise<{ + denseIndex: SearchableIndex; + sparseIndex: SearchableIndex; + }> { + throw new Error('no client'); + } +} + +describe('PineconeIndexSession', () => { + describe('listNamespacesFromKeywordIndex', () => { + it('returns namespace rows when describeIndexStats succeeds', async () => { + const sparse = { + describeIndexStats: vi.fn().mockResolvedValue({ + namespaces: { papers: { recordCount: 42 } }, + }), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ + dense: {} as SearchableIndex, + sparse, + }); + + const result = await session.listNamespacesFromKeywordIndex(); + + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.namespaces).toEqual([{ namespace: 'papers', recordCount: 42 }]); + } + }); + + it('returns ok false when describeIndexStats throws', async () => { + const sparse = { + describeIndexStats: vi.fn().mockRejectedValue(new Error('stats unavailable')), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ + dense: {} as SearchableIndex, + sparse, + }); + + const result = await session.listNamespacesFromKeywordIndex(); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error).toContain('stats unavailable'); + } + }); + }); + + describe('listNamespacesWithMetadata', () => { + it('returns empty when dense stats have no namespaces', async () => { + const dense = { + describeIndexStats: vi.fn().mockResolvedValue({ namespaces: {} }), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ + dense, + sparse: {} as SearchableIndex, + }); + + const rows = await session.listNamespacesWithMetadata(); + expect(rows).toEqual([]); + }); + + it('returns row with empty metadata when recordCount is zero', async () => { + const dense = { + describeIndexStats: vi.fn().mockResolvedValue({ + namespaces: { ns1: { recordCount: 0 } }, + }), + namespace: vi.fn(), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ + dense, + sparse: {} as SearchableIndex, + }); + + const rows = await session.listNamespacesWithMetadata(); + expect(rows).toHaveLength(1); + expect(rows[0]).toEqual({ namespace: 'ns1', recordCount: 0, metadata: {} }); + }); + + it('samples metadata when records exist and namespace.query returns matches', async () => { + const dense = { + describeIndexStats: vi.fn().mockResolvedValue({ + namespaces: { ns1: { recordCount: 2 } }, + dimension: 4, + }), + namespace: () => ({ + query: vi.fn().mockResolvedValue({ + matches: [ + { + metadata: { + title: 'T', + tags: ['a', 'b'], + emptyArr: [], + nested: { x: 1 }, + }, + }, + ], + }), + }), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ + dense, + sparse: {} as SearchableIndex, + }); + + const rows = await session.listNamespacesWithMetadata(); + expect(rows).toHaveLength(1); + expect(rows[0]?.namespace).toBe('ns1'); + expect(rows[0]?.metadata['title']).toBe('string'); + expect(rows[0]?.metadata['tags']).toBe('string[]'); + expect(rows[0]?.metadata['emptyArr']).toBe('array'); + expect(rows[0]?.metadata['nested']).toBe('object'); + }); + }); + + describe('checkIndexes', () => { + it('returns ok when describeIndexStats succeeds for dense and sparse', async () => { + const dense = { + describeIndexStats: vi.fn().mockResolvedValue({}), + } as unknown as SearchableIndex; + const sparse = { + describeIndexStats: vi.fn().mockResolvedValue({}), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ dense, sparse }); + + const result = await session.checkIndexes(); + expect(result.ok).toBe(true); + expect(result.errors).toHaveLength(0); + }); + + it('returns ok false when dense describeIndexStats throws', async () => { + const dense = { + describeIndexStats: vi.fn().mockRejectedValue(new Error('dense down')), + } as unknown as SearchableIndex; + const sparse = { + describeIndexStats: vi.fn().mockResolvedValue({}), + } as unknown as SearchableIndex; + const session = new PineconeIndexSessionTestDouble({ dense, sparse }); + + const result = await session.checkIndexes(); + expect(result.ok).toBe(false); + expect(result.errors.some((e) => e.includes('dense down'))).toBe(true); + }); + + it('returns ok false when ensureIndexes fails', async () => { + const session = new ThrowingEnsureSession(); + const result = await session.checkIndexes(); + expect(result.ok).toBe(false); + expect(result.errors.some((e) => e.includes('no client'))).toBe(true); + }); + }); +}); diff --git a/src/pinecone/indexes.ts b/src/pinecone/indexes.ts new file mode 100644 index 0000000..e1af568 --- /dev/null +++ b/src/pinecone/indexes.ts @@ -0,0 +1,240 @@ +/** + * Lazy Pinecone client and index handles; namespace discovery on dense/sparse indexes. + */ + +import { Pinecone } from '@pinecone-database/pinecone'; +import { error as logError, info as logInfo } from '../logger.js'; +import type { KeywordIndexNamespacesResult, NamespaceHandle, SearchableIndex } from '../types.js'; + +/** + * Infers a human-readable metadata field type for namespace discovery. + * Distinguishes Pinecone-supported list type (string[]) from other arrays. + */ +function inferMetadataFieldType(value: unknown): string { + if (value === null || value === undefined) { + return 'unknown'; + } + if (Array.isArray(value)) { + if (value.length === 0) return 'array'; + if (value.every((item) => typeof item === 'string')) return 'string[]'; + return 'array'; + } + const t = typeof value; + if (t === 'string' || t === 'number' || t === 'boolean') return t; + return 'object'; +} + +/** Holds lazy Pinecone SDK client and dense/sparse index references. */ +export class PineconeIndexSession { + private pc: Pinecone | null = null; + private denseIndex: SearchableIndex | null = null; + private sparseIndex: SearchableIndex | null = null; + private initialized = false; + + constructor( + private readonly apiKey: string, + private readonly indexName: string + ) {} + + /** Same as hybrid sparse index name: `{indexName}-sparse`. */ + getSparseIndexName(): string { + return `${this.indexName}-sparse`; + } + + /** Ensure Pinecone client is initialized */ + ensureClient(): Pinecone { + if (!this.pc) { + if (!this.apiKey) { + throw new Error( + 'Pinecone API key is required. Set PINECONE_API_KEY environment variable or pass apiKey parameter.' + ); + } + this.pc = new Pinecone({ apiKey: this.apiKey }); + logInfo('Pinecone client initialized'); + } + return this.pc; + } + + /** + * Ensure Pinecone indexes are initialized and return them + */ + async ensureIndexes(): Promise<{ + denseIndex: SearchableIndex; + sparseIndex: SearchableIndex; + }> { + if (this.initialized && this.denseIndex !== null && this.sparseIndex !== null) { + return { denseIndex: this.denseIndex, sparseIndex: this.sparseIndex }; + } + + const pc = this.ensureClient(); + const denseName = this.indexName; + const sparseName = this.getSparseIndexName(); + + const dense = pc.index(denseName) as unknown as SearchableIndex; + const sparse = pc.index(sparseName) as unknown as SearchableIndex; + this.denseIndex = dense; + this.sparseIndex = sparse; + this.initialized = true; + + logInfo(`Connected to indexes: ${denseName} and ${sparseName}`); + return { denseIndex: dense, sparseIndex: sparse }; + } + + /** + * List namespaces present on the sparse index (same index used for hybrid sparse and keyword_search). + * Use this to choose a namespace for sparse-only queries instead of the dense index list. + */ + async listNamespacesFromKeywordIndex(): Promise { + try { + const { sparseIndex } = await this.ensureIndexes(); + const stats = sparseIndex.describeIndexStats + ? await sparseIndex.describeIndexStats() + : undefined; + const namespaces = stats?.namespaces ?? {}; + const rows = Object.entries(namespaces).map(([namespace, info]) => ({ + namespace, + recordCount: info?.recordCount ?? 0, + })); + return { ok: true, namespaces: rows }; + } catch (error) { + logError('Error listing namespaces from keyword index', error); + const msg = error instanceof Error ? error.message : String(error); + return { ok: false, error: msg }; + } + } + + /** + * List all available namespaces with their metadata information + * + * Fetches namespaces from the index stats and samples records to discover + * available metadata fields and their types. + */ + async listNamespacesWithMetadata(): Promise< + Array<{ + namespace: string; + recordCount: number; + metadata: Record; + }> + > { + try { + const { denseIndex } = await this.ensureIndexes(); + + // Get index stats to find namespaces + const stats = denseIndex.describeIndexStats + ? await denseIndex.describeIndexStats() + : undefined; + const namespaces = stats?.namespaces ? Object.keys(stats.namespaces) : []; + + logInfo(`Found ${namespaces.length} namespace(s)`); + + // Get metadata info for each namespace by sampling records + const namespacesInfo = await Promise.all( + namespaces.map(async (ns: string) => { + try { + const recordCount = stats?.namespaces?.[ns]?.recordCount || 0; + const metadataFields: Record = {}; + + // Sample a few records to discover metadata fields + if (recordCount > 0 && denseIndex.namespace) { + try { + const nsObj: NamespaceHandle = denseIndex.namespace(ns); + const sampleQuery = + typeof nsObj.query === 'function' + ? await nsObj.query({ + topK: 5, + vector: Array(stats?.dimension ?? 1536).fill(0), + includeMetadata: true, + }) + : { matches: undefined }; + + // Collect unique metadata fields and infer types (including string[]) + if (sampleQuery?.matches) { + sampleQuery.matches.forEach((match: { metadata?: Record }) => { + if (match.metadata) { + Object.entries(match.metadata).forEach(([key, value]) => { + const inferredType = inferMetadataFieldType(value); + if (!(key in metadataFields)) { + metadataFields[key] = inferredType; + } else if ( + (metadataFields[key] === 'object' || metadataFields[key] === 'array') && + inferredType === 'string[]' + ) { + // Prefer array type over generic object when we see it in another sample + metadataFields[key] = inferredType; + } + }); + } + }); + } + } catch (queryError) { + logError(`Error sampling records for namespace ${ns}`, queryError); + } + } + + return { + namespace: ns, + recordCount, + metadata: metadataFields, + }; + } catch (error) { + logError(`Error processing namespace ${ns}`, error); + return { + namespace: ns, + recordCount: 0, + metadata: {}, + }; + } + }) + ); + + return namespacesInfo; + } catch (error) { + logError('Error listing namespaces', error); + return []; + } + } + + /** + * Verify dense and sparse indexes are reachable (describeIndexStats). + * Used by `--check-indexes` / `PINECONE_CHECK_INDEXES` before the server starts. + */ + async checkIndexes(): Promise<{ ok: boolean; errors: string[] }> { + const errors: string[] = []; + const denseName = this.indexName; + const sparseName = this.getSparseIndexName(); + try { + const { denseIndex, sparseIndex } = await this.ensureIndexes(); + + if (typeof denseIndex.describeIndexStats !== 'function') { + errors.push( + `Dense index "${denseName}": describeIndexStats is not available on this SDK surface` + ); + } else { + try { + await denseIndex.describeIndexStats(); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + errors.push(`Dense index "${denseName}": ${msg}`); + } + } + + if (typeof sparseIndex.describeIndexStats !== 'function') { + errors.push( + `Sparse index "${sparseName}": describeIndexStats is not available on this SDK surface` + ); + } else { + try { + await sparseIndex.describeIndexStats(); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + errors.push(`Sparse index "${sparseName}": ${msg}`); + } + } + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + errors.push(`Failed to connect to Pinecone indexes: ${msg}`); + } + + return { ok: errors.length === 0, errors }; + } +} diff --git a/src/pinecone/rerank.test.ts b/src/pinecone/rerank.test.ts new file mode 100644 index 0000000..17c66a5 --- /dev/null +++ b/src/pinecone/rerank.test.ts @@ -0,0 +1,46 @@ +import { describe, it, expect, vi } from 'vitest'; +import { rerankResults } from './rerank.js'; +import type { MergedHit } from '../types.js'; + +const sampleMerged: MergedHit[] = [ + { _id: '1', _score: 0.5, chunk_text: 'hello', metadata: { k: 'v' } }, +]; + +describe('rerankResults', () => { + it('returns empty array when there are no merged hits', async () => { + const pc = {} as Parameters[0]; + const out = await rerankResults(pc, 'any-model', 'q', [], 5); + expect(out).toEqual([]); + }); + + it('maps successful inference.rerank response', async () => { + const rerank = vi.fn().mockResolvedValue({ + data: [ + { + score: 0.99, + document: { _id: '1', chunk_text: 'hello', metadata: { k: 'v' } }, + }, + ], + }); + const pc = { inference: { rerank } } as Parameters[0]; + + 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); + }); + + it('returns unreranked slice when rerank throws', async () => { + const rerank = vi.fn().mockRejectedValue(new Error('rerank unavailable')); + const pc = { inference: { rerank } } as Parameters[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'); + }); +}); diff --git a/src/pinecone/rerank.ts b/src/pinecone/rerank.ts new file mode 100644 index 0000000..90b6c85 --- /dev/null +++ b/src/pinecone/rerank.ts @@ -0,0 +1,63 @@ +/** + * Pinecone inference reranking and mapping to SearchResult. + */ + +import type { Pinecone } from '@pinecone-database/pinecone'; +import { error as logError } from '../logger.js'; +import type { MergedHit, SearchResult } from '../types.js'; + +/** + * Rerank merged hits using Pinecone's reranking model; on failure returns unreranked slice. + */ +export async function rerankResults( + pc: Pinecone, + rerankModel: string, + query: string, + results: MergedHit[], + topN: number +): Promise { + if (!results || results.length === 0) { + return []; + } + + try { + const rerankResult = await pc.inference.rerank({ + model: rerankModel, + query, + // The Pinecone SDK types constrain document values to `Record`, + // but the underlying HTTP API accepts any JSON value. We pass MergedHit objects + // (metadata may contain number/boolean/string[]) and only `chunk_text` — which is + // always a string — is accessed via rankFields. The double cast via `as unknown` + // is intentional: it bypasses the SDK's over-narrow type without stringifying + // metadata values that we need to read back from the returned documents. + documents: results as unknown as (string | Record)[], + topN, + rankFields: ['chunk_text'], + returnDocuments: true, + parameters: { truncate: 'END' }, + }); + + const reranked: SearchResult[] = []; + for (const item of rerankResult.data || []) { + const document = (item.document || {}) as MergedHit; + reranked.push({ + id: document['_id'] || '', + content: document['chunk_text'] || '', + score: parseFloat(String(item.score || 0)), + metadata: document['metadata'] || {}, + reranked: true, + }); + } + return reranked; + } catch (error) { + logError('Error reranking results', error); + // Fall back to returning unreranked results + return results.slice(0, topN).map((result) => ({ + id: result._id || '', + content: result.chunk_text || '', + score: result._score || 0, + metadata: result.metadata || {}, + reranked: false, + })); + } +} diff --git a/src/pinecone/search.test.ts b/src/pinecone/search.test.ts new file mode 100644 index 0000000..98bdbe9 --- /dev/null +++ b/src/pinecone/search.test.ts @@ -0,0 +1,111 @@ +import { describe, it, expect, vi } from 'vitest'; +import { + searchIndex, + mergeResults, + sliceMergedHitsToSearchResults, + mapSparseHitsToSearchResults, + countUniqueDocumentsFromHits, +} from './search.js'; +import type { PineconeHit, SearchableIndex } from '../types.js'; + +describe('searchIndex', () => { + it('uses index.search when available and passes fields', async () => { + const search = vi.fn().mockResolvedValue({ + result: { hits: [{ _id: '1', _score: 1, fields: { chunk_text: 'x' } }] }, + }); + const index = { search } as unknown as SearchableIndex; + const hits = await searchIndex(index, 'hi', 5, 'ns', { k: 'v' }, { fields: ['chunk_text'] }); + expect(hits).toHaveLength(1); + expect(search).toHaveBeenCalledWith( + expect.objectContaining({ + namespace: 'ns', + fields: ['chunk_text'], + query: expect.objectContaining({ filter: { k: 'v' } }), + }) + ); + }); + + it('uses namespace-scoped searchRecords when search is absent', async () => { + const searchRecords = vi.fn().mockResolvedValue({ + result: { hits: [{ _id: 'r1', _score: 0.5, fields: {} }] }, + }); + const index = { + namespace: vi.fn().mockReturnValue({ searchRecords }), + } as unknown as SearchableIndex; + const hits = await searchIndex(index, 'hi', 3, 'ns'); + expect(hits).toHaveLength(1); + expect(searchRecords).toHaveBeenCalled(); + }); + + it('uses top-level searchRecords when there is no namespace', async () => { + const searchRecords = vi.fn().mockResolvedValue({ result: { hits: [] } }); + const index = { searchRecords } as unknown as SearchableIndex; + const hits = await searchIndex(index, 'hi', 2, undefined); + expect(hits).toEqual([]); + expect(searchRecords).toHaveBeenCalled(); + }); + + it('returns empty hits when searchRecords is missing on fallback target', async () => { + const index = { namespace: vi.fn().mockReturnValue({}) } as unknown as SearchableIndex; + const hits = await searchIndex(index, 'hi', 2, 'ns'); + expect(hits).toEqual([]); + }); + + it('wraps errors with namespace context', async () => { + const index = { + search: vi.fn().mockRejectedValue(new Error('boom')), + } as unknown as SearchableIndex; + await expect(searchIndex(index, 'hi', 5, 'my-ns')).rejects.toThrow( + 'Pinecone search failed for namespace "my-ns": boom' + ); + }); +}); + +describe('mergeResults', () => { + it('keeps higher score when duplicate _id appears in dense and sparse', () => { + const merged = mergeResults( + [{ _id: '1', _score: 0.9, fields: { chunk_text: 'a' } }], + [{ _id: '1', _score: 0.1, fields: { chunk_text: 'b' } }] + ); + expect(merged).toHaveLength(1); + expect(merged[0]?.chunk_text).toBe('a'); + }); +}); + +describe('sliceMergedHitsToSearchResults', () => { + it('maps merged hits to SearchResult rows', () => { + const out = sliceMergedHitsToSearchResults( + [ + { _id: 'a', _score: 1, chunk_text: 'c', metadata: {} }, + { _id: 'b', _score: 0.5, chunk_text: 'd', metadata: { x: 1 } }, + ], + 1 + ); + expect(out).toHaveLength(1); + expect(out[0]?.id).toBe('a'); + expect(out[0]?.reranked).toBe(false); + }); +}); + +describe('mapSparseHitsToSearchResults', () => { + it('splits chunk_text from other fields', () => { + const hits: PineconeHit[] = [ + { _id: 'z', _score: 0.2, fields: { chunk_text: 'body', author: 'me' } }, + ]; + const out = mapSparseHitsToSearchResults(hits); + expect(out[0]?.content).toBe('body'); + expect(out[0]?.metadata['author']).toBe('me'); + }); +}); + +describe('countUniqueDocumentsFromHits', () => { + it('dedupes by document_number', () => { + const hits: PineconeHit[] = [ + { _id: 'c1', _score: 1, fields: { document_number: 'D1' } }, + { _id: 'c2', _score: 0.9, fields: { document_number: 'D1' } }, + ]; + const r = countUniqueDocumentsFromHits(hits, 'ns'); + expect(r.count).toBe(1); + expect(r.truncated).toBe(false); + }); +}); diff --git a/src/pinecone/search.ts b/src/pinecone/search.ts new file mode 100644 index 0000000..030631e --- /dev/null +++ b/src/pinecone/search.ts @@ -0,0 +1,191 @@ +/** + * Pinecone search pipeline: single-index text query and dense/sparse merge. + */ + +import { COUNT_FIELDS, COUNT_TOP_K } from '../constants.js'; +import { debug as logDebug, warn as logWarn } from '../logger.js'; +import type { + CountResult, + MergedHit, + PineconeHit, + PineconeMetadataValue, + SearchResult, + SearchableIndex, +} from '../types.js'; + +/** + * Search a Pinecone index using text query with optional metadata filtering. + * When options.fields is set, only those fields are requested (e.g. for count: no chunk_text). + */ +export async function searchIndex( + index: SearchableIndex, + query: string, + topK: number, + namespace?: string, + metadataFilter?: Record, + options?: { fields?: string[] } +): Promise { + // Build query payload in the same shape as Python implementation. + const queryPayload: Record = { + top_k: topK, + inputs: { text: query }, + }; + + // Include filter when explicitly provided (matches Python behavior). + if (metadataFilter !== undefined) { + queryPayload['filter'] = metadataFilter; + logDebug('Applying metadata filter', metadataFilter); + } + + try { + // Preferred path: Pinecone search API. + if (typeof index.search === 'function') { + const searchOpts: { + namespace?: string; + query: Record; + fields?: string[]; + } = { + namespace, + query: queryPayload, + }; + if (options?.fields?.length) { + searchOpts.fields = options.fields; + } + const result = await index.search(searchOpts); + return result?.result?.hits || []; + } + + // Backward-compatible fallback for older API shapes. + const target = namespace && index.namespace ? index.namespace(namespace) : index; + const queryParams: { query: Record; fields?: string[] } = { + query: { + topK, + inputs: { text: query }, + }, + }; + if (metadataFilter !== undefined) { + queryParams.query['filter'] = metadataFilter; + } + if (options?.fields?.length) { + queryParams.fields = options.fields; + } + const result = target.searchRecords + ? await target.searchRecords(queryParams) + : { result: { hits: [] as PineconeHit[] } }; + return result?.result?.hits || []; + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + throw new Error( + `Pinecone search failed for namespace "${namespace ?? 'default'}": ${errorMessage}` + ); + } +} + +/** + * Merge and deduplicate results from dense and sparse searches + * + * Uses the higher score when duplicates are found. + */ +export function mergeResults(denseHits: PineconeHit[], sparseHits: PineconeHit[]): MergedHit[] { + const deduped: Record = {}; + + for (const hit of [...denseHits, ...sparseHits]) { + const hitId = hit._id || ''; + const hitScore = hit._score || 0; + + const existing = deduped[hitId]; + if (existing !== undefined && (existing._score || 0) >= hitScore) { + continue; + } + + const hitMetadata: Record = {}; + let content = ''; + + for (const [key, value] of Object.entries(hit.fields || {})) { + if (key === 'chunk_text') { + content = typeof value === 'string' ? value : ''; + } else { + hitMetadata[key] = value as PineconeMetadataValue; + } + } + + deduped[hitId] = { + _id: hitId, + _score: hitScore, + chunk_text: content, + metadata: hitMetadata, + }; + } + + return Object.values(deduped).sort((a, b) => (b._score || 0) - (a._score || 0)); +} + +/** Top merged hits as SearchResult without calling rerank API. */ +export function sliceMergedHitsToSearchResults(merged: MergedHit[], topK: number): SearchResult[] { + return merged.slice(0, topK).map((result) => ({ + id: result._id || '', + content: result.chunk_text || '', + score: result._score || 0, + metadata: result.metadata || {}, + reranked: false, + })); +} + +/** Map sparse-index hits to SearchResult rows (keyword search; no reranking). */ +export function mapSparseHitsToSearchResults(hits: PineconeHit[]): SearchResult[] { + return hits.map((hit) => { + const fields = hit.fields || {}; + let content = ''; + const metadata: Record = {}; + for (const [key, value] of Object.entries(fields)) { + if (key === 'chunk_text') { + content = typeof value === 'string' ? value : ''; + } else { + metadata[key] = value as PineconeMetadataValue; + } + } + return { + id: hit._id || '', + content, + score: hit._score || 0, + metadata, + reranked: false, + }; + }); +} + +/** Deduplicate semantic search hits by document identifiers for count(). */ +export function countUniqueDocumentsFromHits( + hits: PineconeHit[], + namespace: string | undefined +): CountResult { + const docKeys = new Set(); + let idFallbackCount = 0; + for (const hit of hits) { + const fields = hit.fields || {}; + const docNumber = fields['document_number']; + const url = fields['url']; + const docId = fields['doc_id']; + const docKey = + (typeof docNumber === 'string' ? docNumber : undefined) ?? + (typeof url === 'string' ? url : undefined) ?? + (typeof docId === 'string' ? docId : undefined); + if (docKey !== undefined) { + docKeys.add(docKey); + } else { + // Fall back to chunk ID — this yields a chunk count, not a document count + idFallbackCount++; + docKeys.add(hit._id ?? ''); + } + } + if (idFallbackCount > 0) { + logWarn( + `count(): ${idFallbackCount} hit(s) in namespace "${namespace}" had none of the ` + + `identifier fields (${COUNT_FIELDS.join(', ')}); fell back to chunk ID — result may overcount documents` + ); + } + return { + count: docKeys.size, + truncated: hits.length >= COUNT_TOP_K, + }; +}