diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e5e991c8..11f67490 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -35,6 +35,18 @@ jobs: run: | echo "lsp_version=$(cat refact_lsp)" >> $GITHUB_OUTPUT + - name: Download refact-chat-js + id: download-refact-chat-js + uses: dawidd6/action-download-artifact@v3 + with: + github_token: ${{secrets.GITHUB_TOKEN}} + workflow: node.js.yml + workflow_search: true + repo: smallcloudai/refact-chat-js + branch: alpha + name: refact-chat-js-latest + path: ./src/main/resources/webview/dist + - uses: convictional/trigger-workflow-and-wait@v1.6.1 name: "Build refact-lsp" with: @@ -54,6 +66,7 @@ jobs: branch: ${{ steps.setupvars.outputs.lsp_version }} path: ./src/main/resources/bin + # Validate wrapper - name: Gradle Wrapper Validation uses: gradle/wrapper-validation-action@v3 diff --git a/.gitignore b/.gitignore index 6c319f7f..4b249c25 100644 --- a/.gitignore +++ b/.gitignore @@ -38,4 +38,5 @@ bin/ ### Mac OS ### .DS_Store -src/main/resources/bin/ \ No newline at end of file +src/main/resources/bin/ +src/main/resources/webview/dist \ No newline at end of file diff --git a/build.gradle.kts b/build.gradle.kts index 18abf0b0..ba582c20 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -15,11 +15,15 @@ dependencies { implementation("org.jetbrains.kotlin:kotlin-reflect:1.8.10") implementation("com.vladsch.flexmark:flexmark-all:0.64.8") implementation("io.github.kezhenxu94:cache-lite:0.2.0") + + // test libraries + testImplementation(kotlin("test")) } + group = "com.smallcloud" -version = getVersionString("1.2.25") +version = getVersionString("1.3.0") repositories { mavenCentral() @@ -27,7 +31,7 @@ repositories { // Configure Gradle IntelliJ Plugin -// Read more: https://plugins.jetbrains.com/docs/intellij/tools-gradle-intellij-plugin.html +// Read more: https://plugins.jetbrains.com/docs/intellij/tools-gradle-intellij-plugin.html#intellij-extension-type intellij { // version.set("LATEST-EAP-SNAPSHOT") version.set("2022.3.1") @@ -65,6 +69,10 @@ tasks { channels.set(listOf(System.getenv("PUBLISH_CHANNEL"))) token.set(System.getenv("PUBLISH_TOKEN")) } + + test { + useJUnitPlatform() + } } fun String.runCommand( diff --git a/refact_lsp b/refact_lsp index b19b5211..fe5b6a9d 100644 --- a/refact_lsp +++ b/refact_lsp @@ -1 +1 @@ -v0.8.0 +v0.8.2 \ No newline at end of file diff --git a/src/main/kotlin/com/smallcloud/refactai/io/AsyncConnection.kt b/src/main/kotlin/com/smallcloud/refactai/io/AsyncConnection.kt index 284179be..e2242799 100644 --- a/src/main/kotlin/com/smallcloud/refactai/io/AsyncConnection.kt +++ b/src/main/kotlin/com/smallcloud/refactai/io/AsyncConnection.kt @@ -152,6 +152,7 @@ class AsyncConnection : Disposable { failedDataReceiveEnded: (Throwable?) -> Unit = {}, requestId: String = "" ): CompletableFuture> { + var canceled = false; return CompletableFuture.supplyAsync { return@supplyAsync client.execute( requestProducer, @@ -173,6 +174,8 @@ class AsyncConnection : Disposable { override fun data(src: ByteBuffer?, endOfStream: Boolean) { src ?: return + if(canceled) return; + val part = Charset.forName("UTF-8").decode(src) bufferStr += part if (part.startsWith(STREAMING_PREFIX)) { @@ -223,7 +226,11 @@ class AsyncConnection : Disposable { failedDataReceiveEnded(ex) } - override fun cancelled() {} + override fun cancelled() { + // TODO: figure out how to stop the stream fro the lsp + canceled = true + dataReceiveEnded("Canceled") + } } ) } @@ -253,4 +260,5 @@ class AsyncConnection : Disposable { override fun dispose() { client.close() } + } diff --git a/src/main/kotlin/com/smallcloud/refactai/io/InferenceGlobalContext.kt b/src/main/kotlin/com/smallcloud/refactai/io/InferenceGlobalContext.kt index ef0666f4..025c16ae 100644 --- a/src/main/kotlin/com/smallcloud/refactai/io/InferenceGlobalContext.kt +++ b/src/main/kotlin/com/smallcloud/refactai/io/InferenceGlobalContext.kt @@ -97,7 +97,7 @@ class InferenceGlobalContext : Disposable { get() { return AppSettingsState.inferenceUri?.let { URI(it) } } - + var isNewChatStyle: Boolean = false diff --git a/src/main/kotlin/com/smallcloud/refactai/lsp/LSPProcessHolder.kt b/src/main/kotlin/com/smallcloud/refactai/lsp/LSPProcessHolder.kt index 3675b96e..5b342e82 100644 --- a/src/main/kotlin/com/smallcloud/refactai/lsp/LSPProcessHolder.kt +++ b/src/main/kotlin/com/smallcloud/refactai/lsp/LSPProcessHolder.kt @@ -1,6 +1,7 @@ package com.smallcloud.refactai.lsp import com.google.gson.Gson +import com.google.gson.JsonObject import com.intellij.execution.configurations.GeneralCommandLine import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationInfo @@ -16,18 +17,20 @@ import com.intellij.util.messages.MessageBus import com.intellij.util.messages.Topic import com.smallcloud.refactai.Resources import com.smallcloud.refactai.Resources.binPrefix -import com.smallcloud.refactai.account.AccountManager.Companion.instance import com.smallcloud.refactai.account.AccountManagerChangedNotifier import com.smallcloud.refactai.io.InferenceGlobalContextChangedNotifier import com.smallcloud.refactai.notifications.emitError +import com.smallcloud.refactai.panes.sharedchat.* import org.apache.hc.core5.concurrent.ComplexFuture import java.net.URI import java.nio.file.Files import java.nio.file.Paths import java.nio.file.StandardCopyOption +import java.util.concurrent.CompletableFuture import java.util.concurrent.Future import java.util.concurrent.TimeUnit import kotlin.io.path.Path +import com.smallcloud.refactai.account.AccountManager.Companion.instance as AccountManager import com.smallcloud.refactai.io.InferenceGlobalContext.Companion.instance as InferenceGlobalContext @@ -159,7 +162,7 @@ class LSPProcessHolder(val project: Project): Disposable { InferenceGlobalContext.inferenceUri val newConfig = LSPConfig( address = address, - apiKey = instance.apiKey, + apiKey = AccountManager.apiKey, port = 0, clientVersion = "${Resources.client}-${Resources.version}/${Resources.jbBuildVersion}", useTelemetry = true, @@ -246,6 +249,7 @@ class LSPProcessHolder(val project: Project): Disposable { private val BIN_PATH = Path(getTempDirectory(), ApplicationInfo.getInstance().build.toString().replace(Regex("[^A-Za-z0-9 ]"), "_") + "_refact_lsp${getExeSuffix()}").toString() + // here ? @JvmStatic fun getInstance(project: Project): LSPProcessHolder = project.service() @@ -288,7 +292,7 @@ class LSPProcessHolder(val project: Project): Disposable { try { requestFuture = it.get() as ComplexFuture val out = requestFuture.get() - logger.warn("LSP caps_received " + out) + // logger.warn("LSP caps_received " + out) val gson = Gson() res = gson.fromJson(out as String, LSPCapabilities::class.java) logger.debug("caps_received request finished") @@ -298,4 +302,116 @@ class LSPProcessHolder(val project: Project): Disposable { return res } } + + fun fetchCaps(): Future { +// // This causes the ide to crash :/ +// if(this.capabilities.codeChatModels.isNotEmpty()) { +// println("caps_cached") +// return FutureTask { this.capabilities } +// } + + val res = InferenceGlobalContext.connection.get( + url.resolve("/v1/caps"), + dataReceiveEnded = {}, + errorDataReceived = {} + ) + + return res.thenApply { + val body = it.get() as String + val caps = Gson().fromJson(body, LSPCapabilities::class.java) + caps + } + } + + fun fetchSystemPrompts(): Future { + val res = InferenceGlobalContext.connection.get( + url.resolve("/v1/customization"), + dataReceiveEnded = {}, + errorDataReceived = {} + ) + val json = res.thenApply { + val body = it.get() as String + val prompts = Gson().fromJson(body, CustomPromptsResponse::class.java) + prompts.systemPrompts + } + + return json + } + + fun fetchCommandCompletion(query: String, cursor: Int, count: Int = 5): Future { + + val requestBody = Gson().toJson(mapOf("query" to query, "cursor" to cursor, "top_n" to count)) + + val res = InferenceGlobalContext.connection.post( + url.resolve("/v1/at-command-completion"), + requestBody, + ) + // TODO: could this have a detail message? + val json = res.thenApply { + val body = it.get() as String + // handle error + // if(body.startsWith("detail")) + Gson().fromJson(body, CommandCompletionResponse::class.java) + } + + return json + } + + fun fetchCommandPreview(query: String): Future { + val requestBody = Gson().toJson(mapOf("query" to query)) + + val response = InferenceGlobalContext.connection.post( + url.resolve("/v1/at-command-preview"), + requestBody + ) + + val json: Future = response.thenApply { + val responseBody = it.get() as String + if (responseBody.startsWith("detail")) { + Events.AtCommands.Preview.Response(emptyArray()) + } else { + Events.gson.fromJson(responseBody, Events.AtCommands.Preview.Response::class.java) + } + } + + return json + } + + fun sendChat( + id: String, + messages: ChatMessages, + model: String, + dataReceived: (String, String) -> Unit, + dataReceiveEnded: (String) -> Unit, + errorDataReceived: (JsonObject) -> Unit, + failedDataReceiveEnded: (Throwable?) -> Unit, + ): CompletableFuture> { + + val parameters = mapOf("max_new_tokens" to 1000) + + val requestBody = Gson().toJson(mapOf( + "messages" to messages.map { + val content = if(it.content is String) { it.content } else { Gson().toJson(it.content) } + mapOf("role" to it.role, "content" to content) + }, + "model" to model, + "parameters" to parameters, + "stream" to true + )) + + val headers = mapOf("Authorization" to "Bearer ${AccountManager.apiKey}") + val request = InferenceGlobalContext.connection.post( + url.resolve("/v1/chat"), + requestBody, + headers = headers, + dataReceived = dataReceived, + dataReceiveEnded = dataReceiveEnded, + errorDataReceived = errorDataReceived, + failedDataReceiveEnded = failedDataReceiveEnded, + requestId = id, + ) + + return request + + } } \ No newline at end of file diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/RefactAIToolboxPaneFactory.kt b/src/main/kotlin/com/smallcloud/refactai/panes/RefactAIToolboxPaneFactory.kt index 8f2120f4..b015775e 100644 --- a/src/main/kotlin/com/smallcloud/refactai/panes/RefactAIToolboxPaneFactory.kt +++ b/src/main/kotlin/com/smallcloud/refactai/panes/RefactAIToolboxPaneFactory.kt @@ -7,30 +7,74 @@ import com.intellij.openapi.wm.ToolWindowFactory import com.intellij.openapi.wm.ToolWindowManager import com.intellij.ui.content.Content import com.intellij.ui.content.ContentFactory +import com.intellij.ui.jcef.JBCefApp import com.smallcloud.refactai.Resources import com.smallcloud.refactai.panes.gptchat.ChatGPTPanes +import com.smallcloud.refactai.panes.sharedchat.ChatPanes import com.smallcloud.refactai.utils.getLastUsedProject + class RefactAIToolboxPaneFactory : ToolWindowFactory { override fun init(toolWindow: ToolWindow) { toolWindow.setIcon(Resources.Icons.LOGO_RED_13x13) super.init(toolWindow) } + override fun isApplicable(project: Project): Boolean { + return try { + JBCefApp.isSupported() && JBCefApp.isStarted() + JBCefApp.isSupported() + } catch (_: Exception) { + false + } + } + + override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) { val contentFactory = ContentFactory.getInstance() - val gptChatPanes = ChatGPTPanes(project, toolWindow.disposable) - val content: Content = contentFactory.createContent( - gptChatPanes.getComponent(), - "Chat", - false - ) +// val gptChatPanes = ChatGPTPanes(project, toolWindow.disposable) +// val content: Content = contentFactory.createContent( +// gptChatPanes.getComponent(), +// "Chat", +// false +// ) +// content.isCloseable = false +// content.putUserData(panesKey, gptChatPanes) +// toolWindow.contentManager.addContent(content) + + val chatPanes = ChatPanes(project, toolWindow.disposable) + val content: Content = contentFactory.createContent(chatPanes.getComponent(), "Chat", false) content.isCloseable = false - content.putUserData(panesKey, gptChatPanes) toolWindow.contentManager.addContent(content) + +// val sp = SharedChatPane(project) +// +// val chatIframeContent: Content = contentFactory.createContent( +// sp.getComponent(), +// "Shared Chat", +// false +// ) +// chatIframeContent.isCloseable = false +// +// toolWindow.contentManager.addContent(chatIframeContent) + + + // Uncomment to enable dev tools +// val devToolsBrowser = JBCefBrowser.createBuilder() +// .setCefBrowser(sp.webView.cefBrowser.devTools) +// .setClient(sp.webView.jbCefClient) +// .build(); +// +// val c = contentFactory.createContent(devToolsBrowser.component, "Shared Chat Dev", false) +// toolWindow.contentManager.addContent(c) +// devToolsBrowser.openDevtools() + + } + + companion object { private val panesKey = Key.create("refact.panes") val chat: ChatGPTPanes? diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatHistory.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatHistory.kt new file mode 100644 index 00000000..7bb4770e --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatHistory.kt @@ -0,0 +1,117 @@ +package com.smallcloud.refactai.panes.sharedchat + +import com.google.gson.GsonBuilder +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.* +import com.intellij.util.xmlb.XmlSerializerUtil + +import com.intellij.util.xmlb.annotations.MapAnnotation +import com.intellij.util.xmlb.annotations.Property + + +// @Service + + +data class ChatHistoryItem( + @Property + val id: String, + // @Property // maybe @Collection or maybe a custom class? + @Property + val messages: ChatMessages, + @Property + val model: String, + @Property + val title: String? = null, + @Property + val updatedAt: Long = System.currentTimeMillis(), + @Property + val createdAt: Long = System.currentTimeMillis(), +) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as ChatHistoryItem + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + if (model != other.model) return false + if (title != other.title) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + result = 31 * result + model.hashCode() + result = 31 * result + (title?.hashCode() ?: 0) + result = 31 * result + updatedAt.hashCode() + result = 31 * result + createdAt.hashCode() + return result + } +} + +@State(name = "com.smallcloud.refactai.panes.sharedchat.RefactChatHistory", storages = [ + Storage("refactaiChatHistory.xml"), +]) +class ChatHistory: PersistentStateComponent { + @MapAnnotation + var chatHistory: MutableMap = emptyMap().toMutableMap() + + var gson = GsonBuilder() + .registerTypeAdapter(ChatMessage::class.java, ChatHistorySerializer()) + .registerTypeAdapter(ChatMessage::class.java, ChatMessageDeserializer()) + .create() + + fun setItem(item: ChatHistoryItem) { + val json = gson.toJson(item) + chatHistory[item.id] = json + } + + fun save(id: String, messages: ChatMessages, model: String) { + val maybeItem = this.getItem(id); + if (maybeItem == null) { + val title = messages.first { it.role == ChatRole.USER }.content.toString().let { + val end = it.length.coerceAtMost(16) + it.substring(0, end) + } + + val newItem = ChatHistoryItem(id, messages, model, title) + setItem(newItem) + } else { + val item = maybeItem.copy( + messages = messages, + updatedAt = System.currentTimeMillis()) + this.setItem(item) + } + + } + fun removeItem(id: String?) { + chatHistory.remove(id) + } + + fun getItem(id: String): ChatHistoryItem? { + val json = chatHistory[id] + return gson.fromJson(json, ChatHistoryItem::class.java) + } + + fun getAll(): List { + return chatHistory.map { gson.fromJson(it.value, ChatHistoryItem::class.java) } + } + + override fun getState(): ChatHistory { + return this + } + + override fun loadState(state: ChatHistory) { + XmlSerializerUtil.copyBean(state, this) + } + + + companion object { + @JvmStatic + val instance: ChatHistory + get() = ApplicationManager.getApplication().getService(ChatHistory::class.java) + } +} diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatPanes.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatPanes.kt new file mode 100644 index 00000000..98488bc7 --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatPanes.kt @@ -0,0 +1,157 @@ +package com.smallcloud.refactai.panes.sharedchat + +import com.google.gson.GsonBuilder +import com.intellij.execution.ui.layout.impl.JBRunnerTabs +import com.intellij.icons.AllIcons +import com.intellij.openapi.Disposable +import com.intellij.openapi.actionSystem.* +import com.intellij.openapi.application.invokeLater +import java.awt.BorderLayout +import javax.swing.JPanel +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.Disposer +import com.intellij.ui.components.JBLabel +import com.intellij.ui.jcef.JBCefBrowser +import com.intellij.ui.tabs.TabInfo +import com.intellij.util.containers.ContainerUtil +import com.smallcloud.refactai.RefactAIBundle +import javax.swing.JComponent + +class ChatPanes(val project: Project, private val parent: Disposable) { + private val paneBaseName = "Chat" + private val panes = JBRunnerTabs(project, parent) + private val holder = JPanel().also { + it.layout = BorderLayout() + } + + private fun setupPanes() { + invokeLater { + holder.removeAll() + holder.add(panes) + restoreOrAddNew() + } + } + + private fun restoreOrAddNew() { + ChatHistory.instance.state.let { + if (it.chatHistory.isNotEmpty()) { + it.getAll().forEach { item -> + restoreTab(item) + } + } else { + addTab() + } + } + } + + private fun restoreTab(item: ChatHistoryItem) { + val newPane = SharedChatPane(project) + val component: JComponent = newPane.webView.component + val info = TabInfo(component) + info.text = item.title ?: "Chat" + + newPane.restoreWhenReady(item.id, item.messages, item.model) + Disposer.register(parent, newPane) + // Add button + info.setActions(DefaultActionGroup(object : AnAction(AllIcons.General.Add) { + override fun actionPerformed(e: AnActionEvent) { + addTab() + } + + override fun getActionUpdateThread(): ActionUpdateThread { + return ActionUpdateThread.EDT + } + }), ActionPlaces.EDITOR_TAB) + + // Delete button + info.setTabLabelActions(DefaultActionGroup(object : AnAction(AllIcons.Actions.Close) { + override fun actionPerformed(e: AnActionEvent) { + // TODO: cancel requests + newPane.handleChatStop(item.id) + panes.removeTab(info) + ChatHistory.instance.state.removeItem(item.id) + if (getVisibleTabs().isEmpty()) { + addTab() + } + } + + override fun getActionUpdateThread(): ActionUpdateThread { + return ActionUpdateThread.EDT + } + }), ActionPlaces.EDITOR_TAB) + + panes.addTab(info) + +// val devToolsBrowser = JBCefBrowser.createBuilder() +// .setCefBrowser(newPane.webView.cefBrowser.devTools) +// .setClient(newPane.webView.jbCefClient) +// .build(); +// +// val devInfo = TabInfo(devToolsBrowser.component) +// devInfo.text = "DevTools" +// panes.addTab(devInfo) +// devToolsBrowser.openDevtools() + } + + init { + setupPanes() + } + + fun addTab(title: String? = null) { + val newPane = SharedChatPane(project) + val component: JComponent = newPane.webView.component + val info = TabInfo(component) + info.text = title ?: "Chat" + Disposer.register(parent, newPane) + + // Add button + info.setActions(DefaultActionGroup(object : AnAction(AllIcons.General.Add) { + override fun actionPerformed(e: AnActionEvent) { + addTab() + } + + override fun getActionUpdateThread(): ActionUpdateThread { + return ActionUpdateThread.EDT + } + }), ActionPlaces.EDITOR_TAB) + + // Delete button + info.setTabLabelActions(DefaultActionGroup(object : AnAction(AllIcons.Actions.Close) { + override fun actionPerformed(e: AnActionEvent) { + // TODO: cancel requests + // (info.component as SharedChatPane).cancelRequest() + panes.removeTab(info) + ChatHistory.instance.state.removeItem(newPane.id) + Disposer.dispose(newPane) + if (getVisibleTabs().isEmpty()) { + addTab() + } + } + + override fun getActionUpdateThread(): ActionUpdateThread { + return ActionUpdateThread.EDT + } + }), ActionPlaces.EDITOR_TAB) + + panes.addTab(info) + panes.select(info, true) + +// val devToolsBrowser = JBCefBrowser.createBuilder() +// .setCefBrowser(newPane.webView.cefBrowser.devTools) +// .setClient(newPane.webView.jbCefClient) +// .build(); +// +// val devInfo = TabInfo(devToolsBrowser.component) +// devInfo.text = "DevTools" +// panes.addTab(devInfo) +// devToolsBrowser.openDevtools() + } + + fun getComponent(): JComponent { + return holder + } + + fun getVisibleTabs(): List { + return ContainerUtil.filter(panes.tabs) { tabInfo -> !tabInfo.isHidden } + } +} \ No newline at end of file diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/Events.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/Events.kt new file mode 100644 index 00000000..65c726cd --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/Events.kt @@ -0,0 +1,983 @@ +package com.smallcloud.refactai.panes.sharedchat + +import com.google.gson.annotations.SerializedName +import com.google.gson.* +import com.smallcloud.refactai.lsp.LSPCapabilities +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.Response.ResponsePayload +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.ResponseDeserializer +import java.io.Serializable +import java.lang.reflect.Type + +// Lsp responses + +enum class ChatRole(val value: String): Serializable { + @SerializedName("user") USER("user"), + @SerializedName("assistant") ASSISTANT("assistant"), + @SerializedName("context_file") CONTEXT_FILE("context_file"), + @SerializedName("system") SYSTEM("system"), +} + +data class ChatContextFile( + @SerializedName("file_name") val fileName: String, + @SerializedName("file_content") val fileContent: String, + val line1: Int, + val line2: Int, + val usefulness: Double? = null +) + + + +abstract class ChatMessage( + @SerializedName("role") + val role: ChatRole, + @Transient + @SerializedName("content") + open val content: T +): Serializable { +} +data class UserMessage(override val content: String): ChatMessage(ChatRole.USER, content) + +data class AssistantMessage(override val content: String): ChatMessage(ChatRole.ASSISTANT, content) + +data class SystemMessage(override val content: String): ChatMessage(ChatRole.SYSTEM, content) + +data class ContentFileMessage(override val content: Array): ChatMessage>( + ChatRole.CONTEXT_FILE, content) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as ContentFileMessage + + return content.contentEquals(other.content) + } + + override fun hashCode(): Int { + return content.contentHashCode() + } +} + + +typealias ChatMessages = Array> + +class ChatMessageDeserializer: JsonDeserializer> { + override fun deserialize(p0: JsonElement?, p1: Type?, p2: JsonDeserializationContext?): ChatMessage<*>? { + val role = p0?.asJsonObject?.get("role")?.asString + return when(role) { + ChatRole.USER.value -> p2?.deserialize(p0, UserMessage::class.java) + ChatRole.ASSISTANT.value -> p2?.deserialize(p0, AssistantMessage::class.java) + ChatRole.CONTEXT_FILE.value -> p2?.deserialize(p0, ContentFileMessage::class.java) + ChatRole.SYSTEM.value -> p2?.deserialize(p0, SystemMessage::class.java) + else -> null + } + } +} + +class ChatHistorySerializer: JsonSerializer> { + override fun serialize(p0: ChatMessage<*>, p1: Type?, p2: JsonSerializationContext?): JsonElement? { + return when (p0) { + is UserMessage -> p2?.serialize(p0, UserMessage::class.java) + is AssistantMessage -> p2?.serialize(p0, AssistantMessage::class.java) + is ContentFileMessage -> p2?.serialize(p0, ContentFileMessage::class.java) + is SystemMessage -> p2?.serialize(p0, SystemMessage::class.java) + else -> JsonNull.INSTANCE + } + } +} + + +abstract class Delta(val role: ChatRole, val content: T) +class AssistantDelta(content: String): Delta(ChatRole.ASSISTANT, content) +class ContextFileDelta(content: ChatMessages): Delta(ChatRole.CONTEXT_FILE, content) + +class DeltaDeserializer: JsonDeserializer> { + override fun deserialize(p0: JsonElement?, p1: Type?, p2: JsonDeserializationContext?): Delta<*>? { + val role = p0?.asJsonObject?.get("role")?.asString + return when(role) { + ChatRole.ASSISTANT.value -> p2?.deserialize(p0, AssistantDelta::class.java) + ChatRole.CONTEXT_FILE.value -> p2?.deserialize(p0, ContextFileDelta::class.java) + else -> null + } + } +} + + + +data class ScratchPad( + @SerializedName("default_system_message") val defaultSystemMessage: String, +) + +data class ChatModel( + @SerializedName("default_scratchpad") val defaultScratchPad: String, + @SerializedName("n_ctx") val nCtx: Int, + @SerializedName("similar_models") val similarModels: List, + @SerializedName("supports_scratchpads") val supportsScratchpads: Map +) + +data class CodeCompletionModel( + @SerializedName("default_scratchpad") val defaultScratchPad: String, + @SerializedName("n_ctx") val nCtx: Int, + @SerializedName("similar_models") val similarModels: List, + // Might be wrong + @SerializedName("supports_scratchpads") val supportsScratchpads: Map +) + +data class CapsResponse( + @SerializedName("caps_version") val capsVersion: Int, + @SerializedName("cloud_name") val cloudName: String, + @SerializedName("code_chat_default_model") val codeChatDefaultModel: String, + @SerializedName("code_chat_models") val codeChatModels: Map, + @SerializedName("code_completion_default_model") val codeCompletionDefaultModel: String, + @SerializedName("code_completion_models") val codeCompletionModels: Map, + @SerializedName("code_completion_n_ctx") val codeCompletionNCtx: Int, + @SerializedName("endpoint_chat_passthrough") val endpointChatPassthrough: String, + @SerializedName("endpoint_style") val endpointStyle: String, + @SerializedName("endpoint_template") val endpointTemplate: String, + @SerializedName("running_models") val runningModels: List, + @SerializedName("telemetry_basic_dest") val telemetryBasicDest: String, + @SerializedName("telemetry_corrected_snippets_dest") val telemetryCorrectedSnippetsDest: String, + @SerializedName("tokenizer_path_template") val tokenizerPathTemplate: String, + @SerializedName("tokenizer_rewrite_path") val tokenizerRewritePath: Map> +) + +data class CommandCompletionResponse( + @SerializedName("completions") val completions: Array, + val replace: Array, + @SerializedName("is_cmd_executable") val isCmdExecutable: Boolean +) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as CommandCompletionResponse + + if (!completions.contentEquals(other.completions)) return false + if (!replace.contentEquals(other.replace)) return false + if (isCmdExecutable != other.isCmdExecutable) return false + + return true + } + + override fun hashCode(): Int { + var result = completions.contentHashCode() + result = 31 * result + replace.contentHashCode() + result = 31 * result + isCmdExecutable.hashCode() + return result + } +} + +data class DetailMessage( + val detail: String +) + + + +data class SystemPrompt( + val text: String, + val description: String, +) + +typealias SystemPromptMap = Map + +data class CustomPromptsResponse( + @SerializedName("system_prompts") val systemPrompts: SystemPromptMap, + // Might need to update this + @SerializedName("toolbox_commands") val toolboxCommands: Map +) + +// Events + +class EventNames { + enum class FromChat(val value: String) { + SAVE_CHAT("save_chat_to_history"), + ASK_QUESTION("chat_question"), + REQUEST_CAPS("chat_request_caps"), + STOP_STREAMING("chat_stop_streaming"), + BACK_FROM_CHAT("chat_back_from_chat"), + OPEN_IN_CHAT_IN_TAB("open_chat_in_new_tab"), + SEND_TO_SIDE_BAR("chat_send_to_sidebar"), + READY("chat_ready"), + NEW_FILE("chat_create_new_file"), + PASTE_DIFF("chat_paste_diff"), + REQUEST_AT_COMMAND_COMPLETION("chat_request_at_command_completion"), + REQUEST_PREVIEW_FILES("chat_request_preview_files"), + REQUEST_PROMPTS("chat_request_prompts") + } + + enum class ToChat(val value: String) { + CLEAR_ERROR("chat_clear_error"), + @SerializedName("restore_chat_from_history") RESTORE_CHAT("restore_chat_from_history"), + @SerializedName("chat_response") CHAT_RESPONSE("chat_response"), + BACKUP_MESSAGES("back_up_messages"), + @SerializedName("chat_done_streaming") DONE_STREAMING("chat_done_streaming"), + @SerializedName("chat_error_streaming") ERROR_STREAMING("chat_error_streaming"), + NEW_CHAT("create_new_chat"), + @SerializedName("receive_caps") RECEIVE_CAPS("receive_caps"), + @SerializedName("receive_caps_error") RECEIVE_CAPS_ERROR("receive_caps_error"), + SET_CHAT_MODEL("chat_set_chat_model"), + SET_DISABLE_CHAT("set_disable_chat"), + @SerializedName("chat_active_file_info") ACTIVE_FILE_INFO("chat_active_file_info"), + TOGGLE_ACTIVE_FILE("chat_toggle_active_file"), + @SerializedName("chat_receive_at_command_completion") RECEIVE_AT_COMMAND_COMPLETION("chat_receive_at_command_completion"), + @SerializedName("chat_receive_at_command_preview") RECEIVE_AT_COMMAND_PREVIEW("chat_receive_at_command_preview"), + SET_SELECTED_AT_COMMAND("chat_set_selected_command"), + SET_LAST_MODEL_USED("chat_set_last_model_used"), + + @SerializedName("chat_set_selected_snippet") SET_SELECTED_SNIPPET("chat_set_selected_snippet"), + REMOVE_PREVIEW_FILE_BY_NAME("chat_remove_file_from_preview"), + SET_PREVIOUS_MESSAGES_LENGTH("chat_set_previous_messages_length"), + RECEIVE_TOKEN_COUNT("chat_set_tokens"), + @SerializedName("chat_receive_prompts") RECEIVE_PROMPTS("chat_receive_prompts"), + @SerializedName("chat_receive_prompts_error") RECEIVE_PROMPTS_ERROR("chat_receive_prompts_error"), + SET_SELECTED_SYSTEM_PROMPT("chat_set_selected_system_prompt"), + @SerializedName("receive_config_update") RECEIVE_CONFIG_UPDATE("chat_receive_config_update"), + } +} + + + +class Events { + + open class Payload( + @Transient + @SerializedName("id") + open val id: String + ): Serializable + + abstract class FromChat(val type: EventNames.FromChat, open val payload: Payload): Serializable + + private class FromChatDeserializer : JsonDeserializer { + override fun deserialize(p0: JsonElement?, p1: Type?, p2: JsonDeserializationContext?): FromChat? { + val type = p0?.asJsonObject?.get("type")?.asString + val payload = p0?.asJsonObject?.get("payload") + if(type == null || payload == null) return null + + return when(type) { + EventNames.FromChat.READY.value -> p2?.deserialize(payload, Ready::class.java) + EventNames.FromChat.REQUEST_PROMPTS.value -> p2?.deserialize(payload, SystemPrompts.Request::class.java) + EventNames.FromChat.REQUEST_AT_COMMAND_COMPLETION.value -> p2?.deserialize(payload, AtCommands.Completion.Request::class.java) + EventNames.FromChat.REQUEST_PREVIEW_FILES.value -> p2?.deserialize(payload, AtCommands.Preview.Request::class.java) + EventNames.FromChat.SAVE_CHAT.value -> { + val messages = JsonArray() + payload.asJsonObject.get("messages").asJsonArray.forEach { + val pair = it.asJsonArray + val role = pair.get(0) + val content = pair.get(1) + val obj = JsonObject() + obj.add("role", role) + obj.add("content", content) + messages.add(obj) + } + + payload.asJsonObject.add("messages", messages) + + return p2?.deserialize(payload, Chat.Save::class.java) + } + EventNames.FromChat.ASK_QUESTION.value -> { + val messages = JsonArray() + payload.asJsonObject.get("messages").asJsonArray.forEach { + val pair = it.asJsonArray + val role = pair.get(0) + val content = pair.get(1) + val obj = JsonObject() + obj.add("role", role) + obj.add("content", content) + messages.add(obj) + } + + payload.asJsonObject.add("messages", messages) + return p2?.deserialize(payload, Chat.AskQuestion::class.java) + } + EventNames.FromChat.STOP_STREAMING.value -> p2?.deserialize(payload, Chat.Stop::class.java) + EventNames.FromChat.NEW_FILE.value -> p2?.deserialize(payload, Editor.NewFile::class.java) + EventNames.FromChat.PASTE_DIFF.value -> p2?.deserialize(payload, Editor.Paste::class.java) + EventNames.FromChat.REQUEST_CAPS.value -> p2?.deserialize(payload, Caps.Request::class.java) + else -> null + } + } + + } + + abstract class ToChat( + @SerializedName("type") + val type: EventNames.ToChat, + @SerializedName("payload") + open val payload: Payload + ): Serializable + + + data class Ready(val id: String): FromChat(EventNames.FromChat.READY, Payload(id)) + + class SystemPrompts() { + data class Request(val id: String): FromChat(EventNames.FromChat.REQUEST_PROMPTS, Payload(id)) + + data class SystemPromptsPayload(override val id: String, val prompts: SystemPromptMap): Payload(id) + class Receive(payload: SystemPromptsPayload): ToChat(EventNames.ToChat.RECEIVE_PROMPTS, payload) + + data class SystemPromptsErrorPayload(override val id: String, val error: String): Payload(id) + data class Error(val id: String, val error: String): ToChat(EventNames.ToChat.RECEIVE_PROMPTS_ERROR, SystemPromptsErrorPayload(id, error)) + // set? + } + + class AtCommands { + + class Completion { + + data class RequestPayload( + override val id: String, + val query: String, + val cursor: Int, + val number: Int = 5, + val trigger: String? = null, + ) : Payload(id) + + data class Request( + val id: String, + val query: String, + val cursor: Int, + val number: Int = 5, + val trigger: String? = null, + ) : FromChat(EventNames.FromChat.REQUEST_AT_COMMAND_COMPLETION, RequestPayload(id, query, cursor, number, trigger)) + + data class CompletionPayload( + override val id: String, + val completions: Array, + val replace: Array, + @SerializedName("is_cmd_executable") + val isCmdExecutable: Boolean = false + ) : Payload(id) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as CompletionPayload + + if (id != other.id) return false + if (!completions.contentEquals(other.completions)) return false + if (!replace.contentEquals(other.replace)) return false + if (isCmdExecutable != other.isCmdExecutable) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + completions.contentHashCode() + result = 31 * result + replace.contentHashCode() + result = 31 * result + isCmdExecutable.hashCode() + return result + } + } + + class Receive( payload: CompletionPayload ) : ToChat(EventNames.ToChat.RECEIVE_AT_COMMAND_COMPLETION, payload) + } + + class Preview { + data class RequestPayload(override val id: String, val query: String): Payload(id) + data class Request( val id: String, val query: String): FromChat(EventNames.FromChat.REQUEST_PREVIEW_FILES, RequestPayload(id, query)) + + data class Response( + val messages: Array, + ) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as Response + + return messages.contentEquals(other.messages) + } + + override fun hashCode(): Int { + return messages.contentHashCode() + } + } + + class ResponseDeserializer: JsonDeserializer { + override fun deserialize( + p0: JsonElement?, + p1: Type?, + p2: JsonDeserializationContext? + ): Events.AtCommands.Preview.Response? { + val messages = p0?.asJsonObject?.get("messages")?.asJsonArray + val arr = JsonArray() + messages?.forEach { + val contentAsString = it.asJsonObject.get("content").asString + val a = Gson().fromJson(contentAsString, JsonArray::class.java) + it.asJsonObject.add("content", a) + arr.add(it) + } + + p0?.asJsonObject?.add("messages", arr) + return Gson().fromJson(p0, Events.AtCommands.Preview.Response::class.java) + } + } + + + data class PreviewPayload( + override val id: String, + val preview: Array, + ): Payload(id) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as PreviewPayload + + if (id != other.id) return false + if (!preview.contentEquals(other.preview)) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + preview.contentHashCode() + return result + } + } + + class Receive( + payload: PreviewPayload + ) : ToChat(EventNames.ToChat.RECEIVE_AT_COMMAND_PREVIEW, payload) + } + } + + class Chat { + data class ThreadPayload( + override val id: String, + val messages: ChatMessages, + val model: String, + val title: String? = null, + @SerializedName("attach_file") val attachFile: Boolean = false, + ): Payload(id) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as ThreadPayload + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + if (model != other.model) return false + if (title != other.title) return false + if (attachFile != other.attachFile) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + result = 31 * result + model.hashCode() + result = 31 * result + (title?.hashCode() ?: 0) + result = 31 * result + attachFile.hashCode() + return result + } + } + + data class Save( + val id: String, + val messages: ChatMessages, + val model: String, + val title: String, + val attachFile: Boolean = false, + ): FromChat(EventNames.FromChat.SAVE_CHAT, ThreadPayload(id, messages, model, title, attachFile)) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as Save + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + if (model != other.model) return false + if (title != other.title) return false + if (attachFile != other.attachFile) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + result = 31 * result + model.hashCode() + result = 31 * result + title.hashCode() + result = 31 * result + attachFile.hashCode() + return result + } + } + + data class AskQuestion( + val id: String, + val messages: ChatMessages, + val model: String, + val title: String? = null, + val attachFile: Boolean = false, + ): FromChat(EventNames.FromChat.ASK_QUESTION, ThreadPayload(id, messages, model, title, attachFile)) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as AskQuestion + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + if (model != other.model) return false + if (title != other.title) return false + if (attachFile != other.attachFile) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + result = 31 * result + model.hashCode() + result = 31 * result + (title?.hashCode() ?: 0) + result = 31 * result + attachFile.hashCode() + return result + } + } + + data class Stop(val id: String): FromChat(EventNames.FromChat.STOP_STREAMING, Payload(id)) + + // receive + class Response { + enum class Roles(value: String) { + @SerializedName("user") USER("user"), + @SerializedName("context_file") CONTEXT_FILE("context_file") + } + + abstract class ResponsePayload() + + data class UserMessage( + @SerializedName("role") + val role: Roles, + @SerializedName("content") + val content: String + ): ResponsePayload() + + class UserMessagePayload( + override val id: String, + val role: Roles, + val content: String, + ): Payload(id) + + class UserMessageToChat(payload: UserMessagePayload): ToChat(EventNames.ToChat.CHAT_RESPONSE, payload) + + enum class FinishReasons(value: String) { + STOP("stop"), + ABORT("abort"), + } + + data class Choice( + val delta: AssistantDelta, + val index: Int, + @SerializedName("finish_reason") val finishReason: FinishReasons?, + ) + + data class Choices( + val choices: Array, + val created: String, + val model: String, + ): ResponsePayload() { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as Choices + + if (!choices.contentEquals(other.choices)) return false + if (created != other.created) return false + if (model != other.model) return false + + return true + } + + override fun hashCode(): Int { + var result = choices.contentHashCode() + result = 31 * result + created.hashCode() + result = 31 * result + model.hashCode() + return result + } + } + + class ChoicesPayload( + override val id: String, + val choices: Array, + val created: String, + val model: String, + ): Payload(id) + + class ChoicesToChat(payload: ChoicesPayload): ToChat(EventNames.ToChat.CHAT_RESPONSE, payload) + + data class ChatDone(val message: String? = null): ResponsePayload() + + class ChatDonePayload( + override val id: String, + val message: String?, + ): Payload(id) + + class ChatDoneToChat(payload: ChatDonePayload): ToChat(EventNames.ToChat.DONE_STREAMING, payload) + + data class ChatError(val message: JsonElement): ResponsePayload() + + class ChatErrorPayload( + override val id: String, + val message: String, + ): Payload(id) + + class ChatErrorStreamingToChat(payload: ChatErrorPayload): ToChat(EventNames.ToChat.ERROR_STREAMING, payload) + + data class ChatFailedStream(val message: Throwable?): ResponsePayload() + + // detail + data class DetailMessage(val detail: String): ResponsePayload() + + companion object { + private val gson = GsonBuilder() + .registerTypeAdapter(ResponsePayload::class.java, ResponseDeserializer()) + .registerTypeAdapter(Delta::class.java, DeltaDeserializer()) + .create() + + fun parse(str: String): ResponsePayload { + return gson.fromJson(str, ResponsePayload::class.java) + } + + + fun formatToChat(response: ResponsePayload, id: String): ToChat? { + return when (response) { + is Response.UserMessage -> { + val payload = UserMessagePayload(id, response.role, response.content) + return UserMessageToChat(payload) + } + + is Response.Choices -> { + val payload = ChoicesPayload(id, response.choices, response.created, response.model) + return ChoicesToChat(payload) + } + + is Response.ChatDone -> { + val payload = ChatDonePayload(id, response.message) + return ChatDoneToChat(payload) + } + + is Response.DetailMessage -> { + val payload = ChatErrorPayload(id, response.detail) + return ChatErrorStreamingToChat(payload) + } + + is Response.ChatError -> { + val maybeDetail = response.message.asJsonObject.get("detail").asString + val message = maybeDetail ?: response.message.toString() + val payload = ChatErrorPayload(id, message) + return ChatErrorStreamingToChat(payload) + } + + is Response.ChatFailedStream -> { + val message = "Failed during stream: ${response.message?.message}" + val payload = ChatErrorPayload(id, message) + return ChatErrorStreamingToChat(payload) + } + + else -> null + } + } + } + } + + class ResponseDeserializer : JsonDeserializer { + override fun deserialize(p0: JsonElement?, p1: Type?, p2: JsonDeserializationContext?): Response.ResponsePayload? { + + val role = p0?.asJsonObject?.get("role")?.asString + + if (role == "user" || role == "context_file") { + return p2?.deserialize(p0, Response.UserMessage::class.java) + } + + val choices = p0?.asJsonObject?.get("choices")?.asJsonArray + + if (choices !== null) { + return p2?.deserialize(p0, Response.Choices::class.java) + } + + val detail = p0?.asJsonObject?.has("detail") + if(detail == true) { + return p2?.deserialize(p0, Response.DetailMessage::class.java) + } + + return p2?.deserialize(p0, Response.ResponsePayload::class.java) + + } + + } + + + data class BackupPayload( + override val id: String, + val messages: ChatMessages + ): Payload(id) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as BackupPayload + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + return result + } + + } + + data class Backup( + val id: String, + val messages: ChatMessages, + ): ToChat(EventNames.ToChat.BACKUP_MESSAGES, BackupPayload(id, messages)) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as Backup + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + return result + } + } + + + // last model used + data class LastModelUsedPayload( + override val id: String, + val model: String, + ): Payload(id) + + data class LastModelUsed( + val id: String, + val model: String, + ): ToChat(EventNames.ToChat.SET_LAST_MODEL_USED, LastModelUsedPayload(id, model)) + + // restore + + data class Thread( + val id: String, + val messages: ChatMessages, + val model: String, + val title: String? = null, + @SerializedName("attach_file") val attachFile: Boolean = false, + ) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as Thread + + if (id != other.id) return false + if (!messages.contentEquals(other.messages)) return false + if (model != other.model) return false + if (title != other.title) return false + if (attachFile != other.attachFile) return false + + return true + } + + override fun hashCode(): Int { + var result = id.hashCode() + result = 31 * result + messages.contentHashCode() + result = 31 * result + model.hashCode() + result = 31 * result + (title?.hashCode() ?: 0) + result = 31 * result + attachFile.hashCode() + return result + } + } + + data class RestorePayload( + override val id: String, + val chat: Thread, + val snippet: Editor.Snippet? = null, + ): Payload(id) + + class RestoreToChat( + payload: RestorePayload + ): ToChat(EventNames.ToChat.RESTORE_CHAT, payload) + + // new ? + data class NewChatPayload( + override val id: String, + val snippet: Editor.Snippet? + ): Payload(id) + + data class NewChat( + val id: String, + val snippet: Editor.Snippet?, + ): ToChat(EventNames.ToChat.NEW_CHAT, NewChatPayload(id, snippet)) + + } + + class ActiveFile { + data class FileInfo( + val name: String = "", + val path: String = "", + @SerializedName("can_paste") val canPaste: Boolean = false, + val attach: Boolean = false, + val line1: Int? = null, + val line2: Int? = null, + val cursor: Int? = null, + val content: String? = null, + val usefulness: Int? = null, + ) + + data class FileInfoPayload( + override val id: String, + val file: FileInfo, + ) : Payload(id) + + class ActiveFileToChat(payload: FileInfoPayload): ToChat(EventNames.ToChat.ACTIVE_FILE_INFO, payload) + + } + + class Editor { + data class ContentPayload( + override val id: String, + val content: String + ): Payload(id) + + data class NewFile( + val id: String, + val content: String, + ): FromChat(EventNames.FromChat.NEW_FILE, ContentPayload(id, content)) + + data class Paste( + val id: String, + val content: String + ): FromChat(EventNames.FromChat.PASTE_DIFF, ContentPayload(id, content)) + + data class Snippet( + val language: String = "", + val code: String = "", + val path: String = "", + val basename: String = "", + ) + + data class SetSnippetPayload( + override val id: String, + val snippet: Snippet + ): Payload(id) + + class SetSnippetToChat(payload: SetSnippetPayload): ToChat(EventNames.ToChat.SET_SELECTED_SNIPPET, payload) + + } + + class Caps { + data class Request( + val id: String + ): FromChat(EventNames.FromChat.REQUEST_CAPS, Payload(id)) { + constructor() : this("") + } + + data class CapsPayload( + override val id: String, + val caps: LSPCapabilities + ): Payload(id) + + class Receive( + id: String, + caps: LSPCapabilities + ): ToChat(EventNames.ToChat.RECEIVE_CAPS, CapsPayload(id, caps)) + + data class ErrorPayload( + override val id: String, + val message: String + ): Payload(id) + + data class Error( + val id: String, + val error: String + ): ToChat(EventNames.ToChat.RECEIVE_CAPS_ERROR, ErrorPayload(id, error)) + } + + class Config { + abstract class BaseFeatures() + + data class Features(val ast: Boolean, val vecdb: Boolean): BaseFeatures() + + data class ThemeProps(val mode: String, val hasBackground: Boolean = false, val scale: String = "90%", val accentColor: String ="gray") + + data class UpdatePayload(override val id: String, val features: BaseFeatures, val themeProps: ThemeProps?): Payload(id) + + class Update(id: String, features: BaseFeatures, themeProps: ThemeProps?): ToChat(EventNames.ToChat.RECEIVE_CONFIG_UPDATE, UpdatePayload(id, features, themeProps)) + + } + companion object { + + private class MessageSerializer: JsonSerializer> { + override fun serialize(src: ChatMessage<*>, typeOfSrc: Type, context: JsonSerializationContext): JsonElement { + return when(src) { + is UserMessage -> { + val role = context.serialize(src.role, ChatRole::class.java) + val arr = JsonArray() + arr.add(role) + arr.add(src.content) + return arr + } + is SystemMessage -> { + val role = context.serialize(src.role, ChatRole::class.java) + val arr = JsonArray() + arr.add(role) + arr.add(src.content) + return arr + } + is AssistantMessage-> { + val role = context.serialize(src.role, ChatRole::class.java) + val arr = JsonArray() + arr.add(role) + arr.add(src.content) + return arr + } + is ContentFileMessage -> { + val role = context.serialize(src.role, ChatRole::class.java) + val arr = JsonArray() + arr.add(role) + val fileArray = arrayOf(ChatContextFile("", "", 0, 0)) + val contextFile = context.serialize(src.content, fileArray::class.java) + arr.add(contextFile) + return arr + } + + else -> JsonArray() + } + } + } + + val gson = GsonBuilder() + .registerTypeAdapter(FromChat::class.java, FromChatDeserializer()) + .registerTypeAdapter(AtCommands.Preview.Response::class.java, AtCommands.Preview.ResponseDeserializer()) + .registerTypeAdapter(ChatMessage::class.java, ChatMessageDeserializer()) + .registerTypeHierarchyAdapter(ChatMessage::class.java, MessageSerializer()) + .create() + + fun parse(msg: String?): FromChat? { + return gson.fromJson(msg, FromChat::class.java) + } + + fun stringify(event: ToChat): String { + return gson.toJson(event) + } + } +} diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/SharedChatPane.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/SharedChatPane.kt new file mode 100644 index 00000000..0430fb56 --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/SharedChatPane.kt @@ -0,0 +1,397 @@ +package com.smallcloud.refactai.panes.sharedchat + +import com.intellij.lang.Language +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.command.WriteCommandAction +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.editor.event.SelectionEvent +import com.intellij.openapi.editor.event.SelectionListener +import com.intellij.openapi.fileEditor.FileEditorManager +import com.intellij.openapi.fileEditor.FileEditorManagerEvent +import com.intellij.openapi.fileEditor.FileEditorManagerListener +import com.intellij.openapi.fileEditor.OpenFileDescriptor +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.Disposer +import com.intellij.openapi.util.TextRange +import com.intellij.openapi.vfs.VirtualFile +import com.intellij.psi.PsiDocumentManager +import com.intellij.testFramework.LightVirtualFile +import com.intellij.util.concurrency.AppExecutorUtil +import com.intellij.util.ui.UIUtil +import com.smallcloud.refactai.io.InferenceGlobalContextChangedNotifier +import com.smallcloud.refactai.lsp.LSPProcessHolder +import com.smallcloud.refactai.panes.sharedchat.Events.ActiveFile.ActiveFileToChat +import com.smallcloud.refactai.panes.sharedchat.Events.ActiveFile.FileInfoPayload +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.RestorePayload +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.RestoreToChat +import com.smallcloud.refactai.panes.sharedchat.Events.Editor +import com.smallcloud.refactai.panes.sharedchat.browser.ChatWebView +import com.smallcloud.refactai.settings.AppSettingsState +import org.jetbrains.annotations.NotNull +import java.beans.PropertyChangeListener +import java.util.concurrent.CompletableFuture +import java.util.concurrent.Future +import javax.swing.JPanel +import javax.swing.UIManager + + +class SharedChatPane(val project: Project) : JPanel(), Disposable { + + private val lsp: LSPProcessHolder = LSPProcessHolder.getInstance(project) + + var id: String? = null; + private var defaultChatModel: String? = null + private var chatThreadToRestore: Events.Chat.Thread? = null + private var lastProcess: CompletableFuture>? = null; + + + private fun getLanguage(fm: FileEditorManager): Language? { + val editor = fm.selectedTextEditor + val language = editor?.document?.let { + PsiDocumentManager.getInstance(project).getPsiFile(it)?.language + } + + return language + } + + private fun sendSelectedSnippet(id: String) { + this.getSelectedSnippet { snippet -> + val payload = Editor.SetSnippetPayload(id, snippet) + val message = Editor.SetSnippetToChat(payload) + this.postMessage(message) + } + } + + private fun sendUserConfig(id: String) { + val hasAst = AppSettingsState.instance.astIsEnabled + val hasVecdb = AppSettingsState.instance.vecdbIsEnabled + val features = Events.Config.Features(hasAst, hasVecdb) + val isDarkMode = UIUtil.isUnderDarcula() + val mode = if (isDarkMode) "dark" else "light" + val themeProps = Events.Config.ThemeProps(mode) + val message = Events.Config.Update(id, features, themeProps) + this.postMessage(message) + } + + private fun getSelectedSnippet(cb: (Events.Editor.Snippet) -> Unit) { + ApplicationManager.getApplication().invokeLater { + if (!project.isDisposed && FileEditorManager.getInstance(project).selectedFiles.isNotEmpty()) { + val fileEditorManager = FileEditorManager.getInstance(project) + val editor = fileEditorManager.selectedTextEditor + val file = fileEditorManager.selectedFiles[0] + val path = file.path + val name = file.name + val language = this.getLanguage(fileEditorManager)?.id + val caretModel = editor?.caretModel + + val selection = caretModel?.currentCaret?.selectionRange + val range = TextRange(selection?.startOffset ?: 0, selection?.endOffset ?: 0) + + val code = editor?.document?.getText(range) + if (language == null || code == null) { + cb(Events.Editor.Snippet()) + } else { + val snippet = Events.Editor.Snippet(language, code, path, name) + cb(snippet) + } + } + } + } + + private fun getActiveFileInfo(cb: (Events.ActiveFile.FileInfo) -> Unit) { + ApplicationManager.getApplication().invokeLater { + if (!project.isDisposed && FileEditorManager.getInstance(project).selectedFiles.isNotEmpty()) { + val fileEditorManager = FileEditorManager.getInstance(project) + val editor = fileEditorManager.selectedTextEditor + + val cursor = editor?.caretModel?.offset + val virtualFile = fileEditorManager.selectedFiles[0] + val filePath = virtualFile.path + val fileName = virtualFile.name + + val selection = editor?.caretModel?.currentCaret?.selectionRange + val range = TextRange(selection?.startOffset ?: 0, selection?.endOffset ?: 0) + + val code = editor?.document?.getText(range) + + val canPaste = selection != null && !selection.isEmpty + + val fileInfo = Events.ActiveFile.FileInfo( + fileName, + filePath, + canPaste, + cursor = cursor, + line1 = selection?.startOffset, + line2 = selection?.endOffset, + content = code, + ) + cb(fileInfo) + + } else { + val fileInfo = Events.ActiveFile.FileInfo() + cb(fileInfo) + } + } + } + + private fun sendActiveFileInfo(id: String) { + this.getActiveFileInfo { file -> + val payload = FileInfoPayload(id, file) + val message = ActiveFileToChat(payload) + this.postMessage(message) + } + } + + private fun handleCaps(id: String) { + this.lsp.fetchCaps().also { caps -> + val res = caps.get() + val message: Events.Caps.Receive = Events.Caps.Receive(id, res) + this.defaultChatModel = res.codeChatDefaultModel + this.postMessage(message) + } + } + + private fun handleSystemPrompts(id: String) { + this.lsp.fetchSystemPrompts().also { res -> + val prompts: SystemPromptMap = res.get() + val payload = Events.SystemPrompts.SystemPromptsPayload(id, prompts) + val message: Events.SystemPrompts.Receive = Events.SystemPrompts.Receive(payload) + this.postMessage(message) + } + } + + private fun handleCompletion( + id: String, + query: String, + cursor: Int, + number: Int = 5, + ) { + AppExecutorUtil.getAppExecutorService().submit { + try { + this.lsp.fetchCommandCompletion(query, cursor, number).also { res -> + val completions = res.get() + val payload = Events.AtCommands.Completion.CompletionPayload( + id, completions.completions, completions.replace, completions.isCmdExecutable + ) + val message = Events.AtCommands.Completion.Receive(payload) + this.postMessage(message) + } + } catch (e: Exception) { + println("Commands error") + println(e) + } + } + } + + private fun handleChat(id: String, messages: ChatMessages, model: String, title: String? = null) { + + val future = this.lsp.sendChat(id, messages, model, dataReceived = { str, requestId -> + when (val res = Events.Chat.Response.parse(str)) { + is Events.Chat.Response.Choices -> { + val message = Events.Chat.Response.formatToChat(res, requestId) + this.postMessage(message) + } + + is Events.Chat.Response.UserMessage -> { + val message = Events.Chat.Response.formatToChat(res, requestId) + this.postMessage(message) + } + + is Events.Chat.Response.DetailMessage -> { + val message = Events.Chat.Response.formatToChat(res, requestId) + this.postMessage(message) + } + } + }, dataReceiveEnded = { str -> + val res = Events.Chat.Response.ChatDone(str) + val message = Events.Chat.Response.formatToChat(res, id) + this.postMessage(message) + }, errorDataReceived = { json -> + val res = Events.Chat.Response.ChatError(json) + val message = Events.Chat.Response.formatToChat(res, id) + this.postMessage(message) + }, failedDataReceiveEnded = { e -> + val res = Events.Chat.Response.ChatFailedStream(e) + val message = Events.Chat.Response.formatToChat(res, id) + this.postMessage(message) + }) + + this.lastProcess = future + + } + + private fun handleChatSave(id: String, messages: ChatMessages, maybeModel: String) { + val model = maybeModel.ifEmpty { this.defaultChatModel ?: "" } + ChatHistory.instance.state.save(id, messages, model) + } + + fun handleChatStop(id: String) { + // TODO: stop the stream from the lsp + this.lastProcess?.get()?.cancel(true) + } + + private fun handlePaste(id: String, content: String) { + ApplicationManager.getApplication().invokeLater { + val editor = FileEditorManager.getInstance(project).selectedTextEditor + val selection = editor?.caretModel?.currentCaret?.selectionRange + if (selection != null) { + WriteCommandAction.runWriteCommandAction(project) { + editor.document.replaceString(selection.startOffset, selection.endOffset, content) + } + } + } + } + + private fun handleNewFile(id: String, content: String) { + // TODO: file type? + val vf = LightVirtualFile("Untitled", content) + + val fileDescriptor = OpenFileDescriptor(project, vf) + + ApplicationManager.getApplication().invokeLater { + val e = FileEditorManager.getInstance(project).openTextEditor(fileDescriptor, true) + } + } + + private fun addEventListeners(id: String) { + println("Adding ide event listeners") + val listener: FileEditorManagerListener = object : FileEditorManagerListener { + override fun fileOpened(@NotNull source: FileEditorManager, @NotNull file: VirtualFile) { + this@SharedChatPane.sendActiveFileInfo(id) + } + + override fun fileClosed(@NotNull source: FileEditorManager, @NotNull file: VirtualFile) { + this@SharedChatPane.sendActiveFileInfo(id) + } + + override fun selectionChanged(@NotNull event: FileEditorManagerEvent) { + this@SharedChatPane.sendActiveFileInfo(id) + this@SharedChatPane.sendSelectedSnippet(id) + } + + } + + val selectionListener = object : SelectionListener { + override fun selectionChanged(event: SelectionEvent) { + this@SharedChatPane.sendActiveFileInfo(id) + this@SharedChatPane.sendSelectedSnippet(id) + } + + } + + project.messageBus.connect().subscribe(FileEditorManagerListener.FILE_EDITOR_MANAGER, listener) + + val ef = EditorFactory.getInstance() + ef.eventMulticaster.addSelectionListener(selectionListener, this) + + UIManager.addPropertyChangeListener(uiChangeListener) + + // ast and vecdb settings change + project.messageBus.connect() + .subscribe(InferenceGlobalContextChangedNotifier.TOPIC, object : InferenceGlobalContextChangedNotifier { + override fun astFlagChanged(newValue: Boolean) { + println("ast changed to: $newValue") + this@SharedChatPane.sendUserConfig(id) + } + + override fun vecdbFlagChanged(newValue: Boolean) { + println("vecdb changed to: $newValue") + this@SharedChatPane.sendUserConfig(id) + } + }) + } + + private fun setLookAndFeel() { + this.browser.setStyle() + if (this.id != null) { + this.sendUserConfig(this.id!!) + } + } + + private val uiChangeListener = PropertyChangeListener { event -> + if (event.propertyName == "lookAndFeel") { + this.setLookAndFeel() + } + } + + fun restoreWhenReady(id: String, messages: ChatMessages, model: String) { + val chatThread = Events.Chat.Thread(id, messages, model) + this.chatThreadToRestore = chatThread + } + + private fun maybeRestore(id: String) { + if (this.chatThreadToRestore != null) { + val payload = RestorePayload(id, this.chatThreadToRestore!!) + val event = RestoreToChat(payload) + this.id = payload.id + this.postMessage(event) + this.chatThreadToRestore = null + } + } + + private fun handlePreviewFileRequest(id: String, query: String) { + try { + this.lsp.fetchCommandPreview(query).also { res -> + val preview = res.get() + val payload = Events.AtCommands.Preview.PreviewPayload(id, preview.messages) + val message = Events.AtCommands.Preview.Receive(payload) + this.postMessage(message) + } + } catch (e: Exception) { + println("Command preview error") + println(e) + } + } + + private fun handleReadyMessage(id: String) { + this.id = id; + this.sendActiveFileInfo(id) + this.sendSelectedSnippet(id) + this.addEventListeners(id) + this.sendUserConfig(id) + this.maybeRestore(id) + } + + private fun handleEvent(event: Events.FromChat) { + // println("Event received: ${event}") + when (event) { + is Events.Ready -> this.handleReadyMessage(event.id) + is Events.Caps.Request -> this.handleCaps(event.id) + is Events.SystemPrompts.Request -> this.handleSystemPrompts(event.id) + is Events.AtCommands.Completion.Request -> this.handleCompletion( + event.id, event.query, event.cursor, event.number + ) + + is Events.AtCommands.Preview.Request -> this.handlePreviewFileRequest(event.id, event.query) + is Events.Chat.AskQuestion -> this.handleChat(event.id, event.messages, event.model, event.title) + is Events.Chat.Save -> this.handleChatSave(event.id, event.messages, event.model) + is Events.Chat.Stop -> this.handleChatStop(event.id) + is Events.Editor.Paste -> this.handlePaste(event.id, event.content) + is Events.Editor.NewFile -> this.handleNewFile(event.id, event.content) + else -> Unit + } + } + + private val browser by lazy { + ChatWebView { event -> + this.handleEvent(event) + } + } + + val webView by lazy { + browser.webView + } + + private fun postMessage(message: Events.ToChat?) { + this.browser.postMessage(message) + } + + + override fun dispose() { + UIManager.removePropertyChangeListener(uiChangeListener) + webView.dispose() + Disposer.dispose(this) + } +} + diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/ChatWebView.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/ChatWebView.kt new file mode 100644 index 00000000..dac4f918 --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/ChatWebView.kt @@ -0,0 +1,133 @@ +package com.smallcloud.refactai.panes.sharedchat.browser + +import com.intellij.ide.BrowserUtil +import com.intellij.openapi.Disposable +import com.intellij.ui.jcef.* +import com.intellij.util.ui.UIUtil +import com.smallcloud.refactai.panes.sharedchat.Events +import org.cef.CefApp +import org.cef.browser.CefBrowser +import org.cef.handler.CefLoadHandlerAdapter +import javax.swing.JComponent + +class ChatWebView(val messageHandler: (event: Events.FromChat) -> Unit): Disposable { + private val jsPoolSize = "200" + + init { + System.setProperty("ide.browser.jcef.jsQueryPoolSize", jsPoolSize) + } + + fun setStyle() { + val isDarkMode = UIUtil.isUnderDarcula() + val mode = if (isDarkMode) { "dark" } else { "light" } + val bodyClass = if (isDarkMode) { "vscode-dark" } else { "vscode-light" } + val backgroundColour = UIUtil.getPanelBackground() + val red = backgroundColour.red + val green = backgroundColour.green + val blue = backgroundColour.blue + this.webView.executeJavaScriptAsync("""document.body.style.setProperty("background-color", "rgb($red, $green, $blue");""") + this.webView.executeJavaScriptAsync("""document.body.class = "$bodyClass";""") + this.webView.executeJavaScriptAsync("""document.documentElement.className = "$mode";""") + + } + + val webView by lazy { + // TODO: handle JBCef not being available + val browser = JBCefBrowser() + browser.jbCefClient.setProperty( + JBCefClient.Properties.JS_QUERY_POOL_SIZE, + jsPoolSize, + ) + browser.setProperty(JBCefBrowserBase.Properties.NO_CONTEXT_MENU, true) + + CefApp.getInstance().registerSchemeHandlerFactory("http", "refactai", RequestHandlerFactory()) + + browser.loadURL("http://refactai/index.html") + + val myJSQueryOpenInBrowser = JBCefJSQuery.create((browser as JBCefBrowserBase?)!!) + myJSQueryOpenInBrowser.addHandler { msg -> + val event = Events.parse(msg) + if(event != null) { + this.messageHandler(event) + } + null + } + + val myJSQueryOpenInBrowserRedirectHyperlink = JBCefJSQuery.create((browser as JBCefBrowserBase?)!!) + myJSQueryOpenInBrowserRedirectHyperlink.addHandler { href -> + BrowserUtil.browse(href) + null + } + + var installedScript = false + + browser.jbCefClient.addLoadHandler(object: CefLoadHandlerAdapter() { + override fun onLoadingStateChange( + browser: CefBrowser, + isLoading: Boolean, + canGoBack: Boolean, + canGoForward: Boolean + ) { + if(!installedScript) { + installedScript = setUpJavaScriptMessageBus(browser, myJSQueryOpenInBrowser) + } + if(!isLoading) { + setUpJavaScriptMessageBusRedirectHyperlink(browser, myJSQueryOpenInBrowserRedirectHyperlink) + setStyle() + } + } + + }, browser.cefBrowser) + + browser + } + + fun setUpJavaScriptMessageBusRedirectHyperlink(browser: CefBrowser?, myJSQueryOpenInBrowser: JBCefJSQuery) { + val script = """window.openLink = function(href) { + ${myJSQueryOpenInBrowser.inject("href")} + } + document.addEventListener('click', function(event) { + if (event.target.tagName.toLowerCase() === 'a') { + event.preventDefault(); + window.openLink(event.target.href); + } + });""".trimIndent() + browser?.executeJavaScript(script, browser?.url, 0) + } + + fun setUpJavaScriptMessageBus(browser: CefBrowser?, myJSQueryOpenInBrowser: JBCefJSQuery): Boolean { + + val script = """window.postIntellijMessage = function(event) { + const msg = JSON.stringify(event); + ${myJSQueryOpenInBrowser.inject("msg")} + }""".trimIndent() + if(browser != null) { + browser.executeJavaScript(script, browser.url, 0); + return true + } + return false + } + + fun postMessage(message: Events.ToChat?) { + if(message != null) { + val json = Events.stringify(message) + this.postMessage(json) + } + } + + fun postMessage(message: String) { + // println("postMessage: $message") + val script = """window.postMessage($message, "*");""" + webView.cefBrowser.executeJavaScript(script, webView.cefBrowser.url, 0) + } + + fun getComponent(): JComponent { + return webView.component + } + + override fun dispose() { + this.webView.dispose() + } + + +} \ No newline at end of file diff --git a/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/RequestHandler.kt b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/RequestHandler.kt new file mode 100644 index 00000000..fc2e7ad7 --- /dev/null +++ b/src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/browser/RequestHandler.kt @@ -0,0 +1,170 @@ +package com.smallcloud.refactai.panes.sharedchat.browser + +import com.intellij.openapi.project.DumbAware +import org.cef.browser.CefBrowser +import org.cef.browser.CefFrame +import org.cef.callback.CefCallback +import org.cef.callback.CefSchemeHandlerFactory +import org.cef.handler.CefLoadHandler +import org.cef.handler.CefResourceHandler +import org.cef.misc.IntRef +import org.cef.misc.StringRef +import org.cef.network.CefRequest +import org.cef.network.CefResponse +import java.io.IOException +import java.io.InputStream +import java.net.URLConnection + +class RequestHandlerFactory : CefSchemeHandlerFactory { + override fun create( + browser: CefBrowser?, + frame: CefFrame?, + schemeName: String, + request: CefRequest + ): CefResourceHandler { + return ResourceHandler() + } +} + +data object ClosedConnection : ResourceHandlerState() { + override fun getResponseHeaders( + cefResponse: CefResponse, + responseLength: IntRef, + redirectUrl: StringRef + ) { + cefResponse.status = 404 + } +} + +sealed class ResourceHandlerState { + open fun getResponseHeaders( + cefResponse: CefResponse, + responseLength: IntRef, + redirectUrl: StringRef + ) { + } + + open fun readResponse( + dataOut: ByteArray, + bytesToRead: Int, + bytesRead: IntRef, + callback: CefCallback + ): Boolean = false + + open fun close() {} +} + +class OpenedConnection(private val connection: URLConnection?) : + ResourceHandlerState() { + + private val inputStream: InputStream? by lazy { + connection?.inputStream + } + + override fun getResponseHeaders( + cefResponse: CefResponse, + responseLength: IntRef, + redirectUrl: StringRef + ) { + try { + if (connection != null) { + val url = connection.url.toString() + when { + url.contains("css") -> cefResponse.mimeType = "text/css" + url.contains("js") -> cefResponse.mimeType = "text/javascript" + url.contains("html") -> cefResponse.mimeType = "text/html" + else -> cefResponse.mimeType = connection.contentType + } + responseLength.set(inputStream?.available() ?: 0) + cefResponse.status = 200 + } else { + // Handle the case where connection is null + cefResponse.error = CefLoadHandler.ErrorCode.ERR_FAILED + cefResponse.statusText = "Connection is null" + cefResponse.status = 500 + } + } catch (e: IOException) { + println("Error: $e"); + cefResponse.error = CefLoadHandler.ErrorCode.ERR_FILE_NOT_FOUND + cefResponse.statusText = e.localizedMessage + cefResponse.status = 404 + } + } + + + override fun readResponse( + dataOut: ByteArray, + bytesToRead: Int, + bytesRead: IntRef, + callback: CefCallback + ): Boolean { + return inputStream?.let {inputStream -> + val availableSize = inputStream.available() + return if (availableSize > 0) { + val maxBytesToRead = minOf(availableSize, bytesToRead) + val realBytesRead = inputStream.read(dataOut, 0, maxBytesToRead) + bytesRead.set(realBytesRead) + true + } else { + inputStream.close() + false + } + } ?: false + } + + override fun close() { + inputStream?.close() + } +} + +class ResourceHandler : CefResourceHandler, DumbAware { + private var state: ResourceHandlerState = ClosedConnection + private var currentUrl: String? = null + override fun processRequest( + cefRequest: CefRequest, + cefCallback: CefCallback + ): Boolean { + val url = cefRequest.url + return if (url != null) { + val pathToResource = url.replace("http://refactai/", "webview/") + val newUrl = javaClass.classLoader.getResource(pathToResource) + state = OpenedConnection(newUrl?.openConnection()) + currentUrl = url + cefCallback.Continue() + true + } else { + false + } + } + + override fun getResponseHeaders( + cefResponse: CefResponse, + responseLength: IntRef, + redirectUrl: StringRef + ) { + + if (currentUrl !== null){ + when { + currentUrl!!.contains("css") -> cefResponse.mimeType = "text/css" + currentUrl!!.contains("js") -> cefResponse.mimeType = "text/javascript" + currentUrl!!.contains("html") -> cefResponse.mimeType = "text/html" + else -> {} + } + } + state.getResponseHeaders(cefResponse, responseLength, redirectUrl) + } + + override fun readResponse( + dataOut: ByteArray, + bytesToRead: Int, + bytesRead: IntRef, + callback: CefCallback + ): Boolean { + return state.readResponse(dataOut, bytesToRead, bytesRead, callback) + } + + override fun cancel() { + state.close() + state = ClosedConnection + } +} diff --git a/src/main/kotlin/com/smallcloud/refactai/settings/AppSettingsComponent.kt b/src/main/kotlin/com/smallcloud/refactai/settings/AppSettingsComponent.kt index dc09a5ba..7ceef1b9 100644 --- a/src/main/kotlin/com/smallcloud/refactai/settings/AppSettingsComponent.kt +++ b/src/main/kotlin/com/smallcloud/refactai/settings/AppSettingsComponent.kt @@ -33,6 +33,7 @@ class AppSettingsComponent { myXDebugLSPPortLabel.isVisible = true myStagingVersionText.isVisible = true myStagingVersionLabel.isVisible = true + astCheckbox.isVisible = true vecdbCheckbox.isVisible = true } } @@ -58,8 +59,9 @@ class AppSettingsComponent { isVisible = false } private val defaultSystemPromptTextArea = JBTextArea() + private val vecdbCheckbox = JCheckBox("VECDB").apply { - isVisible = false + isVisible = true } @@ -107,6 +109,7 @@ class AppSettingsComponent { addComponent(developerModeCheckBox, UIUtil.LARGE_VGAP) addLabeledComponent(myXDebugLSPPortLabel, myXDebugLSPPort, UIUtil.LARGE_VGAP) addLabeledComponent(myStagingVersionLabel, myStagingVersionText, UIUtil.LARGE_VGAP) + addComponent(astCheckbox, UIUtil.LARGE_VGAP) addComponent(vecdbCheckbox, UIUtil.LARGE_VGAP) addComponentFillVertically(JPanel(), 0) }.panel @@ -119,7 +122,6 @@ class AppSettingsComponent { val preferredFocusedComponent: JComponent get() = myTokenText - var tokenText: String get() = myTokenText.text set(newText) { diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index e4376a3a..566f2c98 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -136,6 +136,7 @@ integrated into a single package that follows your privacy settings.

+ diff --git a/src/main/resources/webview/index.html b/src/main/resources/webview/index.html new file mode 100644 index 00000000..16be679b --- /dev/null +++ b/src/main/resources/webview/index.html @@ -0,0 +1,53 @@ + + + + + Refact.ai + + + + + +
+ + + + \ No newline at end of file diff --git a/src/test/kotlin/com/smallcloud/refactai/panes/sharedchat/EventsTest.kt b/src/test/kotlin/com/smallcloud/refactai/panes/sharedchat/EventsTest.kt new file mode 100644 index 00000000..ed60b1aa --- /dev/null +++ b/src/test/kotlin/com/smallcloud/refactai/panes/sharedchat/EventsTest.kt @@ -0,0 +1,273 @@ +package com.smallcloud.refactai.panes.sharedchat + +import com.google.gson.* +import com.intellij.util.containers.toArray +import com.smallcloud.refactai.lsp.LSPCapabilities +import com.smallcloud.refactai.panes.sharedchat.Events.ActiveFile.ActiveFileToChat +import com.smallcloud.refactai.panes.sharedchat.Events.ActiveFile.FileInfoPayload +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.RestorePayload +import com.smallcloud.refactai.panes.sharedchat.Events.Chat.RestoreToChat +import com.smallcloud.refactai.panes.sharedchat.Events.Editor +import kotlin.test.Test +import org.junit.jupiter.api.Assertions.* +import java.lang.reflect.Type + +class EventsTest { + @Test + fun parseReadyMessage() { + val msg = """{ + type: "${EventNames.FromChat.READY.value}", + payload: { + id: "foo" + } + }""".trimIndent() + val expected = Events.Ready("foo") + val result = Events.parse(msg) + + assertEquals(result, expected) + } + + + @Test + fun parseResponseContextTest() { + // TODO: move to lsp + val msg = """{"role":"context_file", "content":"[{\"file_name\":\"/main.py\",\"file_content\":\"foo\",\"line1\":1,\"line2\":15,\"symbol\":\"00000000-0000-0000-0000-000000000000\",\"gradient_type\":-1,\"usefulness\":0.0}]"}""" + val result = Events.Chat.Response.parse(msg) + val expected = Events.Chat.Response.UserMessage(Events.Chat.Response.Roles.CONTEXT_FILE,"""[{"file_name":"/main.py","file_content":"foo","line1":1,"line2":15,"symbol":"00000000-0000-0000-0000-000000000000","gradient_type":-1,"usefulness":0.0}]""") + assertEquals(expected, result) + } + + @Test + fun formatResponseContextTest() { + val message = Events.Chat.Response.UserMessage(Events.Chat.Response.Roles.CONTEXT_FILE,"""[]""") + val id = "foo"; + val toChat = Events.Chat.Response.formatToChat(message, id) + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_response","payload":{"id":"foo","role":"context_file","content":"[]"}}""" + + assertEquals(result, expected) + } + + @Test + fun parseResponseDetail() { + val str = """{"detail":"JSON problem: invalid type: sequence, expected a string at line 1 column 46"}""" + val res = Events.Chat.Response.parse(str) + val toChat = Events.Chat.Response.formatToChat(res, "foo") + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_error_streaming","payload":{"id":"foo","message":"JSON problem: invalid type: sequence, expected a string at line 1 column 46"}}""" + assertEquals(expected, result) + } + + @Test + fun parseAtCommandCompletionMessageTest() { + val message = """{"type":"chat_request_at_command_completion","payload":{"id":"foo","query":"@","cursor":1,"trigger":"@","number":5}}""" + val result = Events.parse(message) + val expected = Events.AtCommands.Completion.Request("foo", "@", 1, 5, "@") + assertEquals(expected, result) + } + + @Test + fun parsePreviewFileRequest() { + val message = """{"type":"chat_request_preview_files","payload":{"id":"foo","query": "hello"}}""" + val result = Events.parse(message) + val expected = Events.AtCommands.Preview.Request("foo", "hello") + assertEquals(expected, result) + } + + @Test + fun parsePreviewResponse() { + val response = """{"highlight":[],"messages":[{"content":"[]","role":"context_file"}],"model":"deepseek-coder/6.7b/instruct-finetune"}""" + val result = Events.gson.fromJson(response, Events.AtCommands.Preview.Response::class.java) + val expected = Events.AtCommands.Preview.Response(arrayOf(ContentFileMessage(emptyArray()))) + assertEquals(expected, result) + + val payload = Events.AtCommands.Preview.PreviewPayload("foo", result.messages) + val message = Events.stringify(Events.AtCommands.Preview.Receive(payload)) + val json = """{"type":"chat_receive_at_command_preview","payload":{"id":"foo","preview":[["context_file",[]]]}}""" + assertEquals(json, message) + } + + @Test + fun parseSaveChatMessage() { + val message = """{"type":"save_chat_to_history","payload":{"id":"foo","messages":[["context_file",[{"file_name":"/main.py","file_content":"hello\n","line1":1,"line2":1,"symbol":"00000000-0000-0000-0000-000000000000","gradient_type":-1,"usefulness":0}]],["user","hello"],["assistant","Hello there"]],"title":"","model":""}}""" + val result = Events.parse(message) + val messages: ChatMessages = arrayOf( + ContentFileMessage(arrayOf(ChatContextFile("/main.py","hello\n",1,1,0.0 ))), + UserMessage("hello"), + AssistantMessage("Hello there") + ) + val expected = Events.Chat.Save("foo", messages, "", "") + assertEquals(expected, result) + } + + + @Test + fun parseCompletionsResponse() { + val input = """{"completions":["@file","@workspace","@symbols-at","@references","@definition"],"replace":[0,1],"is_cmd_executable":false}""" + val result = Gson().fromJson(input, CommandCompletionResponse::class.java) + val expected = CommandCompletionResponse( + arrayOf("@file", "@workspace", "@symbols-at","@references","@definition"), + arrayOf(0, 1), + false + ) + + assertEquals(expected, result) + + } + + @Test + fun parseCapsRequest() { + val input = """{"type":"chat_request_caps","payload":{"id":"foo"}}""" + val result = Events.parse(input) + val expected = Events.Caps.Request("foo") + assertEquals(expected, result) + } + + @Test + fun stringifyCompletions() { + val payload = Events.AtCommands.Completion.CompletionPayload("foo", emptyArray(), arrayOf(0, 1), false) + val input = Events.AtCommands.Completion.Receive(payload) + val result = Events.stringify(input) + val expected = """{"type":"chat_receive_at_command_completion","payload":{"id":"foo","completions":[],"replace":[0,1],"is_cmd_executable":false}}""" + assertEquals(expected, result) + } + + + @Test + fun stringifyCompletionPreview() { + val preview = ChatContextFile("test.py", "foo", 0, 1, 100.0) + val payload = Events.AtCommands.Preview.PreviewPayload("foo", arrayOf(ContentFileMessage(arrayOf(preview)))) + val input = Events.AtCommands.Preview.Receive(payload) + val result = Events.stringify(input) + val expected = """{"type":"chat_receive_at_command_preview","payload":{"id":"foo","preview":[["context_file",[{"file_name":"test.py","file_content":"foo","line1":0,"line2":1,"usefulness":100.0}]]]}}""" + assertEquals(expected, result) + } + + + @Test + fun formatResponseChoicesTest() { + val delta = AssistantDelta("hello") + val choice = Events.Chat.Response.Choice(delta, 0, null) + val choices = arrayOf(choice) + val message = Events.Chat.Response.Choices(choices, "0", "refact" ) + val toChat = Events.Chat.Response.formatToChat(message, "foo") + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_response","payload":{"id":"foo","choices":[{"delta":{"role":"assistant","content":"hello"},"index":0}],"created":"0","model":"refact"}}""" + + assertEquals(expected, result) + } + + @Test + fun formatResponseDoneTest() { + val msg = "data: [DONE]" + val done = Events.Chat.Response.ChatDone(msg) + val toChat = Events.Chat.Response.formatToChat(done, "foo") + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_done_streaming","payload":{"id":"foo","message":"data: [DONE]"}}""" + + assertEquals(expected, result) + + } + + @Test + fun formatResponseErrorMessage() { + val msg = JsonObject() + msg.addProperty("detail", "test error") + val err = Events.Chat.Response.ChatError(msg) + val toChat = Events.Chat.Response.formatToChat(err, "foo") + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_error_streaming","payload":{"id":"foo","message":"test error"}}""" + + assertEquals(expected, result) + } + + @Test + fun formatResponseFailedStreamingTest() { + + val e = Throwable("test error") + val err = Events.Chat.Response.ChatFailedStream(e) + val toChat = Events.Chat.Response.formatToChat(err, "foo") + val result = Events.stringify(toChat!!) + val expected = """{"type":"chat_error_streaming","payload":{"id":"foo","message":"Failed during stream: test error"}}""" + assertEquals(expected, result) + } + + @Test + fun formatSnippetToChatTest() { + val id = "foo" + val snippet = Events.Editor.Snippet() + val payload = Editor.SetSnippetPayload(id, snippet) + val message = Editor.SetSnippetToChat(payload) + val result = Events.stringify(message) + val expected = """{"type":"chat_set_selected_snippet","payload":{"id":"foo","snippet":{"language":"","code":"","path":"","basename":""}}}""" + assertEquals(expected, result) + } + + @Test + fun formatActiveFileToChatTest() { + val file = Events.ActiveFile.FileInfo() + val id = "foo" + val payload = FileInfoPayload(id, file) + val message = ActiveFileToChat(payload) + val result = Events.stringify(message) + val expected = """{"type":"chat_active_file_info","payload":{"id":"foo","file":{"name":"","path":"","can_paste":false,"attach":false}}}""" + + assertEquals(expected, result) + } + + @Test + fun formatRestoreChatToChat() { + val chatMessages: ChatMessages = arrayOf( + ContentFileMessage(arrayOf( + ChatContextFile("/main.py", "hello", 1, 15, 0.0) + )), + UserMessage("hello"), + AssistantMessage("hello") + ) + + val currentChatId = "foo" + val chat = Events.Chat.Thread("bar", chatMessages, "refact") + val payload = RestorePayload(currentChatId, chat) + val event = RestoreToChat(payload) + val result = Events.stringify(event) + + val expected = """{"type":"restore_chat_from_history","payload":{"id":"foo","chat":{"id":"bar","messages":[["context_file",[{"file_name":"/main.py","file_content":"hello","line1":1,"line2":15,"usefulness":0.0}]],["user","hello"],["assistant","hello"]],"model":"refact","attach_file":false}}}""" + + assertEquals(expected, result) + } + + @Test + fun systemPromptsMessage() { + val prompts: SystemPromptMap = mapOf("default" to SystemPrompt(text="Use backquotes for code blocks.\nPay close attention to indent when editing code blocks: indent must be exactly the same as in the original code block.\n", description=""),) + val payload = Events.SystemPrompts.SystemPromptsPayload("foo", prompts) + val message: Events.SystemPrompts.Receive = Events.SystemPrompts.Receive(payload) + val result = Events.stringify(message) + val expected = """{"type":"chat_receive_prompts","payload":{"id":"foo","prompts":{"default":{"text":"Use backquotes for code blocks.\nPay close attention to indent when editing code blocks: indent must be exactly the same as in the original code block.\n","description":""}}}}""" + assertEquals(expected, result) + + } + + @Test + fun receiveCapsMessage() { + val caps = LSPCapabilities() + val message = Events.Caps.Receive("foo", caps) + val result = Events.stringify(message) + val expected = """{"type":"receive_caps","payload":{"id":"foo","caps":{"cloud_name":"","code_chat_default_model":"","code_chat_models":{},"code_completion_default_model":"","code_completion_models":{},"endpoint_style":"","endpoint_template":"","running_models":[],"telemetry_basic_dest":"","tokenizer_path_template":"","tokenizer_rewrite_path":{}}}}""" + assertEquals(expected, result) + + } + + @Test + fun configMessage() { + val message = Events.Config.Update( + "chat-id", + Events.Config.Features(true, false), + Events.Config.ThemeProps("light") + ) + val result = Events.stringify(message) + val expected = """{"type":"receive_config_update","payload":{"id":"chat-id","features":{"ast":true,"vecdb":false},"themeProps":{"mode":"light","hasBackground":false,"scale":"90%","accentColor":"gray"}}}""" + + assertEquals(expected, result) + } + +} \ No newline at end of file