(nvim) Add script that injects AI-generated docstring

private_dot_config/nvim/after/plugin/neoai-extended.lua

@@ -17,4 +17,54 @@ local inject_commit_message = function()
   )
 end
 
+local inject_docstring = function()
+  local docstring_formats = { 
+	["python"] =  "reST" ,
+	["javascript"] =  "jsdoc" ,
+	["typescript"] =  "jsdoc" ,
+	["typescriptreact"] =  "jsdoc" ,
+	["javascriptreact"] =  "jsdoc" ,
+	["c"] =  "doxygen" ,
+	["cpp"] =  "doxygen" ,
+	["lua"] =  "luadoc" ,
+	["rust"] =  "rustdoc" ,
+	["go"] =  "godoc" ,
+	["java"] =  "javadoc" ,
+	["php"] =  "phpdoc" ,
+	["ruby"] =  "yard" ,
+	["elixir"] =  "exdoc" ,
+	["erlang"] =  "edoc" ,
+	["clojure"] =  "codox" ,
+	["haskell"] =  "haddock" ,
+	["scala"] =  "scaladoc" ,
+	["kotlin"] =  "kdoc" ,
+	["swift"] =  "swiftdoc" ,
+	["julia"] =  "juliadoc" ,
+	["dart"] =  "dartdoc" ,
+	["vim"] =  "vimdoc" ,
+	["viml"] =  "vimdoc" ,
+  }
+
+  vim.ui.select( 
+  	{"korean", "english"},
+	{
+	  prompt = "Select language: ",
+	},
+	function(language)
+	  if language ~= nil then 
+		local filetype = vim.bo.filetype
+		local docstring_format = docstring_formats[filetype]
+		local prompt = require('utilities.prompt-engineering').generate_docstring(language, docstring_format)
+		local current_position = vim.api.nvim_win_get_cursor(0)
+		local start_line = vim.fn.line("'<")
+		vim.api.nvim_win_set_cursor(0, {start_line - 2, current_position[2]})
+		require('neoai').context_inject(prompt, nil, start_line - 1, start_line - 1)
+	  end
+	end
+  )
+end
+
 vim.api.nvim_create_user_command("InjectCommitMessage", inject_commit_message, {})
+vim.api.nvim_create_user_command("InjectDocstring", inject_docstring, { range = true })
+
+

private_dot_config/nvim/lua/utilities/prompt-engineering.lua

@@ -18,4 +18,73 @@ M.generate_commit_message = function(language)
   return prompt
 end
 
+M.generate_docstring = function(language, docstring_format)
+  local buffer = vim.api.nvim_get_current_buf()
+  local filetype = vim.api.nvim_buf_get_option(buffer, "filetype")
+  local lines = vim.api.nvim_buf_get_lines(buffer, 0, -1, false)
+
+  print(language, docstring_format);
+  local docstring_format_references = {
+    ["yard"] = "https://www.rubydoc.info/gems/yard/file/docs/Tags.md#tag-reference",
+    ["numpy"] = "https://numpydoc.readthedocs.io/en/latest/format.html",
+    ["pdoc"] = "https://pdoc3.github.io/pdoc/doc/pdoc/#docstring-formats",
+    ["restructuredtext"] = "https://www.python.org/dev/peps/pep-0287/",
+    ["javadoc"] = "https://www.oracle.com/technical-resources/articles/java/javadoc-tool.html",
+    ["phpdoc"] = "https://docs.phpdoc.org/latest/references/phpdoc/basic-syntax.html",
+    ["jsdoc"] = "https://jsdoc.app/tags-example.html",
+  }
+
+  local docstring_format_reference = docstring_format_references[docstring_format];
+
+  local additional_prompt_about_docstring_format = ""
+
+  if docstring_format_reference ~= nil then
+    additional_prompt_about_docstring_format = [[
+      And I strongly recommend you to use tag notation if it exists.
+      For more information about the docstring format, please refer to:
+      ]] .. docstring_format_reference .. [[
+    ]]
+  end
+
+  local start_line = vim.fn.line("'<")
+  local end_line = vim.fn.line("'>")
+
+  local lines_to_format = {}
+  for i = start_line, end_line do
+	table.insert(lines_to_format, lines[i])
+  end
+
+  local code_snippet = table.concat(lines_to_format, "\n")
+  local prompt = [[
+      This code snippet is written in ]] .. filetype .. [[.
+      Using the following code snippets generate a consise and clear docstring comment that consists of these:
+
+	  1. A identifier of the function/class/method (Essential)
+	  2. A short summary of the function/class/method (Essential)
+	  3. A description of the function/class/method (Essential)
+	  4. A list of parameters and their types (Essential)
+	  5. A list of return values and their types (Essential)
+	  6. A list of exceptions that can be raised (Optional)
+	  7. A list of side effects (Optional)
+	  8. A list of restrictions (Optional)
+	  9. A list of examples (Optional)
+	  10. A list of references (Optional)
+
+  ]] .. " and you should generate docstring in " .. language .. " and its docstring format should be " .. docstring_format .. additional_prompt_about_docstring_format .. [[:
+  ]] .. [[
+  ```
+  ]] .. code_snippet .. [[
+  ```
+  And you shoud give me docstring immediately.
+  And you don't need to explain unnecessary things.
+  And don't use qoute to answer this question.
+  **AND DON'T INCLUDE CODE SNIPPET IN YOUR ANSWER.**
+  YOU SHOULD ANSWER THIS QUESTION IN DOCSTRING ONLY.
+  PLEASE DON'T INCLUDE YOUR COMMENT IN YOUR ANSWER.
+  **IF THERE IS LINE BREAK, YOU SHOULD FORMAT IT AS DOCSTRING FORMAT.**
+  ]]
+
+  return prompt
+end
+
 return M