feat(prompting): add PromptOptimizer to trim code spaces #317

Introduce `PromptOptimizer` to remove leading/trailing spaces and empty lines from code prompts. Integrate it into `CustomLLMProvider` for optional prompt trimming based on settings. Include comprehensive unit tests for the new functionality.

Author: phodal

core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt

@@ -2,6 +2,7 @@ package cc.unitmesh.devti.llms.custom
 
 import cc.unitmesh.devti.gui.chat.message.ChatRole
 import cc.unitmesh.devti.llms.LLMProvider
+import cc.unitmesh.devti.prompting.optimizer.PromptOptimizer
 import cc.unitmesh.devti.settings.AutoDevSettingsState
 import cc.unitmesh.devti.settings.coder.coderSetting
 import com.intellij.openapi.diagnostic.logger
@@ -92,7 +93,13 @@ class CustomLLMProvider(val project: Project) : LLMProvider, CustomSSEProcessor(
     }
 
     fun prompt(instruction: String, input: String): String {
-        messages += Message("user", instruction)
+        val prompt = if (project.coderSetting.state.trimCodeBeforeSend) {
+            PromptOptimizer.trimCodeSpace(instruction)
+        } else {
+            instruction
+        }
+
+        messages += Message("user", prompt)
         val customRequest = CustomRequest(messages)
         val requestContent = Json.encodeToString<CustomRequest>(customRequest)
 

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

@@ -0,0 +1,16 @@
+package cc.unitmesh.devti.prompting.optimizer
+
+object PromptOptimizer {
+    /**
+     * 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
+        return prompt.lines()
+            .filter { it.isNotBlank() }
+            .joinToString("\n") { it.trim() }
+    }
+}

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

@@ -0,0 +1,180 @@
+package cc.unitmesh.devti.prompting.optimizer
+
+import org.junit.Test
+import org.assertj.core.api.Assertions.assertThat
+import org.intellij.lang.annotations.Language
+
+class PromptOptimizerTest {
+    @Test
+    fun should_trim_leading_and_trailing_spaces_from_each_line() {
+        // given
+        val input = """
+            // Leading space in first line
+             // Leading tab in second line
+            Third line with trailing space 
+            Fourth line with trailing tab	
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(input)
+
+        // then
+        val expected = """
+            // Leading space in first line
+            // Leading tab in second line
+            Third line with trailing space
+            Fourth line with trailing tab
+        """.trimIndent()
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun should_remove_empty_lines_from_the_prompt() {
+        // given
+        val input = """
+            First line
+            
+            Second line
+            
+            Third line
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(input)
+
+        // then
+        val expected = """
+            First line
+            Second line
+            Third line
+        """.trimIndent()
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun should_handle_prompt_with_only_empty_lines() {
+        // given
+        val input = """
+            
+            
+            
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(input)
+
+        // then
+        val expected = ""
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun should_return_original_prompt_if_no_spaces_or_empty_lines() {
+        // given
+        val input = """
+            First line
+            Second line
+            Third line
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(input)
+
+        // then
+        val expected = """
+            First line
+            Second line
+            Third line
+        """.trimIndent()
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun should_handle_for_rust_code_in_issue() {
+        @Language("Rust")
+        val code = """
+use crate::{find_target, Plot};
+use anyhow::{Context, Result};
+use std::{env, process};
+
+impl Plot {
+    pub fn generate_plot(&mut self) -> Result<(), anyhow::Error> {
+        eprintln!("Generating plot");
+
+        self.target =
+            find_target(&self.target).context("⚠️  couldn't find the target for plotting")?;
+
+        // The cargo executable
+        let cargo = env::var("CARGO").unwrap_or_else(|_| String::from("cargo"));
+
+        let fuzzer_data_dir = format!(
+            "{}/{}/afl/{}/",
+            &self.ziggy_output.display(),
+            &self.target,
+            &self.input
+        );
+
+        let plot_dir = self
+            .output
+            .display()
+            .to_string()
+            .replace("{ziggy_output}", &self.ziggy_output.display().to_string())
+            .replace("{target_name}", &self.target);
+        println!("{plot_dir}");
+        println!("{}", self.target);
+
+        // We run the afl-plot command
+        process::Command::new(cargo)
+            .args(["afl", "plot", &fuzzer_data_dir, &plot_dir])
+            .spawn()
+            .context("⚠️  couldn't spawn afl plot")?
+            .wait()
+            .context("⚠️  couldn't wait for the afl plot")?;
+
+        Ok(())
+    }
+}
+        """.trimIndent()
+
+        // when
+        val result = PromptOptimizer.trimCodeSpace(code)
+
+        // then
+        val expected = """
+use crate::{find_target, Plot};
+use anyhow::{Context, Result};
+use std::{env, process};
+impl Plot {
+pub fn generate_plot(&mut self) -> Result<(), anyhow::Error> {
+eprintln!("Generating plot");
+self.target =
+find_target(&self.target).context("⚠️  couldn't find the target for plotting")?;
+// The cargo executable
+let cargo = env::var("CARGO").unwrap_or_else(|_| String::from("cargo"));
+let fuzzer_data_dir = format!(
+"{}/{}/afl/{}/",
+&self.ziggy_output.display(),
+&self.target,
+&self.input
+);
+let plot_dir = self
+.output
+.display()
+.to_string()
+.replace("{ziggy_output}", &self.ziggy_output.display().to_string())
+.replace("{target_name}", &self.target);
+println!("{plot_dir}");
+println!("{}", self.target);
+// We run the afl-plot command
+process::Command::new(cargo)
+.args(["afl", "plot", &fuzzer_data_dir, &plot_dir])
+.spawn()
+.context("⚠️  couldn't spawn afl plot")?
+.wait()
+.context("⚠️  couldn't wait for the afl plot")?;
+Ok(())
+}
+}
+"""
+    }
+}