diff --git a/src/control/__tests__/indexOperationsBuilder.test.ts b/src/control/__tests__/indexOperationsBuilder.test.ts new file mode 100644 index 00000000..7aab28fa --- /dev/null +++ b/src/control/__tests__/indexOperationsBuilder.test.ts @@ -0,0 +1,40 @@ +import { indexOperationsBuilder } from '../indexOperationsBuilder'; +import { Configuration } from '../../pinecone-generated-ts-fetch'; + +jest.mock('../../pinecone-generated-ts-fetch', () => ({ + ...jest.requireActual('../../pinecone-generated-ts-fetch'), + Configuration: jest.fn(), +})); + +describe('indexOperationsBuilder', () => { + test('API Configuration basePath is set to api.pinecone.io by default', () => { + const config = { apiKey: 'test-api-key' }; + indexOperationsBuilder(config); + expect(Configuration).toHaveBeenCalledWith( + expect.objectContaining({ basePath: 'https://api.pinecone.io' }) + ); + }); + + test('controllerHostUrl overwrites the basePath in API Configuration', () => { + const controllerHostUrl = 'https://test-controller-host-url.io'; + const config = { + apiKey: 'test-api-key', + controllerHostUrl, + }; + indexOperationsBuilder(config); + expect(Configuration).toHaveBeenCalledWith( + expect.objectContaining({ basePath: controllerHostUrl }) + ); + }); + + test('additionalHeaders are passed to the API Configuration', () => { + const additionalHeaders = { 'x-test-header': 'test-value' }; + const config = { apiKey: 'test-api-key', additionalHeaders }; + indexOperationsBuilder(config); + expect(Configuration).toHaveBeenCalledWith( + expect.objectContaining({ + headers: expect.objectContaining(additionalHeaders), + }) + ); + }); +}); diff --git a/src/control/indexOperationsBuilder.ts b/src/control/indexOperationsBuilder.ts index 7b0f5065..b8cb71aa 100644 --- a/src/control/indexOperationsBuilder.ts +++ b/src/control/indexOperationsBuilder.ts @@ -24,7 +24,7 @@ export const indexOperationsBuilder = ( apiKey, queryParamsStringify, headers: { - 'User-Agent': buildUserAgent(), + 'User-Agent': buildUserAgent(config), ...headers, }, fetchApi: getFetch(config), diff --git a/src/data/__tests__/dataOperationsProvider.test.ts b/src/data/__tests__/dataOperationsProvider.test.ts index da5d7e4e..c21f4386 100644 --- a/src/data/__tests__/dataOperationsProvider.test.ts +++ b/src/data/__tests__/dataOperationsProvider.test.ts @@ -1,8 +1,17 @@ import { DataOperationsProvider } from '../dataOperationsProvider'; import { IndexHostSingleton } from '../indexHostSingleton'; +import { Configuration } from '../../pinecone-generated-ts-fetch'; + +jest.mock('../../pinecone-generated-ts-fetch', () => ({ + ...jest.requireActual('../../pinecone-generated-ts-fetch'), + Configuration: jest.fn(), +})); describe('DataOperationsProvider', () => { let real; + const config = { + apiKey: 'test-api-key', + }; beforeAll(() => { real = IndexHostSingleton.getHostUrl; @@ -41,9 +50,6 @@ describe('DataOperationsProvider', () => { }); test('passing indexHostUrl skips hostUrl resolution', async () => { - const config = { - apiKey: 'test-api-key', - }; const indexHostUrl = 'http://index-host-url'; const provider = new DataOperationsProvider( config, @@ -58,4 +64,21 @@ describe('DataOperationsProvider', () => { expect(IndexHostSingleton.getHostUrl).not.toHaveBeenCalled(); expect(provider.buildDataOperationsConfig).toHaveBeenCalled(); }); + + test('passing additionalHeaders applies them to the API Configuration', async () => { + const additionalHeaders = { 'x-custom-header': 'custom-value' }; + const provider = new DataOperationsProvider( + config, + 'index-name', + undefined, + additionalHeaders + ); + + await provider.provide(); + expect(Configuration).toHaveBeenCalledWith( + expect.objectContaining({ + headers: expect.objectContaining(additionalHeaders), + }) + ); + }); }); diff --git a/src/data/__tests__/index.test.ts b/src/data/__tests__/index.test.ts index 952a27fb..bd2b95cf 100644 --- a/src/data/__tests__/index.test.ts +++ b/src/data/__tests__/index.test.ts @@ -2,6 +2,7 @@ import { FetchCommand } from '../fetch'; import { QueryCommand } from '../query'; import { UpdateCommand } from '../update'; import { UpsertCommand } from '../upsert'; +import { DataOperationsProvider } from '../dataOperationsProvider'; import { Index } from '../index'; import type { ScoredPineconeRecord } from '../query'; @@ -9,6 +10,7 @@ jest.mock('../fetch'); jest.mock('../query'); jest.mock('../update'); jest.mock('../upsert'); +jest.mock('../dataOperationsProvider'); describe('Index', () => { let config; @@ -24,165 +26,188 @@ describe('Index', () => { }; }); - test('can write functions that take types with generic params', () => { - // eslint-disable-next-line @typescript-eslint/no-unused-vars - const fn1 = (record: ScoredPineconeRecord) => { - // no type errors on this because typescript doesn't know anything about what keys are defined - // ScoredPineconeRecord without specifying the generic type param - console.log(record.metadata && record.metadata.yolo); - }; - type MyMeta = { - name: string; - }; - - // eslint-disable-next-line @typescript-eslint/no-unused-vars - const fn2 = (record: ScoredPineconeRecord) => { - console.log(record.metadata && record.metadata.name); - - // @ts-expect-error because bogus not in MyMeta - console.log(record.metadata && record.metadata.bogus); - }; + describe('index initialization', () => { + test('passes config, indexName, indexHostUrl, and additionalHeaders to DataOperationsProvider', () => { + const indexHostUrl = 'https://test-api-pinecone.io'; + const additionalHeaders = { 'x-custom-header': 'custom-value' }; + new Index( + 'index-name', + config, + undefined, + indexHostUrl, + additionalHeaders + ); + expect(DataOperationsProvider).toHaveBeenCalledTimes(1); + expect(DataOperationsProvider).toHaveBeenCalledWith( + config, + 'index-name', + indexHostUrl, + additionalHeaders + ); + }); }); - test('can be used without generic types param', async () => { - const index = new Index('index-name', config, 'namespace'); + describe('metadata', () => { + test('can write functions that take types with generic params', () => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + const fn1 = (record: ScoredPineconeRecord) => { + // no type errors on this because typescript doesn't know anything about what keys are defined + // ScoredPineconeRecord without specifying the generic type param + console.log(record.metadata && record.metadata.yolo); + }; + type MyMeta = { + name: string; + }; + + // eslint-disable-next-line @typescript-eslint/no-unused-vars + const fn2 = (record: ScoredPineconeRecord) => { + console.log(record.metadata && record.metadata.name); + + // @ts-expect-error because bogus not in MyMeta + console.log(record.metadata && record.metadata.bogus); + }; + }); - // You can use the index class without passing the generic type for metadata, - // but you lose type safety in that case. - await index.update({ id: '1', metadata: { foo: 'bar' } }); - await index.update({ id: '1', metadata: { baz: 'quux' } }); + test('can be used without generic types param', async () => { + const index = new Index('index-name', config, 'namespace'); - // Same thing with upsert. You can upsert anything in metadata field without type. - await index.upsert([ - { id: '2', values: [0.1, 0.2], metadata: { hello: 'world' } }, - ]); + // You can use the index class without passing the generic type for metadata, + // but you lose type safety in that case. + await index.update({ id: '1', metadata: { foo: 'bar' } }); + await index.update({ id: '1', metadata: { baz: 'quux' } }); - // @ts-expect-error even when you haven't passed a generic type, it enforces the expected shape of RecordMetadata - await index.upsert([{ id: '2', values: [0.1, 0.2], metadata: 2 }]); - }); + // Same thing with upsert. You can upsert anything in metadata field without type. + await index.upsert([ + { id: '2', values: [0.1, 0.2], metadata: { hello: 'world' } }, + ]); - test('preserves metadata typing through chained namespace calls', async () => { - const index = new Index('index-name', config, 'namespace'); - const ns1 = index.namespace('ns1'); + // @ts-expect-error even when you haven't passed a generic type, it enforces the expected shape of RecordMetadata + await index.upsert([{ id: '2', values: [0.1, 0.2], metadata: 2 }]); + }); - // @ts-expect-error because MovieMetadata metadata still expected after chained namespace call - await ns1.update({ id: '1', metadata: { title: 'Vertigo', rating: 5 } }); - }); + test('preserves metadata typing through chained namespace calls', async () => { + const index = new Index('index-name', config, 'namespace'); + const ns1 = index.namespace('ns1'); - test('upsert: has type errors when passing malformed metadata', async () => { - const index = new Index('index-name', config, 'namespace'); - expect(UpsertCommand).toHaveBeenCalledTimes(1); + // @ts-expect-error because MovieMetadata metadata still expected after chained namespace call + await ns1.update({ id: '1', metadata: { title: 'Vertigo', rating: 5 } }); + }); - // No ts errors when upserting with proper MovieMetadata - await index.upsert([ - { - id: '1', - values: [0.1, 0.1, 0.1], - metadata: { - genre: 'romance', - runtime: 120, - }, - }, - ]); - - // No ts errors when upserting with no metadata - await index.upsert([ - { - id: '2', - values: [0.1, 0.1, 0.1], - }, - ]); - - // ts error expected when passing metadata that doesn't match MovieMetadata - await index.upsert([ - { - id: '3', - values: [0.1, 0.1, 0.1], - metadata: { - // @ts-expect-error - somethingElse: 'foo', + test('upsert: has type errors when passing malformed metadata', async () => { + const index = new Index('index-name', config, 'namespace'); + expect(UpsertCommand).toHaveBeenCalledTimes(1); + + // No ts errors when upserting with proper MovieMetadata + await index.upsert([ + { + id: '1', + values: [0.1, 0.1, 0.1], + metadata: { + genre: 'romance', + runtime: 120, + }, }, - }, - ]); - }); + ]); - test('fetch: response is typed with generic metadata type', async () => { - const index = new Index('index-name', config, 'namespace'); - expect(FetchCommand).toHaveBeenCalledTimes(1); - - const response = await index.fetch(['1']); - if (response && response.records) { - // eslint-disable-next-line - Object.entries(response.records).forEach(([key, value]) => { - // No errors on these because they are properties from MovieMetadata - console.log(value.metadata?.genre); - console.log(value.metadata?.runtime); + // No ts errors when upserting with no metadata + await index.upsert([ + { + id: '2', + values: [0.1, 0.1, 0.1], + }, + ]); + + // ts error expected when passing metadata that doesn't match MovieMetadata + await index.upsert([ + { + id: '3', + values: [0.1, 0.1, 0.1], + metadata: { + // @ts-expect-error + somethingElse: 'foo', + }, + }, + ]); + }); - // @ts-expect-error because result is expecting metadata to be MovieMetadata - console.log(value.metadata?.bogus); - }); - } - }); + test('fetch: response is typed with generic metadata type', async () => { + const index = new Index('index-name', config, 'namespace'); + expect(FetchCommand).toHaveBeenCalledTimes(1); + + const response = await index.fetch(['1']); + if (response && response.records) { + // eslint-disable-next-line + Object.entries(response.records).forEach(([key, value]) => { + // No errors on these because they are properties from MovieMetadata + console.log(value.metadata?.genre); + console.log(value.metadata?.runtime); + + // @ts-expect-error because result is expecting metadata to be MovieMetadata + console.log(value.metadata?.bogus); + }); + } + }); - test('query: returns typed results', async () => { - const index = new Index('index-name', config, 'namespace'); - expect(QueryCommand).toHaveBeenCalledTimes(1); + test('query: returns typed results', async () => { + const index = new Index('index-name', config, 'namespace'); + expect(QueryCommand).toHaveBeenCalledTimes(1); - const results = await index.query({ id: '1', topK: 5 }); - if (results && results.matches) { - if (results.matches.length > 0) { - const firstResult = results.matches[0]; + const results = await index.query({ id: '1', topK: 5 }); + if (results && results.matches) { + if (results.matches.length > 0) { + const firstResult = results.matches[0]; - // no ts error because score is part of ScoredPineconeRecord - console.log(firstResult.score); + // no ts error because score is part of ScoredPineconeRecord + console.log(firstResult.score); - // no ts error because genre and runtime part of MovieMetadata - console.log(firstResult.metadata?.genre); - console.log(firstResult.metadata?.runtime); + // no ts error because genre and runtime part of MovieMetadata + console.log(firstResult.metadata?.genre); + console.log(firstResult.metadata?.runtime); - // @ts-expect-error because bogus not part of MovieMetadata - console.log(firstResult.metadata?.bogus); + // @ts-expect-error because bogus not part of MovieMetadata + console.log(firstResult.metadata?.bogus); + } } - } - }); + }); - test('update: has typed arguments', async () => { - const index = new Index('index-name', config, 'namespace'); - expect(UpdateCommand).toHaveBeenCalledTimes(1); + test('update: has typed arguments', async () => { + const index = new Index('index-name', config, 'namespace'); + expect(UpdateCommand).toHaveBeenCalledTimes(1); - // Can update metadata only without ts errors - await index.update({ - id: '1', - metadata: { genre: 'romance', runtime: 90 }, - }); + // Can update metadata only without ts errors + await index.update({ + id: '1', + metadata: { genre: 'romance', runtime: 90 }, + }); - // Can update values only without ts errors - await index.update({ id: '2', values: [0.1, 0.2, 0.3] }); + // Can update values only without ts errors + await index.update({ id: '2', values: [0.1, 0.2, 0.3] }); - // Can update sparseValues only without ts errors - await index.update({ - id: '3', - sparseValues: { indices: [0, 3], values: [0.2, 0.5] }, - }); + // Can update sparseValues only without ts errors + await index.update({ + id: '3', + sparseValues: { indices: [0, 3], values: [0.2, 0.5] }, + }); - // Can update all fields without ts errors - await index.update({ - id: '4', - values: [0.1, 0.2, 0.3], - sparseValues: { indices: [0], values: [0.789] }, - metadata: { genre: 'horror', runtime: 10 }, - }); + // Can update all fields without ts errors + await index.update({ + id: '4', + values: [0.1, 0.2, 0.3], + sparseValues: { indices: [0], values: [0.789] }, + metadata: { genre: 'horror', runtime: 10 }, + }); - // @ts-expect-error when id is missing - await index.update({ metadata: { genre: 'drama', runtime: 97 } }); + // @ts-expect-error when id is missing + await index.update({ metadata: { genre: 'drama', runtime: 97 } }); - // @ts-expect-error when metadata has unexpected fields - await index.update({ id: '5', metadata: { title: 'Vertigo' } }); + // @ts-expect-error when metadata has unexpected fields + await index.update({ id: '5', metadata: { title: 'Vertigo' } }); - await index.update({ - id: '6', - // @ts-expect-error when metadata has extra properties - metadata: { genre: 'comedy', runtime: 80, title: 'Miss Congeniality' }, + await index.update({ + id: '6', + // @ts-expect-error when metadata has extra properties + metadata: { genre: 'comedy', runtime: 80, title: 'Miss Congeniality' }, + }); }); }); }); diff --git a/src/data/dataOperationsProvider.ts b/src/data/dataOperationsProvider.ts index d013bc3b..71d8c74a 100644 --- a/src/data/dataOperationsProvider.ts +++ b/src/data/dataOperationsProvider.ts @@ -4,6 +4,7 @@ import { ConfigurationParameters, DataPlaneApi, } from '../pinecone-generated-ts-fetch'; +import type { HTTPHeaders } from '../pinecone-generated-ts-fetch'; import { queryParamsStringify, buildUserAgent, @@ -18,15 +19,18 @@ export class DataOperationsProvider { private indexName: string; private indexHostUrl?: string; private dataOperations?: DataPlaneApi; + private additionalHeaders?: HTTPHeaders; constructor( config: PineconeConfiguration, indexName: string, - indexHostUrl?: string + indexHostUrl?: string, + additionalHeaders?: HTTPHeaders ) { this.config = config; this.indexName = indexName; this.indexHostUrl = normalizeUrl(indexHostUrl); + this.additionalHeaders = additionalHeaders; } async provide() { @@ -51,12 +55,15 @@ export class DataOperationsProvider { } buildDataOperationsConfig() { + const headers = this.additionalHeaders || null; + const indexConfigurationParameters: ConfigurationParameters = { basePath: this.indexHostUrl, apiKey: this.config.apiKey, queryParamsStringify, headers: { - 'User-Agent': buildUserAgent(), + 'User-Agent': buildUserAgent(this.config), + ...headers, }, fetchApi: getFetch(this.config), middleware, diff --git a/src/data/index.ts b/src/data/index.ts index 157cc61a..789ab007 100644 --- a/src/data/index.ts +++ b/src/data/index.ts @@ -14,6 +14,7 @@ import { describeIndexStats } from './describeIndexStats'; import { DataOperationsProvider } from './dataOperationsProvider'; import { listPaginated } from './list'; import type { ListOptions } from './list'; +import type { HTTPHeaders } from '../pinecone-generated-ts-fetch'; import type { PineconeConfiguration, RecordMetadata, @@ -333,12 +334,14 @@ export class Index { * @param config - The configuration from the Pinecone client. * @param namespace - The namespace for the index. * @param indexHostUrl - An optional override for the host address used for data operations. + * @param additionalHeaders - An optional object of additional header to send with each request. */ constructor( indexName: string, config: PineconeConfiguration, namespace = '', - indexHostUrl?: string + indexHostUrl?: string, + additionalHeaders?: HTTPHeaders ) { this.config = config; this.target = { @@ -350,7 +353,8 @@ export class Index { const apiProvider = new DataOperationsProvider( config, indexName, - indexHostUrl + indexHostUrl, + additionalHeaders ); this._deleteAll = deleteAll(apiProvider, namespace); diff --git a/src/data/types.ts b/src/data/types.ts index 6f00a58a..215bd082 100644 --- a/src/data/types.ts +++ b/src/data/types.ts @@ -41,6 +41,11 @@ export type PineconeConfiguration = { * Optional headers to be included in all requests. */ additionalHeaders?: HTTPHeaders; + + /** + * Optional sourceTag that is applied to the User-Agent header with all requests. + */ + sourceTag?: string; }; export const RecordIdSchema = Type.String({ minLength: 1 }); diff --git a/src/pinecone.ts b/src/pinecone.ts index 34da677f..9888887a 100644 --- a/src/pinecone.ts +++ b/src/pinecone.ts @@ -17,6 +17,7 @@ import { ConfigureIndexRequestSpecPod, CreateCollectionRequest, } from './pinecone-generated-ts-fetch'; +import type { HTTPHeaders } from './pinecone-generated-ts-fetch'; import { IndexHostSingleton } from './data/indexHostSingleton'; import { PineconeConfigurationError, @@ -619,14 +620,22 @@ export class Pinecone { * @typeParam T - The type of metadata associated with each record. * @param indexName - The name of the index to target. * @param indexHostUrl - An optional host url to use for operations against this index. If not provided, the host url will be resolved by calling {@link describeIndex}. + * @param additionalHeaders - An optional object containing additional headers to pass with each index request. * @typeParam T - The type of the metadata object associated with each record. * @returns An {@link Index} object that can be used to perform data operations. */ index( indexName: string, - indexHostUrl?: string + indexHostUrl?: string, + additionalHeaders?: HTTPHeaders ) { - return new Index(indexName, this.config, undefined, indexHostUrl); + return new Index( + indexName, + this.config, + undefined, + indexHostUrl, + additionalHeaders + ); } /** @@ -635,8 +644,9 @@ export class Pinecone { // Alias method to match the Python SDK capitalization Index( indexName: string, - indexHostUrl?: string + indexHostUrl?: string, + additionalHeaders?: HTTPHeaders ) { - return this.index(indexName, indexHostUrl); + return this.index(indexName, indexHostUrl, additionalHeaders); } } diff --git a/src/utils/__tests__/user-agent.test.ts b/src/utils/__tests__/user-agent.test.ts new file mode 100644 index 00000000..20aa63ce --- /dev/null +++ b/src/utils/__tests__/user-agent.test.ts @@ -0,0 +1,47 @@ +import { buildUserAgent } from '../user-agent'; +import * as EnvironmentModule from '../environment'; + +describe('user-agent', () => { + describe('buildUserAgent', () => { + test('applies Edge Runtime when running in an edge environment', () => { + jest.spyOn(EnvironmentModule, 'isEdge').mockReturnValue(true); + const config = { apiKey: 'test-api-key' }; + const userAgent = buildUserAgent(config); + + expect(userAgent).toContain('Edge Runtime'); + }); + + test('applies source_tag when provided via PineconeConfiguration', () => { + const config = { + apiKey: 'test-api-key', + sourceTag: 'test source tag', + }; + + const userAgent = buildUserAgent(config); + expect(userAgent).toContain('source_tag=test_source_tag'); + }); + }); + + describe('normalizeSourceTag', () => { + test('normalizes variations of sourceTag', () => { + const config = { + apiKey: 'test-api-key', + sourceTag: 'my source tag!!!', + }; + let userAgent = buildUserAgent(config); + expect(userAgent).toContain('source_tag=my_source_tag'); + + config.sourceTag = 'My Source Tag'; + userAgent = buildUserAgent(config); + expect(userAgent).toContain('source_tag=my_source_tag'); + + config.sourceTag = ' My Source Tag 123 '; + userAgent = buildUserAgent(config); + expect(userAgent).toContain('source_tag=my_source_tag_123'); + + config.sourceTag = ' MY SOURCE TAG 1234 ##### !!!!!!'; + userAgent = buildUserAgent(config); + expect(userAgent).toContain('source_tag=my_source_tag_1234'); + }); + }); +}); diff --git a/src/utils/user-agent.ts b/src/utils/user-agent.ts index a5ffd513..82c08458 100644 --- a/src/utils/user-agent.ts +++ b/src/utils/user-agent.ts @@ -1,7 +1,8 @@ import { isEdge } from './environment'; +import type { PineconeConfiguration } from '../data/types'; import * as packageInfo from '../version.json'; -export const buildUserAgent = () => { +export const buildUserAgent = (config: PineconeConfiguration) => { // We always want to include the package name and version // along with the langauge name to help distinguish these // requests from those emitted by other clients @@ -19,5 +20,28 @@ export const buildUserAgent = () => { userAgentParts.push(`node ${process.version}`); } + if (config.sourceTag) { + userAgentParts.push(`source_tag=${normalizeSourceTag(config.sourceTag)}`); + } + return userAgentParts.join('; '); }; + +const normalizeSourceTag = (sourceTag: string) => { + if (!sourceTag) { + return; + } + + /** + * normalize sourceTag + * 1. Lowercase + * 2. Limit charset to [a-z0-9_ ] + * 3. Trim left/right spaces + * 4. Condense multiple spaces to one, and replace with underscore + */ + return sourceTag + .toLowerCase() + .replace(/[^a-z0-9_ ]/g, '') + .trim() + .replace(/[ ]+/g, '_'); +};