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,
  Settings,
} 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 Settings to define the rules (eg. LLM API key, chunk size, etc.):

Settings.llm = new OpenAI();
Settings.nodeParser = new SimpleNodeParser({
  chunkSize: 1024,
});

Creating Indices

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

const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
const summaryIndex = await SummaryIndex.fromDocuments(documents);

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",
    },
  ],
});

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,
  Settings,
} from "llamaindex";

Settings.llm = new OpenAI();
Settings.nodeParser = new SimpleNodeParser({
  chunkSize: 1024,
});

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

  // Create indices
  const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
  const summaryIndex = await SummaryIndex.fromDocuments(documents);

  // 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",
      },
    ],
  });

  // 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"));