Skip to content

Commit

Permalink
Merge pull request #151 from FlowiseAI/bugfix/ChromaURL
Browse files Browse the repository at this point in the history
Bugfix/Add url to chroma
  • Loading branch information
HenryHengZJ authored May 22, 2023
2 parents 0fb16c7 + f8e081f commit d3de0d3
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand All @@ -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()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand All @@ -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() : []
Expand All @@ -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()
Expand Down

0 comments on commit d3de0d3

Please sign in to comment.