Skip to content

Commit 7ce823a

Browse files
committed
feat(llm): add token usage tracking with event system
- Add TokenUsageEvent data class and TokenUsageListener interface - Make Usage data class serializable for event transmission - Standardize topic naming convention across message bus topics - Include session tracking and notification helper for token usage events
1 parent afbe257 commit 7ce823a

File tree

3 files changed

+44
-1
lines changed

3 files changed

+44
-1
lines changed
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package cc.unitmesh.devti.llm2
2+
3+
import cc.unitmesh.devti.llms.custom.Usage
4+
import com.intellij.util.messages.Topic
5+
import kotlinx.serialization.Serializable
6+
7+
/**
8+
* Token usage event data that will be sent when SSE stream finishes
9+
*/
10+
@Serializable
11+
data class TokenUsageEvent(
12+
val usage: Usage,
13+
val model: String? = null,
14+
val sessionId: String? = null,
15+
val timestamp: Long = System.currentTimeMillis()
16+
)
17+
18+
/**
19+
* Listener interface for token usage events
20+
*/
21+
interface TokenUsageListener {
22+
/**
23+
* Called when token usage data is available from LLM response
24+
*
25+
* @param event The token usage event containing usage statistics
26+
*/
27+
fun onTokenUsage(event: TokenUsageEvent)
28+
29+
companion object {
30+
val TOPIC = Topic.create("autodev.llm.token.usage", TokenUsageListener::class.java)
31+
32+
/**
33+
* Notify all subscribers about token usage
34+
*/
35+
fun notify(event: TokenUsageEvent) {
36+
com.intellij.openapi.application.ApplicationManager.getApplication().messageBus
37+
.syncPublisher(TOPIC)
38+
.onTokenUsage(event)
39+
}
40+
}
41+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import com.fasterxml.jackson.annotation.JsonAlias
44
import com.fasterxml.jackson.annotation.JsonInclude
55
import com.fasterxml.jackson.annotation.JsonProperty
66
import com.fasterxml.jackson.databind.JsonNode
7+
import kotlinx.serialization.Serializable
78

89
data class ChatFunctionCall(
910
val name: String? = null,
@@ -27,6 +28,7 @@ data class ChatCompletionChoice(
2728
val finishReason: String? = null,
2829
)
2930

31+
@Serializable
3032
data class Usage(
3133
@JsonProperty("prompt_tokens")
3234
val promptTokens: Long = 0,

core/src/main/kotlin/cc/unitmesh/devti/mcp/ui/McpConfigService.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ interface SketchConfigListener {
1919
fun onSelectedToolsChanged(tools: Map<String, Set<Tool>>)
2020

2121
companion object {
22-
val TOPIC = Topic.create("SketchConfigChanged", SketchConfigListener::class.java)
22+
val TOPIC = Topic.create("autodev.sketch.configChanged", SketchConfigListener::class.java)
2323
}
2424
}
2525

0 commit comments

Comments
 (0)