add url to chroma

This commit is contained in:
Henry 2023-05-22 18:19:29 +01:00
parent 0fb16c7849
commit f8e081f1df
2 changed files with 28 additions and 6 deletions

View File

@ -32,6 +32,12 @@ class Chroma_Existing_VectorStores implements INode {
label: 'Collection Name', label: 'Collection Name',
name: 'collectionName', name: 'collectionName',
type: 'string' type: 'string'
},
{
label: 'Chroma URL',
name: 'chromaURL',
type: 'string',
optional: true
} }
] ]
this.outputs = [ this.outputs = [
@ -51,11 +57,16 @@ class Chroma_Existing_VectorStores implements INode {
async init(nodeData: INodeData): Promise<any> { async init(nodeData: INodeData): Promise<any> {
const collectionName = nodeData.inputs?.collectionName as string const collectionName = nodeData.inputs?.collectionName as string
const embeddings = nodeData.inputs?.embeddings as Embeddings const embeddings = nodeData.inputs?.embeddings as Embeddings
const chromaURL = nodeData.inputs?.chromaURL as string
const output = nodeData.outputs?.output as string const output = nodeData.outputs?.output as string
const vectorStore = await Chroma.fromExistingCollection(embeddings, { const obj: {
collectionName collectionName: string
}) url?: string
} = { collectionName }
if (chromaURL) obj.url = chromaURL
const vectorStore = await Chroma.fromExistingCollection(embeddings, obj)
if (output === 'retriever') { if (output === 'retriever') {
const retriever = vectorStore.asRetriever() const retriever = vectorStore.asRetriever()

View File

@ -39,6 +39,12 @@ class ChromaUpsert_VectorStores implements INode {
label: 'Collection Name', label: 'Collection Name',
name: 'collectionName', name: 'collectionName',
type: 'string' type: 'string'
},
{
label: 'Chroma URL',
name: 'chromaURL',
type: 'string',
optional: true
} }
] ]
this.outputs = [ this.outputs = [
@ -59,6 +65,7 @@ class ChromaUpsert_VectorStores implements INode {
const collectionName = nodeData.inputs?.collectionName as string const collectionName = nodeData.inputs?.collectionName as string
const docs = nodeData.inputs?.document as Document[] const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings const embeddings = nodeData.inputs?.embeddings as Embeddings
const chromaURL = nodeData.inputs?.chromaURL as string
const output = nodeData.outputs?.output as string const output = nodeData.outputs?.output as string
const flattenDocs = docs && docs.length ? docs.flat() : [] const flattenDocs = docs && docs.length ? docs.flat() : []
@ -67,9 +74,13 @@ class ChromaUpsert_VectorStores implements INode {
finalDocs.push(new Document(flattenDocs[i])) finalDocs.push(new Document(flattenDocs[i]))
} }
const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, { const obj: {
collectionName collectionName: string
}) url?: string
} = { collectionName }
if (chromaURL) obj.url = chromaURL
const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, obj)
if (output === 'retriever') { if (output === 'retriever') {
const retriever = vectorStore.asRetriever() const retriever = vectorStore.asRetriever()