diff --git a/packages/components/credentials/ChromaApi.credential.ts b/packages/components/credentials/ChromaApi.credential.ts new file mode 100644 index 00000000..759c113c --- /dev/null +++ b/packages/components/credentials/ChromaApi.credential.ts @@ -0,0 +1,24 @@ +import { INodeParams, INodeCredential } from '../src/Interface' + +class ChromaApi implements INodeCredential { + label: string + name: string + description: string + version: number + inputs: INodeParams[] + + constructor() { + this.label = 'Chroma API' + this.name = 'chromaApi' + this.version = 1.0 + this.inputs = [ + { + label: 'Chroma Api Key', + name: 'chromaApiKey', + type: 'password' + } + ] + } +} + +module.exports = { credClass: ChromaApi } diff --git a/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts b/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts index 8b84d33e..47266fda 100644 --- a/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts +++ b/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts @@ -1,7 +1,8 @@ -import { INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface' -import { Chroma } from 'langchain/vectorstores/chroma' +import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface' +import { Chroma, ChromaLibArgs } from 'langchain/vectorstores/chroma' import { Embeddings } from 'langchain/embeddings/base' -import { getBaseClasses } from '../../../src/utils' +import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils' +import type { Collection } from 'chromadb' class Chroma_Existing_VectorStores implements INode { label: string @@ -13,6 +14,7 @@ class Chroma_Existing_VectorStores implements INode { category: string baseClasses: string[] inputs: INodeParams[] + credential: INodeParams outputs: INodeOutputsValue[] constructor() { @@ -24,6 +26,14 @@ class Chroma_Existing_VectorStores implements INode { this.category = 'Vector Stores' this.description = 'Load existing index from Chroma (i.e: Document has been upserted)' this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever'] + this.credential = { + label: 'Connect Credential', + name: 'credential', + type: 'credential', + description: 'Only needed if you have chroma on cloud services with X-Api-key', + optional: true, + credentialNames: ['chromaApi'] + } this.inputs = [ { label: 'Embeddings', @@ -65,7 +75,7 @@ class Chroma_Existing_VectorStores implements INode { ] } - async init(nodeData: INodeData): Promise { + async init(nodeData: INodeData, _: string, options: ICommonObject): Promise { const collectionName = nodeData.inputs?.collectionName as string const embeddings = nodeData.inputs?.embeddings as Embeddings const chromaURL = nodeData.inputs?.chromaURL as string @@ -73,13 +83,18 @@ class Chroma_Existing_VectorStores implements INode { const topK = nodeData.inputs?.topK as string const k = topK ? parseFloat(topK) : 4 + const credentialData = await getCredentialData(nodeData.credential ?? '', options) + const chromaApiKey = getCredentialParam('chromaApiKey', credentialData, nodeData) + const obj: { collectionName: string url?: string + chromaApiKey?: string } = { collectionName } if (chromaURL) obj.url = chromaURL + if (chromaApiKey) obj.chromaApiKey = chromaApiKey - const vectorStore = await Chroma.fromExistingCollection(embeddings, obj) + const vectorStore = await ChromaExisting.fromExistingCollection(embeddings, obj) if (output === 'retriever') { const retriever = vectorStore.asRetriever(k) @@ -92,4 +107,50 @@ class Chroma_Existing_VectorStores implements INode { } } +interface ChromaAuth { + chromaApiKey?: string +} + +class ChromaExisting extends Chroma { + chromaApiKey?: string + + constructor(embeddings: Embeddings, args: ChromaLibArgs & Partial) { + super(embeddings, args) + this.chromaApiKey = args.chromaApiKey + } + + static async fromExistingCollection(embeddings: Embeddings, dbConfig: ChromaLibArgs & Partial): Promise { + const instance = new this(embeddings, dbConfig) + await instance.ensureCollection() + return instance + } + + async ensureCollection(): Promise { + if (!this.collection) { + if (!this.index) { + const { ChromaClient } = await Chroma.imports() + const obj: any = { + path: this.url + } + if (this.chromaApiKey) { + obj.fetchOptions = { + headers: { + 'X-Api-Key': this.chromaApiKey + } + } + } + this.index = new ChromaClient(obj) + } + try { + this.collection = await this.index.getOrCreateCollection({ + name: this.collectionName + }) + } catch (err) { + throw new Error(`Chroma getOrCreateCollection error: ${err}`) + } + } + return this.collection + } +} + module.exports = { nodeClass: Chroma_Existing_VectorStores } diff --git a/packages/components/package.json b/packages/components/package.json index 85267d38..374e7949 100644 --- a/packages/components/package.json +++ b/packages/components/package.json @@ -29,7 +29,7 @@ "@types/jsdom": "^21.1.1", "axios": "^0.27.2", "cheerio": "^1.0.0-rc.12", - "chromadb": "^1.4.2", + "chromadb": "^1.5.3", "cohere-ai": "^6.2.0", "d3-dsv": "2", "dotenv": "^16.0.0",