File size: 4,018 Bytes
eb1202c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import type { Message, Part } from "@opencode-ai/sdk/v2/client"
export type SessionContextBreakdownKey = "system" | "user" | "assistant" | "tool" | "other"
export type SessionContextBreakdownSegment = {
key: SessionContextBreakdownKey
tokens: number
width: number
percent: number
}
const estimateTokens = (chars: number) => Math.ceil(chars / 4)
const toPercent = (tokens: number, input: number) => (tokens / input) * 100
const toPercentLabel = (tokens: number, input: number) => Math.round(toPercent(tokens, input) * 10) / 10
const charsFromUserPart = (part: Part) => {
if (part.type === "text") return part.text.length
if (part.type === "file") return part.source?.text.value.length ?? 0
if (part.type === "agent") return part.source?.value.length ?? 0
return 0
}
const charsFromAssistantPart = (part: Part) => {
if (part.type === "text") return { assistant: part.text.length, tool: 0 }
if (part.type === "reasoning") return { assistant: part.text.length, tool: 0 }
if (part.type !== "tool") return { assistant: 0, tool: 0 }
const input = Object.keys(part.state.input).length * 16
if (part.state.status === "pending") return { assistant: 0, tool: input + part.state.raw.length }
if (part.state.status === "completed") return { assistant: 0, tool: input + part.state.output.length }
if (part.state.status === "error") return { assistant: 0, tool: input + part.state.error.length }
return { assistant: 0, tool: input }
}
const build = (
tokens: { system: number; user: number; assistant: number; tool: number; other: number },
input: number,
) => {
return [
{
key: "system",
tokens: tokens.system,
},
{
key: "user",
tokens: tokens.user,
},
{
key: "assistant",
tokens: tokens.assistant,
},
{
key: "tool",
tokens: tokens.tool,
},
{
key: "other",
tokens: tokens.other,
},
]
.filter((x) => x.tokens > 0)
.map((x) => ({
key: x.key,
tokens: x.tokens,
width: toPercent(x.tokens, input),
percent: toPercentLabel(x.tokens, input),
})) as SessionContextBreakdownSegment[]
}
export function estimateSessionContextBreakdown(args: {
messages: Message[]
parts: Record<string, Part[] | undefined>
input: number
systemPrompt?: string
}) {
if (!args.input) return []
const counts = args.messages.reduce(
(acc, msg) => {
const parts = args.parts[msg.id] ?? []
if (msg.role === "user") {
const user = parts.reduce((sum, part) => sum + charsFromUserPart(part), 0)
return { ...acc, user: acc.user + user }
}
if (msg.role !== "assistant") return acc
const assistant = parts.reduce(
(sum, part) => {
const next = charsFromAssistantPart(part)
return {
assistant: sum.assistant + next.assistant,
tool: sum.tool + next.tool,
}
},
{ assistant: 0, tool: 0 },
)
return {
...acc,
assistant: acc.assistant + assistant.assistant,
tool: acc.tool + assistant.tool,
}
},
{
system: args.systemPrompt?.length ?? 0,
user: 0,
assistant: 0,
tool: 0,
},
)
const tokens = {
system: estimateTokens(counts.system),
user: estimateTokens(counts.user),
assistant: estimateTokens(counts.assistant),
tool: estimateTokens(counts.tool),
}
const estimated = tokens.system + tokens.user + tokens.assistant + tokens.tool
if (estimated <= args.input) {
return build({ ...tokens, other: args.input - estimated }, args.input)
}
const scale = args.input / estimated
const scaled = {
system: Math.floor(tokens.system * scale),
user: Math.floor(tokens.user * scale),
assistant: Math.floor(tokens.assistant * scale),
tool: Math.floor(tokens.tool * scale),
}
const total = scaled.system + scaled.user + scaled.assistant + scaled.tool
return build({ ...scaled, other: Math.max(0, args.input - total) }, args.input)
}
|