Creating a Retrieval Chain

Now that you have a vector store, you can use it to retrieve chunks of text that are semantically similar to a user’s question.

In this challenge, you will create a chain that will use the vector search index to find movies with similar plots.

You must first:

  1. Use the initVectorStore() function implemented in the previous lesson to create a vector store and retriever

  2. Create an instance of the Answer Generation chain.

Then, create a chain that will:

  1. Takes the string input and assigns it the input variable

  2. Uses the input to retrieve similar movie plots

  3. Uses the answer generation chain to generate an answer

  4. Use the saveHistory() function to save the response and context to the database

  5. Returns the output as a string.

Existing function

The modules/agent/tools/vector-retrieval.chain.ts file contains the following placeholder functions for saving and retrieving history.

typescript
export default async function initVectorRetrievalChain(
  llm: BaseLanguageModel,
  embeddings: Embeddings
): Promise<Runnable<AgentToolInput, string>> {
  // TODO: Create vector store instance
  // const vectorStore = ...
  // TODO: Initialize a retriever wrapper around the vector store
  // const vectorStoreRetriever = ...
  // TODO: Initialize Answer chain
  // const answerChain = ...
  // TODO: Return chain
  // return RunnablePassthrough.assign( ... )
}
Open vector-retrieval.chain.ts

Instantiate Tools

Inside the initVectorRetrievalChain() function, replace the // TODO comments to create an instance of the vector store using the initVectorStore() function from the previous lesson.

typescript
Vector Store
//  Create vector store instance
const vectorStore = await initVectorStore(embeddings);

Next, call the .asRetriever() method on the vectorStore object to create a new VectorStoreRetriever instance.

typescript
Vector Store
// Initialize a retriever wrapper around the vector store
const vectorStoreRetriever = vectorStore.asRetriever(5);

Finally, create an answer generation chain using the initGenerateAnswerChain() function.

typescript
Answer Generation Chain
// Initialize  Answer Chain
const answerChain = initGenerateAnswerChain(llm);

Building the Chain

As this chain will be called by an agent, it will receive a structured input containing an input and rephrasedQuestion.

typescript
Agent to Tool Input
export interface AgentToolInput {
  input: string;
  rephrasedQuestion: string;
}

Because the chain will receive an object as the input, you can use RunnablePassthrough.assign() to modify the input directly rather than the RunnableSequence.from() method used in the previous lessons.

This should be used to collect relevant context using the retriever.

typescript
Get Documents
// Get the rephrased question and generate context
return (
  RunnablePassthrough.assign({
    documents: new RunnablePick("rephrasedQuestion").pipe(
      vectorStoreRetriever
    ),
  })

Next, the elementIds of the document must be extracted from the to create the :CONTEXT relationship between the (:Response) and (:Movie) nodes. At the same time, the context needs to be converted to a string so it can be used in the Answer Generation Chain.

Helper Functions

These functions in vector-retrieval.chain.ts are used to extract information to modify the context.

typescript
Helper Functions
// Helper function to extract document IDs from Movie node metadata
const extractDocumentIds = (
  documents: DocumentInterface<{ _id: string; [key: string]: any }>[]
): string[] => documents.map((document) => document.metadata._id);

// Convert documents to string to be included in the prompt
const docsToJson = (documents: DocumentInterface[]) =>
  JSON.stringify(documents);

The RunnablePassthrough is a fluent interface, so the .assign() method can be called to chain the steps together.

typescript
Mutate State
// Get the rephrased question and generate context
return (
  RunnablePassthrough.assign({
    documents: new RunnablePick("rephrasedQuestion").pipe(
      vectorStoreRetriever
    ),
  })
    .assign({
      // Extract the IDs
      ids: new RunnablePick("documents").pipe(extractDocumentIds),
      // convert documents to string
      context: new RunnablePick("documents").pipe(docsToJson),
    })

The rephrased question and context can then be passed to the answerChain to generate an output.

typescript
Generate an answer
.assign({
  output: (input: RetrievalChainThroughput) =>
    answerChain.invoke({
      question: input.rephrasedQuestion,
      context: input.context,
    }),
})

Then, the input, rephrased question and output can be saved to the database using the saveHistory() function created in Conversation Memory module.

typescript
Generate an answer
.assign({
  responseId: async (input: RetrievalChainThroughput, options) =>
    saveHistory(
      options?.config.configurable.sessionId,
      "vector",
      input.input,
      input.rephrasedQuestion,
      input.output,
      input.ids
    ),
})

Before, finally picking the output as a string.

typescript
Return the output
.assign({
  responseId: async (input: RetrievalChainThroughput, options) =>
    saveHistory(
      options?.config.configurable.sessionId,
      "vector",
      input.input,
      input.rephrasedQuestion,
      input.output,
      input.ids
    ),
})
    .pick("output")
);

If you have followed the instructions correctly, your code should resemble the following:

typescript
Full Function
export default async function initVectorRetrievalChain(
  llm: BaseLanguageModel,
  embeddings: Embeddings
): Promise<Runnable<AgentToolInput, string>> {
  //  Create vector store instance
  const vectorStore = await initVectorStore(embeddings);

  // Initialize a retriever wrapper around the vector store
  const vectorStoreRetriever = vectorStore.asRetriever(5);

  // Initialize  Answer Chain
  const answerChain = initGenerateAnswerChain(llm);

  // Get the rephrased question and generate context
  return (
    RunnablePassthrough.assign({
      documents: new RunnablePick("rephrasedQuestion").pipe(
        vectorStoreRetriever
      ),
    })
      .assign({
        // Extract the IDs
        ids: new RunnablePick("documents").pipe(extractDocumentIds),
        // convert documents to string
        context: new RunnablePick("documents").pipe(docsToJson),
      })
      .assign({
        output: (input: RetrievalChainThroughput) =>
          answerChain.invoke({
            question: input.rephrasedQuestion,
            context: input.context,
          }),
      })
      .assign({
        responseId: async (input: RetrievalChainThroughput, options) =>
          saveHistory(
            options?.config.configurable.sessionId,
            "vector",
            input.input,
            input.rephrasedQuestion,
            input.output,
            input.ids
          ),
      })
      .pick("output")
  );
}

Testing your changes

If you have followed the instructions, you should be able to run the following unit test to verify the response using the npm run test command.

sh
Running the Test
npm run test vector-retrieval.chain.test.ts
View Unit Test
typescript
vector-retrieval.chain.test.ts
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { config } from "dotenv";
import { BaseChatModel } from "langchain/chat_models/base";
import { Embeddings } from "langchain/embeddings/base";
import { Runnable } from "@langchain/core/runnables";
import initVectorRetrievalChain from "./vector-retrieval.chain";
import { Neo4jGraph } from "@langchain/community/graphs/neo4j_graph";
import { AgentToolInput } from "../agent.types";
import { close } from "../../graph";

describe("Vector Retrieval Chain", () => {
  let graph: Neo4jGraph;
  let llm: BaseChatModel;
  let embeddings: Embeddings;
  let chain: Runnable<AgentToolInput, string>;

  beforeAll(async () => {
    config({ path: ".env.local" });

    graph = await Neo4jGraph.initialize({
      url: process.env.NEO4J_URI as string,
      username: process.env.NEO4J_USERNAME as string,
      password: process.env.NEO4J_PASSWORD as string,
      database: process.env.NEO4J_DATABASE as string | undefined,
    });

    llm = new ChatOpenAI({
      openAIApiKey: process.env.OPENAI_API_KEY,
      modelName: "gpt-3.5-turbo",
      temperature: 0,
      configuration: {
        baseURL: process.env.OPENAI_API_BASE,
      },
    });

    embeddings = new OpenAIEmbeddings({
      openAIApiKey: process.env.OPENAI_API_KEY as string,
      configuration: {
        baseURL: process.env.OPENAI_API_BASE,
      },
    });

    chain = await initVectorRetrievalChain(llm, embeddings);
  });

  afterAll(async () => {
    await graph.close();
    await close();
  });

  it("should provide a recommendation", async () => {
    const sessionId = "vector-retriever-1";
    const input = "[redacted]";
    const rephrasedQuestion = "Recommend a movie about ghosts";

    const output = await chain.invoke(
      {
        input,
        rephrasedQuestion,
      },
      { configurable: { sessionId } }
    );

    // Should generate an answer
    expect(output).toBeDefined();

    // Should save to the database
    const res = await graph.query(
      `
        MATCH (s:Session {id: $sessionId})-[:LAST_RESPONSE]->(r)
        RETURN s.id AS session, r.input AS input, r.output AS output,
          r.rephrasedQuestion AS rephrasedQuestion,
          [ (r)-[:CONTEXT]->(m) | m.title ] AS context
        ORDER BY r.createdAt DESC LIMIT 1
    `,
      { sessionId }
    );

    expect(res).toBeDefined();

    // Should have properties set
    const [first] = res!;

    expect(first.input).toEqual(input);
    expect(first.rephrasedQuestion).toEqual(rephrasedQuestion);
    expect(first.output).toEqual(output);
    expect(first.input).toEqual(input);

    // Should save with context
    expect(first.context.length).toBeGreaterThanOrEqual(1);

    // Any of the movies in the context should be mentioned
    let found = false;

    for (const title of first.context) {
      if (output.includes(title.replace(", The", ""))) {
        found = true;
      }
    }

    expect(found).toBe(true);
  });
});

Verifying the Test

If every test in the test suite has passed, a new (:Session) node with a .id property of vector-retriever-1 will have been created in your database.

The session should have atleast one (:Response) node, linked with a :CONTEXT relationship to at least one movie.

Click the Check Database button below to verify the tests have succeeded.

Hint

You can compare your code with the solution in src/solutions/modules/agent/tools/vector-retrieval.chain.ts and double-check that the conditions have been met in the test suite.

Solution

You can compare your code with the solution in src/solutions/modules/agent/tools/vector-retrieval.chain.ts and double-check that the conditions have been met in the test suite.

You can also run the following Cypher statement to double-check that the index has been created in your database.

cypher
MATCH (s:Session {id: 'vector-retriever-1'})
RETURN s,
    [ (s)-[:HAS_RESPONSE]->(r)
        | [ r,
            [ (r)-[:CONTEXT]->(c) | c ]
        ]
    ]

Once you have verified your code and re-ran the tests, click Try again…​* to complete the challenge.

Summary

In this lesson, you built a chain that takes the components built in the course so far to build a chain that retrieves documents from the vector search index and uses them to answer a question.

The chain then saves the response to the database.

In the next lesson, you will see how the response can be used to filter out documents that have been used to provide unhelpful responses in the past.