Update gradle.properties

core/src/main/kotlin/com/simiacryptus/cognotik/platform/ClientManager.kt

@@ -12,6 +12,8 @@ import com.simiacryptus.cognotik.util.ImmediateExecutorService
 import com.simiacryptus.jopenai.ChatClient
 import com.simiacryptus.jopenai.OpenAIClient
 import com.simiacryptus.jopenai.models.APIProvider
+import com.simiacryptus.jopenai.models.ApiModel
+import com.simiacryptus.jopenai.models.OpenAIModel
 import com.simiacryptus.jopenai.util.ClientUtil
 import org.slf4j.LoggerFactory
 import java.util.concurrent.ScheduledThreadPoolExecutor
@@ -64,11 +66,19 @@ open class ClientManager {
             val userSettings = userSettingsManager.getUserSettings(user)
             val userApi =
                 if (userSettings.apiKeys.isNotEmpty()) {
-                    ChatClient(
+                    object : ChatClient(
                         key = userSettings.apiKeys,
                         apiBase = userSettings.apiBase,
                         workPool = getPool(session, user),
-                    ).apply {
+                    ){
+                        override fun onUsage(
+                            model: OpenAIModel?,
+                            tokens: ApiModel.Usage
+                        ) {
+                            super.onUsage(model, tokens)
+                            ApplicationServices.usageManager.incrementUsage(session, user, model!!, tokens)
+                        }
+                    }.apply {
                         this.session = session
                         this.user = user
                         logStreams += sessionDir.resolve("openai.log").outputStream().buffered()
@@ -81,10 +91,18 @@ open class ClientManager {
         )
         if (!canUseGlobalKey) throw RuntimeException("No API key")
         return (if (ClientUtil.keyMap.isNotEmpty()) {
-            ChatClient(
+            object : ChatClient(
                 key = ClientUtil.keyMap.mapKeys { APIProvider.valueOf(it.key) },
                 workPool = getPool(session, user),
-            ).apply {
+            ){
+                override fun onUsage(
+                    model: OpenAIModel?,
+                    tokens: ApiModel.Usage
+                ) {
+                    super.onUsage(model, tokens)
+                    ApplicationServices.usageManager.incrementUsage(session, user, model!!, tokens)
+                }
+            }.apply {
                 this.session = session
                 this.user = user
                 logStreams += sessionDir.resolve("openai.log").outputStream().buffered()

gradle.properties

@@ -1,7 +1,7 @@
 pluginName=Cognotik
 pluginRepositoryUrl=https://github.com/SimiaCryptus/Cognotik
 libraryGroup=com.simiacryptus
-libraryVersion=2.0.1
+libraryVersion=2.0.2
 
 gradleVersion=8.13
 org.gradle.caching=true

intellij/src/main/kotlin/com/simiacryptus/cognotik/util/IdeaChatClient.kt

@@ -72,8 +72,8 @@ open class IdeaChatClient(
     private val isInRequest = AtomicBoolean(false)
 
     override fun onUsage(model: OpenAIModel?, tokens: Usage) {
-
         ApplicationServices.usageManager.incrementUsage(currentSession, localUser, model!!, tokens)
+        super.onUsage(model, tokens)
     }
 
     override fun authorize(request: HttpRequest, apiProvider: APIProvider) {

jo-penai/src/main/kotlin/com/simiacryptus/jopenai/ChatClient.kt

@@ -71,14 +71,20 @@ open class ChatClient(
             super.log(level, msg)
             [email protected](level, msg)
         }
+
+        override fun onUsage(
+            model: OpenAIModel?,
+            tokens: Usage
+        ) {
+            [email protected](model, tokens)
+            super.onUsage(model, tokens)
+        }
     }
 
     open fun getChildClient(): ChatClient = ChildClient()
 
     protected open fun onUsage(model: OpenAIModel?, tokens: Usage) {
-        log.debug(
-            "Usage recorded for session: {}, user: {}, model: {}, tokens: {}", session, user, model, tokens
-        )
+        log.debug("Usage recorded for session: {}, user: {}, model: {}, tokens: {}", session, user, model, tokens)
         if (null != budget) budget = budget!!.toDouble() - (tokens.cost ?: 0.0)
     }