Merge pull request #292 from FlowiseAI/feature/TopK

Feature/Add top K to vector stores
This commit is contained in:
Henry Heng
2023-06-10 18:03:36 +01:00
committed by GitHub
36 changed files with 925 additions and 238 deletions
@@ -45,8 +45,9 @@ class BabyAGI_Agents implements INode {
const model = nodeData.inputs?.model as BaseChatModel
const vectorStore = nodeData.inputs?.vectorStore as VectorStore
const taskLoop = nodeData.inputs?.taskLoop as string
const k = (vectorStore as any)?.k ?? 4
const babyAgi = BabyAGI.fromLLM(model, vectorStore, parseInt(taskLoop, 10))
const babyAgi = BabyAGI.fromLLM(model, vectorStore, parseInt(taskLoop, 10), k)
return babyAgi
}
@@ -154,18 +154,22 @@ export class BabyAGI {
maxIterations = 3
topK = 4
constructor(
taskCreationChain: TaskCreationChain,
taskPrioritizationChain: TaskPrioritizationChain,
executionChain: ExecutionChain,
vectorStore: VectorStore,
maxIterations: number
maxIterations: number,
topK: number
) {
this.taskCreationChain = taskCreationChain
this.taskPrioritizationChain = taskPrioritizationChain
this.executionChain = executionChain
this.vectorStore = vectorStore
this.maxIterations = maxIterations
this.topK = topK
}
addTask(task: Task) {
@@ -219,7 +223,7 @@ export class BabyAGI {
this.printNextTask(task)
// Step 2: Execute the task
const result = await executeTask(this.vectorStore, this.executionChain, objective, task.task_name)
const result = await executeTask(this.vectorStore, this.executionChain, objective, task.task_name, this.topK)
const thisTaskId = task.task_id
finalResult = result
this.printTaskResult(result)
@@ -257,10 +261,10 @@ export class BabyAGI {
return finalResult
}
static fromLLM(llm: BaseChatModel, vectorstore: VectorStore, maxIterations = 3): BabyAGI {
static fromLLM(llm: BaseChatModel, vectorstore: VectorStore, maxIterations = 3, topK = 4): BabyAGI {
const taskCreationChain = TaskCreationChain.from_llm(llm)
const taskPrioritizationChain = TaskPrioritizationChain.from_llm(llm)
const executionChain = ExecutionChain.from_llm(llm)
return new BabyAGI(taskCreationChain, taskPrioritizationChain, executionChain, vectorstore, maxIterations)
return new BabyAGI(taskCreationChain, taskPrioritizationChain, executionChain, vectorstore, maxIterations, topK)
}
}
@@ -46,7 +46,7 @@ class MultiRetrievalQAChain_Chains implements INode {
for (const vs of vectorStoreRetriever) {
retrieverNames.push(vs.name)
retrieverDescriptions.push(vs.description)
retrievers.push(vs.vectorStore.asRetriever())
retrievers.push(vs.vectorStore.asRetriever((vs.vectorStore as any).k ?? 4))
}
const chain = MultiRetrievalQAChain.fromRetrievers(model, retrieverNames, retrieverDescriptions, retrievers, undefined, {
@@ -40,7 +40,10 @@ class VectorDBQAChain_Chains implements INode {
const model = nodeData.inputs?.model as BaseLanguageModel
const vectorStore = nodeData.inputs?.vectorStore as VectorStore
const chain = VectorDBQAChain.fromLLM(model, vectorStore, { verbose: process.env.DEBUG === 'true' ? true : false })
const chain = VectorDBQAChain.fromLLM(model, vectorStore, {
k: (vectorStore as any)?.k ?? 4,
verbose: process.env.DEBUG === 'true' ? true : false
})
return chain
}
@@ -38,6 +38,15 @@ class Chroma_Existing_VectorStores implements INode {
name: 'chromaURL',
type: 'string',
optional: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -59,6 +68,8 @@ class Chroma_Existing_VectorStores implements INode {
const embeddings = nodeData.inputs?.embeddings as Embeddings
const chromaURL = nodeData.inputs?.chromaURL as string
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const obj: {
collectionName: string
@@ -69,9 +80,10 @@ class Chroma_Existing_VectorStores implements INode {
const vectorStore = await Chroma.fromExistingCollection(embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -46,6 +46,15 @@ class ChromaUpsert_VectorStores implements INode {
name: 'chromaURL',
type: 'string',
optional: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -68,6 +77,8 @@ class ChromaUpsert_VectorStores implements INode {
const embeddings = nodeData.inputs?.embeddings as Embeddings
const chromaURL = nodeData.inputs?.chromaURL as string
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const flattenDocs = docs && docs.length ? flatten(docs) : []
const finalDocs = []
@@ -84,9 +95,10 @@ class ChromaUpsert_VectorStores implements INode {
const vectorStore = await Chroma.fromDocuments(finalDocs, embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -34,6 +34,15 @@ class Faiss_Existing_VectorStores implements INode {
description: 'Path to load faiss.index file',
placeholder: `C:\\Users\\User\\Desktop`,
type: 'string'
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -54,13 +63,16 @@ class Faiss_Existing_VectorStores implements INode {
const embeddings = nodeData.inputs?.embeddings as Embeddings
const basePath = nodeData.inputs?.basePath as string
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const vectorStore = await FaissStore.load(basePath, embeddings)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -42,6 +42,15 @@ class FaissUpsert_VectorStores implements INode {
description: 'Path to store faiss.index file',
placeholder: `C:\\Users\\User\\Desktop`,
type: 'string'
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -63,6 +72,8 @@ class FaissUpsert_VectorStores implements INode {
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const basePath = nodeData.inputs?.basePath as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const flattenDocs = docs && docs.length ? flatten(docs) : []
const finalDocs = []
@@ -74,9 +85,10 @@ class FaissUpsert_VectorStores implements INode {
await vectorStore.save(basePath)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -35,6 +35,14 @@ class InMemoryVectorStore_VectorStores implements INode {
label: 'Embeddings',
name: 'embeddings',
type: 'Embeddings'
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
optional: true
}
]
this.outputs = [
@@ -55,6 +63,8 @@ class InMemoryVectorStore_VectorStores implements INode {
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const flattenDocs = docs && docs.length ? flatten(docs) : []
const finalDocs = []
@@ -65,9 +75,10 @@ class InMemoryVectorStore_VectorStores implements INode {
const vectorStore = await MemoryVectorStore.fromDocuments(finalDocs, embeddings)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -49,6 +49,7 @@ class Pinecone_Existing_VectorStores implements INode {
name: 'pineconeNamespace',
type: 'string',
placeholder: 'my-first-namespace',
additionalParams: true,
optional: true
},
{
@@ -57,6 +58,15 @@ class Pinecone_Existing_VectorStores implements INode {
type: 'json',
optional: true,
additionalParams: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -79,9 +89,10 @@ class Pinecone_Existing_VectorStores implements INode {
const index = nodeData.inputs?.pineconeIndex as string
const pineconeNamespace = nodeData.inputs?.pineconeNamespace as string
const pineconeMetadataFilter = nodeData.inputs?.pineconeMetadataFilter
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const client = new PineconeClient()
await client.init({
@@ -104,9 +115,10 @@ class Pinecone_Existing_VectorStores implements INode {
const vectorStore = await PineconeStore.fromExistingIndex(embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -57,6 +57,16 @@ class PineconeUpsert_VectorStores implements INode {
name: 'pineconeNamespace',
type: 'string',
placeholder: 'my-first-namespace',
additionalParams: true,
optional: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
@@ -82,6 +92,8 @@ class PineconeUpsert_VectorStores implements INode {
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const client = new PineconeClient()
await client.init({
@@ -106,9 +118,10 @@ class PineconeUpsert_VectorStores implements INode {
const vectorStore = await PineconeStore.fromDocuments(finalDocs, embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -55,6 +55,15 @@ class Supabase_Existing_VectorStores implements INode {
type: 'json',
optional: true,
additionalParams: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -79,6 +88,8 @@ class Supabase_Existing_VectorStores implements INode {
const embeddings = nodeData.inputs?.embeddings as Embeddings
const supabaseMetadataFilter = nodeData.inputs?.supabaseMetadataFilter
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const client = createClient(supabaseProjUrl, supabaseApiKey)
@@ -96,9 +107,10 @@ class Supabase_Existing_VectorStores implements INode {
const vectorStore = await SupabaseVectorStore.fromExistingIndex(embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -56,6 +56,15 @@ class SupabaseUpsert_VectorStores implements INode {
label: 'Query Name',
name: 'queryName',
type: 'string'
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -80,6 +89,8 @@ class SupabaseUpsert_VectorStores implements INode {
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const client = createClient(supabaseProjUrl, supabaseApiKey)
@@ -96,9 +107,10 @@ class SupabaseUpsert_VectorStores implements INode {
})
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -79,6 +79,15 @@ class Weaviate_Existing_VectorStores implements INode {
placeholder: `["foo"]`,
optional: true,
additionalParams: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -102,9 +111,10 @@ class Weaviate_Existing_VectorStores implements INode {
const weaviateApiKey = nodeData.inputs?.weaviateApiKey as string
const weaviateTextKey = nodeData.inputs?.weaviateTextKey as string
const weaviateMetadataKeys = nodeData.inputs?.weaviateMetadataKeys as string
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const clientConfig: any = {
scheme: weaviateScheme,
@@ -125,9 +135,10 @@ class Weaviate_Existing_VectorStores implements INode {
const vectorStore = await WeaviateStore.fromExistingIndex(embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore
@@ -87,6 +87,15 @@ class WeaviateUpsert_VectorStores implements INode {
placeholder: `["foo"]`,
optional: true,
additionalParams: true
},
{
label: 'Top K',
name: 'topK',
description: 'Number of top results to fetch. Default to 4',
placeholder: '4',
type: 'number',
additionalParams: true,
optional: true
}
]
this.outputs = [
@@ -110,10 +119,11 @@ class WeaviateUpsert_VectorStores implements INode {
const weaviateApiKey = nodeData.inputs?.weaviateApiKey as string
const weaviateTextKey = nodeData.inputs?.weaviateTextKey as string
const weaviateMetadataKeys = nodeData.inputs?.weaviateMetadataKeys as string
const docs = nodeData.inputs?.document as Document[]
const embeddings = nodeData.inputs?.embeddings as Embeddings
const output = nodeData.outputs?.output as string
const topK = nodeData.inputs?.topK as string
const k = topK ? parseInt(topK, 10) : 4
const clientConfig: any = {
scheme: weaviateScheme,
@@ -140,9 +150,10 @@ class WeaviateUpsert_VectorStores implements INode {
const vectorStore = await WeaviateStore.fromDocuments(finalDocs, embeddings, obj)
if (output === 'retriever') {
const retriever = vectorStore.asRetriever()
const retriever = vectorStore.asRetriever(k)
return retriever
} else if (output === 'vectorStore') {
;(vectorStore as any).k = k
return vectorStore
}
return vectorStore