diff --git a/webui/src/main/kotlin/com/simiacryptus/skyenet/servers/AppServerBase.kt b/webui/src/main/kotlin/com/simiacryptus/skyenet/servers/AppServerBase.kt index 0dbcf050..0e1f6ad2 100644 --- a/webui/src/main/kotlin/com/simiacryptus/skyenet/servers/AppServerBase.kt +++ b/webui/src/main/kotlin/com/simiacryptus/skyenet/servers/AppServerBase.kt @@ -5,6 +5,7 @@ import com.simiacryptus.openai.OpenAIClient import com.simiacryptus.skyenet.OutputInterceptor import com.simiacryptus.skyenet.util.AwsUtil.decryptResource import com.simiacryptus.skyenet.servlet.AuthenticatedWebsite +import com.simiacryptus.skyenet.servlet.UsageServlet import com.simiacryptus.skyenet.webui.ApplicationBase import com.simiacryptus.skyenet.servlet.UserInfoServlet import com.simiacryptus.skyenet.servlet.UserSettingsServlet @@ -46,6 +47,7 @@ abstract class AppServerBase( val welcomeResources = Resource.newResource(javaClass.classLoader.getResource("welcome")) val userInfoServlet = UserInfoServlet() val userSettingsServlet = UserSettingsServlet() + val usageServlet = UsageServlet() protected fun _main(args: Array) { try { @@ -65,6 +67,7 @@ abstract class AppServerBase( *(arrayOf( newWebAppContext("/userInfo", userInfoServlet), newWebAppContext("/userSettings", userSettingsServlet), + newWebAppContext("/usage", usageServlet), newWebAppContext("/proxy", ProxyHttpServlet()), authentication.configure( newWebAppContext( @@ -110,6 +113,7 @@ abstract class AppServerBase( requestURI == "/index.html" -> resp?.writer?.write(homepage().trimIndent()) requestURI.startsWith("/userInfo") -> userInfoServlet.doGet(req!!, resp!!) requestURI.startsWith("/userSettings") -> userSettingsServlet.doGet(req!!, resp!!) + requestURI.startsWith("/usage") -> usageServlet.doGet(req!!, resp!!) else -> try { val inputStream = welcomeResources.addPath(requestURI)?.inputStream inputStream?.copyTo(resp?.outputStream!!) diff --git a/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/SessionSettingsServlet.kt b/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/SessionSettingsServlet.kt index b9480006..28c20d18 100644 --- a/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/SessionSettingsServlet.kt +++ b/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/SessionSettingsServlet.kt @@ -13,31 +13,31 @@ class SessionSettingsServlet( resp.contentType = "text/html" resp.status = HttpServletResponse.SC_OK val sessionId = req.getParameter("sessionId") - if (null == sessionId) { - resp.status = HttpServletResponse.SC_BAD_REQUEST - resp.writer.write("Session ID is required") - } else { + if (null != sessionId) { val settings = server.getSettings(sessionId) val json = if(settings != null) JsonUtil.toJson(settings) else "" //language=HTML resp.writer.write( """ - | - | - | Settings - | - | - | - |
- | - | - | - | - |
- | - | - """.trimMargin() + | + | + | Settings + | + | + | + |
+ | + | + | + | + |
+ | + | + """.trimMargin() ) + } else { + resp.status = HttpServletResponse.SC_BAD_REQUEST + resp.writer.write("Session ID is required") } } diff --git a/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/UsageServlet.kt b/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/UsageServlet.kt new file mode 100644 index 00000000..0ffdaaf2 --- /dev/null +++ b/webui/src/main/kotlin/com/simiacryptus/skyenet/servlet/UsageServlet.kt @@ -0,0 +1,107 @@ +package com.simiacryptus.skyenet.servlet + +import com.google.api.services.oauth2.model.Userinfo +import com.simiacryptus.openai.OpenAIClient +import jakarta.servlet.http.HttpServlet +import jakarta.servlet.http.HttpServletRequest +import jakarta.servlet.http.HttpServletResponse +import java.util.concurrent.atomic.AtomicInteger + +class UsageServlet : HttpServlet() { + public override fun doGet(req: HttpServletRequest, resp: HttpServletResponse) { + resp.contentType = "text/html" + resp.status = HttpServletResponse.SC_OK + + val sessionId = req.getParameter("sessionId") + if (null != sessionId) { + serve(resp, getSessionUsageSummary(sessionId)) + } else { + val userinfo = AuthenticatedWebsite.getUser(req) + if (null == userinfo) { + resp.status = HttpServletResponse.SC_BAD_REQUEST + } else { + val usage = getUserUsageSummary(userinfo) + serve(resp, usage) + } + } + } + + private fun serve( + resp: HttpServletResponse, + usage: Map + ) { + resp.writer.write( + """ + | + | + | Usage + | + | + | + | + | + | + | + | + | ${ + usage.entries.joinToString("\n") { (model, count) -> + """ + | + | + | + | + """.trimMargin() + } + } + |
ModelUsage
$model$count
+ | + | + """.trimMargin() + ) + } + + data class UsageCounters( + val tokensPerModel: HashMap = HashMap(), + ) + + companion object { + val log = org.slf4j.LoggerFactory.getLogger(UsageServlet::class.java) + + private val usagePerSession = HashMap() + val sessionsByUser = HashMap>() + + fun incrementUsage(sessionId: String, userinfo: Userinfo?, model: OpenAIClient.Model, tokens: Int) { + val usage = usagePerSession.getOrPut(sessionId) { + UsageCounters() + } + val tokensPerModel = usage.tokensPerModel.getOrPut(model) { + AtomicInteger() + } + tokensPerModel.addAndGet(tokens) + if (userinfo != null) { + val sessions = sessionsByUser.getOrPut(userinfo.id) { + ArrayList() + } + sessions.add(sessionId) + } + } + + fun getUserUsageSummary(userinfo: Userinfo): Map { + val sessions = sessionsByUser[userinfo.id] + return sessions?.flatMap { sessionId -> + val usage = usagePerSession[sessionId] + usage?.tokensPerModel?.entries?.map { (model, counter) -> + model to counter.get() + } ?: emptyList() + }?.groupBy { it.first }?.mapValues { it.value.map { it.second }.sum() } ?: emptyMap() + } + + fun getSessionUsageSummary(sessionId: String): Map { + val usage = usagePerSession[sessionId] + return usage?.tokensPerModel?.entries?.map { (model, counter) -> + model to counter.get() + }?.groupBy { it.first }?.mapValues { it.value.map { it.second }.sum() } ?: emptyMap() + } + + } +} diff --git a/webui/src/main/kotlin/com/simiacryptus/skyenet/webui/MessageWebSocket.kt b/webui/src/main/kotlin/com/simiacryptus/skyenet/webui/MessageWebSocket.kt index 482cbfb4..9ea688c9 100644 --- a/webui/src/main/kotlin/com/simiacryptus/skyenet/webui/MessageWebSocket.kt +++ b/webui/src/main/kotlin/com/simiacryptus/skyenet/webui/MessageWebSocket.kt @@ -2,6 +2,7 @@ package com.simiacryptus.skyenet.webui import com.simiacryptus.openai.OpenAIClient import com.simiacryptus.skyenet.servlet.AuthenticatedWebsite +import com.simiacryptus.skyenet.servlet.UsageServlet.Companion.incrementUsage import com.simiacryptus.skyenet.servlet.UserSettingsServlet.Companion.getUserSettings import org.eclipse.jetty.websocket.api.Session import org.eclipse.jetty.websocket.api.WebSocketAdapter @@ -19,20 +20,30 @@ class MessageWebSocket( val userinfo = AuthenticatedWebsite.users[authId] val userApi: OpenAIClient? = if (userinfo == null) null else { val userSettings = getUserSettings(userinfo) - if (userSettings.apiKey.isBlank()) null else OpenAIClient( + if (userSettings.apiKey.isBlank()) null else object : OpenAIClient( key = userSettings.apiKey, logLevel = Level.DEBUG, logStreams = mutableListOf( sessionDataStorage.getSessionDir(sessionId).resolve("openai.log").outputStream().buffered() ), - ) + ) { + override fun incrementTokens(model: Model?, tokens: Int) { + incrementUsage(sessionId, userinfo, model!!, tokens) + super.incrementTokens(model, tokens) + } + } } - return userApi ?: OpenAIClient( + return userApi ?: object : OpenAIClient( logLevel = Level.DEBUG, logStreams = mutableListOf( sessionDataStorage.getSessionDir(sessionId).resolve("openai.log").outputStream().buffered() ) - ) + ) { + override fun incrementTokens(model: Model?, tokens: Int) { + incrementUsage(sessionId, userinfo, model!!, tokens) + super.incrementTokens(model, tokens) + } + } } override fun onWebSocketConnect(session: Session) { diff --git a/webui/src/main/resources/simpleSession/index.html b/webui/src/main/resources/simpleSession/index.html index 15b8cc64..ce8c5d67 100644 --- a/webui/src/main/resources/simpleSession/index.html +++ b/webui/src/main/resources/simpleSession/index.html @@ -18,6 +18,7 @@ Session List Session Settings Files + Usage
diff --git a/webui/src/main/resources/simpleSession/main.js b/webui/src/main/resources/simpleSession/main.js index 9aef3f05..78b0e4a8 100644 --- a/webui/src/main/resources/simpleSession/main.js +++ b/webui/src/main/resources/simpleSession/main.js @@ -73,6 +73,7 @@ document.addEventListener('DOMContentLoaded', () => { const form = document.getElementById('form'); const messageInput = document.getElementById('message'); + const usage = document.getElementById('usage'); form.addEventListener('submit', (event) => { event.preventDefault(); @@ -90,6 +91,7 @@ document.addEventListener('DOMContentLoaded', () => { const sessionId = getSessionId(); if (sessionId) { connect(sessionId, onWebSocketText); + usage.href = '/usage/?sessionId=' + sessionId; } else { connect(undefined, onWebSocketText); }