Skip to content

Commit 1f30963

Browse files
authored
Merge pull request #691 from FlowiseAI/feature/ChromaAuth
Feature/chroma auth
2 parents 4436467 + 3669eee commit 1f30963

File tree

7 files changed

+112
-16
lines changed

7 files changed

+112
-16
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { INodeParams, INodeCredential } from '../src/Interface'
2+
3+
class ChromaApi implements INodeCredential {
4+
label: string
5+
name: string
6+
description: string
7+
version: number
8+
inputs: INodeParams[]
9+
10+
constructor() {
11+
this.label = 'Chroma API'
12+
this.name = 'chromaApi'
13+
this.version = 1.0
14+
this.inputs = [
15+
{
16+
label: 'Chroma Api Key',
17+
name: 'chromaApiKey',
18+
type: 'password'
19+
}
20+
]
21+
}
22+
}
23+
24+
module.exports = { credClass: ChromaApi }

packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts packages/components/nodes/vectorstores/Chroma/Chroma_Existing.ts

+19-4
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
import { INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
1+
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
22
import { Chroma } from 'langchain/vectorstores/chroma'
33
import { Embeddings } from 'langchain/embeddings/base'
4-
import { getBaseClasses } from '../../../src/utils'
4+
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
5+
import { ChromaExtended } from './core'
56

67
class Chroma_Existing_VectorStores implements INode {
78
label: string
@@ -13,6 +14,7 @@ class Chroma_Existing_VectorStores implements INode {
1314
category: string
1415
baseClasses: string[]
1516
inputs: INodeParams[]
17+
credential: INodeParams
1618
outputs: INodeOutputsValue[]
1719

1820
constructor() {
@@ -24,6 +26,14 @@ class Chroma_Existing_VectorStores implements INode {
2426
this.category = 'Vector Stores'
2527
this.description = 'Load existing index from Chroma (i.e: Document has been upserted)'
2628
this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever']
29+
this.credential = {
30+
label: 'Connect Credential',
31+
name: 'credential',
32+
type: 'credential',
33+
description: 'Only needed if you have chroma on cloud services with X-Api-key',
34+
optional: true,
35+
credentialNames: ['chromaApi']
36+
}
2737
this.inputs = [
2838
{
2939
label: 'Embeddings',
@@ -65,21 +75,26 @@ class Chroma_Existing_VectorStores implements INode {
6575
]
6676
}
6777

68-
async init(nodeData: INodeData): Promise<any> {
78+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
6979
const collectionName = nodeData.inputs?.collectionName as string
7080
const embeddings = nodeData.inputs?.embeddings as Embeddings
7181
const chromaURL = nodeData.inputs?.chromaURL as string
7282
const output = nodeData.outputs?.output as string
7383
const topK = nodeData.inputs?.topK as string
7484
const k = topK ? parseFloat(topK) : 4
7585

86+
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
87+
const chromaApiKey = getCredentialParam('chromaApiKey', credentialData, nodeData)
88+
7689
const obj: {
7790
collectionName: string
7891
url?: string
92+
chromaApiKey?: string
7993
} = { collectionName }
8094
if (chromaURL) obj.url = chromaURL
95+
if (chromaApiKey) obj.chromaApiKey = chromaApiKey
8196

82-
const vectorStore = await Chroma.fromExistingCollection(embeddings, obj)
97+
const vectorStore = await ChromaExtended.fromExistingCollection(embeddings, obj)
8398

8499
if (output === 'retriever') {
85100
const retriever = vectorStore.asRetriever(k)

packages/components/nodes/vectorstores/Chroma_Upsert/Chroma_Upsert.ts packages/components/nodes/vectorstores/Chroma/Chroma_Upsert.ts

+19-4
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
import { INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
1+
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
22
import { Chroma } from 'langchain/vectorstores/chroma'
33
import { Embeddings } from 'langchain/embeddings/base'
44
import { Document } from 'langchain/document'
5-
import { getBaseClasses } from '../../../src/utils'
5+
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
66
import { flatten } from 'lodash'
7+
import { ChromaExtended } from './core'
78

89
class ChromaUpsert_VectorStores implements INode {
910
label: string
@@ -15,6 +16,7 @@ class ChromaUpsert_VectorStores implements INode {
1516
category: string
1617
baseClasses: string[]
1718
inputs: INodeParams[]
19+
credential: INodeParams
1820
outputs: INodeOutputsValue[]
1921

2022
constructor() {
@@ -26,6 +28,14 @@ class ChromaUpsert_VectorStores implements INode {
2628
this.category = 'Vector Stores'
2729
this.description = 'Upsert documents to Chroma'
2830
this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever']
31+
this.credential = {
32+
label: 'Connect Credential',
33+
name: 'credential',
34+
type: 'credential',
35+
description: 'Only needed if you have chroma on cloud services with X-Api-key',
36+
optional: true,
37+
credentialNames: ['chromaApi']
38+
}
2939
this.inputs = [
3040
{
3141
label: 'Document',
@@ -73,7 +83,7 @@ class ChromaUpsert_VectorStores implements INode {
7383
]
7484
}
7585

76-
async init(nodeData: INodeData): Promise<any> {
86+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
7787
const collectionName = nodeData.inputs?.collectionName as string
7888
const docs = nodeData.inputs?.document as Document[]
7989
const embeddings = nodeData.inputs?.embeddings as Embeddings
@@ -82,6 +92,9 @@ class ChromaUpsert_VectorStores implements INode {
8292
const topK = nodeData.inputs?.topK as string
8393
const k = topK ? parseFloat(topK) : 4
8494

95+
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
96+
const chromaApiKey = getCredentialParam('chromaApiKey', credentialData, nodeData)
97+
8598
const flattenDocs = docs && docs.length ? flatten(docs) : []
8699
const finalDocs = []
87100
for (let i = 0; i < flattenDocs.length; i += 1) {
@@ -91,10 +104,12 @@ class ChromaUpsert_VectorStores implements INode {
91104
const obj: {
92105
collectionName: string
93106
url?: string
107+
chromaApiKey?: string
94108
} = { collectionName }
95109
if (chromaURL) obj.url = chromaURL
110+
if (chromaApiKey) obj.chromaApiKey = chromaApiKey
96111

97-
const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, obj)
112+
const vectorStore = await ChromaExtended.fromDocuments(finalDocs, embeddings, obj)
98113

99114
if (output === 'retriever') {
100115
const retriever = vectorStore.asRetriever(k)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import { Chroma, ChromaLibArgs } from 'langchain/vectorstores/chroma'
2+
import { Embeddings } from 'langchain/embeddings/base'
3+
import type { Collection } from 'chromadb'
4+
5+
interface ChromaAuth {
6+
chromaApiKey?: string
7+
}
8+
9+
export class ChromaExtended extends Chroma {
10+
chromaApiKey?: string
11+
12+
constructor(embeddings: Embeddings, args: ChromaLibArgs & Partial<ChromaAuth>) {
13+
super(embeddings, args)
14+
this.chromaApiKey = args.chromaApiKey
15+
}
16+
17+
static async fromExistingCollection(embeddings: Embeddings, dbConfig: ChromaLibArgs & Partial<ChromaAuth>): Promise<Chroma> {
18+
const instance = new this(embeddings, dbConfig)
19+
await instance.ensureCollection()
20+
return instance
21+
}
22+
23+
async ensureCollection(): Promise<Collection> {
24+
if (!this.collection) {
25+
if (!this.index) {
26+
const { ChromaClient } = await Chroma.imports()
27+
const obj: any = {
28+
path: this.url
29+
}
30+
if (this.chromaApiKey) {
31+
obj.fetchOptions = {
32+
headers: {
33+
'X-Api-Key': this.chromaApiKey
34+
}
35+
}
36+
}
37+
this.index = new ChromaClient(obj)
38+
}
39+
try {
40+
this.collection = await this.index.getOrCreateCollection({
41+
name: this.collectionName
42+
})
43+
} catch (err) {
44+
throw new Error(`Chroma getOrCreateCollection error: ${err}`)
45+
}
46+
}
47+
return this.collection
48+
}
49+
}

packages/components/nodes/vectorstores/Chroma_Upsert/chroma.svg

-7
This file was deleted.

packages/components/package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
"@types/jsdom": "^21.1.1",
3030
"axios": "^0.27.2",
3131
"cheerio": "^1.0.0-rc.12",
32-
"chromadb": "^1.4.2",
32+
"chromadb": "^1.5.3",
3333
"cohere-ai": "^6.2.0",
3434
"d3-dsv": "2",
3535
"dotenv": "^16.0.0",

0 commit comments

Comments
 (0)