Fix sliding window taking into account message order

This commit is contained in:
Saoud Rizwan
2024-08-25 00:53:00 -04:00
parent 650283a11f
commit cfd4cbfa36

View File

@@ -7,7 +7,7 @@ export function slidingWindowContextManagement(
messages: Anthropic.Messages.MessageParam[],
tools: Anthropic.Messages.Tool[]
): Anthropic.Messages.MessageParam[] {
const adjustedContextWindow = contextWindow - 10000 // Buffer to account for tokenizer differences
const adjustedContextWindow = contextWindow - 10_000 // Buffer to account for tokenizer differences
const systemPromptTokens = countTokens(systemPrompt)
const toolsTokens = countTokens(JSON.stringify(tools))
let availableTokens = adjustedContextWindow - systemPromptTokens - toolsTokens
@@ -19,17 +19,50 @@ export function slidingWindowContextManagement(
// If over limit, remove messages starting from the third message onwards (task and claude's step-by-step thought process are important to keep in context)
const newMessages = [...messages]
while (totalMessageTokens > availableTokens && newMessages.length > 2) {
const removedMessage = newMessages.splice(2, 1)[0]
const removedTokens = countMessageTokens(removedMessage)
totalMessageTokens -= removedTokens
let index = 2
while (totalMessageTokens > availableTokens && index < newMessages.length) {
const messageToEmpty = newMessages[index]
const originalTokens = countMessageTokens(messageToEmpty)
// Empty the content of the message (messages must be in a specific order so we can't just remove)
if (typeof messageToEmpty.content === "string") {
messageToEmpty.content = ""
} else if (Array.isArray(messageToEmpty.content)) {
messageToEmpty.content = messageToEmpty.content.map((item) => {
if (typeof item === "string") {
return {
type: "text",
text: "(truncated due to context window)",
} as Anthropic.Messages.TextBlockParam
} else if (item.type === "text") {
return {
type: "text",
text: "(truncated due to context window)",
} as Anthropic.Messages.TextBlockParam
} else if (item.type === "image") {
return { ...item, source: { type: "base64", data: "" } } as Anthropic.Messages.ImageBlockParam
} else if (item.type === "tool_use") {
return { ...item, input: {} } as Anthropic.Messages.ToolUseBlockParam
} else if (item.type === "tool_result") {
return {
...item,
content: Array.isArray(item.content)
? item.content.map((contentItem) =>
contentItem.type === "text"
? { ...contentItem, text: "(truncated due to context window)" }
: contentItem.type === "image"
? { ...contentItem, source: { type: "base64", data: "" } }
: contentItem
)
: "",
} as Anthropic.Messages.ToolResultBlockParam
}
return item
})
}
const newTokens = countMessageTokens(messageToEmpty)
totalMessageTokens -= originalTokens - newTokens
index++
}
if (totalMessageTokens > availableTokens) {
// Over the limit due to the first two messages
throw new Error("Task exceeds available context window")
}
return newMessages
}