@smythos/sre 1.7.18 → 1.7.40

package/src/subsystems/IO/VectorDB.service/connectors/MilvusVectorDB.class.ts

@@ -1,4 +1,3 @@
-//==[ SRE: S3Storage ]======================
 import { ConnectorService } from '@sre/Core/ConnectorsService';
 import { JSONContentHelper } from '@sre/helpers/JsonContent.helper';
 import { Logger } from '@sre/helpers/Log.helper';
@@ -17,12 +16,21 @@ import {
     VectorDBResult,
     VectorsResultData,
 } from '@sre/types/VectorDB.types';
+import { calcSizeMb } from '@sre/utils/string.utils';
 import { CreateIndexSimpleReq, DataType, ErrorCode, FieldType, MilvusClient } from '@zilliz/milvus2-sdk-node';
 import crypto from 'crypto';
 import { jsonrepair } from 'jsonrepair';
 import { EmbeddingsFactory } from '../embed';
 import { BaseEmbedding, TEmbeddings } from '../embed/BaseEmbedding';
 import { DeleteTarget, VectorDBConnector } from '../VectorDBConnector';
+import { NKVConnector } from '@sre/IO/NKV.service/NKVConnector';
+
+//! Due to current bug (still investigating), any consumer of this connector
+//! needs to install @zilliz/milvus2-sdk-node in the package.json of the project so it can work
+
+//* Note, we are storing Datasources info inside both NKV and Milvus (using some quirks).
+//* The connector favors NKV as storage number 1, and Milvus acting as fallback storage.
+//* This is because Milvus operations are heavy and can be slow the more u add big datasources
 
 const console = Logger('Milvus');
 
@@ -42,7 +50,7 @@ export type MilvusConfig = {
 };
 
 // Define schema field names as a type for strong typing
-type SchemaFieldNames = 'id' | 'text' | 'namespaceId' | 'datasourceId' | 'datasourceLabel' | 'vector' | 'acl' | 'user_metadata';
+type SchemaFieldNames = 'id' | 'text' | 'namespaceId' | 'datasourceId' | 'datasourceLabel' | 'vector' | 'acl' | 'user_metadata' | 'smyth_metadata';
 
 type SchemaField = FieldType & { name: SchemaFieldNames };
 
@@ -55,6 +63,7 @@ export class MilvusVectorDB extends VectorDBConnector {
     public embedder: BaseEmbedding;
     private SCHEMA_DEFINITION: SchemaField[];
     private INDEX_PARAMS: IndexParams;
+    private nkvConnector: NKVConnector;
 
     constructor(protected _settings: MilvusConfig) {
         super(_settings);
@@ -72,16 +81,22 @@ export class MilvusVectorDB extends VectorDBConnector {
 
         console.log('clientConfig', clientConfig);
 
-        this.client = new MilvusClient(clientConfig);
+        this.client = new MilvusClient(clientConfig, undefined, undefined, undefined, {
+            'grpc.max_receive_message_length': 50 * 1024 * 1024,
+            'grpc.max_send_message_length': 50 * 1024 * 1024,
+            max_receive_message_length: 50 * 1024 * 1024,
+            max_send_message_length: 50 * 1024 * 1024,
+        });
         console.info('Milvus client initialized');
         this.accountConnector = ConnectorService.getAccountConnector();
         this.cache = ConnectorService.getCacheConnector();
+        this.nkvConnector = ConnectorService.getNKVConnector();
 
         if (!_settings.embeddings) {
             _settings.embeddings = { provider: 'OpenAI', model: 'text-embedding-3-large', dimensions: 3072 };
         }
 
-        if (!_settings.embeddings.dimensions) _settings.embeddings.dimensions = 3072;
+        if (!_settings.embeddings?.dimensions) _settings.embeddings.dimensions = 3072;
 
         this.embedder = EmbeddingsFactory.create(_settings.embeddings.provider, _settings.embeddings);
 
@@ -128,6 +143,11 @@ export class MilvusVectorDB extends VectorDBConnector {
                 data_type: DataType.VarChar,
                 max_length: 2048,
             },
+            {
+                name: 'smyth_metadata',
+                data_type: DataType.VarChar,
+                max_length: 65535,
+            },
         ];
         this.INDEX_PARAMS = {
             index_type: 'AUTOINDEX',
@@ -192,6 +212,11 @@ export class MilvusVectorDB extends VectorDBConnector {
             throw new Error(`Error dropping collection: ${res}`);
         }
 
+        // delete the linked datasources from nkv
+        await this.nkvConnector
+            .requester(acRequest.candidate as AccessCandidate)
+            .deleteAll(`vectorDB:${this.id}:namespaces:${preparedNs}:datasources`);
+
         await this.deleteACL(AccessCandidate.clone(acRequest.candidate), namespace);
     }
 
@@ -213,10 +238,13 @@ export class MilvusVectorDB extends VectorDBConnector {
         const result = await this.client.search({
             data: _vector as number[],
             collection_name: preparedNs,
-            output_fields: ['id', 'text', this.USER_METADATA_KEY, 'namespaceId', 'datasourceId', 'datasourceLabel', 'vector'],
             limit: options.topK || 10,
         });
 
+        if (result.status.error_code !== ErrorCode.SUCCESS) {
+            throw new Error(`Error searching data: ${result.status.detail}`);
+        }
+
         return result.results.map((match) => {
             let _record = match;
             if (match?.[this.USER_METADATA_KEY]) {
@@ -249,17 +277,45 @@ export class MilvusVectorDB extends VectorDBConnector {
 
         const sourceType = this.embedder.detectSourceType(sourceWrapper[0].source);
         if (sourceType === 'unknown' || sourceType === 'url') throw new Error('Unsupported source type');
+
         const transformedSource = await this.embedder.transformSource(sourceWrapper, sourceType, acRequest.candidate as AccessCandidate);
-        const preparedSource: Record<SchemaFieldNames, any>[] = transformedSource.map((s) => ({
-            id: s.id,
-            text: s.metadata?.text,
-            user_metadata: s.metadata?.[this.USER_METADATA_KEY],
-            namespaceId: preparedNs,
-            datasourceId: s.metadata?.datasourceId,
-            datasourceLabel: s.metadata?.datasourceLabel,
-            vector: s.source,
-            acl: s.metadata?.acl,
-        }));
+
+        const liveSchema = await this.client.describeCollection({
+            collection_name: preparedNs,
+        });
+
+        // only incl the fields that are in the live schema
+
+        const preparedSource: Partial<Record<SchemaFieldNames, any>>[] = transformedSource.map((s) => {
+            function schemaFieldExists(field: SchemaFieldNames) {
+                return liveSchema.schema.fields.some((f) => f.name === field);
+            }
+            return {
+                id: s.id,
+                text: s.metadata?.text,
+                user_metadata: s.metadata?.[this.USER_METADATA_KEY],
+                namespaceId: preparedNs,
+                datasourceId: s.metadata?.datasourceId, // legacy field
+                datasourceLabel: s.metadata?.datasourceLabel, // legacy field
+                vector: s.source,
+                acl: s.metadata?.acl,
+                ...(schemaFieldExists('smyth_metadata')
+                    ? {
+                          smyth_metadata: JSON.stringify({
+                              datasource: {
+                                  id: s.metadata?.datasourceId,
+                                  label: s.metadata?.datasourceLabel,
+                                  chunkSize: s.metadata?.chunkSize,
+                                  chunkOverlap: s.metadata?.chunkOverlap,
+                                  createdAt: s.metadata?.createdAt,
+                                  datasourceSizeMb: s.metadata?.datasourceSizeMb,
+                                  chunkIndex: s.metadata?.chunkIndex,
+                              },
+                          }),
+                      }
+                    : {}),
+            };
+        });
 
         const res = await this.client.insert({
             collection_name: preparedNs,
@@ -322,6 +378,9 @@ export class MilvusVectorDB extends VectorDBConnector {
         const acl = new ACL().addAccess(acRequest.candidate.role, acRequest.candidate.id, TAccessLevel.Owner);
         const dsId = datasource.id || crypto.randomUUID();
 
+        if (!datasource.chunkSize) datasource.chunkSize = 1000;
+        if (!datasource.chunkOverlap) datasource.chunkOverlap = 200;
+
         const formattedNs = this.constructNsName(acRequest.candidate as AccessCandidate, namespace);
         const chunkedText = this.embedder.chunkText(datasource.text, {
             chunkSize: datasource.chunkSize,
@@ -329,6 +388,7 @@ export class MilvusVectorDB extends VectorDBConnector {
         });
         const ids = Array.from({ length: chunkedText.length }, (_, i) => crypto.randomUUID());
         const label = datasource.label || 'Untitled';
+        const totalSizeMb = calcSizeMb(datasource.text);
         const source: IVectorDataSourceDto[] = chunkedText.map<IVectorDataSourceDto>((doc, i) => {
             return {
                 id: ids[i],
@@ -338,6 +398,11 @@ export class MilvusVectorDB extends VectorDBConnector {
                     namespaceId: formattedNs,
                     datasourceId: dsId,
                     datasourceLabel: label,
+                    chunkSize: datasource?.chunkSize?.toString(),
+                    chunkOverlap: datasource?.chunkOverlap?.toString(),
+                    createdAt: new Date().getTime().toString(),
+                    datasourceSizeMb: totalSizeMb.toString(),
+                    chunkIndex: i,
                     user_metadata: datasource.metadata ? jsonrepair(JSON.stringify(datasource.metadata)) : JSON.stringify({}),
                 },
             };
@@ -354,7 +419,16 @@ export class MilvusVectorDB extends VectorDBConnector {
             text: datasource.text,
             vectorIds: _vIds.map((v) => v.id),
             id: dsId,
+            chunkSize: datasource.chunkSize,
+            chunkOverlap: datasource.chunkOverlap,
+            createdAt: new Date(),
+            datasourceSizeMb: totalSizeMb,
         };
+
+        await this.nkvConnector
+            .requester(acRequest.candidate as AccessCandidate)
+            .set(`vectorDB:${this.id}:namespaces:${formattedNs}:datasources`, dsId, JSON.stringify(dsData));
+
         if (datasource.returnFullVectorInfo) {
             dsData.vectorInfo = _vIds;
         }
@@ -367,6 +441,10 @@ export class MilvusVectorDB extends VectorDBConnector {
         const formattedNs = this.constructNsName(acRequest.candidate as AccessCandidate, namespace);
 
         await this.delete(acRequest, namespace, { datasourceId });
+
+        await this.nkvConnector
+            .requester(acRequest.candidate as AccessCandidate)
+            .delete(`vectorDB:${this.id}:namespaces:${formattedNs}:datasources`, datasourceId);
     }
 
     @SecureConnector.AccessControl
@@ -374,12 +452,26 @@ export class MilvusVectorDB extends VectorDBConnector {
         //const teamId = await this.accountConnector.getCandidateTeam(acRequest.candidate);
         const formattedNs = this.constructNsName(acRequest.candidate as AccessCandidate, namespace);
 
+        //* [1] Get datasources from NKV
+        try {
+            const nkvDatasources = await this.nkvConnector
+                .requester(acRequest.candidate as AccessCandidate)
+                .list(`vectorDB:${this.id}:namespaces:${formattedNs}:datasources`)
+                .then((ds) => ds.map((d) => JSONContentHelper.create(d.data?.toString()).tryParse() as IStorageVectorDataSource));
+
+            return nkvDatasources;
+        } catch (error) {
+            console.error('[NKV] Error listing datasources: ', error);
+        }
+
+        console.info('Trying to get datasources from Milvus');
+        //* [2] Get datasources from Milvus [EXPENSIVE OPERATION] that may exhaust rpc channel
         // Use queryIterator for memory-efficient pagination
         const batchSize = 1000; // Process 1000 records at a time
         const iterator = await this.client.queryIterator({
             collection_name: formattedNs,
             batchSize: batchSize,
-            output_fields: ['id', 'text', this.USER_METADATA_KEY, 'namespaceId', 'datasourceId', 'datasourceLabel', 'vector'],
+            // output_fields: ['id', 'text', this.USER_METADATA_KEY, 'namespaceId', 'datasourceId', 'datasourceLabel', 'vector'],
         });
 
         // Group records by datasourceId using Map for efficient lookups
@@ -395,16 +487,19 @@ export class MilvusVectorDB extends VectorDBConnector {
                             namespaceId: formattedNs,
                             candidateId: acRequest.candidate.id,
                             candidateRole: acRequest.candidate.role,
-                            text: record.text,
+                            // text: record.text,
                             name: record.datasourceLabel,
                             metadata: record[this.USER_METADATA_KEY]
                                 ? JSONContentHelper.create(record[this.USER_METADATA_KEY].toString()).tryParse()
                                 : undefined,
-                            vectorIds: [],
                             id: datasourceId,
+                            vectorIds: [], // to be filled iteratively
+                            text: '', // to be filled iteratively
                         });
                     }
                     datasourceMap.get(datasourceId)!.vectorIds.push(record.id);
+                    // the text here represents the total text of the datasource (not a vector-stored text)
+                    datasourceMap.get(datasourceId)!.text += record.text;
                 }
             }
         } finally {
@@ -418,10 +513,41 @@ export class MilvusVectorDB extends VectorDBConnector {
     protected async getDatasource(acRequest: AccessRequest, namespace: string, datasourceId: string): Promise<IStorageVectorDataSource | undefined> {
         //const teamId = await this.accountConnector.getCandidateTeam(acRequest.candidate);
         const formattedNs = this.constructNsName(acRequest.candidate as AccessCandidate, namespace);
+
+        //* [1] Get datasource from NKV
+        try {
+            const nkvDatasource = await this.nkvConnector
+                .requester(acRequest.candidate as AccessCandidate)
+                .get(`vectorDB:${this.id}:namespaces:${formattedNs}:datasources`, datasourceId)
+                .then((ds) => JSONContentHelper.create(ds?.toString()).tryParse() as IStorageVectorDataSource);
+
+            if (nkvDatasource) {
+                return nkvDatasource;
+            } else {
+                console.info('Datasource not found in NKV');
+            }
+        } catch (error) {
+            console.error('[NKV] Error getting datasource: ', error);
+        }
+
+        console.info('Trying to get datasource from Milvus');
+
+        //* [2] Get datasource from Milvus
         const res = await this.client.query({
             collection_name: formattedNs,
             expr: `datasourceId == "${datasourceId}"`,
-            output_fields: ['id', 'text', this.USER_METADATA_KEY, 'namespaceId', 'datasourceId', 'datasourceLabel', 'vector'],
+            // output_fields: [
+            //     'id',
+            //     'text',
+            //     this.USER_METADATA_KEY,
+            //     'namespaceId',
+            //     'datasourceId',
+            //     'datasourceLabel',
+            //     'vector',
+            //     'chunkSize',
+            //     'chunkOverlap',
+            //     'createdAt',
+            // ],
         });
         // if 0 results, throw error
         if (res.data.length === 0) {

package/src/subsystems/IO/VectorDB.service/connectors/PineconeVectorDB.class.ts

@@ -24,20 +24,29 @@ import { CacheConnector } from '@sre/MemoryManager/Cache.service/CacheConnector'
 import crypto from 'crypto';
 import { BaseEmbedding, TEmbeddings } from '../embed/BaseEmbedding';
 import { EmbeddingsFactory, SupportedProviders, SupportedModels } from '../embed';
-
+import { calcSizeMb } from '@sre/utils/string.utils';
 import { jsonrepair } from 'jsonrepair';
+import { chunkArr } from '@sre/utils/array.utils';
 
 const console = Logger('Pinecone VectorDB');
 
 export type PineconeConfig = {
     /**
-     * The Pinecone API key
+     * The Pinecone API key [LEGACY]
      */
-    apiKey: string;
+    apiKey?: string;
     /**
-     * The Pinecone index name
+     * The Pinecone index name [LEGACY]
      */
-    indexName: string;
+    indexName?: string;
+
+    /**
+     * The Pinecone credentials [New unified format]
+     */
+    credentials?: {
+        apiKey: string;
+        indexName: string;
+    };
     /**
      * The embeddings model to use
      */
@@ -55,28 +64,28 @@ export class PineconeVectorDB extends VectorDBConnector {
 
     constructor(protected _settings: PineconeConfig) {
         super(_settings);
-        if (!_settings.apiKey) {
+        if (!_settings.apiKey && !_settings?.credentials?.apiKey) {
             console.warn('Missing Pinecone API key : returning empty Pinecone connector');
             return;
         }
-        if (!_settings.indexName) {
+        if (!_settings.indexName && !_settings?.credentials?.indexName) {
             console.warn('Missing Pinecone index name : returning empty Pinecone connector');
             return;
         }
 
         this.client = new Pinecone({
-            apiKey: _settings.apiKey,
+            apiKey: _settings.apiKey || _settings.credentials?.apiKey,
         });
         console.info('Pinecone client initialized');
-        console.info('Pinecone index name:', _settings.indexName);
-        this.indexName = _settings.indexName;
+        console.info('Pinecone index name:', _settings.indexName || _settings.credentials?.indexName);
+        this.indexName = _settings.indexName || _settings.credentials?.indexName;
         this.accountConnector = ConnectorService.getAccountConnector();
         this.cache = ConnectorService.getCacheConnector();
         this.nkvConnector = ConnectorService.getNKVConnector();
         if (!_settings.embeddings) {
             _settings.embeddings = { provider: 'OpenAI', model: 'text-embedding-3-large', dimensions: 3072 };
         }
-        if (!_settings.embeddings.dimensions) _settings.embeddings.dimensions = 3072;
+        if (!_settings.embeddings?.dimensions) _settings.embeddings.dimensions = 3072;
 
         this.embedder = EmbeddingsFactory.create(_settings.embeddings.provider, _settings.embeddings);
     }
@@ -229,17 +238,28 @@ export class PineconeVectorDB extends VectorDBConnector {
         const sourceType = this.embedder.detectSourceType(sourceWrapper[0].source);
         if (sourceType === 'unknown' || sourceType === 'url') throw new Error('Invalid source type');
         const transformedSource = await this.embedder.transformSource(sourceWrapper, sourceType, acRequest.candidate as AccessCandidate);
-        const preparedSource = transformedSource.map((s) => ({
+        const preparedSources = transformedSource.map((s) => ({
             id: s.id,
             values: s.source as number[],
             metadata: s.metadata,
         }));
 
-        // await pineconeStore.addDocuments(chunks, ids);
-        await this.client
-            .Index(this.indexName)
-            .namespace(this.constructNsName(acRequest.candidate as AccessCandidate, namespace))
-            .upsert(preparedSource);
+        // pinecone advices to use batches of 100 at a time
+        const batchSize = 100;
+        const chunkedVectors = chunkArr(preparedSources, batchSize);
+
+        // await this.client
+        //     .Index(this.indexName)
+        //     .namespace(this.constructNsName(acRequest.candidate as AccessCandidate, namespace))
+        //     .upsert(preparedSources);
+
+        const promises = chunkedVectors.map(async (chunk) => {
+            await this.client
+                .Index(this.indexName)
+                .namespace(this.constructNsName(acRequest.candidate as AccessCandidate, namespace))
+                .upsert(chunk);
+        });
+        await Promise.all(promises);
 
         const accessCandidate = acRequest.candidate;
 
@@ -249,7 +269,7 @@ export class PineconeVectorDB extends VectorDBConnector {
             await this.setACL(acRequest, namespace, acl);
         }
 
-        return preparedSource.map((s) => {
+        return preparedSources.map((s) => {
             const { text, acl, user_metadata, ...restMetadata } = s.metadata || {};
             return {
                 id: s.id,
@@ -272,10 +292,16 @@ export class PineconeVectorDB extends VectorDBConnector {
         } else {
             const _ids = Array.isArray(deleteTarget) ? deleteTarget : [deleteTarget];
 
-            const res = await this.client
-                .Index(this.indexName)
-                .namespace(this.constructNsName(acRequest.candidate as AccessCandidate, namespace))
-                .deleteMany(_ids);
+            const batchSize = 800;
+            const chunkedIds = chunkArr(_ids, batchSize);
+
+            const promises = chunkedIds.map(async (chunk) => {
+                await this.client
+                    .Index(this.indexName)
+                    .namespace(this.constructNsName(acRequest.candidate as AccessCandidate, namespace))
+                    .deleteMany(chunk);
+            });
+            await Promise.all(promises);
         }
     }
 
@@ -285,6 +311,9 @@ export class PineconeVectorDB extends VectorDBConnector {
         const acl = new ACL().addAccess(acRequest.candidate.role, acRequest.candidate.id, TAccessLevel.Owner);
         const dsId = datasource.id || crypto.randomUUID();
 
+        if (!datasource.chunkSize) datasource.chunkSize = 2000;
+        if (!datasource.chunkOverlap) datasource.chunkOverlap = 200;
+
         const formattedNs = this.constructNsName(acRequest.candidate as AccessCandidate, namespace);
         const chunkedText = this.embedder.chunkText(datasource.text, {
             chunkSize: datasource.chunkSize,
@@ -301,6 +330,7 @@ export class PineconeVectorDB extends VectorDBConnector {
                     namespaceId: formattedNs,
                     datasourceId: dsId,
                     datasourceLabel: label,
+                    chunkIndex: i,
                     user_metadata: datasource.metadata ? jsonrepair(JSON.stringify(datasource.metadata)) : undefined,
                 },
             };
@@ -317,6 +347,10 @@ export class PineconeVectorDB extends VectorDBConnector {
             text: datasource.text,
             vectorIds: _vIds.map((v) => v.id),
             id: dsId,
+            datasourceSizeMb: calcSizeMb(datasource.text),
+            chunkSize: datasource.chunkSize,
+            chunkOverlap: datasource.chunkOverlap,
+            createdAt: new Date(),
         };
         if (datasource.returnFullVectorInfo) {
             dsData.vectorInfo = _vIds;

package/src/subsystems/IO/VectorDB.service/embed/GoogleEmbedding.ts

@@ -18,6 +18,7 @@ export class GoogleEmbeds extends BaseEmbedding {
     }
 
     async embedTexts(texts: string[], candidate: AccessCandidate): Promise<number[][]> {
+        // we split into batches to avoid provider limits
         const batches = this.chunkArr(this.processTexts(texts), this.batchSize);
 
         const batchRequests = batches.map((batch) => {

package/src/subsystems/IO/VectorDB.service/embed/OpenAIEmbedding.ts

@@ -24,7 +24,7 @@ export class OpenAIEmbeds extends BaseEmbedding {
     protected client: OpenAIClient;
     protected clientConfig: ClientOptions;
 
-    public static models = ['text-embedding-ada-002', 'text-embedding-3-large'];
+    public static models = ['text-embedding-3-large', 'text-embedding-ada-002'];
     public canSpecifyDimensions = true;
 
     constructor(private settings?: Partial<TEmbeddings>) {
@@ -41,6 +41,7 @@ export class OpenAIEmbeds extends BaseEmbedding {
     }
 
     async embedTexts(texts: string[], candidate: AccessCandidate): Promise<number[][]> {
+        // we split into batches to avoid provider limits
         const batches = this.chunkArr(this.processTexts(texts), this.batchSize);
 
         const batchRequests = batches.map((batch) => {

package/src/subsystems/IO/VectorDB.service/embed/index.ts

@@ -34,4 +34,22 @@ export class EmbeddingsFactory {
 
         return new supportedProviders[provider as SupportedProviders].embedder(config);
     }
+
+    public static getProviderByModel(model: SupportedModels): SupportedProviders {
+        return Object.keys(supportedProviders).find((provider) => supportedProviders[provider].models.includes(model)) as SupportedProviders;
+    }
+
+    public static getModels() {
+        return Object.keys(supportedProviders)
+            .reduce((acc, provider) => {
+                acc.push(
+                    ...supportedProviders[provider].models.map((model) => ({
+                        provider,
+                        model,
+                    }))
+                );
+                return acc;
+            }, [] as { provider: SupportedProviders; model: SupportedModels[SupportedProviders] }[])
+            .filter((item) => item.model !== 'text-embedding-ada-002'); //! SPECIAL case for ada-002, it doesn't support dimensions passing
+    }
 }