Handling Full Context Windows

Handling Full Context Windows

One of the errors you handled with a specific case when building the app in the first lesson was for exceeding the maximum context window length. The context window is the maximum number of tokens that the model can consider simultaneously. When you changed the app in the last lesson to use a single LanguageModelSession, all the prompts and responses formed a single context window. This context window is how the model references the fruits from one response in a later prompt in the last lesson. Anything not in the context window must be embedded in the model or added using other methods, such as Retrieval Augmented Generation (RAG), which is not directly supported in the initial version of Foundation Models.

The current size of the context window is 4,096 tokens. A single token typically corresponds to three or four letters in most Western languages, such as English, and one character in languages like Chinese, Japanese, and Korean. Every prompt and response takes up part of the context window. Instructions also take up space in the context window. What happens when you exceed the context length? Find the text folder in the starter project, which contains several text files. Copy and paste the contents of macbeth-act-one.txt, which includes the text of Act One of Shakespeare’s play Macbeth, into a prompt. You will trigger a LanguageModelSession.GenerationError.exceededContextWindowSize error.

You have several options to handle exceeding the context window length. In some cases, it may be possible to throw away the context and create a new session. This would lose all information in the session; however, for tasks that only require the included information, it may be an acceptable option. If you need to process text that is longer than the context window, your best option will be to split the text into smaller chunks, each sized to fit into the context window. You can then combine the results. You’ll need to modify the current app so that if the context window is exceeded, the app will try to summarize the current conversation and initiate a new session before adding the summary to the context window.

The limited size of the context window means you may often encounter similar situations. You have several options to handle exceeding the context window length:

• Start a new session:
  In some cases, it may be possible to throw away the context and create a new session. This would lose all information in the session; however, for tasks that only require the included information, it may be an acceptable option.

• Split text into smaller chunks:
  If you need to process text that is longer than the context window, your best option will be to split the text into smaller chunks, each sized to fit into the context window. You can then combine the results.

You’ll need to modify the current app so that if the context window is exceeded, the app will try to summarize the current conversation and initiate a new session before adding the summary to the context window.

Open ChatView.swift from the starter project. Add the following method to the end of the view, after the sendMessage() method:

func summarizeChat() {
}

This new method will summarize the current chat and place that summary in the context window of a new session. First, find the catch LanguageModelSession.GenerationError.exceededContextWindowSize case in the sendMessage() method, which now just shows an error. Delete the addMessage(_:isFromUser:animate:) method call and replace it with the following code to call your new method:

summarizeChat()

Now add the following code to the summarizeChat() method:

// 1
var allText = ""
// 2
for entry in session.transcript {
  // 3
  switch entry {
  case .prompt(let prompt):
    allText += prompt.description + "\n"
  case .response(let response):
    allText += response.description + "\n"
  default:
    allText += "\n"
  }
}
// 4
addMessage("Context windows exceeded. Summarizing Chat", isFromUser: false)

This code will gather the current prompts and responses in the context window. To do so:

1. First, you set an empty string to hold the gathered text.
2. Each session has a transcript property which documents all interactions with the model. This code creates a loop through all records in the transcript.
3. Each entry in the transcript is of type Transcript.Entry, an enumerable of the different possible entry types. The closure creates a switch on the type and handles the .prompt and .response entries, which contain the prompts and responses that make up the chat. You add the contents of any of these entries, along with a newline, to the allText string.
4. This will add a message to the chat that the summarization is in progress, since it will likely take at least a few seconds

When the loop completes, allText will contain the contents of the chat. The next step will be to summarize the text. Continue the method with the following code:

let summarySession = LanguageModelSession(instructions: "Summarize all text presented to the model.")
let summarizedText = try? await summarySession.respond(to: allText).content

This creates a new LanguageModelSession specifying instructions to summarize any text sent to the model. It then uses the respond(to:options:) method on the session and gets the content property, which holds the summary. The try? await pattern will set summarizedText to nil if any errors occur.

Now you can create a new session with the summarized text.

// 1
if let summarizedText = summarizedText {
  // 2
  resetChatHistory()
  addMessage(summarizedText, isFromUser: false)
  // 3
  session = LanguageModelSession(instructions: promptInstructions)
  // 4
  let response = try? await session.respond(to: summarizedText)
  addMessage(response?.content ?? "", isFromUser: false)
} else {
  // 5
  resetChatHistory()
}

Here’s how this code finishes the summarization process.

1. First, you attempt to unwrap the summarizedText value. If anything went wrong summarizing the text, summarizedText will be nil. You’ll deal with that later in this code block.
2. If you could unwrap summarizedText, then you now have a valid summary. You add a new message containing the summarized text of the previous chat.
3. First, you clear the existing messages and create a new LanguageModelSession with the current promptInstructions.
4. Now, you prompt the newly created session with the summarized text to add it to the context window. You then present any reply back to the chat window or add an empty chat bubble if an error happens.
5. In the case where an error occurs attempting to summarize the text, meaning summarizedText will be nil, you reset the chat.

Run the app. Enter the following prompt:

Produce a brief summary of the following text.

Now find the macbeth-act-one-s12.txt and macbeth-act-one-s34.txt files in the text folder in the Starter project. These contain scenes one and two and scenes three and four of the first act of Macbeth, respectively. Paste macbeth-act-one-s12.txt as a prompt. Then paste macbeth-act-one-s34.txt into the prompt. The two combined will be too large for the context window, triggering your new summarization code.

You will notice that the summarized text does not include any information from the prompt that triggered the error, as the error prevents the prompt from appearing in the transcript. You can paste macbeth-act-one-s34.txt again to continue the chat.

In a real app, you’d likely want to perform summarization in the background instead of pausing the app while it takes place. It does, however, provide a good starting point for handling overloading of the context window in your apps.