UNPKG

@convo-lang/convo-lang-pinecone

Version:
203 lines 8.27 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.ConvoPineconeRagService = void 0; const convo_lang_1 = require("@convo-lang/convo-lang"); const pinecone_1 = require("@pinecone-database/pinecone"); const convo_pinecone_deps_1 = require("./convo-pinecone-deps"); const pinecone_const_1 = require("./pinecone-const"); class ConvoPineconeRagService { static fromScope(scope) { const mnt = scope.to(convo_pinecone_deps_1.pineconeIndexMountsParam).get(); let indexMounts; if (mnt) { indexMounts = mnt.split(/[,;]/g).map(v => { const i = v.indexOf(':'); if (i === -1) { throw new Error(`Invalid ConvoPineconeIndexMount expression - ${mnt}`); } const path = v.substring(0, i).trim(); const index = v.substring(i + 1).trim(); if (!path || !index) { throw new Error(`Invalid ConvoPineconeIndexMount expression - ${mnt}`); } return { path, index }; }); } const index = (0, convo_pinecone_deps_1.pineconeIndexParam)(scope); return new ConvoPineconeRagService({ apiKey: (0, convo_pinecone_deps_1.pineconeApiKeyParam)(scope), index, indexKey: (0, convo_pinecone_deps_1.pineconeIndexKeyParam)(scope), cloud: (0, convo_pinecone_deps_1.pineconeCloudParam)(scope), region: (0, convo_pinecone_deps_1.pineconeRegionParam)(scope), model: (0, convo_pinecone_deps_1.pineconeModelParam)(scope), autoCreateIndex: (0, convo_pinecone_deps_1.pineconeAutoCreateIndexParam)(scope), allowedReadIndexes: scope.to(convo_pinecone_deps_1.pineconeAllowedReadIndexesParam).get()?.split(',').map(i => i.trim()).filter(i => i) ?? indexMounts?.map(i => i.index) ?? [index], allowedWriteIndexes: scope.to(convo_pinecone_deps_1.pineconeAllowedWriteIndexesParam).get()?.split(',').map(i => i.trim()).filter(i => i), indexMounts }); } constructor({ apiKey, index = pinecone_const_1.defaultPineconeIndex, indexKey = pinecone_const_1.defaultPineconeIndexKey, cloud = pinecone_const_1.defaultPineconeCloud, region = pinecone_const_1.defaultPineconeRegion, model = pinecone_const_1.defaultPineconeModel, autoCreateIndex = pinecone_const_1.defaultPineconeAutoCreateIndex, allowedReadIndexes, allowedWriteIndexes, indexMounts, }) { this._clientPromise = null; this.indexPromises = {}; this.options = { apiKey, index, indexKey, cloud, region, model, autoCreateIndex, allowedReadIndexes, allowedWriteIndexes, indexMounts, }; } async getClientAsync() { return await this._clientPromise ?? (this._clientPromise = this.createClient()); } async createClient() { const client = new pinecone_1.Pinecone({ apiKey: this.options.apiKey, }); return client; } async getIndexAsync(index, forRead) { if (forRead ? !this.isReadIndexAllowed(index) : !this.isWriteIndexAllowed(index)) { throw new Error(`${forRead ? 'Reading' : 'Writing to'} Pinecone index ${index} not allowed`); } const client = await this.getClientAsync(); return await (this.indexPromises[index] ?? (this.indexPromises[index] = this.createIndexAsync(client, index))); } async createIndexAsync(client, index) { try { const idx = await client.describeIndex(index); if (idx) { return client.index(index); } } catch (ex) { console.error(`Failed to check status of Pinecone index - ${index}`, ex); } if (!this.options.autoCreateIndex) { throw new Error('Auto create Pinecone index not enabled'); } await client.createIndexForModel({ name: index, cloud: this.options.cloud, region: this.options.region, embed: { model: this.options.model, fieldMap: { text: this.options.indexKey } }, waitUntilReady: true, }); return client.index(index); } isReadIndexAllowed(index) { return this.options.allowedReadIndexes ? this.options.allowedReadIndexes.includes(index) : this.options.index === index; } isWriteIndexAllowed(index) { return this.options.allowedWriteIndexes?.includes(index) ?? false; } getMountedIndex(path) { if (!path || !this.options.indexMounts) { return { index: this.options.index, path }; } if (path.startsWith('/')) { path = path.substring(1); } const rootedPath = '/' + path; for (const mnt of this.options.indexMounts) { if (path.startsWith(mnt.path) || rootedPath.startsWith(mnt.path)) { return { index: mnt.index, path, }; } } return undefined; } async searchAsync(search) { if (!search.content) { return { items: [], usage: (0, convo_lang_1.createEmptyConvoTokenUsage)() }; } if (search.paths) { const indexMap = {}; for (const p of search.paths) { const mnt = this.getMountedIndex(p); if (!mnt) { continue; } (indexMap[mnt.index] ?? (indexMap[mnt.index] = [])).push(mnt.path); } const keys = Object.keys(indexMap); if (!keys.length) { return { items: [], usage: (0, convo_lang_1.createEmptyConvoTokenUsage)() }; } const r = await Promise.all(keys.map(k => this.searchIndexAsync(search, k, indexMap[k]?.filter(p => p)))); return (0, convo_lang_1.mergeConvoRagResults)(r); } else { return await this.searchIndexAsync(search, this.options.index); } } async searchIndexAsync(search, index, paths) { if (!search.content) { return { items: [], usage: (0, convo_lang_1.createEmptyConvoTokenUsage)() }; } const idx = await this.getIndexAsync(index, true); const r = await idx.searchRecords({ query: { topK: search.limit ?? convo_lang_1.defaultConvoRagSearchLimit, inputs: { text: search.content }, } }); return { usage: (0, convo_lang_1.createEmptyConvoTokenUsage)(), items: r.result.hits.filter(h => h._score <= search.tolerance).map(h => { const content = h.fields?.[this.options.indexKey] ?? ''; return { id: h._id, distance: h._score, document: { sourceId: h._id, content, } }; }) }; } async upsertAsync(documents) { const indexMap = {}; for (const doc of documents) { const mnt = this.getMountedIndex(doc.path); if (!mnt) { continue; } (indexMap[mnt.index] ?? (indexMap[mnt.index] = [])).push({ mntPath: mnt.path, doc, }); } const keys = Object.keys(indexMap); await Promise.all(keys.map(async (index) => { const docs = indexMap[index]; if (!docs) { return; } const idx = await this.getIndexAsync(index, false); await idx.upsertRecords(docs.map(d => { const doc = { ...d.doc }; delete doc.vector; if (doc.sourceId) { doc._id = doc.sourceId; delete doc.sourceId; } return doc; })); })); } } exports.ConvoPineconeRagService = ConvoPineconeRagService; //# sourceMappingURL=ConvoPineconeRagService.js.map