Making Long Context LLMs Usable with Context Caching

Google recently added support for context caching to their Gemini API, and with this they are trying to address some of the major limitations of long context LLMs. In theory, long context LLMs can hold a lot of information. However, they suffer from some major issue because with every query you are going to be sending the whole context to the LLM. Since the LLM has to process all the tokens with every query, that will result in increased processing time because you're processing a lot of data, and it results in higher latency because of the large data transfers that are happening. And since the API providers charge you per token basis, that will increase your cost as well. With context caching, Google is trying to address some of these issues. Okay, so here's a quick representation of how the system works based on the details that Google has provided. This is my own representation, so it may not be exactly how they have implemented it, but as a first step, you need to initialize your cache, and that's basically you provide a system instruction. Or it could be a large context, such as a huge amount of documents, video files, audio files, but these are the files or documents or data that you want to cache. Now, with each cache there is a unique identifier. Think of this as the name of the cache and another parameter which they are calling time to live. This is basically for how long you want to cache the content of this specific data set. Then Gemini uses its internal storage system, which is optimized for quick retrieval, to store your cache. When the Gemini API receives a user query, then it simply analyzes all the cache data sets, since each cache has its own unique identifier. So it can pick the appropriate cache and can also verify the availability or validity of the cache data by looking at the time to live. Then it combines both the cache data plus the original query from the user, and that becomes another input to the system which is processed by the Gemini API and you get a response. In this system, both the cached data plus the queries or the new user query are treated in exactly the same way. So this in theory will reduce both your processing time as well as the cost associated with it. Because now you are just sending the your data set only once and it's cached and you are just sending the user queries, which are going to be a lot shorter when it comes to the number of tokens. Google actually says that if you use the caching for up to 128,000 tokens, this is almost four times less than if you don't use caching and you Just send all your contacts with every query. Now there is another element that you need to consider and that is the storage cost, which is $1 per million tokens per hour of storage for the cash. And there are a number of different factors that can impact the cost. One is the cash token count for how long you are keeping that cache, which is basically determined by time to live. That's the parameter that you can set. And we're going to look at a couple of examples there. And there are some other factors such as non cached input tokens and output tokens. Currently, caching is supported by both the Pro as well as the Flash models. So this is pretty great. You can use the faster Flash model if you need to. And for each cache you will need to define the time to live. If you don't provide any value, by default it's going to be set to one hour. And the minimum value that you can set is 60 seconds or a minute. Okay. Before looking at an example implementation of how this works, there are a couple of other things that you need to keep in mind. The minimum input token counts for context caching is about 32,000 tokens, so you will need to cache this many tokens for it to work. I hope that they will reduce it to something like 4 or 8,000 tokens, which is going to be a lot more manageable and useful. But I think those improvements are going to come with time. And when it comes to the maximum number of tokens that you can cache is basically the max context window of a given model. So I believe both in case of Gemini Pro as well as Flash, that is about 2 million tokens. So let's look at an example. So for this example, I uploaded a PDF course material that I'm working on. This is an implementation of a RAC system and it's in the form of a book, which is about 34,000 tokens. So it's just beyond the 32,000 tokens context window that Gemini Caching currently support. When I process it through the API, I think it was close to 40,000 tokens. So I'm not sure like what those extra 6,000 tokens are coming from. But first let's look at all the different packages that we'll need to set it up. So first we need to install the Google Generative AI client or Python package. We will also install the PYPDF package for reading data from PDF files. Then we import the caching from the Generative AI client. This is the one that we're going to be using to create our cache, we also are going to look at or use the markdown for displaying the text that we get. And then we create the Gemini client. Now we will need the Google API key. You can get this from AI Studio. Now in my case that's set as a secret in the Google Colab notebook that I'm using. And then I just pointing to that specific API key with my Google generative AI client. Next I am pointing to that PDF file. We load all the pages of that PDF file, read all of those pages together and convert it into a single string. We look at the number of pages. I think the total number of pages in this case is about 48 pages. This is basically a material that I'm working on for exploring different RAC techniques. I have a course which talks about techniques beyond basics when it comes to Rack. Link is going to be in the video description if you are interested in that. Next we take all the pages, put them together into a single string and here are the contents of this string. Now when we look at the total number of elements, it has about 150,000 words, right? So it's a pretty large file. And the number of tokens comes out to be around 40,000 tokens. And we're going to look at that later in the video. Now I want to do retrieval on this. So I was going to use this system prompt. You are an expert on RAG systems since the document actually talks about rag. So your task is to answer questions based on the provided text. Now before caching this huge file, I wanted to see how long it's going to take if we use the Gemini Flash model as it is. So in this case we created instance of the generative AI model from Google. We provide the model name which is Gemini 1.5 Flash in this case. And then here are three different queries that I'm going to run in parallel. So the first one is what is rag? What are the main components? How to use hide in LangChain, provide code examples and how does routing work in LangChain? So these are queries that are very specific to the contents of that PDF file. When I ran this, here are the Ming's information that it took. This is the CPU time. So the actual time you're going to see is going to be a lot more than this. But we kind of have a baseline that we can compare against when we look at the cache content. Now here is the response and basically we got a response regarding our first question, what is rag? Then how to use hide? And it actually was able to give us a proper code example based on the documents that we have provided. And then it also was able to get the information related to how routing works in LangChain. So pretty neat stuff here. And this is actually one of the great things about in context learning. So you can provide these models with new information and since these are relatively bigger models, so they will be able to use that information to actually give you much more reliable results. Let's look at how do we cache this information so that we are just sending a query one at a time. In the previous case, with every query it was going to send this whole context. Now in order to do that, we are going to be using the caching object from the generative AI package. Then we need to tell it to create a cache. So first we need to tell it which model to use. We're going to be again using Gemini 1.5 flash. Then if you want to pass on any system instruction. In the previous example I didn't do anything, but here we're just providing the system instruction then the actual contents. So basically this is the content that is going to be cached by the system. You also need to tell it how long to keep it. So that's basically the time to live parameter that you need to provide. If you don't provide any value here, it's going to assume that you want to keep it for one hour in this case when you're creating the model object. So you're going to be using this method from the generative table AI model class, which is from cached content and you provide your cached content here. You can also provide a name for for each different cache and then you can use that name to actually retrieve different cache. I'll just show you a few examples of how that will differ. So now we created the model, but in this case it has access to this cached content. If we run the same queries again, you will see that it takes about the same time. And the reason is that it seems like Google has so far only implemented the cost reduction, but not the the reduction in latency. That update probably is going to come in near future. So it's definitely something you need to keep in mind if you're expecting a lot lower latency. Okay, next we can get the response from the model. So this is very similar response to what we have seen before. So it talks about understanding rag and its component, then it talks about height in LangChain and it has a code example and then there's a section on routing in LangChain as well. Now, currently it's not possible to retrieve or view cached content, but you can retrieve some metadata related to your cache. So from the response you can actually see that how many tokens were used in total. So this is basically the total number of tokens that were used. These are the candidates token count, which I believe is the number of tokens generated by the model. Now in this case, the total number of tokens that were cached were about 40,800 or 900, and the number of tokens that are coming in from the prompt are probably around 30 tokens. So when you add those, that becomes token, prompt, token count. So basically it will contain both your prompt as well as your cached content. Okay, now if you want to update your cached time to live, right? So let's say if you want to increase it, you can do that by simply calling the update function and then providing a new value. So now in this case that's basically updated to two hours. I've been playing around with this, so I, I wanted to look at all the different caches that I have created. So you can actually see that so far I have created three different ones and here are the names of those. So this was the identifier that I was talking about that you can actually provide different names to your each of your cache and then you can use that if you're trying to retrieve information from a specific one. Now, with each one of them there is time at which it was created, then the updated. So basically in most of the cases, I simply updated the cache duration or time to live to two hours. So you can see that what the expiration time is going to be. Now, for some reason, if you want to simply delete the cache before the expiration time, you can do that as well just by calling the delete function on top of the cached content. Okay, so this was a quick tutorial on how to do caching using the Gemini API. It's definitely a feature that is going to be extremely helpful for developers, especially if you're building on top of the Gemini API. And I hope to see that the other API providers such as OpenAI or Anthropic or Cohere will be able to implement something like this. I think this will significantly improve the quality of life for developers, the cost as well as the processing time. So I highly recommend to check it out. You can play with both the Flash and the Pro version of Gemini API for free. Let me know what you think and let me know if you have any questions in the comment section below. I hope you found this video useful. Thanks for watching and as always, see you in the next one.