fix: ensure LLM outputs only code, not plain text

Improved prompts and code extraction to ensure generated code: - Matches the file type and existing code structure - Contains NO markdown code blocks (```) - Contains NO explanations or comments about what was done - Is ready to be directly inserted into the target file Changes: - Enhanced system prompts with stricter "ABSOLUTE RULES" - Added explicit instructions about no markdown in user prompts - Improved extract_code() to strip markdown and explanations - Added processing notification with prompt preview
2026-01-11 15:31:23 -05:00
parent bba0647b47
commit fe04e624db
3 changed files with 109 additions and 47 deletions
--- a/lua/codetyper/commands.lua
+++ b/lua/codetyper/commands.lua
@@ -88,6 +88,23 @@ local function cmd_toggle()
  window.toggle_split(target_path, coder_path)
 end

+--- Build enhanced user prompt with context
+---@param clean_prompt string The cleaned user prompt
+---@param context table Context information
+---@return string Enhanced prompt
+local function build_user_prompt(clean_prompt, context)
+  local enhanced = "TASK: " .. clean_prompt .. "\n\n"
+  
+  enhanced = enhanced .. "REQUIREMENTS:\n"
+  enhanced = enhanced .. "- Generate ONLY " .. (context.language or "code") .. " code\n"
+  enhanced = enhanced .. "- NO markdown code blocks (no ```)\n"
+  enhanced = enhanced .. "- NO explanations or comments about what you did\n"
+  enhanced = enhanced .. "- Match the coding style of the existing file exactly\n"
+  enhanced = enhanced .. "- Output must be ready to insert directly into the file\n"
+  
+  return enhanced
+end
+
 --- Process prompt at cursor and generate code
 local function cmd_process()
  local parser = require("codetyper.parser")
@@ -111,8 +128,13 @@ local function cmd_process()
  local prompt_type = parser.detect_prompt_type(prompt.content)
  local context = llm.build_context(target_path, prompt_type)
  local clean_prompt = parser.clean_prompt(prompt.content)
+  
+  -- Build enhanced prompt with explicit instructions
+  local enhanced_prompt = build_user_prompt(clean_prompt, context)

-  llm.generate(clean_prompt, context, function(response, err)
+  utils.notify("Processing: " .. clean_prompt:sub(1, 50) .. "...", vim.log.levels.INFO)
+
+  llm.generate(enhanced_prompt, context, function(response, err)
    if err then
      utils.notify("Generation failed: " .. err, vim.log.levels.ERROR)
      return
@@ -122,6 +144,7 @@ local function cmd_process()
      -- Inject code into target file
      local inject = require("codetyper.inject")
      inject.inject_code(target_path, response, prompt_type)
+      utils.notify("Code generated and injected!", vim.log.levels.INFO)
    end
  end)
 end
--- a/lua/codetyper/llm/init.lua
+++ b/lua/codetyper/llm/init.lua
@@ -44,8 +44,14 @@ function M.build_system_prompt(context)
  system = system:gsub("{{language}}", context.language or "unknown")
  system = system:gsub("{{filepath}}", context.file_path or "unknown")

-  if context.file_content then
-    system = system .. "\n\nExisting file content:\n```\n" .. context.file_content .. "\n```"
+  -- Add file content with analysis hints
+  if context.file_content and context.file_content ~= "" then
+    system = system .. "\n\n===== EXISTING FILE CONTENT (analyze and match this style) =====\n"
+    system = system .. context.file_content
+    system = system .. "\n===== END OF EXISTING FILE =====\n"
+    system = system .. "\nYour generated code MUST follow the exact patterns shown above."
+  else
+    system = system .. "\n\nThis is a new/empty file. Generate clean, idiomatic " .. (context.language or "code") .. " following best practices."
  end

  return system
@@ -89,11 +95,31 @@ end
 ---@param response string Raw LLM response
 ---@return string Extracted code
 function M.extract_code(response)
-  -- Remove markdown code blocks if present
-  local code = response:gsub("```%w*\n?", ""):gsub("\n?```", "")
-
-  -- Trim whitespace
-  code = code:match("^%s*(.-)%s*$")
+  local code = response
+  
+  -- Remove markdown code blocks with language tags (```typescript, ```javascript, etc.)
+  code = code:gsub("```%w+%s*\n", "")
+  code = code:gsub("```%w+%s*$", "")
+  code = code:gsub("^```%w*\n?", "")
+  code = code:gsub("\n?```%s*$", "")
+  code = code:gsub("\n```\n", "\n")
+  code = code:gsub("```", "")
+  
+  -- Remove common explanation prefixes that LLMs sometimes add
+  code = code:gsub("^Here.-:\n", "")
+  code = code:gsub("^Here's.-:\n", "")
+  code = code:gsub("^This.-:\n", "")
+  code = code:gsub("^The following.-:\n", "")
+  code = code:gsub("^Below.-:\n", "")
+  
+  -- Remove common explanation suffixes
+  code = code:gsub("\n\nThis code.-$", "")
+  code = code:gsub("\n\nThe above.-$", "")
+  code = code:gsub("\n\nNote:.-$", "")
+  code = code:gsub("\n\nExplanation:.-$", "")
+  
+  -- Trim leading/trailing whitespace but preserve internal formatting
+  code = code:match("^%s*(.-)%s*$") or code

  return code
 end
--- a/lua/codetyper/prompts/system.lua
+++ b/lua/codetyper/prompts/system.lua
@@ -5,20 +5,28 @@
 local M = {}

 --- Base system prompt for code generation
-M.code_generation = [[You are an expert code generation assistant integrated into Neovim via Codetyper.nvim.
-Your task is to generate high-quality, production-ready code based on the user's prompt.
+M.code_generation = [[You are an expert code generation assistant integrated into Neovim.
+Your task is to generate production-ready code that EXACTLY matches the style of the existing file.

-CRITICAL RULES:
-1. Output ONLY the code - no explanations, no markdown code blocks, no comments about what you did
-2. Match the coding style, conventions, and patterns of the existing file
-3. Use proper indentation and formatting for the language
-4. Follow best practices for the specific language/framework
-5. Preserve existing functionality unless explicitly asked to change it
-6. Use meaningful variable and function names
-7. Handle edge cases and errors appropriately
+ABSOLUTE RULES - FOLLOW STRICTLY:
+1. Output ONLY raw code - NO explanations, NO markdown, NO code fences (```), NO comments about what you did
+2. DO NOT wrap output in ```typescript``` or ```javascript``` or any markdown
+3. The output must be valid {{language}} code that can be directly inserted into the file
+4. MATCH the existing code patterns:
+   - Same indentation style (spaces/tabs)
+   - Same naming conventions (camelCase, snake_case, etc.)
+   - Same import style
+   - Same comment style
+   - Same function/class patterns used in the file
+5. If the file has existing exports, follow the same export pattern
+6. If the file uses certain libraries/frameworks, use the same ones
+7. Include proper TypeScript types if the file uses TypeScript
+8. Include proper error handling following the file's patterns

 Language: {{language}}
 File: {{filepath}}
+
+REMEMBER: Output ONLY the code. No markdown. No explanations. Just the code.
 ]]

 --- System prompt for code explanation/ask
@@ -38,59 +46,64 @@ IMPORTANT: When file contents are provided, analyze them carefully and base your
 ]]

 --- System prompt for refactoring
-M.refactor = [[You are an expert code refactoring assistant integrated into Neovim via Codetyper.nvim.
+M.refactor = [[You are an expert code refactoring assistant integrated into Neovim.
 Your task is to refactor code while maintaining its functionality.

-CRITICAL RULES:
-1. Output ONLY the refactored code - no explanations
-2. Preserve ALL existing functionality
-3. Improve code quality, readability, and maintainability
-4. Follow SOLID principles and best practices
-5. Keep the same coding style as the original
+ABSOLUTE RULES - FOLLOW STRICTLY:
+1. Output ONLY the refactored code - NO explanations, NO markdown, NO code fences (```)
+2. DO NOT wrap output in ```typescript``` or any markdown
+3. Preserve ALL existing functionality
+4. Improve code quality, readability, and maintainability
+5. Keep the EXACT same coding style as the original file
 6. Do not add new features unless explicitly requested
-7. Optimize performance where possible without sacrificing readability
+7. Output must be valid {{language}} code ready to replace the original

 Language: {{language}}
+
+REMEMBER: Output ONLY the code. No markdown. No explanations.
 ]]

 --- System prompt for documentation
-M.document = [[You are a documentation expert integrated into Neovim via Codetyper.nvim.
-Your task is to generate clear, comprehensive documentation for code.
+M.document = [[You are a documentation expert integrated into Neovim.
+Your task is to generate documentation comments for code.

-CRITICAL RULES:
-1. Output ONLY the documentation/comments - ready to be inserted into code
-2. Use the appropriate documentation format for the language:
-   - JavaScript/TypeScript: JSDoc
-   - Python: Docstrings (Google or NumPy style)
-   - Lua: LuaDoc/EmmyLua
+ABSOLUTE RULES - FOLLOW STRICTLY:
+1. Output ONLY the documentation comments - NO explanations, NO markdown
+2. DO NOT wrap output in code fences (```)
+3. Use the appropriate format for {{language}}:
+   - JavaScript/TypeScript: JSDoc (/** ... */)
+   - Python: Docstrings (triple quotes)
+   - Lua: LuaDoc/EmmyLua (---)
   - Go: GoDoc
-   - Rust: RustDoc
-   - Java: Javadoc
-3. Document all parameters, return values, and exceptions
-4. Include usage examples where helpful
-5. Be concise but complete
+   - Rust: RustDoc (///)
+4. Document all parameters, return values, and exceptions
+5. Output must be valid comment syntax for {{language}}

 Language: {{language}}
+
+REMEMBER: Output ONLY the documentation comments. No markdown. No explanations.
 ]]

 --- System prompt for test generation
-M.test = [[You are a test generation expert integrated into Neovim via Codetyper.nvim.
-Your task is to generate comprehensive unit tests for the provided code.
+M.test = [[You are a test generation expert integrated into Neovim.
+Your task is to generate unit tests for the provided code.

-CRITICAL RULES:
-1. Output ONLY the test code - no explanations
-2. Use the appropriate testing framework for the language:
+ABSOLUTE RULES - FOLLOW STRICTLY:
+1. Output ONLY the test code - NO explanations, NO markdown, NO code fences (```)
+2. DO NOT wrap output in ```typescript``` or any markdown
+3. Use the appropriate testing framework for {{language}}:
   - JavaScript/TypeScript: Jest or Vitest
   - Python: pytest
   - Lua: busted or plenary
   - Go: testing package
   - Rust: built-in tests
-3. Cover happy paths, edge cases, and error scenarios
-4. Use descriptive test names
+4. Cover happy paths, edge cases, and error scenarios
 5. Follow AAA pattern: Arrange, Act, Assert
-6. Mock external dependencies appropriately
+6. Output must be valid {{language}} test code

 Language: {{language}}
+
+REMEMBER: Output ONLY the test code. No markdown. No explanations.
 ]]

 return M