From 8cd0de045de5e3f2b6765d0fa0e04170a018abc3 Mon Sep 17 00:00:00 2001
From: Andrew Charneski <acharneski@gmail.com>
Date: Sun, 12 Nov 2023 02:01:09 -0500
Subject: [PATCH] wip

---
 .../aicoder/actions/dev/AppServer.kt          |  2 +-
 .../aicoder/actions/dev/CodeChatAction.kt     |  2 +-
 .../aicoder/actions/dev/CodeChatServer.kt     | 93 +++++++++----------
 .../aicoder/util/IdeaOpenAIClient.kt          |  4 +-
 4 files changed, 49 insertions(+), 52 deletions(-)

diff --git a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/AppServer.kt b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/AppServer.kt
index a7546105..bd1d389f 100644
--- a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/AppServer.kt
+++ b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/AppServer.kt
@@ -4,7 +4,7 @@ import com.github.simiacryptus.aicoder.config.AppSettingsState
 import com.github.simiacryptus.aicoder.util.UITools
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.project.Project
-import com.simiacryptus.skyenet.webui.WebSocketServer
+import com.simiacryptus.skyenet.sessions.WebSocketServer
 import org.eclipse.jetty.server.Server
 import org.eclipse.jetty.server.handler.ContextHandlerCollection
 import org.eclipse.jetty.webapp.WebAppContext
diff --git a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatAction.kt b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatAction.kt
index e3305e48..a7b5f841 100644
--- a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatAction.kt
+++ b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatAction.kt
@@ -21,7 +21,7 @@ class CodeChatAction : BaseAction() {
         val language = ComputerLanguage.getComputerLanguage(e)?.name ?: return
         val server = AppServer.getServer(e.project)
         val uuid = UUID.randomUUID().toString()
-        server.addApp("/$uuid", CodeChatServer(e.project!!, language, selectedText))
+        server.addApp("/$uuid", CodeChatServer(e.project!!, language, selectedText, api = api))
         Thread {
             Thread.sleep(500)
             try {
diff --git a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatServer.kt b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatServer.kt
index bc21ded3..b3bac850 100644
--- a/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatServer.kt
+++ b/src/main/kotlin/com/github/simiacryptus/aicoder/actions/dev/CodeChatServer.kt
@@ -3,53 +3,37 @@ package com.github.simiacryptus.aicoder.actions.dev
 import com.github.simiacryptus.aicoder.config.AppSettingsState
 import com.intellij.openapi.project.Project
 import com.simiacryptus.openai.OpenAIClient
-import com.simiacryptus.skyenet.webui.*
+import com.simiacryptus.skyenet.sessions.*
+import com.simiacryptus.skyenet.util.ClasspathResource
+import com.simiacryptus.skyenet.util.MarkdownUtil
 import org.eclipse.jetty.util.resource.Resource
+import org.intellij.lang.annotations.Language
 
 class CodeChatServer(
     val project: Project,
     val language: String,
     val codeSelection: String,
-) : BasicChat(
+    val api: OpenAIClient,
+) : ChatApplicationBase(
     applicationName = "Code Chat",
-    model = AppSettingsState.instance.defaultChatModel()
 ) {
 
     private val rootOperationID = (0..5).map { ('a'..'z').random() }.joinToString("")
     private var rootMessageTrail: String = ""
 
-    override fun newSession(sessionId: String): CodeChatSession {
-        val newSession = CodeChatSession(sessionId)
-rootMessageTrail =
-"""$rootOperationID,<div><h3>Code:</h3><pre><code class="language-$language">${htmlEscape(codeSelection)}</code></pre></div>"""
-        newSession.send(rootMessageTrail)
-        return newSession
-    }
-
-    private fun htmlEscape(html: String): String {
-        return html.replace("&", "&amp;").replace("<", "&lt;")
-            .replace(">", "&gt;").replace("\"", "&quot;")
-            .replace("'", "&#039;")
-    }
-
-    open inner class CodeChatSession(sessionId: String) : BasicChatSession(
-        parent = this@CodeChatServer,
-        model = model,
-        sessionId = sessionId
-    ) {
-        override fun run(userMessage: String) {
-            var messageTrail = divInitializer()
-            send("""$messageTrail<div>$userMessage</div><div>$spinner</div>""")
-            messages += OpenAIClient.ChatMessage(OpenAIClient.ChatMessage.Role.user, userMessage)
-            val response = api.chat(chatRequest, model).choices.first().message?.content.orEmpty()
-            messages += OpenAIClient.ChatMessage(OpenAIClient.ChatMessage.Role.assistant, response)
-            messageTrail += MarkdownUtil.renderMarkdown(response)
-            send(messageTrail)
-        }
-
-        override val messages = listOf(
-            OpenAIClient.ChatMessage(
-                OpenAIClient.ChatMessage.Role.system, """
+    override fun newSession(sessionId: String): ChatSession {
+        val newSession = ChatSession(
+            sessionId = sessionId,
+            parent = this@CodeChatServer,
+            model = AppSettingsState.instance.defaultChatModel(),
+            api = api,
+            visiblePrompt = """
+            <div><h3>Code:</h3>
+                <pre><code class="language-$language">${htmlEscape(codeSelection)}</code></pre>
+            </div>
+        """.trimIndent().trim(),
+            hiddenPrompt = "",
+            systemPrompt = """
                     |You are a helpful AI that helps people with coding.
                     |
                     |You will be answering questions about the following code:
@@ -59,22 +43,35 @@ rootMessageTrail =
                     |```
                     |
                     |Responses may use markdown formatting.
-                    """.trimMargin()
-            )
-        ).toMutableList()
+                    """.trimMargin(),
+        )
+        @Language("HTML") val html = """
+            <div><h3>Code:</h3>
+                <pre><code class="language-$language">${htmlEscape(codeSelection)}</code></pre>
+            </div>
+            """.trimIndent().trim()
+        rootMessageTrail = """$rootOperationID,$html"""
+        newSession.send(rootMessageTrail)
+        return newSession
+    }
 
-        private val chatRequest: OpenAIClient.ChatRequest
-            get() {
-                val chatRequest = OpenAIClient.ChatRequest()
-                val model = AppSettingsState.instance.defaultChatModel()
-                chatRequest.model = model.modelName
-                chatRequest.temperature = AppSettingsState.instance.temperature
-                chatRequest.messages = messages.toTypedArray()
-                return chatRequest
-            }
+    override fun processMessage(
+        sessionId: String,
+        userMessage: String,
+        session: PersistentSessionBase,
+        sessionDiv: SessionDiv,
+        socket: MessageWebSocket
+    ) {
+        TODO("Not yet implemented")
+    }
+
+    private fun htmlEscape(html: String): String {
+        return html.replace("&", "&amp;").replace("<", "&lt;")
+            .replace(">", "&gt;").replace("\"", "&quot;")
+            .replace("'", "&#039;")
     }
 
     override val baseResource: Resource
-        get() = ClasspathResource(javaClass.classLoader.getResource(resourceBase))
+        get() = ClasspathResource(javaClass.classLoader.getResource(resourceBase)!!)
 
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/github/simiacryptus/aicoder/util/IdeaOpenAIClient.kt b/src/main/kotlin/com/github/simiacryptus/aicoder/util/IdeaOpenAIClient.kt
index bbd81917..7dbad08a 100644
--- a/src/main/kotlin/com/github/simiacryptus/aicoder/util/IdeaOpenAIClient.kt
+++ b/src/main/kotlin/com/github/simiacryptus/aicoder/util/IdeaOpenAIClient.kt
@@ -21,8 +21,8 @@ class IdeaOpenAIClient : OpenAIClient(
 ) {
     private val isInRequest = AtomicBoolean(false)
 
-    override fun incrementTokens(totalTokens: Int) {
-        AppSettingsState.instance.tokenCounter += totalTokens
+    override fun incrementTokens(model: Model?, tokens: Int) {
+        AppSettingsState.instance.tokenCounter += tokens
     }
 
     override fun authorize(request: HttpRequestBase) {