feat(agent): add local mode support and logging for LLM responses #379

Author: phodal

core/src/main/resources/messages/AutoDevBundle_en.properties

@@ -267,3 +267,6 @@ mcp.editor.preview.title=Preview
 mcp.editor.preview.tooltip=Preview panel
 mcp.editor.refresh.title=Refresh
 shire.toolchain.function.not.found=Toolchain Not Found
+shire.run.local.mode=Run Local Model
+devins.llm.notfound=LLM not found
+devins.llm.done=Done

core/src/main/resources/messages/AutoDevBundle_zh.properties

@@ -257,3 +257,6 @@ mcp.editor.preview.title=预览
 mcp.editor.preview.tooltip=预览面板
 mcp.editor.refresh.title=刷新
 shire.toolchain.function.not.found=Toolchain Not Found
+shire.run.local.mode=Run Local Model
+devins.llm.notfound=LLM not found
+devins.llm.done=Done

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/ShireConstants.kt

@@ -0,0 +1,24 @@
+package cc.unitmesh.devti.language
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.guessProjectDir
+import com.intellij.openapi.vfs.VirtualFile
+
+public const val SHIRE_CHAT_BOX_FILE = "shire-chatbox.default.shire"
+public const val SHIRE_TEMP_OUTPUT = ".shire-output"
+public const val LLM_LOGGING_JSONL = "logging.jsonl"
+public const val LLM_LOGGING = "logging.log"
+public const val SHIRE_MKT_HOST = "https://shire.run/packages.json"
+
+object ShireConstants {
+    fun outputDir(project: Project): VirtualFile? {
+        val baseDir = project.guessProjectDir() ?: throw IllegalStateException("Project directory not found")
+        val virtualFile = baseDir.findFileByRelativePath(SHIRE_TEMP_OUTPUT)
+        if (virtualFile == null) {
+            baseDir.createChildDirectory(this, SHIRE_TEMP_OUTPUT)
+        }
+
+        return virtualFile
+    }
+}
+

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/ast/HobbitHole.kt

@@ -19,7 +19,7 @@ import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
-import com.phodal.shirelang.compiler.ast.action.PatternAction
+import cc.unitmesh.devti.language.ast.action.PatternAction
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/ast/action/PatternAction.kt

@@ -1,11 +1,8 @@
-package com.phodal.shirelang.compiler.ast.action
+package cc.unitmesh.devti.language.ast.action
 
 import cc.unitmesh.devti.language.ast.FrontMatterType
-import cc.unitmesh.devti.language.ast.action.DirectAction
-import cc.unitmesh.devti.language.ast.action.RuleBasedPatternAction
 import com.intellij.openapi.diagnostic.logger
 import cc.unitmesh.devti.language.ast.ShirePsiQueryStatement
-import cc.unitmesh.devti.language.ast.action.PatternActionFunc
 
 /**
  * PatternFun is a sealed class in Kotlin representing different pattern processing functions.
@@ -14,7 +11,7 @@ import cc.unitmesh.devti.language.ast.action.PatternActionFunc
  *
  * @property funcName The name of the pattern processing function.
  */
-data class PatternAction(
+class PatternAction(
     val pattern: String,
     val patternFuncs: List<PatternActionFunc>,
     val isQueryStatement: Boolean = false

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/streaming/LoggingStreamingService.kt

@@ -0,0 +1,88 @@
+package cc.unitmesh.devti.language.compiler.streaming
+
+import cc.unitmesh.devti.gui.chat.message.ChatRole
+import cc.unitmesh.devti.language.LLM_LOGGING
+import cc.unitmesh.devti.language.LLM_LOGGING_JSONL
+import cc.unitmesh.devti.language.ShireConstants
+import cc.unitmesh.devti.language.console.DevInConsoleViewBase
+import cc.unitmesh.devti.llms.custom.ChatMessage
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.WriteAction
+import com.intellij.openapi.application.runInEdt
+import com.intellij.openapi.application.runWriteAction
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.testFramework.LightVirtualFile
+import kotlinx.serialization.encodeToString
+import kotlinx.serialization.json.Json
+import java.io.File
+
+/**
+ * The `LoggingStreamingService` class is an implementation of the `StreamingServiceProvider` interface.
+ * It provides functionality to log streaming data to a file within a project's directory.
+ *
+ * ### Properties:
+ * - `name`: A string that represents the name of the streaming service, initialized to "logging".
+ * - `result`: A private string that accumulates the streaming data received.
+ */
+class LoggingStreamingService : StreamingServiceProvider {
+    private var outputDir: VirtualFile = LightVirtualFile()
+    override var name: String = "logging"
+
+    private var result: String = ""
+    private var userPrompt: String = ""
+
+    override fun onBeforeStreaming(project: Project, userPrompt: String, console: DevInConsoleViewBase?) {
+        this.userPrompt = userPrompt
+        this.outputDir = ShireConstants.outputDir(project) ?: throw IllegalStateException("Project directory not found")
+        if (outputDir.findChild(LLM_LOGGING) == null) {
+            ApplicationManager.getApplication().invokeAndWait {
+                WriteAction.compute<VirtualFile, Throwable> {
+                    outputDir.createChildData(this, LLM_LOGGING)
+                }
+            }
+        } else {
+            runInEdt {
+                val file = outputDir.findChild(LLM_LOGGING)
+                runWriteAction {
+                    file?.setBinaryContent(ByteArray(0))
+                }
+            }
+        }
+
+        if (outputDir.findChild(LLM_LOGGING_JSONL) == null) {
+            ApplicationManager.getApplication().invokeAndWait {
+                WriteAction.compute<VirtualFile, Throwable> {
+                    outputDir.createChildData(this, LLM_LOGGING_JSONL)
+                }
+            }
+        }
+    }
+
+    override fun onStreaming(project: Project, flow: String, args: List<Any>) {
+        result += flow
+
+        val virtualFile = outputDir.findChild(LLM_LOGGING)
+        val file = virtualFile?.path?.let { File(it) }
+        file?.appendText(flow)
+    }
+
+    override fun afterStreamingDone(project: Project) {
+        ApplicationManager.getApplication().invokeAndWait {
+            WriteAction.compute<VirtualFile, Throwable> {
+                val virtualFile = outputDir.createChildData(this, LLM_LOGGING_JSONL)
+                val file = File(virtualFile.path)
+                val value: List<ChatMessage> = listOf(
+                    ChatMessage(ChatRole.User.name, userPrompt),
+                    ChatMessage(ChatRole.System.name, result)
+                )
+
+                val result = Json.encodeToString<List<ChatMessage>>(value)
+
+                file.appendText(result)
+                file.appendText("\n")
+                virtualFile
+            }
+        }
+    }
+}

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/streaming/OnStreamingService.kt

@@ -0,0 +1,75 @@
+package cc.unitmesh.devti.language.compiler.streaming
+
+import cc.unitmesh.devti.AutoDevNotifications
+import cc.unitmesh.devti.language.console.DevInConsoleViewBase
+import cc.unitmesh.devti.language.middleware.post.LifecycleProcessorSignature
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.project.Project
+
+
+/**
+ * The OnStreamingService class is responsible for managing all the [StreamingServiceProvider] instances related to streaming services.
+ * It offers methods for registering, clearing, and initiating streaming services within the application.
+ *
+ * This class is annotated with the @Service annotation at the project level, indicating its role in the service management infrastructure.
+ *
+ * The class maintains a mutable map to associate [LifecycleProcessorSignature] objects with corresponding [StreamingServiceProvider] instances.
+ * It also holds an optional reference to a console view object that can be used for outputting information to the user.
+ */
+@Service(Service.Level.PROJECT)
+class OnStreamingService {
+    val map = mutableMapOf<LifecycleProcessorSignature, StreamingServiceProvider>()
+    var console: DevInConsoleViewBase? = null
+
+    fun registerStreamingService(sign: LifecycleProcessorSignature, console: DevInConsoleViewBase?) {
+        this.console = console
+        val streamingService = StreamingServiceProvider.Companion.getStreamingService(sign.funcName)
+        if (streamingService != null) {
+            map[sign] = streamingService
+            streamingService.onCreated(console)
+        }
+    }
+
+    fun clearStreamingService() {
+        map.clear()
+    }
+
+    fun all(): List<StreamingServiceProvider> {
+        return StreamingServiceProvider.Companion.all()
+    }
+
+    fun onStart(project: Project, userPrompt: String) {
+        map.forEach { (_, service) ->
+            try {
+                service.onBeforeStreaming(project, userPrompt, console)
+            } catch (e: Exception) {
+                AutoDevNotifications.error(project, "Error on start streaming service: ${e.message}")
+            }
+        }
+    }
+
+    fun onStreaming(project: Project, chunk: String) {
+        map.forEach { (sign, service) ->
+            try {
+                service.
+                onStreaming(project, chunk, sign.args)
+            } catch (e: Exception) {
+                AutoDevNotifications.error(project, "Error on streaming service: ${e.message}")
+            }
+        }
+    }
+
+    fun onDone(project: Project) {
+        map.forEach { (_, service) ->
+            try {
+                service.afterStreamingDone(project)
+            } catch (e: Exception) {
+                AutoDevNotifications.error(project, "Error on done streaming service: ${e.message}")
+            }
+        }
+    }
+
+    fun onStreamingError() {
+        // todo
+    }
+}

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/streaming/ProfilingStreamingService.kt

@@ -0,0 +1,39 @@
+package cc.unitmesh.devti.language.compiler.streaming
+
+import cc.unitmesh.devti.AutoDevNotifications
+import cc.unitmesh.devti.language.console.DevInConsoleViewBase
+import com.intellij.execution.ui.ConsoleViewContentType
+import com.intellij.openapi.project.Project
+import com.intellij.util.io.IOUtil
+
+
+/**
+ * The ProfilingStreamingService class is a concrete implementation of the StreamingServiceProvider interface.
+ * It provides profiling capabilities during the streaming process, outputting memory usage information to the console.
+ */
+class ProfilingStreamingService : StreamingServiceProvider {
+    override var name: String = "profiling"
+    private var console: DevInConsoleViewBase? = null
+
+    override fun onBeforeStreaming(project: Project, userPrompt: String, console: DevInConsoleViewBase?) {
+        this.console = console
+        console?.print("Start profiling: ${getMemory()}", ConsoleViewContentType.SYSTEM_OUTPUT)
+
+    }
+
+    override fun afterStreamingDone(project: Project) {
+        console?.print("End profiling: ${getMemory()}", ConsoleViewContentType.SYSTEM_OUTPUT)
+        AutoDevNotifications.warn(project, "Memory: ${getMemory()}MB")
+    }
+
+    private fun getMemory(): Long {
+        val runtime = Runtime.getRuntime()
+        val allocatedMem = runtime.totalMemory()
+        val usedMem = allocatedMem - runtime.freeMemory()
+        return toMb(usedMem)
+    }
+
+    private fun toMb(value: Long): Long {
+        return value / IOUtil.MiB
+    }
+}

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/streaming/StreamingServiceProvider.kt

@@ -0,0 +1,60 @@
+package cc.unitmesh.devti.language.compiler.streaming
+
+import cc.unitmesh.devti.language.console.DevInConsoleViewBase
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.project.Project
+
+/**
+ * all - Returns a list of all registered StreamingServiceProvider implementations.
+ *
+ * @return A list of StreamingServiceProvider instances.
+ */
+interface StreamingServiceProvider : Disposable {
+    var name: String
+
+    /**
+     * When create the service, you can do some initialization here, like start timer, etc.
+     */
+    fun onCreated(console: DevInConsoleViewBase?) {
+        /// do nothing
+    }
+
+    /**
+     * For the start of the LLM streaming, you can do some initialization here, for example, you can create a file to log the data
+     */
+    fun onBeforeStreaming(project: Project, userPrompt: String, console: DevInConsoleViewBase?) {
+        /// do nothing
+    }
+
+    /**
+     * For the streaming data, you can do some processing here, for example, you can log the data to a file
+     */
+    fun onStreaming(project: Project, flow: String, args: List<Any>) {
+        /// do nothing
+    }
+
+    /**
+     * For the end of the streaming, for example, you can do some cleanup here, or show some notification
+     */
+    fun afterStreamingDone(project: Project) {
+        /// do nothing
+    }
+
+    override fun dispose() {
+        /// do nothing
+    }
+
+    companion object {
+        val EP_NAME =
+            ExtensionPointName.create<StreamingServiceProvider>("com.phodal.shireStreamingService")
+
+        fun getStreamingService(name: String): StreamingServiceProvider? {
+            return EP_NAME.extensions.firstOrNull { it.name == name }
+        }
+
+        fun all(): List<StreamingServiceProvider> {
+            return EP_NAME.extensions.toList()
+        }
+    }
+}

exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/streaming/TimingStreamingService.kt

@@ -0,0 +1,34 @@
+package cc.unitmesh.devti.language.compiler.streaming
+
+import cc.unitmesh.devti.AutoDevNotifications
+import cc.unitmesh.devti.language.console.DevInConsoleViewBase
+import com.intellij.execution.ui.ConsoleViewContentType
+import com.intellij.openapi.project.Project
+
+/**
+ * Logging start time and end time for each lifecycle
+ */
+class TimingStreamingService : StreamingServiceProvider {
+    override var name: String = "timing"
+
+    private var time: Long = 0
+    private var console: DevInConsoleViewBase? = null
+
+    override fun onCreated(console: DevInConsoleViewBase?) {
+        this.console = console
+        val currentTime = System.currentTimeMillis()
+        time = currentTime
+        // new line
+        console?.print("\n", ConsoleViewContentType.SYSTEM_OUTPUT)
+        console?.print("Start timing: $currentTime \n", ConsoleViewContentType.SYSTEM_OUTPUT)
+    }
+
+    override fun afterStreamingDone(project: Project) {
+        val currentTime = System.currentTimeMillis()
+
+        console?.print("\n", ConsoleViewContentType.SYSTEM_OUTPUT)
+        console?.print("End timing: $currentTime \n", ConsoleViewContentType.SYSTEM_OUTPUT)
+
+        AutoDevNotifications.warn(project, "Timing: ${currentTime - time}ms")
+    }
+}