From f8e081f1df0fb51192e84a57ee6f269f977e4ef1 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Mon, 22 May 2023 18:19:29 +0100
Subject: [PATCH] add url to chroma

---
 .../Chroma_Existing/Chroma_Existing.ts          | 17 ++++++++++++++---
 .../vectorstores/Chroma_Upsert/Chroma_Upsert.ts | 17 ++++++++++++++---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts b/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts
index ec9d28a176a..fbaa5cbb45a 100644
--- a/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts
+++ b/packages/components/nodes/vectorstores/Chroma_Existing/Chroma_Existing.ts
@@ -32,6 +32,12 @@ class Chroma_Existing_VectorStores implements INode {
                 label: 'Collection Name',
                 name: 'collectionName',
                 type: 'string'
+            },
+            {
+                label: 'Chroma URL',
+                name: 'chromaURL',
+                type: 'string',
+                optional: true
             }
         ]
         this.outputs = [
@@ -51,11 +57,16 @@ class Chroma_Existing_VectorStores implements INode {
     async init(nodeData: INodeData): Promise<any> {
         const collectionName = nodeData.inputs?.collectionName as string
         const embeddings = nodeData.inputs?.embeddings as Embeddings
+        const chromaURL = nodeData.inputs?.chromaURL as string
         const output = nodeData.outputs?.output as string
 
-        const vectorStore = await Chroma.fromExistingCollection(embeddings, {
-            collectionName
-        })
+        const obj: {
+            collectionName: string
+            url?: string
+        } = { collectionName }
+        if (chromaURL) obj.url = chromaURL
+
+        const vectorStore = await Chroma.fromExistingCollection(embeddings, obj)
 
         if (output === 'retriever') {
             const retriever = vectorStore.asRetriever()
diff --git a/packages/components/nodes/vectorstores/Chroma_Upsert/Chroma_Upsert.ts b/packages/components/nodes/vectorstores/Chroma_Upsert/Chroma_Upsert.ts
index b63a9845855..fb7c404e347 100644
--- a/packages/components/nodes/vectorstores/Chroma_Upsert/Chroma_Upsert.ts
+++ b/packages/components/nodes/vectorstores/Chroma_Upsert/Chroma_Upsert.ts
@@ -39,6 +39,12 @@ class ChromaUpsert_VectorStores implements INode {
                 label: 'Collection Name',
                 name: 'collectionName',
                 type: 'string'
+            },
+            {
+                label: 'Chroma URL',
+                name: 'chromaURL',
+                type: 'string',
+                optional: true
             }
         ]
         this.outputs = [
@@ -59,6 +65,7 @@ class ChromaUpsert_VectorStores implements INode {
         const collectionName = nodeData.inputs?.collectionName as string
         const docs = nodeData.inputs?.document as Document[]
         const embeddings = nodeData.inputs?.embeddings as Embeddings
+        const chromaURL = nodeData.inputs?.chromaURL as string
         const output = nodeData.outputs?.output as string
 
         const flattenDocs = docs && docs.length ? docs.flat() : []
@@ -67,9 +74,13 @@ class ChromaUpsert_VectorStores implements INode {
             finalDocs.push(new Document(flattenDocs[i]))
         }
 
-        const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, {
-            collectionName
-        })
+        const obj: {
+            collectionName: string
+            url?: string
+        } = { collectionName }
+        if (chromaURL) obj.url = chromaURL
+
+        const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, obj)
 
         if (output === 'retriever') {
             const retriever = vectorStore.asRetriever()