|
| 1 | +package cc.unitmesh.devti.gui.toolbar |
| 2 | + |
| 3 | +import cc.unitmesh.devti.AutoDevIcons |
| 4 | +import cc.unitmesh.devti.indexer.DomainDictGenerateAction |
| 5 | +import cc.unitmesh.devti.indexer.DomainDictGenerateContext |
| 6 | +import cc.unitmesh.devti.indexer.provider.LangDictProvider |
| 7 | +import cc.unitmesh.devti.indexer.usage.PromptEnhancer |
| 8 | +import cc.unitmesh.devti.llms.LlmFactory |
| 9 | +import cc.unitmesh.devti.observer.agent.AgentStateService |
| 10 | +import cc.unitmesh.devti.settings.coder.coderSetting |
| 11 | +import cc.unitmesh.devti.statusbar.AutoDevStatus |
| 12 | +import cc.unitmesh.devti.statusbar.AutoDevStatusService |
| 13 | +import cc.unitmesh.devti.template.GENIUS_CODE |
| 14 | +import cc.unitmesh.devti.template.TemplateRender |
| 15 | +import cc.unitmesh.devti.util.AutoDevCoroutineScope |
| 16 | +import com.intellij.icons.AllIcons |
| 17 | +import com.intellij.openapi.actionSystem.* |
| 18 | +import com.intellij.openapi.actionSystem.ex.CustomComponentAction |
| 19 | +import com.intellij.openapi.application.ApplicationManager |
| 20 | +import com.intellij.openapi.command.WriteCommandAction |
| 21 | +import com.intellij.openapi.diagnostic.logger |
| 22 | +import com.intellij.openapi.editor.ScrollType |
| 23 | +import com.intellij.openapi.fileEditor.FileEditorManager |
| 24 | +import com.intellij.openapi.project.Project |
| 25 | +import com.intellij.openapi.project.guessProjectDir |
| 26 | +import com.intellij.openapi.vfs.LocalFileSystem |
| 27 | +import com.intellij.ui.components.panels.Wrapper |
| 28 | +import com.intellij.util.ui.JBInsets |
| 29 | +import com.intellij.util.ui.JBUI |
| 30 | +import kotlinx.coroutines.flow.Flow |
| 31 | +import kotlinx.coroutines.flow.cancellable |
| 32 | +import kotlinx.coroutines.launch |
| 33 | +import javax.swing.JButton |
| 34 | +import javax.swing.JComponent |
| 35 | +import kotlin.io.path.createDirectories |
| 36 | +import kotlin.io.path.exists |
| 37 | + |
| 38 | +class SummaryMessagesAction : AnAction("Summary Messages", "Summary all current messages to memorize.md", AllIcons.Nodes.Target), |
| 39 | + CustomComponentAction { |
| 40 | + override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.EDT |
| 41 | + override fun actionPerformed(e: AnActionEvent) { |
| 42 | + val project = e.project ?: return |
| 43 | + val baseDir = project.coderSetting.state.teamPromptsDir |
| 44 | + val presentation = e.presentation |
| 45 | + |
| 46 | + copyMessages(project) |
| 47 | + |
| 48 | + AutoDevCoroutineScope.scope(project).launch { |
| 49 | + try { |
| 50 | + updatePresentation(presentation, AutoDevIcons.LOADING, false) |
| 51 | + AutoDevStatusService.notifyApplication(AutoDevStatus.InProgress) |
| 52 | + val promptDir = project.guessProjectDir()!!.toNioPath().resolve(baseDir) |
| 53 | + if (!promptDir.exists()) { |
| 54 | + promptDir.createDirectories() |
| 55 | + } |
| 56 | + |
| 57 | + val file = promptDir.resolve("memories.md").toFile() |
| 58 | + if (!file.exists()) { |
| 59 | + file.createNewFile() |
| 60 | + } |
| 61 | + |
| 62 | + val systemPrompt = buildPrompt(project) |
| 63 | + val userPrompt = copyMessages(project) |
| 64 | + |
| 65 | + val fileEditorManager = FileEditorManager.getInstance(project) |
| 66 | + ApplicationManager.getApplication().invokeAndWait { |
| 67 | + val virtualFile = LocalFileSystem.getInstance().refreshAndFindFileByIoFile(file) |
| 68 | + if (virtualFile != null) { |
| 69 | + fileEditorManager.setSelectedEditor(virtualFile, "text-editor") |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + val editor = fileEditorManager.selectedTextEditor |
| 74 | + val stream: Flow<String> = LlmFactory.create(project).stream(systemPrompt, userPrompt) |
| 75 | + val result = StringBuilder() |
| 76 | + |
| 77 | + stream.cancellable().collect { chunk -> |
| 78 | + result.append(chunk) |
| 79 | + WriteCommandAction.writeCommandAction(project).compute<Any, RuntimeException> { |
| 80 | + editor?.document?.setText(result.toString()) |
| 81 | + editor?.caretModel?.moveToOffset(editor?.document?.textLength ?: 0) |
| 82 | + editor?.scrollingModel?.scrollToCaret(ScrollType.RELATIVE) |
| 83 | + } |
| 84 | + } |
| 85 | + |
| 86 | + AutoDevStatusService.notifyApplication(AutoDevStatus.Done) |
| 87 | + } catch (e: Exception) { |
| 88 | + AutoDevStatusService.notifyApplication(AutoDevStatus.Error) |
| 89 | + e.printStackTrace() |
| 90 | + } finally { |
| 91 | + updatePresentation(presentation, AutoDevIcons.AI_COPILOT, true) |
| 92 | + } |
| 93 | + } |
| 94 | + } |
| 95 | + |
| 96 | + override fun createCustomComponent(presentation: Presentation, place: String): JComponent { |
| 97 | + val button: JButton = object : JButton() { |
| 98 | + init { |
| 99 | + putClientProperty("ActionToolbar.smallVariant", true) |
| 100 | + putClientProperty("customButtonInsets", JBInsets(1, 1, 1, 1).asUIResource()) |
| 101 | + |
| 102 | + setOpaque(false) |
| 103 | + addActionListener { |
| 104 | + copyMessages(ActionToolbar.getDataContextFor(this).getData(CommonDataKeys.PROJECT)) |
| 105 | + } |
| 106 | + } |
| 107 | + } |
| 108 | + |
| 109 | + return Wrapper(button).also { |
| 110 | + it.setBorder(JBUI.Borders.empty(0, 10)) |
| 111 | + } |
| 112 | + } |
| 113 | + |
| 114 | + private fun copyMessages(project: Project?): String { |
| 115 | + val agentStateService = project?.getService(AgentStateService::class.java) ?: return "" |
| 116 | + val allText = agentStateService.getAllMessages().joinToString("\n") { it.content } |
| 117 | + return allText |
| 118 | + } |
| 119 | + |
| 120 | + private fun updatePresentation(presentation: Presentation, icon: javax.swing.Icon, enabled: Boolean) { |
| 121 | + presentation.icon = icon |
| 122 | + presentation.isEnabled = enabled |
| 123 | + } |
| 124 | + |
| 125 | + private suspend fun buildPrompt(project: Project): String { |
| 126 | + val names = LangDictProvider.all(project) |
| 127 | + val templateRender = TemplateRender(GENIUS_CODE) |
| 128 | + val template = templateRender.getTemplate("memory.vm") |
| 129 | + val readmeMe = PromptEnhancer.readmeFile(project) |
| 130 | + |
| 131 | + val context = DomainDictGenerateContext(names.joinToString(", "), readmeMe) |
| 132 | + val prompt = templateRender.renderTemplate(template, context) |
| 133 | + return prompt |
| 134 | + } |
| 135 | +} |
0 commit comments