add auth to chroma

This commit is contained in:
Henry 2023-08-03 11:52:17 +01:00
parent d8749bd495
commit 3669eeee41
5 changed files with 71 additions and 60 deletions

View File

@ -1,8 +1,8 @@
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
import { Chroma, ChromaLibArgs } from 'langchain/vectorstores/chroma'
import { Chroma } from 'langchain/vectorstores/chroma'
import { Embeddings } from 'langchain/embeddings/base'
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
import type { Collection } from 'chromadb'
import { ChromaExtended } from './core'
class Chroma_Existing_VectorStores implements INode {
label: string
@ -94,7 +94,7 @@ class Chroma_Existing_VectorStores implements INode {
if (chromaURL) obj.url = chromaURL
if (chromaApiKey) obj.chromaApiKey = chromaApiKey
const vectorStore = await ChromaExisting.fromExistingCollection(embeddings, obj)
const vectorStore = await ChromaExtended.fromExistingCollection(embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever(k)
@ -107,50 +107,4 @@ class Chroma_Existing_VectorStores implements INode {
}
}
interface ChromaAuth {
chromaApiKey?: string
}
class ChromaExisting extends Chroma {
chromaApiKey?: string
constructor(embeddings: Embeddings, args: ChromaLibArgs & Partial<ChromaAuth>) {
super(embeddings, args)
this.chromaApiKey = args.chromaApiKey
}
static async fromExistingCollection(embeddings: Embeddings, dbConfig: ChromaLibArgs & Partial<ChromaAuth>): Promise<Chroma> {
const instance = new this(embeddings, dbConfig)
await instance.ensureCollection()
return instance
}
async ensureCollection(): Promise<Collection> {
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 }

View File

@ -1,9 +1,10 @@
import { INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
import { Chroma } from 'langchain/vectorstores/chroma'
import { Embeddings } from 'langchain/embeddings/base'
import { Document } from 'langchain/document'
import { getBaseClasses } from '../../../src/utils'
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
import { flatten } from 'lodash'
import { ChromaExtended } from './core'
class ChromaUpsert_VectorStores implements INode {
label: string
@ -15,6 +16,7 @@ class ChromaUpsert_VectorStores implements INode {
category: string
baseClasses: string[]
inputs: INodeParams[]
credential: INodeParams
outputs: INodeOutputsValue[]
constructor() {
@ -26,6 +28,14 @@ class ChromaUpsert_VectorStores implements INode {
this.category = 'Vector Stores'
this.description = 'Upsert documents to Chroma'
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: 'Document',
@ -73,7 +83,7 @@ class ChromaUpsert_VectorStores implements INode {
]
}
async init(nodeData: INodeData): Promise<any> {
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
const collectionName = nodeData.inputs?.collectionName as string
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
@ -82,6 +92,9 @@ class ChromaUpsert_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 flattenDocs = docs && docs.length ? flatten(docs) : []
const finalDocs = []
for (let i = 0; i < flattenDocs.length; i += 1) {
@ -91,10 +104,12 @@ class ChromaUpsert_VectorStores implements INode {
const obj: {
collectionName: string
url?: string
chromaApiKey?: string
} = { collectionName }
if (chromaURL) obj.url = chromaURL
if (chromaApiKey) obj.chromaApiKey = chromaApiKey
const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, obj)
const vectorStore = await ChromaExtended.fromDocuments(finalDocs, embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever(k)

View File

Before

Width:  |  Height:  |  Size: 622 B

After

Width:  |  Height:  |  Size: 622 B

View File

@ -0,0 +1,49 @@
import { Chroma, ChromaLibArgs } from 'langchain/vectorstores/chroma'
import { Embeddings } from 'langchain/embeddings/base'
import type { Collection } from 'chromadb'
interface ChromaAuth {
chromaApiKey?: string
}
export class ChromaExtended extends Chroma {
chromaApiKey?: string
constructor(embeddings: Embeddings, args: ChromaLibArgs & Partial<ChromaAuth>) {
super(embeddings, args)
this.chromaApiKey = args.chromaApiKey
}
static async fromExistingCollection(embeddings: Embeddings, dbConfig: ChromaLibArgs & Partial<ChromaAuth>): Promise<Chroma> {
const instance = new this(embeddings, dbConfig)
await instance.ensureCollection()
return instance
}
async ensureCollection(): Promise<Collection> {
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
}
}

View File

@ -1,7 +0,0 @@
<svg width="209" height="135" viewBox="0 0 209 135" fill="none" xmlns="http://www.w3.org/2000/svg">
<ellipse cx="136.019" cy="67.2304" rx="66.6667" ry="64" fill="#FFDE2D"/>
<ellipse cx="69.352" cy="67.2304" rx="66.6667" ry="64" fill="#327EFF"/>
<path d="M2.68528 67.2304C2.68527 31.8842 32.5329 3.23047 69.3519 3.23047L69.3519 67.2304L2.68528 67.2304Z" fill="#327EFF"/>
<path d="M136.019 67.2305C136.019 102.577 106.171 131.23 69.3519 131.23L69.3519 67.2305L136.019 67.2305Z" fill="#FF6446"/>
<path d="M69.352 67.2304C69.352 31.8842 99.1997 3.23047 136.019 3.23047L136.019 67.2304L69.352 67.2304Z" fill="#FF6446"/>
</svg>

Before

Width:  |  Height:  |  Size: 622 B