add subquery engine

This commit is contained in:
Henry 2024-01-28 23:46:55 +00:00
parent 3e59b9b658
commit 21c47d8049
11 changed files with 453 additions and 21 deletions

View File

@ -41,6 +41,14 @@ class ChatOpenAI_LlamaIndex_LLMs implements INode {
label: 'gpt-4',
name: 'gpt-4'
},
{
label: 'gpt-4-turbo-preview',
name: 'gpt-4-turbo-preview'
},
{
label: 'gpt-4-0125-preview',
name: 'gpt-4-0125-preview'
},
{
label: 'gpt-4-1106-preview',
name: 'gpt-4-1106-preview'

View File

@ -32,6 +32,27 @@ class OpenAIEmbedding_LlamaIndex_Embeddings implements INode {
credentialNames: ['openAIApi']
}
this.inputs = [
{
label: 'Model Name',
name: 'modelName',
type: 'options',
options: [
{
label: 'text-embedding-3-large',
name: 'text-embedding-3-large'
},
{
label: 'text-embedding-3-small',
name: 'text-embedding-3-small'
},
{
label: 'text-embedding-ada-002',
name: 'text-embedding-ada-002'
}
],
default: 'text-embedding-ada-002',
optional: true
},
{
label: 'Timeout',
name: 'timeout',
@ -51,12 +72,14 @@ class OpenAIEmbedding_LlamaIndex_Embeddings implements INode {
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
const timeout = nodeData.inputs?.timeout as string
const modelName = nodeData.inputs?.modelName as string
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const openAIApiKey = getCredentialParam('openAIApiKey', credentialData, nodeData)
const obj: Partial<OpenAIEmbedding> = {
apiKey: openAIApiKey
apiKey: openAIApiKey,
model: modelName
}
if (timeout) obj.timeout = parseInt(timeout, 10)

View File

@ -0,0 +1,193 @@
import { flatten } from 'lodash'
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
import {
TreeSummarize,
SimpleResponseBuilder,
Refine,
BaseEmbedding,
ResponseSynthesizer,
CompactAndRefine,
QueryEngineTool,
LLMQuestionGenerator,
SubQuestionQueryEngine,
BaseNode,
Metadata,
serviceContextFromDefaults
} from 'llamaindex'
import { reformatSourceDocuments } from '../EngineUtils'
class SubQuestionQueryEngine_LlamaIndex implements INode {
label: string
name: string
version: number
description: string
type: string
icon: string
category: string
baseClasses: string[]
tags: string[]
inputs: INodeParams[]
outputs: INodeOutputsValue[]
sessionId?: string
constructor(fields?: { sessionId?: string }) {
this.label = 'Sub Question Query Engine'
this.name = 'subQuestionQueryEngine'
this.version = 1.0
this.type = 'SubQuestionQueryEngine'
this.icon = 'subQueryEngine.svg'
this.category = 'Engine'
this.description =
'Breaks complex query into sub questions for each relevant data source, then gather all the intermediate reponses and synthesizes a final response'
this.baseClasses = [this.type]
this.tags = ['LlamaIndex']
this.inputs = [
{
label: 'QueryEngine Tools',
name: 'queryEngineTools',
type: 'QueryEngineTool',
list: true
},
{
label: 'Chat Model',
name: 'model',
type: 'BaseChatModel_LlamaIndex'
},
{
label: 'Embeddings',
name: 'embeddings',
type: 'BaseEmbedding_LlamaIndex'
},
{
label: 'Response Synthesizer',
name: 'responseSynthesizer',
type: 'ResponseSynthesizer',
description:
'ResponseSynthesizer is responsible for sending the query, nodes, and prompt templates to the LLM to generate a response. See <a target="_blank" href="https://ts.llamaindex.ai/modules/low_level/response_synthesizer">more</a>',
optional: true
},
{
label: 'Return Source Documents',
name: 'returnSourceDocuments',
type: 'boolean',
optional: true
}
]
this.sessionId = fields?.sessionId
}
async init(): Promise<any> {
return null
}
async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string | object> {
const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
const embeddings = nodeData.inputs?.embeddings as BaseEmbedding
const model = nodeData.inputs?.model
const serviceContext = serviceContextFromDefaults({
llm: model,
embedModel: embeddings
})
let queryEngineTools = nodeData.inputs?.queryEngineTools as QueryEngineTool[]
queryEngineTools = flatten(queryEngineTools)
let queryEngine = SubQuestionQueryEngine.fromDefaults({
serviceContext,
queryEngineTools,
questionGen: new LLMQuestionGenerator({ llm: model })
})
const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer
if (responseSynthesizerObj) {
if (responseSynthesizerObj.type === 'TreeSummarize') {
const responseSynthesizer = new ResponseSynthesizer({
responseBuilder: new TreeSummarize(serviceContext, responseSynthesizerObj.textQAPromptTemplate),
serviceContext
})
queryEngine = SubQuestionQueryEngine.fromDefaults({
responseSynthesizer,
serviceContext,
queryEngineTools,
questionGen: new LLMQuestionGenerator({ llm: model })
})
} else if (responseSynthesizerObj.type === 'CompactAndRefine') {
const responseSynthesizer = new ResponseSynthesizer({
responseBuilder: new CompactAndRefine(
serviceContext,
responseSynthesizerObj.textQAPromptTemplate,
responseSynthesizerObj.refinePromptTemplate
),
serviceContext
})
queryEngine = SubQuestionQueryEngine.fromDefaults({
responseSynthesizer,
serviceContext,
queryEngineTools,
questionGen: new LLMQuestionGenerator({ llm: model })
})
} else if (responseSynthesizerObj.type === 'Refine') {
const responseSynthesizer = new ResponseSynthesizer({
responseBuilder: new Refine(
serviceContext,
responseSynthesizerObj.textQAPromptTemplate,
responseSynthesizerObj.refinePromptTemplate
),
serviceContext
})
queryEngine = SubQuestionQueryEngine.fromDefaults({
responseSynthesizer,
serviceContext,
queryEngineTools,
questionGen: new LLMQuestionGenerator({ llm: model })
})
} else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
const responseSynthesizer = new ResponseSynthesizer({
responseBuilder: new SimpleResponseBuilder(serviceContext),
serviceContext
})
queryEngine = SubQuestionQueryEngine.fromDefaults({
responseSynthesizer,
serviceContext,
queryEngineTools,
questionGen: new LLMQuestionGenerator({ llm: model })
})
}
}
let text = ''
let sourceDocuments: ICommonObject[] = []
let sourceNodes: BaseNode<Metadata>[] = []
let isStreamingStarted = false
const isStreamingEnabled = options.socketIO && options.socketIOClientId
if (isStreamingEnabled) {
const stream = await queryEngine.query({ query: input, stream: true })
for await (const chunk of stream) {
text += chunk.response
if (chunk.sourceNodes) sourceNodes = chunk.sourceNodes
if (!isStreamingStarted) {
isStreamingStarted = true
options.socketIO.to(options.socketIOClientId).emit('start', chunk.response)
}
options.socketIO.to(options.socketIOClientId).emit('token', chunk.response)
}
if (returnSourceDocuments) {
sourceDocuments = reformatSourceDocuments(sourceNodes)
options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', sourceDocuments)
}
} else {
const response = await queryEngine.query({ query: input })
text = response?.response
sourceDocuments = reformatSourceDocuments(response?.sourceNodes ?? [])
}
if (returnSourceDocuments) return { text, sourceDocuments }
else return { text }
}
}
module.exports = { nodeClass: SubQuestionQueryEngine_LlamaIndex }

View File

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-filter-question" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round"><path stroke="none" d="M0 0h24v24H0z" fill="none"/><path d="M15 19l-6 2v-8.5l-4.48 -4.928a2 2 0 0 1 -.52 -1.345v-2.227h16v2.172a2 2 0 0 1 -.586 1.414l-4.414 4.414" /><path d="M19 22v.01" /><path d="M19 19a2.003 2.003 0 0 0 .914 -3.782a1.98 1.98 0 0 0 -2.414 .483" /></svg>

After

Width:  |  Height:  |  Size: 506 B

View File

@ -0,0 +1,68 @@
import { INode, INodeData, INodeParams } from '../../../src/Interface'
import { VectorStoreIndex } from 'llamaindex'
class QueryEngine_Tools implements INode {
label: string
name: string
version: number
description: string
type: string
icon: string
category: string
tags: string[]
baseClasses: string[]
inputs?: INodeParams[]
constructor() {
this.label = 'QueryEngine Tool'
this.name = 'queryEngineToolLlamaIndex'
this.version = 1.0
this.type = 'QueryEngineTool'
this.icon = 'queryEngineTool.svg'
this.category = 'Tools'
this.tags = ['LlamaIndex']
this.description = 'Tool used to invoke query engine'
this.baseClasses = [this.type]
this.inputs = [
{
label: 'Vector Store Index',
name: 'vectorStoreIndex',
type: 'VectorStoreIndex'
},
{
label: 'Tool Name',
name: 'toolName',
type: 'string',
description: 'Tool name must be small capital letter with underscore. Ex: my_tool'
},
{
label: 'Tool Description',
name: 'toolDesc',
type: 'string',
rows: 4
}
]
}
async init(nodeData: INodeData): Promise<any> {
const vectorStoreIndex = nodeData.inputs?.vectorStoreIndex as VectorStoreIndex
const toolName = nodeData.inputs?.toolName as string
const toolDesc = nodeData.inputs?.toolDesc as string
const queryEngineTool = {
queryEngine: vectorStoreIndex.asQueryEngine({
preFilters: {
...(vectorStoreIndex as any).metadatafilter
}
}),
metadata: {
name: toolName,
description: toolDesc
},
vectorStoreIndex
}
return queryEngineTool
}
}
module.exports = { nodeClass: QueryEngine_Tools }

View File

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-brand-google-big-query" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round"><path stroke="none" d="M0 0h24v24H0z" fill="none"/><path d="M17.73 19.875a2.225 2.225 0 0 1 -1.948 1.125h-7.283a2.222 2.222 0 0 1 -1.947 -1.158l-4.272 -6.75a2.269 2.269 0 0 1 0 -2.184l4.272 -6.75a2.225 2.225 0 0 1 1.946 -1.158h7.285c.809 0 1.554 .443 1.947 1.158l3.98 6.75a2.33 2.33 0 0 1 0 2.25l-3.98 6.75v-.033z" /><path d="M11.5 11.5m-3.5 0a3.5 3.5 0 1 0 7 0a3.5 3.5 0 1 0 -7 0" /><path d="M14 14l2 2" /></svg>

After

Width:  |  Height:  |  Size: 654 B

View File

@ -13,7 +13,7 @@ import {
import { FetchResponse, Index, Pinecone, ScoredPineconeRecord } from '@pinecone-database/pinecone'
import { flatten } from 'lodash'
import { Document as LCDocument } from 'langchain/document'
import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
import { flattenObject, getCredentialData, getCredentialParam } from '../../../src/utils'
class PineconeLlamaIndex_VectorStores implements INode {
@ -28,6 +28,7 @@ class PineconeLlamaIndex_VectorStores implements INode {
baseClasses: string[]
inputs: INodeParams[]
credential: INodeParams
outputs: INodeOutputsValue[]
constructor() {
this.label = 'Pinecone'
@ -93,6 +94,18 @@ class PineconeLlamaIndex_VectorStores implements INode {
optional: true
}
]
this.outputs = [
{
label: 'Pinecone Retriever',
name: 'retriever',
baseClasses: this.baseClasses
},
{
label: 'Pinecone Vector Store Index',
name: 'vectorStore',
baseClasses: [this.type, 'VectorStoreIndex']
}
]
}
//@ts-ignore
@ -155,8 +168,10 @@ class PineconeLlamaIndex_VectorStores implements INode {
}
if (pineconeNamespace) obj.namespace = pineconeNamespace
let metadatafilter = {}
if (pineconeMetadataFilter) {
const metadatafilter = typeof pineconeMetadataFilter === 'object' ? pineconeMetadataFilter : JSON.parse(pineconeMetadataFilter)
metadatafilter = typeof pineconeMetadataFilter === 'object' ? pineconeMetadataFilter : JSON.parse(pineconeMetadataFilter)
obj.queryFilter = metadatafilter
}
@ -171,11 +186,21 @@ class PineconeLlamaIndex_VectorStores implements INode {
serviceContext
})
const output = nodeData.outputs?.output as string
if (output === 'retriever') {
const retriever = index.asRetriever()
retriever.similarityTopK = k
;(retriever as any).serviceContext = serviceContext
return retriever
} else if (output === 'vectorStore') {
;(index as any).k = k
if (metadatafilter) {
;(index as any).metadatafilter = metadatafilter
}
return index
}
return index
}
}

View File

@ -63,6 +63,18 @@ class SimpleStoreUpsert_LlamaIndex_VectorStores implements INode {
optional: true
}
]
this.outputs = [
{
label: 'SimpleStore Retriever',
name: 'retriever',
baseClasses: this.baseClasses
},
{
label: 'SimpleStore Vector Store Index',
name: 'vectorStore',
baseClasses: [this.type, 'VectorStoreIndex']
}
]
}
//@ts-ignore
@ -114,10 +126,19 @@ class SimpleStoreUpsert_LlamaIndex_VectorStores implements INode {
const storageContext = await storageContextFromDefaults({ persistDir: filePath })
const index = await VectorStoreIndex.init({ storageContext, serviceContext })
const output = nodeData.outputs?.output as string
if (output === 'retriever') {
const retriever = index.asRetriever()
retriever.similarityTopK = k
;(retriever as any).serviceContext = serviceContext
return retriever
} else if (output === 'vectorStore') {
;(index as any).k = k
return index
}
return index
}
}

View File

@ -181,6 +181,28 @@
"credentialNames": ["openAIApi"],
"id": "openAIEmbedding_LlamaIndex_0-input-credential-credential"
},
{
"label": "Model Name",
"name": "modelName",
"type": "options",
"options": [
{
"label": "text-embedding-3-large",
"name": "text-embedding-3-large"
},
{
"label": "text-embedding-3-small",
"name": "text-embedding-3-small"
},
{
"label": "text-embedding-ada-002",
"name": "text-embedding-ada-002"
}
],
"default": "text-embedding-ada-002",
"optional": true,
"id": "openAIEmbedding_LlamaIndex_0-input-modelName-options"
},
{
"label": "Timeout",
"name": "timeout",
@ -315,13 +337,29 @@
},
"outputAnchors": [
{
"id": "pineconeLlamaIndex_0-output-pineconeLlamaIndex-Pinecone|VectorIndexRetriever",
"name": "pineconeLlamaIndex",
"label": "Pinecone",
"name": "output",
"label": "Output",
"type": "options",
"options": [
{
"id": "pineconeLlamaIndex_0-output-retriever-Pinecone|VectorIndexRetriever",
"name": "retriever",
"label": "Pinecone Retriever",
"type": "Pinecone | VectorIndexRetriever"
},
{
"id": "pineconeLlamaIndex_0-output-retriever-Pinecone|VectorStoreIndex",
"name": "vectorStore",
"label": "Pinecone Vector Store Index",
"type": "Pinecone | VectorStoreIndex"
}
],
"outputs": {},
"default": "retriever"
}
],
"outputs": {
"output": "retriever"
},
"selected": false
},
"selected": false,
@ -367,6 +405,14 @@
"label": "gpt-4",
"name": "gpt-4"
},
{
"label": "gpt-4-turbo-preview",
"name": "gpt-4-turbo-preview"
},
{
"label": "gpt-4-0125-preview",
"name": "gpt-4-0125-preview"
},
{
"label": "gpt-4-1106-preview",
"name": "gpt-4-1106-preview"
@ -672,6 +718,14 @@
"label": "gpt-4",
"name": "gpt-4"
},
{
"label": "gpt-4-turbo-preview",
"name": "gpt-4-turbo-preview"
},
{
"label": "gpt-4-0125-preview",
"name": "gpt-4-0125-preview"
},
{
"label": "gpt-4-1106-preview",
"name": "gpt-4-1106-preview"

View File

@ -163,13 +163,29 @@
},
"outputAnchors": [
{
"id": "pineconeLlamaIndex_0-output-pineconeLlamaIndex-Pinecone|VectorIndexRetriever",
"name": "pineconeLlamaIndex",
"label": "Pinecone",
"name": "output",
"label": "Output",
"type": "options",
"options": [
{
"id": "pineconeLlamaIndex_0-output-retriever-Pinecone|VectorIndexRetriever",
"name": "retriever",
"label": "Pinecone Retriever",
"type": "Pinecone | VectorIndexRetriever"
},
{
"id": "pineconeLlamaIndex_0-output-retriever-Pinecone|VectorStoreIndex",
"name": "vectorStore",
"label": "Pinecone Vector Store Index",
"type": "Pinecone | VectorStoreIndex"
}
],
"outputs": {},
"default": "retriever"
}
],
"outputs": {
"output": "retriever"
},
"selected": false
},
"selected": false,
@ -206,6 +222,28 @@
"credentialNames": ["openAIApi"],
"id": "openAIEmbedding_LlamaIndex_0-input-credential-credential"
},
{
"label": "Model Name",
"name": "modelName",
"type": "options",
"options": [
{
"label": "text-embedding-3-large",
"name": "text-embedding-3-large"
},
{
"label": "text-embedding-3-small",
"name": "text-embedding-3-small"
},
{
"label": "text-embedding-ada-002",
"name": "text-embedding-ada-002"
}
],
"default": "text-embedding-ada-002",
"optional": true,
"id": "openAIEmbedding_LlamaIndex_0-input-modelName-options"
},
{
"label": "Timeout",
"name": "timeout",

View File

@ -844,7 +844,7 @@ export const isFlowValidForStream = (reactFlowNodes: IReactFlowNode[], endingNod
const whitelistAgents = ['openAIFunctionAgent', 'csvAgent', 'airtableAgent', 'conversationalRetrievalAgent']
isValidChainOrAgent = whitelistAgents.includes(endingNodeData.name)
} else if (endingNodeData.category === 'Engine') {
const whitelistEngine = ['contextChatEngine', 'simpleChatEngine', 'queryEngine']
const whitelistEngine = ['contextChatEngine', 'simpleChatEngine', 'queryEngine', 'subQuestionQueryEngine']
isValidChainOrAgent = whitelistEngine.includes(endingNodeData.name)
}