diff --git a/packages/components/credentials/AstraApi.credential.ts b/packages/components/credentials/AstraApi.credential.ts new file mode 100644 index 000000000..ad4c65a81 --- /dev/null +++ b/packages/components/credentials/AstraApi.credential.ts @@ -0,0 +1,34 @@ +import { INodeParams, INodeCredential } from '../src/Interface' + +class AstraApi implements INodeCredential { + label: string + name: string + version: number + description: string + inputs: INodeParams[] + + constructor() { + this.label = 'Astra API' + this.name = 'AstraApi' + this.version = 1.0 + this.inputs = [ + { + label: 'Colection Name', + name: 'collectionName', + type: 'string' + }, + { + label: 'Astra DB Application Token', + name: 'applicationToken', + type: 'password' + }, + { + label: 'Astra DB Api Endpoint', + name: 'dbEndPoint', + type: 'string' + } + ] + } +} + +module.exports = { credClass: AstraApi } diff --git a/packages/components/nodes/vectorstores/Astra/Astra.ts b/packages/components/nodes/vectorstores/Astra/Astra.ts new file mode 100644 index 000000000..648a8b491 --- /dev/null +++ b/packages/components/nodes/vectorstores/Astra/Astra.ts @@ -0,0 +1,190 @@ +import { flatten } from 'lodash' +import { Embeddings } from 'langchain/embeddings/base' +import { Document } from 'langchain/document' +import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface' +import { getBaseClasses, getCredentialData } from '../../../src/utils' +import { AstraDBVectorStore, AstraLibArgs } from '@langchain/community/vectorstores/astradb' + +class Astra_VectorStores implements INode { + label: string + name: string + version: number + description: string + type: string + icon: string + category: string + badge: string + baseClasses: string[] + inputs: INodeParams[] + credential: INodeParams + outputs: INodeOutputsValue[] + + constructor() { + this.label = 'Astra' + this.name = 'Astra' + this.version = 1.0 + this.type = 'Astra' + this.icon = 'astra.svg' + this.category = 'Vector Stores' + this.description = `Upsert embedded data and perform similarity search upon query using DataStax Astra DB, a serverless vector database that’s perfect for managing mission-critical AI workloads` + this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever'] + this.badge = 'NEW' + this.credential = { + label: 'Connect Credential', + name: 'credential', + type: 'credential', + credentialNames: ['AstraApi'] + } + this.inputs = [ + { + label: 'Document', + name: 'document', + type: 'Document', + list: true, + optional: true + }, + { + label: 'Embeddings', + name: 'embeddings', + type: 'Embeddings' + }, + { + label: 'Vector Dimension', + name: 'vectorDimension', + type: 'number', + placeholder: '1536', + optional: true, + description: 'Dimension used for storing vector embedding' + }, + { + label: 'Similarity Metric', + name: 'similarityMetric', + type: 'string', + placeholder: 'cosine', + optional: true, + description: 'cosine | euclidean | dot_product' + }, + { + label: 'Top K', + name: 'topK', + description: 'Number of top results to fetch. Default to 4', + placeholder: '4', + type: 'number', + additionalParams: true, + optional: true + } + ] + this.outputs = [ + { + label: 'Astra Retriever', + name: 'retriever', + baseClasses: this.baseClasses + }, + { + label: 'Astra Vector Store', + name: 'vectorStore', + baseClasses: [this.type, ...getBaseClasses(AstraDBVectorStore)] + } + ] + } + + //@ts-ignore + vectorStoreMethods = { + async upsert(nodeData: INodeData, options: ICommonObject): Promise { + const docs = nodeData.inputs?.document as Document[] + const embeddings = nodeData.inputs?.embeddings as Embeddings + const vectorDimension = nodeData.inputs?.vectorDimension as number + const similarityMetric = nodeData.inputs?.similarityMetric as 'cosine' | 'euclidean' | 'dot_product' | undefined + const credentialData = await getCredentialData(nodeData.credential ?? '', options) + + const expectedSimilarityMetric = ['cosine', 'euclidean', 'dot_product'] + if (similarityMetric && !expectedSimilarityMetric.includes(similarityMetric)) { + throw new Error(`Invalid Similarity Metric should be one of 'cosine' | 'euclidean' | 'dot_product'`) + } + + const clientConfig = { + token: credentialData?.applicationToken ?? 'dummy', + endpoint: credentialData?.dbEndPoint ?? 'dummy' + } + + const astraConfig: AstraLibArgs = { + ...clientConfig, + collection: credentialData.collectionName ?? 'flowise_test', + collectionOptions: { + vector: { + dimension: vectorDimension ?? 1536, + metric: similarityMetric ?? 'cosine' + } + } + } + + const flattenDocs = docs && docs.length ? flatten(docs) : [] + const finalDocs = [] + for (let i = 0; i < flattenDocs.length; i += 1) { + if (flattenDocs[i] && flattenDocs[i].pageContent) { + finalDocs.push(new Document(flattenDocs[i])) + } + } + + try { + await AstraDBVectorStore.fromDocuments(finalDocs, embeddings, astraConfig) + } catch (e) { + throw new Error(e) + } + } + } + + async init(nodeData: INodeData, _: string, options: ICommonObject): Promise { + const docs = nodeData.inputs?.document as Document[] + const embeddings = nodeData.inputs?.embeddings as Embeddings + const vectorDimension = nodeData.inputs?.vectorDimension as number + const similarityMetric = nodeData.inputs?.similarityMetric as 'cosine' | 'euclidean' | 'dot_product' | undefined + const output = nodeData.outputs?.output as string + const topK = nodeData.inputs?.topK as string + const k = topK ? parseFloat(topK) : 4 + + const credentialData = await getCredentialData(nodeData.credential ?? '', options) + + const expectedSimilarityMetric = ['cosine', 'euclidean', 'dot_product'] + if (similarityMetric && !expectedSimilarityMetric.includes(similarityMetric)) { + throw new Error(`Invalid Similarity Metric should be one of 'cosine' | 'euclidean' | 'dot_product'`) + } + + const clientConfig = { + token: credentialData?.applicationToken ?? 'dummy', + endpoint: credentialData?.dbEndPoint ?? 'dummy' + } + + const astraConfig: AstraLibArgs = { + ...clientConfig, + collection: credentialData.collectionName ?? 'flowise_test', + collectionOptions: { + vector: { + dimension: vectorDimension ?? 1536, + metric: similarityMetric ?? 'cosine' + } + } + } + + const flattenDocs = docs && docs.length ? flatten(docs) : [] + const finalDocs = [] + for (let i = 0; i < flattenDocs.length; i += 1) { + if (flattenDocs[i] && flattenDocs[i].pageContent) { + finalDocs.push(new Document(flattenDocs[i])) + } + } + + const vectorStore = await AstraDBVectorStore.fromExistingIndex(embeddings, astraConfig) + + if (output === 'retriever') { + const retriever = vectorStore.asRetriever(k) + return retriever + } else if (output === 'vectorStore') { + ;(vectorStore as any).k = k + return vectorStore + } + return vectorStore + } +} + +module.exports = { nodeClass: Astra_VectorStores } diff --git a/packages/components/nodes/vectorstores/Astra/astra.svg b/packages/components/nodes/vectorstores/Astra/astra.svg new file mode 100644 index 000000000..59c2fc3f6 --- /dev/null +++ b/packages/components/nodes/vectorstores/Astra/astra.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/packages/components/package.json b/packages/components/package.json index a77d91e4b..ff0b687b4 100644 --- a/packages/components/package.json +++ b/packages/components/package.json @@ -19,6 +19,7 @@ "@aws-sdk/client-bedrock-runtime": "3.422.0", "@aws-sdk/client-dynamodb": "^3.360.0", "@aws-sdk/client-s3": "^3.427.0", + "@datastax/astra-db-ts": "^0.1.2", "@dqbd/tiktoken": "^1.0.7", "@elastic/elasticsearch": "^8.9.0", "@getzep/zep-js": "^0.9.0", @@ -26,6 +27,7 @@ "@gomomento/sdk-core": "^1.51.1", "@google-ai/generativelanguage": "^0.2.1", "@huggingface/inference": "^2.6.1", + "@langchain/community": "^0.0.16", "@langchain/google-genai": "^0.0.6", "@langchain/mistralai": "^0.0.6", "@notionhq/client": "^2.2.8",