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[], messages: Anthropic.Messages.MessageParam[],
tools: Anthropic.Messages.Tool[] tools: Anthropic.Messages.Tool[]
): Anthropic.Messages.MessageParam[] { ): 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 systemPromptTokens = countTokens(systemPrompt)
const toolsTokens = countTokens(JSON.stringify(tools)) const toolsTokens = countTokens(JSON.stringify(tools))
let availableTokens = adjustedContextWindow - systemPromptTokens - toolsTokens 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) // 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] const newMessages = [...messages]
while (totalMessageTokens > availableTokens && newMessages.length > 2) { let index = 2
const removedMessage = newMessages.splice(2, 1)[0] while (totalMessageTokens > availableTokens && index < newMessages.length) {
const removedTokens = countMessageTokens(removedMessage) const messageToEmpty = newMessages[index]
totalMessageTokens -= removedTokens 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
if (totalMessageTokens > availableTokens) { })
// Over the limit due to the first two messages }
throw new Error("Task exceeds available context window") const newTokens = countMessageTokens(messageToEmpty)
totalMessageTokens -= originalTokens - newTokens
index++
} }
return newMessages return newMessages
} }