UNPKG

remcode

Version:

Turn your AI assistant into a codebase expert. Intelligent code analysis, semantic search, and software engineering guidance through MCP integration.

201 lines (200 loc) 8.68 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.PineconeStorage = void 0; const logger_1 = require("../../utils/logger"); const pinecone_1 = require("@pinecone-database/pinecone"); const uuid_1 = require("uuid"); const logger = (0, logger_1.getLogger)('PineconeStorage'); class PineconeStorage { constructor(options) { this.initialized = false; this.client = null; this.index = null; this.options = { ...options, dimension: options.dimension || 768, metric: options.metric || 'cosine', }; } async initialize() { logger.info(`Initializing Pinecone index: ${this.options.indexName}`); if (!this.options.apiKey) { throw new Error('Pinecone API key is required'); } try { // Initialize the Pinecone client with v6 API this.client = new pinecone_1.Pinecone({ apiKey: this.options.apiKey }); // Check if the index exists const indexList = await this.client.listIndexes(); const existingIndex = indexList.indexes?.find(idx => idx.name === this.options.indexName); if (!existingIndex) { logger.info(`Creating new Pinecone index: ${this.options.indexName}`); // Create index using v6 API format await this.client.createIndex({ name: this.options.indexName, spec: { serverless: { cloud: 'aws', region: 'us-east-1' } }, dimension: this.options.dimension || 768, metric: this.options.metric || 'cosine', waitUntilReady: true }); logger.info(`Created new Pinecone index: ${this.options.indexName}`); } else { logger.info(`Using existing Pinecone index: ${this.options.indexName}`); } // Connect to the index this.index = this.client.index(this.options.indexName); logger.info(`Successfully initialized Pinecone storage`); this.initialized = true; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to initialize Pinecone: ${errorMessage}`); throw new Error(`Pinecone initialization failed: ${errorMessage}`); } } async storeVectors(vectors) { if (!this.initialized || !this.index) { throw new Error('Pinecone storage is not initialized'); } logger.info(`Storing ${vectors.length} vectors in Pinecone`); try { const pineconeVectors = vectors.map(vec => ({ id: vec.id || (0, uuid_1.v4)(), values: vec.embedding, metadata: vec.metadata || {}, })); // Batch upsert (100 vectors per batch) const chunkSize = 100; const namespace = this.options.namespace || ''; for (let i = 0; i < pineconeVectors.length; i += chunkSize) { const chunk = pineconeVectors.slice(i, i + chunkSize); await this.index.namespace(namespace).upsert(chunk); logger.debug(`Upserted batch ${Math.floor(i / chunkSize) + 1}/${Math.ceil(pineconeVectors.length / chunkSize)}`); } logger.info(`Successfully stored ${vectors.length} vectors in Pinecone`); } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to store vectors in Pinecone: ${errorMessage}`); throw new Error(`Vector storage failed: ${errorMessage}`); } } async queryVectors(embeddings, topK = 10, filter, namespace) { if (!this.initialized || !this.index) { throw new Error('Pinecone storage is not initialized'); } logger.info(`Querying Pinecone index for top ${topK} matches`); try { const ns = namespace || this.options.namespace || ''; // Build query object - only include filter if it has valid content const queryOptions = { vector: embeddings, topK, includeMetadata: true, includeValues: false }; // Only add filter if it exists and has at least one key-value pair if (filter && Object.keys(filter).length > 0) { queryOptions.filter = filter; } const queryResponse = await this.index.namespace(ns).query(queryOptions); logger.info(`Found ${queryResponse.matches?.length || 0} matches in Pinecone`); return queryResponse.matches || []; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to query vectors in Pinecone: ${errorMessage}`); throw new Error(`Vector query failed: ${errorMessage}`); } } async deleteVectors(ids, deleteAll = false, filter, namespace) { if (!this.initialized || !this.index) { throw new Error('Pinecone storage is not initialized'); } const ns = namespace || this.options.namespace || ''; const indexNamespace = this.index.namespace(ns); try { if (deleteAll) { logger.info(`Deleting all vectors from namespace: ${ns}`); await indexNamespace.deleteAll(); return -1; // All vectors deleted } else if (ids && ids.length > 0) { logger.info(`Deleting ${ids.length} vectors by ID`); // Delete in batches const batchSize = 100; for (let i = 0; i < ids.length; i += batchSize) { const batch = ids.slice(i, i + batchSize); await indexNamespace.deleteMany(batch); } return ids.length; } else if (filter) { logger.info(`Deleting vectors by filter`); await indexNamespace.deleteMany({ filter }); return -1; // Unknown count } else { throw new Error('Either ids, deleteAll flag, or filter must be provided'); } } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to delete vectors: ${errorMessage}`); throw new Error(`Vector deletion failed: ${errorMessage}`); } } async deleteVectorsByMetadata(metadata, namespace) { if (!metadata || Object.keys(metadata).length === 0) { throw new Error('Metadata filter is required'); } logger.info(`Deleting vectors by metadata`); // Build filter for Pinecone v6 const filter = {}; Object.entries(metadata).forEach(([key, value]) => { filter[key] = { $eq: value }; }); return await this.deleteVectors(undefined, false, filter, namespace); } async listIndexes() { if (!this.client) { throw new Error('Pinecone client is not initialized'); } try { const response = await this.client.listIndexes(); return response.indexes || []; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to list indexes: ${errorMessage}`); throw new Error(`Failed to list indexes: ${errorMessage}`); } } async getIndexStats(namespace) { if (!this.initialized || !this.index) { throw new Error('Pinecone storage is not initialized'); } try { const stats = await this.index.describeIndexStats(); const ns = namespace || this.options.namespace || ''; if (ns && stats.namespaces?.[ns]) { return stats.namespaces[ns]; } return stats; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); logger.error(`Failed to get index stats: ${errorMessage}`); throw new Error(`Failed to get index stats: ${errorMessage}`); } } } exports.PineconeStorage = PineconeStorage;