feat(writing): add prompt type for custom action prompt

unit-mesh · Nov 2, 2023 · 0bae8b0 · 0bae8b0
1 parent 0c16166
commit 0bae8b0
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 10 deletions.
diff --git a/...esh/devti/custom/team/TeamActionPrompt.kt → ...h/devti/custom/team/CustomActionPrompt.kt b/...esh/devti/custom/team/TeamActionPrompt.kt → ...h/devti/custom/team/CustomActionPrompt.kt
@@ -5,15 +5,21 @@ import cc.unitmesh.template.TemplateRoleSplitter
 import com.intellij.openapi.diagnostic.logger
 import org.yaml.snakeyaml.Yaml
 
-data class TeamActionPrompt(
+enum class CustomActionType {
+    Default,
+    QuickAction,
+}
+
+data class CustomActionPrompt(
     var interaction: InteractionType = InteractionType.AppendCursorStream,
     var priority: Int = 0,
+    var type: CustomActionType = CustomActionType.Default,
     var other: Map<String, Any> = mapOf(),
     // the rest of the content is the chat messages
     var msgs: List<LlmMsg.ChatMessage> = listOf(),
 ) {
     companion object {
-        val logger = logger<TeamActionPrompt>()
+        val logger = logger<CustomActionPrompt>()
 
         /**
          * Parses the given content and returns a TeamActionPrompt object.
@@ -61,11 +67,11 @@ data class TeamActionPrompt(
          * )
          * ```
          */
-        fun fromContent(content: String): TeamActionPrompt {
+        fun fromContent(content: String): CustomActionPrompt {
             val regex = """^---\s*\n(.*?)\n---\s*\n(.*)$""".toRegex(RegexOption.DOT_MATCHES_ALL)
             val matchResult = regex.find(content)
 
-            val prompt = TeamActionPrompt()
+            val prompt = CustomActionPrompt()
             if (matchResult != null) {
                 val frontMatter = matchResult.groupValues[1]
                 val yaml = Yaml()
@@ -82,7 +88,15 @@ data class TeamActionPrompt(
                     0
                 }
 
-                prompt.other = frontMatterMap.filterKeys { it != "interaction" && it != "priority" }
+                prompt.type = try {
+                    CustomActionType.valueOf(frontMatterMap["type"] as String)
+                } catch (e: Exception) {
+                    CustomActionType.Default
+                }
+
+                prompt.other = frontMatterMap.filterKeys {
+                    it != "interaction" && it != "priority" && it != "type"
+                }
 
                 val chatContent = matchResult?.groupValues?.get(2) ?: content
                 prompt.msgs = parseChatMessage(chatContent)

diff --git a/src/main/kotlin/cc/unitmesh/devti/custom/team/TeamPromptsBuilder.kt b/src/main/kotlin/cc/unitmesh/devti/custom/team/TeamPromptsBuilder.kt
@@ -18,7 +18,7 @@ class TeamPromptsBuilder(private val project: Project) {
             val promptName = it.nameWithoutExtension.replace("-", " ")
             // load content of the prompt file
             val promptContent = runReadAction { it.inputStream.readBytes().toString(Charsets.UTF_8) }
-            val actionPrompt = TeamActionPrompt.fromContent(promptContent)
+            val actionPrompt = CustomActionPrompt.fromContent(promptContent)
 
             TeamPromptAction(promptName, actionPrompt)
         }
@@ -27,5 +27,5 @@ class TeamPromptsBuilder(private val project: Project) {
 
 data class TeamPromptAction(
     val actionName: String,
-    val actionPrompt: TeamActionPrompt = TeamActionPrompt(),
+    val actionPrompt: CustomActionPrompt = CustomActionPrompt(),
 )
diff --git a/...devti/custom/team/TeamActionPromptTest.kt → ...vti/custom/team/CustomActionPromptTest.kt b/...devti/custom/team/TeamActionPromptTest.kt → ...vti/custom/team/CustomActionPromptTest.kt
@@ -5,7 +5,7 @@ import io.kotest.matchers.shouldBe
 import junit.framework.TestCase.assertEquals
 import org.junit.Test
 
-class TeamActionPromptTest {
+class CustomActionPromptTest {
 
     @Test
     fun `should create TeamActionPrompt object from content with frontmatter`() {
@@ -24,7 +24,7 @@ class TeamActionPromptTest {
         """.trimIndent()
 
         // when
-        val prompt = TeamActionPrompt.fromContent(content)
+        val prompt = CustomActionPrompt.fromContent(content)
 
         // then
         assertEquals(InteractionType.AppendCursorStream, prompt.interaction)
@@ -47,7 +47,7 @@ class TeamActionPromptTest {
         """.trimIndent()
 
         // when
-        val prompt = TeamActionPrompt.fromContent(content)
+        val prompt = CustomActionPrompt.fromContent(content)
 
         // then
         assertEquals(InteractionType.AppendCursorStream, prompt.interaction)
@@ -58,4 +58,29 @@ class TeamActionPromptTest {
             LlmMsg.ChatMessage(LlmMsg.ChatRole.User, "Chat message 2\n", null),
         )
     }
+
+    @Test
+    fun `should_handle_for_notion_like_summarize`() {
+        val content = """
+            ---
+            type: QuickAction
+            name: Summarize
+            category: Generate
+            interaction: AppendCursorStream
+            ---
+
+            ```System```
+
+            You are an assistant helping summarize a document. Use this format, replacing text in brackets with the result. Do not include the brackets in the output: 
+
+            Summary in [Identified language of the document]: 
+
+            [One-paragaph summary of the document using the identified language.].
+
+            ```User```
+        """.trimIndent()
+
+        val prompt = CustomActionPrompt.fromContent(content)
+        prompt.type shouldBe CustomActionType.QuickAction
+    }
 }