Zum Hauptinhalt springen

Router Query Engine

In this tutorial, we define a custom router query engine that selects one out of several candidate query engines to execute a query.

Setup

First, we need to install import the necessary modules from llamaindex:

pnpm i lamaindex
import {
OpenAI,
RouterQueryEngine,
SimpleDirectoryReader,
SimpleNodeParser,
SummaryIndex,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";

Loading Data

Next, we need to load some data. We will use the SimpleDirectoryReader to load documents from a directory:

const documents = await new SimpleDirectoryReader().loadData({
directoryPath: "node_modules/llamaindex/examples",
});

Service Context

Next, we need to define some basic rules and parse the documents into nodes. We will use the SimpleNodeParser to parse the documents into nodes and ServiceContext to define the rules (eg. LLM API key, chunk size, etc.):

const nodeParser = new SimpleNodeParser({
chunkSize: 1024,
});

const serviceContext = serviceContextFromDefaults({
nodeParser,
llm: new OpenAI(),
});

Creating Indices

Next, we need to create some indices. We will create a VectorStoreIndex and a SummaryIndex:

const vectorIndex = await VectorStoreIndex.fromDocuments(documents, {
serviceContext,
});

const summaryIndex = await SummaryIndex.fromDocuments(documents, {
serviceContext,
});

Creating Query Engines

Next, we need to create some query engines. We will create a VectorStoreQueryEngine and a SummaryQueryEngine:

const vectorQueryEngine = vectorIndex.asQueryEngine();
const summaryQueryEngine = summaryIndex.asQueryEngine();

Creating a Router Query Engine

Next, we need to create a router query engine. We will use the RouterQueryEngine to create a router query engine:

We're defining two query engines, one for summarization and one for retrieving specific context. The router query engine will select the most appropriate query engine based on the query.

const queryEngine = RouterQueryEngine.fromDefaults({
queryEngineTools: [
{
queryEngine: vectorQueryEngine,
description: "Useful for summarization questions related to Abramov",
},
{
queryEngine: summaryQueryEngine,
description: "Useful for retrieving specific context from Abramov",
},
],
serviceContext,
});

Querying the Router Query Engine

Finally, we can query the router query engine:

const summaryResponse = await queryEngine.query({
query: "Give me a summary about his past experiences?",
});

console.log({
answer: summaryResponse.response,
metadata: summaryResponse?.metadata?.selectorResult,
});

Full code

import {
OpenAI,
RouterQueryEngine,
SimpleDirectoryReader,
SimpleNodeParser,
SummaryIndex,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";

async function main() {
// Load documents from a directory
const documents = await new SimpleDirectoryReader().loadData({
directoryPath: "node_modules/llamaindex/examples",
});

// Parse the documents into nodes
const nodeParser = new SimpleNodeParser({
chunkSize: 1024,
});

// Create a service context
const serviceContext = serviceContextFromDefaults({
nodeParser,
llm: new OpenAI(),
});

// Create indices
const vectorIndex = await VectorStoreIndex.fromDocuments(documents, {
serviceContext,
});

const summaryIndex = await SummaryIndex.fromDocuments(documents, {
serviceContext,
});

// Create query engines
const vectorQueryEngine = vectorIndex.asQueryEngine();
const summaryQueryEngine = summaryIndex.asQueryEngine();

// Create a router query engine
const queryEngine = RouterQueryEngine.fromDefaults({
queryEngineTools: [
{
queryEngine: vectorQueryEngine,
description: "Useful for summarization questions related to Abramov",
},
{
queryEngine: summaryQueryEngine,
description: "Useful for retrieving specific context from Abramov",
},
],
serviceContext,
});

// Query the router query engine
const summaryResponse = await queryEngine.query({
query: "Give me a summary about his past experiences?",
});

console.log({
answer: summaryResponse.response,
metadata: summaryResponse?.metadata?.selectorResult,
});

const specificResponse = await queryEngine.query({
query: "Tell me about abramov first job?",
});

console.log({
answer: specificResponse.response,
metadata: specificResponse.metadata.selectorResult,
});
}

main().then(() => console.log("Done"));