refactor: improve AI response content handling (#1214)

Author: tangcent

idea-plugin/src/main/kotlin/com/itangcent/ai/DeepSeekService.kt

@@ -8,7 +8,6 @@ import com.itangcent.http.RawContentType
 import com.itangcent.http.contentType
 import com.itangcent.idea.plugin.condition.ConditionOnSetting
 import com.itangcent.idea.plugin.settings.helper.AISettingsHelper
-import com.itangcent.idea.plugin.utils.AIUtils
 import com.itangcent.intellij.extend.sub
 import com.itangcent.intellij.logger.Logger
 import com.itangcent.suv.http.HttpClientProvider
@@ -87,8 +86,7 @@ open class DeepSeekService : AIService {
             val jsonElement = GsonUtils.parseToJsonTree(responseBody)
             val content = jsonElement.sub("choices")?.asJsonArray?.firstOrNull()
                 ?.asJsonObject?.sub("message")?.sub("content")?.asString
-            return content?.let { AIUtils.cleanMarkdownCodeBlocks(it) }
-                ?: throw AIApiException("Could not parse response from DeepSeek API")
+            return content ?: throw AIApiException("Could not parse response from DeepSeek API")
         } catch (e: AIException) {
             // Re-throw AI exceptions
             throw e

idea-plugin/src/main/kotlin/com/itangcent/ai/LocalLLMClient.kt

@@ -4,7 +4,6 @@ import com.itangcent.common.utils.GsonUtils
 import com.itangcent.http.HttpClient
 import com.itangcent.http.RawContentType
 import com.itangcent.http.contentType
-import com.itangcent.idea.plugin.utils.AIUtils
 import com.itangcent.intellij.extend.sub
 
 /**
@@ -61,8 +60,7 @@ class LocalLLMClient(
             val content = jsonElement.sub("choices")?.asJsonArray?.firstOrNull()
                 ?.asJsonObject?.sub("message")?.sub("content")?.asString
             val errMsg = jsonElement.sub("error")?.asString
-            return content?.let { AIUtils.cleanMarkdownCodeBlocks(it) }
-                ?: throw AIApiException(errMsg ?: "Could not parse response from Local LLM server")
+            return content ?: throw AIApiException(errMsg ?: "Could not parse response from Local LLM server")
         } catch (e: AIException) {
             throw e
         } catch (e: Exception) {

idea-plugin/src/main/kotlin/com/itangcent/ai/OpenAIService.kt

@@ -6,7 +6,6 @@ import com.itangcent.common.logger.traceError
 import com.itangcent.idea.plugin.condition.ConditionOnSetting
 import com.itangcent.idea.plugin.settings.helper.AISettingsHelper
 import com.itangcent.idea.plugin.settings.helper.HttpSettingsHelper
-import com.itangcent.idea.plugin.utils.AIUtils
 import com.itangcent.intellij.logger.Logger
 import com.openai.client.OpenAIClient
 import com.openai.client.okhttp.OpenAIOkHttpClient
@@ -64,8 +63,7 @@ open class OpenAIService : AIService {
             val content = response.choices().firstOrNull()?.message()?.content()
                 ?.orElse("")
 
-            return content?.let { AIUtils.cleanMarkdownCodeBlocks(it) }
-                ?: throw AIApiException("Empty response from OpenAI API")
+            return content ?: throw AIApiException("Empty response from OpenAI API")
         } catch (e: OpenAIException) {
             logger.traceError("OpenAI API error: ${e.message}", e)
             throw AIApiException("OpenAI API error: ${e.message}", e)

idea-plugin/src/main/kotlin/com/itangcent/idea/plugin/api/export/translation/APITranslationHelper.kt

@@ -12,6 +12,7 @@ import com.itangcent.common.model.Request
 import com.itangcent.common.utils.GsonUtils
 import com.itangcent.common.utils.notNullOrEmpty
 import com.itangcent.idea.plugin.settings.helper.AISettingsHelper
+import com.itangcent.idea.plugin.utils.AIUtils
 import com.itangcent.intellij.logger.Logger
 import java.util.concurrent.ConcurrentHashMap
 
@@ -375,10 +376,10 @@ class APITranslationHelper {
                     "Return only the translated text without any explanations or additional formatting."
 
             // Send translation request to AI
-            val translatedText = aiService.sendPrompt(systemMessage, text)
+            val translatedText = AIUtils.getGeneralContent(aiService.sendPrompt(systemMessage, text))
 
             // Cache the result
-            if (translatedText.notNullOrEmpty()) {
+            if (translatedText.isNotEmpty()) {
                 translationCache[cacheKey] = translatedText
                 return translatedText
             }
@@ -434,10 +435,10 @@ class APITranslationHelper {
             """.trimIndent()
 
             // Send translation request to AI
-            val translatedJson = aiService.sendPrompt(systemMessage, jsonString)
+            val translatedJson = AIUtils.getGeneralContent(aiService.sendPrompt(systemMessage, jsonString))
 
             // Cache the result if it's valid JSON
-            if (translatedJson.notNullOrEmpty() && isValidJson(translatedJson)) {
+            if (translatedJson.isNotEmpty() && isValidJson(translatedJson)) {
                 return translatedJson
             }
 

idea-plugin/src/main/kotlin/com/itangcent/idea/plugin/api/infer/AIMethodInferHelper.kt

@@ -12,6 +12,7 @@ import com.itangcent.cache.withoutCache
 import com.itangcent.common.logger.Log
 import com.itangcent.common.logger.traceError
 import com.itangcent.common.utils.GsonUtils
+import com.itangcent.idea.plugin.utils.AIUtils
 import com.itangcent.intellij.context.ActionContext
 import com.itangcent.intellij.context.ThreadFlag
 import com.itangcent.intellij.extend.isNotActive
@@ -371,14 +372,18 @@ class AIMethodInferHelper : MethodInferHelper {
                 // Call the AI API with the system message and prompt
                 val aiResponse = if (currentRetry > 0) {
                     cacheSwitcher.withoutCache {
-                        aiService.sendPrompt(AIPromptFormatter.METHOD_RETURN_TYPE_INFERENCE_MESSAGE, prompt)
+                        AIUtils.getGeneralContent(
+                            aiService.sendPrompt(AIPromptFormatter.METHOD_RETURN_TYPE_INFERENCE_MESSAGE, prompt)
+                        )
                     }
                 } else {
-                    aiService.sendPrompt(AIPromptFormatter.METHOD_RETURN_TYPE_INFERENCE_MESSAGE, prompt)
+                    AIUtils.getGeneralContent(
+                        aiService.sendPrompt(AIPromptFormatter.METHOD_RETURN_TYPE_INFERENCE_MESSAGE, prompt)
+                    )
                 }
 
                 // Check if the response is valid before parsing
-                if (aiResponse.isBlank()) {
+                if (aiResponse.isEmpty()) {
                     logger.warn("Empty AI response received for method ${methodInfo.className}.${methodInfo.methodName}, attempt ${currentRetry + 1}/$maxRetries")
                     currentRetry++
                     continue

idea-plugin/src/main/kotlin/com/itangcent/idea/plugin/utils/AIUtils.kt

@@ -13,19 +13,66 @@ object AIUtils {
      */
     fun cleanMarkdownCodeBlocks(content: String): String {
         val trimmedContent = content.trim()
-        
+
         // Check if content is surrounded by code block delimiters
         if (trimmedContent.startsWith("```") && trimmedContent.endsWith("```")) {
             // Remove starting delimiter (with optional language identifier)
             val withoutStart = trimmedContent.replaceFirst(Regex("^```\\w*", RegexOption.MULTILINE), "")
-            
+
             // Remove ending delimiter - handles both with and without newline
             val result = withoutStart.replace(Regex("(\\n)?```$", RegexOption.MULTILINE), "")
-            
+
             // Trim any leading/trailing whitespace
             return result.trim()
         }
-        
+
         return trimmedContent
     }
+
+    /**
+     * Extracts the first code block from content, including the language identifier if present
+     *
+     * @param content The content that may contain one or more code blocks
+     * @return The first code block found, or null if no code block is found
+     */
+    fun extractFirstCodeBlock(content: String): String? {
+        // First try to find a multi-line code block
+        val multiLinePattern = Regex("^```\\w*\\n([\\s\\S]*?)```$", RegexOption.MULTILINE)
+        val multiLineMatch = multiLinePattern.find(content)
+        if (multiLineMatch != null) {
+            return multiLineMatch.groupValues[1].trim()
+        }
+
+        // If no multi-line block found, try single-line
+        val singleLinePattern = Regex("```([^\\n]+?)```", RegexOption.MULTILINE)
+        val singleLineMatch = singleLinePattern.find(content)
+        if (singleLineMatch != null) {
+            return singleLineMatch.groupValues[1].trim()
+        }
+
+        return null
+    }
+
+    /**
+     * Extracts the first code block of a specific language type from content
+     *
+     * @param content The content that may contain one or more code blocks
+     * @param languageType The specific language type to look for (e.g., "json", "java", "kotlin")
+     * @return The first code block found with the specified language type, or null if no matching code block is found
+     */
+    fun extractFirstCodeBlockByType(content: String, languageType: String): String? {
+        val codeBlockPattern = Regex("^```$languageType\\n([\\s\\S]*?)```$", RegexOption.MULTILINE)
+        return codeBlockPattern.find(content)?.groupValues?.get(1)?.trim()
+    }
+
+    /**
+     * Gets the general content by first trying to extract a code block.
+     * If no code block is found, returns the input content.
+     *
+     * @param content The content that may contain code blocks
+     * @return The extracted code block content if found, otherwise the original content
+     */
+    fun getGeneralContent(content: String): String {
+        return extractFirstCodeBlock(content) ?: content
+    }
 }