1+ package org.example.vicky.context
2+
3+ import com.aallam.openai.api.chat.ChatCompletionRequest
4+ import com.aallam.openai.api.chat.ChatMessage
5+ import com.aallam.openai.api.chat.ChatRole
6+ import com.aallam.openai.client.OpenAI
7+ import org.example.vicky.agent.AgentConfig
8+
9+ /* *
10+ * 上下文压缩器，提供两级上下文管理：
11+ *
12+ * 1. **轮数裁剪** ([maxMemoryRounds])：保留最近 N 轮用户消息，截断更早的消息。
13+ * 2. **LLM 压缩** ([maxContextLength])：当估算 token 超限时，调用 LLM 生成结构化摘要，替换旧消息。
14+ *
15+ * 在每次 LLM 请求前调用 [ensureContextBudget] 执行两级检查。
16+ */
17+ class ContextCompactor (
18+ private val config : AgentConfig ,
19+ private val openAi : OpenAI ,
20+ ) {
21+ companion object {
22+ private const val CHARS_PER_TOKEN = 4
23+ private const val KEEP_RECENT_MESSAGES = 4
24+ private const val SUMMARY_MARKER = " [context-summary]"
25+
26+ const val COMPRESSION_PROMPT = """ 请将以下对话历史压缩为结构化摘要，保留关键信息。使用以下格式，空的部分写"N/A"：
27+
28+ ## 初始请求
29+ 用户的原始目标或请求。
30+
31+ ## 已完成步骤
32+ - 按顺序列出每个已执行的操作、调用的工具、做出的决定。
33+
34+ ## 当前状态
35+ 目前已完成的工作，涉及的文件/代码/数据修改。
36+
37+ ## 待办任务
38+ 仍需完成的工作。
39+
40+ ## 技术决策
41+ 关键的架构、设计或实现选择。
42+
43+ ## 涉及文件
44+ 所有读取、写入、引用的文件路径。
45+
46+ ## 遇到的错误
47+ 错误及解决方案（或未解决的情况）。
48+
49+ ## 环境信息
50+ 相关上下文：OS、语言版本、框架版本、项目结构。
51+
52+ ## 下一步建议
53+ 建议的后续操作。
54+
55+ ---
56+
57+ 对话内容：
58+ {conversation}"""
59+ }
60+
61+ /* *
62+ * 主入口：在每次 LLM 请求前调用。
63+ * 依次执行轮数裁剪和 LLM 压缩。
64+ */
65+ suspend fun ensureContextBudget (history : MutableList <ChatMessage >) {
66+ if (config.maxMemoryRounds > 0 ) {
67+ trimMemoryRounds(history)
68+ }
69+ if (config.maxContextLength > 0 ) {
70+ compressToContextBudget(history)
71+ }
72+ }
73+
74+ // ── 机制一：轮数裁剪 ──────────────────────────────────────
75+
76+ private fun trimMemoryRounds (history : MutableList <ChatMessage >) {
77+ val systemMsg = history.firstOrNull { it.role == ChatRole .System }
78+ val nonSystemMessages = history.filter { it.role != ChatRole .System }
79+
80+ val userMessageIndices = nonSystemMessages
81+ .mapIndexedNotNull { i, m -> if (m.role == ChatRole .User ) i else null }
82+
83+ if (userMessageIndices.size <= config.maxMemoryRounds) return
84+
85+ val roundsToTrim = userMessageIndices.size - config.maxMemoryRounds
86+ val cutoffNonSystemIndex = userMessageIndices[roundsToTrim]
87+ val cutoffHistoryIndex = cutoffNonSystemIndex + if (systemMsg != null ) 1 else 0
88+
89+ val kept = history.subList(cutoffHistoryIndex, history.size).toMutableList()
90+ history.clear()
91+ if (systemMsg != null ) history.add(systemMsg)
92+ history.addAll(kept)
93+ }
94+
95+ // ── 机制二：LLM 压缩 ──────────────────────────────────────
96+
97+ private suspend fun compressToContextBudget (history : MutableList <ChatMessage >) {
98+ val totalTokens = history.sumOf { it.estimateChars() } / CHARS_PER_TOKEN
99+ if (totalTokens <= config.maxContextLength) return
100+
101+ val systemMsg = if (history.firstOrNull()?.role == ChatRole .System ) history.first() else null
102+ val preserveEnd = if (systemMsg != null ) 1 else 0
103+ val totalNonSystem = history.size - preserveEnd
104+ if (totalNonSystem <= KEEP_RECENT_MESSAGES ) return
105+
106+ val compressStart = preserveEnd
107+ // 向后推进 compressEnd 跨过任何 Tool 消息，避免它们与压缩区内的 Assistant.toolCalls 失配
108+ var compressEnd = history.size - KEEP_RECENT_MESSAGES
109+ while (compressEnd < history.size && history[compressEnd].role == ChatRole .Tool ) {
110+ compressEnd++
111+ }
112+ if (compressEnd <= compressStart) return
113+
114+ val toCompress = history.subList(compressStart, compressEnd).map { it.deepCopy() }
115+ val toKeepRecent = history.subList(compressEnd, history.size).map { it.deepCopy() }
116+
117+ val summary = callLLMForSummary(toCompress)
118+
119+ history.clear()
120+ if (systemMsg != null ) history.add(systemMsg)
121+ history.add(
122+ ChatMessage (
123+ role = ChatRole .System ,
124+ content = " $SUMMARY_MARKER \n $summary "
125+ ),
126+ )
127+ history.addAll(toKeepRecent)
128+ }
129+
130+ private suspend fun callLLMForSummary (messages : List <ChatMessage >): String {
131+ val formattedMessages = messages.joinToString(" \n\n " ->
132+ val role = when (msg.role) {
133+ ChatRole .User -> " User"
134+ ChatRole .Assistant -> " Assistant"
135+ ChatRole .Tool -> " Tool(${msg.name ? : " ?" )"
136+ else -> msg.role.toString()
137+ }
138+ val content = msg.content?.take(2000 ) ? : " (empty)"
139+ " [$role ]: $content "
140+ }
141+
142+ val prompt = COMPRESSION_PROMPT .replace(" {conversation}"
143+
144+ val request = ChatCompletionRequest (
145+ model = config.model,
146+ messages = listOf (
147+ ChatMessage (role = ChatRole .System , content = " You are a conversation summarizer."
148+ ChatMessage (role = ChatRole .User , content = prompt),
149+ ),
150+ temperature = 0.0 ,
151+ )
152+
153+ return try {
154+ openAi.chatCompletion(request).choices.first().message.content
155+ ? : " (summary generation returned empty)"
156+ } catch (e: Exception ) {
157+ " (summary generation failed: ${e.message} )"
158+ }
159+ }
160+
161+ private fun ChatMessage.estimateChars (): Int {
162+ val contentLen = (content ? : " "
163+ val toolCallsLen = toolCalls?.sumOf { tc ->
164+ when (tc) {
165+ is com.aallam.openai.api.chat.ToolCall .Function -> tc.function.arguments.length
166+ else -> 0
167+ }
168+ } ? : 0
169+ return contentLen + toolCallsLen
170+ }
171+
172+ private fun ChatMessage.deepCopy (): ChatMessage = ChatMessage (
173+ role = role,
174+ content = content,
175+ toolCalls = toolCalls,
176+ toolCallId = toolCallId,
177+ name = name,
178+ )
179+ }
0 commit comments