Add Memory to AI21 Labs Model Using LangChain: A Quick Tutorial

Introduction

You have probably already been exposed to models from AI21 Labs. I am very comfortable working with them, but the big pain has always been having to manually record the history of my interaction with the model. Much has changed, however, and LangChain makes it possible to implement such an operation in the blink of an eye!

In this tutorial, I will explain how to implement this quickly, then you can play with it yourself!

Implementation

Dependencies

First, we need to create a project directory, virtual environment, and install some dependencies. Let’s do it!

Creating a Virtual Environment

Now, we need to create a virtual environment:

Installing Dependencies

Last step of this stage - installing dependencies:

Coding Time!

Now we can start coding! First, we need to create a .env file with our API key. You can get it from AI21 Labs Studio. I will use AI21_API_KEY as the variable name.

Now, we can create main.py file and start coding!

Importing Modules and Loading the API Key

I will import all necessary modules and load the API key.

I decided to build a simple chatbot for this tutorial. It will be very simple and run in a terminal. Feel free to create any other application; I just want to show you how you can do it! Anyway, let’s do it!

Creating a Prompt Template

I will create a prompt template that will make it easier for the model to understand our task. As I will be using a model that by default is not a chat model but a regular LLM, this is necessary.

I will also create a memory object that will store our conversation history.

I will use ConversationBufferMemory, but you can use another type of memory. In this article about conversation memory types, others are well explained! Check it out!

Creating a LLMChain

Now we can create a LLMChain that will take care of our conversation.

I use verbose=True to browse the input from memory to Chain.

Main Loop

Great, now we can create a main loop that iterates through the stages of our conversation.

Results

Let’s run our application!

Conversation

I conducted a simple conversation to test the performance of our program. As you can see, the story is already working in the second stage, as it should be!

As you can see, implementing this is not that difficult and does not take long. With the combination of AI21 and LangChain, we are able to build really powerful applications. I encourage you to check it out during our Plug into AI with AI21 Hackathon! Feel free to join and build with us!