feat(prompting): enhance trimCodeSpace to handle Python code #316

Refactor trimCodeSpace to preserve code fences for Python and improve overall trimming logic.

Author: phodal

core/src/main/kotlin/cc/unitmesh/devti/prompting/optimizer/PromptOptimizer.kt

@@ -1,14 +1,28 @@
 package cc.unitmesh.devti.prompting.optimizer
 
+import cc.unitmesh.devti.util.parser.CodeFence
+
 object PromptOptimizer {
+    fun trimCodeSpace(prompt: String): String {
+        val fences = CodeFence.parseAll(prompt)
+        return fences.joinToString("\n") {
+            if (it.originLanguage == "python") {
+                "```${it.originLanguage}\n${it.text}\n```"
+            } else {
+                trim(it.text)
+            }
+        }
+    }
+
     /**
+     * Related to [#317](https://github.com/unit-mesh/auto-dev/issues/317)
+     *
      * Similar to the following shell command:
      * ```bash
      * grep -Ev '^[ \t]*$ input.rs | sed 's/^[ \t]*\/\/' | sed 's/[ \t]$//'
      * ```
      */
-    fun trimCodeSpace(prompt: String): String {
-        /// check language of CodeFence skip for Python
+    private fun trim(prompt: String): String {
         return prompt.lines()
             .filter { it.isNotBlank() }
             .joinToString("\n") { it.trim() }

core/src/test/kotlin/cc/unitmesh/devti/prompting/optimizer/PromptOptimizerTest.kt

@@ -140,8 +140,7 @@ impl Plot {
         val result = PromptOptimizer.trimCodeSpace(code)
 
         // then
-        val expected = """
-use crate::{find_target, Plot};
+        val expected = """use crate::{find_target, Plot};
 use anyhow::{Context, Result};
 use std::{env, process};
 impl Plot {
@@ -174,7 +173,32 @@ process::Command::new(cargo)
 .context("⚠️  couldn't wait for the afl plot")?;
 Ok(())
 }
-}
-"""
+}"""
+
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun should_handle_for_python_code_in_issue() {
+        @Language("Markdown")
+        val code = """
+Here is the code:       
+```python
+def foo():
+    print("Hello, World!")
+```
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(code)
+
+        // then
+        val expected = """Here is the code:
+```python
+def foo():
+    print("Hello, World!")
+```"""
+
+        assertThat(result).isEqualTo(expected)
     }
 }