Skip to content

Commit 8ac27e1

Browse files
committed
feat(agent): add token-based message compression
- Change `content` in `Message` to mutable (`var`) for in-place updates. - Add `TokenizerFactory` and `MarkdownCodeHelper` to compress messages based on token count. - Compress messages by removing markdown code if token count exceeds `maxToken`.
1 parent 143d075 commit 8ac27e1

File tree

2 files changed

+17
-3
lines changed

2 files changed

+17
-3
lines changed

core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomSSEProcessor.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ open class CustomSSEProcessor(private val project: Project) {
183183
}
184184

185185
@Serializable
186-
data class Message(val role: String, val content: String)
186+
data class Message(val role: String, var content: String)
187187

188188
@Serializable
189189
data class CustomRequest(val messages: List<Message>)

core/src/main/kotlin/cc/unitmesh/devti/observer/agent/AgentStateService.kt

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@ package cc.unitmesh.devti.observer.agent
33
import cc.unitmesh.devti.agent.tool.AgentTool
44
import cc.unitmesh.devti.devin.dataprovider.BuiltinCommand
55
import cc.unitmesh.devti.llms.custom.Message
6+
import cc.unitmesh.devti.llms.tokenizer.TokenizerFactory
67
import cc.unitmesh.devti.observer.plan.AgentTaskEntry
78
import cc.unitmesh.devti.observer.plan.MarkdownPlanParser
89
import cc.unitmesh.devti.observer.plan.PlanUpdateListener
910
import cc.unitmesh.devti.settings.AutoDevSettingsState
1011
import cc.unitmesh.devti.settings.customize.customizeSetting
12+
import cc.unitmesh.devti.util.parser.MarkdownCodeHelper
1113
import com.intellij.openapi.application.ApplicationManager
1214
import com.intellij.openapi.components.Service
1315
import com.intellij.openapi.diagnostic.logger
@@ -18,6 +20,9 @@ import com.intellij.openapi.vcs.changes.Change
1820
class AgentStateService(val project: Project) {
1921
val maxToken = AutoDevSettingsState.maxTokenLength
2022
var state: AgentState = AgentState()
23+
var tokenizer = lazy {
24+
TokenizerFactory.createTokenizer()
25+
}
2126

2227
fun addTools(tools: List<BuiltinCommand>) {
2328
state.usedTools = tools.map {
@@ -52,8 +57,17 @@ class AgentStateService(val project: Project) {
5257
* Call some LLM to compress it or use some other method to compress the history
5358
*/
5459
fun processMessages(messages: List<Message>): List<Message> {
55-
state.messages = messages
56-
return messages
60+
val countLength = tokenizer.value.count(messages.joinToString("\n") { it.content })
61+
if (countLength < maxToken) {
62+
state.messages = messages
63+
return messages
64+
}
65+
66+
state.messages = messages.map {
67+
it.copy(content = MarkdownCodeHelper.removeAllMarkdownCode(it.content))
68+
}
69+
70+
return state.messages
5771
}
5872

5973
fun updatePlan(items: MutableList<AgentTaskEntry>) {

0 commit comments

Comments
 (0)