From 44a8458b97db06ea9582e55524dc14884a8d991d Mon Sep 17 00:00:00 2001 From: Dominik Polakovics Date: Sat, 14 Dec 2024 15:28:42 +0100 Subject: [PATCH] feat: add the possibility to add initial files for the prompt --- .chatgpt_config.yaml | 3 ++ lua/chatgpt_nvim/config.lua | 27 ++++++--------- lua/chatgpt_nvim/init.lua | 68 +++++++++++++++++++++++++++---------- 3 files changed, 64 insertions(+), 34 deletions(-) diff --git a/.chatgpt_config.yaml b/.chatgpt_config.yaml index 247ecb1..5f1d13c 100644 --- a/.chatgpt_config.yaml +++ b/.chatgpt_config.yaml @@ -3,4 +3,7 @@ default_prompt_blocks: - "basic-prompt" directories: - "." +initial_files: + - "README.md" + - "lua" debug: false diff --git a/lua/chatgpt_nvim/config.lua b/lua/chatgpt_nvim/config.lua index e3e370c..83e05d8 100644 --- a/lua/chatgpt_nvim/config.lua +++ b/lua/chatgpt_nvim/config.lua @@ -3,29 +3,28 @@ local uv = vim.loop local ok_yaml, lyaml = pcall(require, "lyaml") +local prompt_blocks = { + ["go-development"] = "You are a coding assistant specialized in Go development. You will receive a project’s context and user instructions related to Go code, and you must return the requested modifications or guidance. When returning modifications, follow the specified YAML structure. Keep your suggestions aligned with Go best practices and idiomatic Go.", + ["typo3-development"] = "You are a coding assistant specialized in TYPO3 development. You have access to the project’s context and the user’s instructions. Your answers should focus on TYPO3 coding guidelines, extension development best practices, and TSconfig or TypoScript recommendations. When returning modifications, follow the YAML structure given.", + ["basic-prompt"] = "You are a coding assistant who receives a project's context and user instructions. The user will provide a prompt, and you will guide them through a workflow:\n1. First, ask them which files from the project are needed to understand and perform the coding task.\n2. If more information or context is needed, ask the user (outside of the YAML) to provide that.\n3. When all necessary information is gathered, provide the final YAML with the project's name and a list of files to be created or modified.\n\nThe final YAML must have a top-level key named 'project_name' that matches the project's configured name, and a top-level key named 'files', which is a list of file changes. Each element in 'files' must be a mapping with:\n- 'path' for the file path relative to the project’s root directory.\n- either 'content' with a multiline string for new content, or 'delete: true' if the file should be deleted.\n\nIf more context is needed at any point before providing the final YAML, request it outside of the YAML." +} + local function get_project_root() - -- Get the directory of the currently opened file. - -- If no file is open, we fallback to current working directory. local current_file = vim.fn.expand("%:p") local root_dir if current_file == "" then - -- No file opened, fallback to cwd root_dir = vim.fn.getcwd() else - -- Extract directory from current file path local file_dir = current_file:match("(.*)/") if not file_dir then - -- If something went wrong extracting the directory, fallback to cwd root_dir = vim.fn.getcwd() else - -- Attempt to find git root from the file's directory local cmd = string.format("cd %s && git rev-parse --show-toplevel 2>/dev/null", vim.fn.shellescape(file_dir)) local git_root = vim.fn.systemlist(cmd) if vim.v.shell_error == 0 and git_root and #git_root > 0 then root_dir = git_root[1] else - -- Not a git repo or failed to find git root, fallback to the file's directory root_dir = file_dir end end @@ -39,12 +38,6 @@ local function get_config_path() return root .. "/.chatgpt_config.yaml" end -local prompt_blocks = { - ["go-development"] = "You are a coding assistant specialized in Go development. You will receive a project’s context and user instructions related to Go code, and you must return the requested modifications or guidance. When returning modifications, follow the specified YAML structure. Keep your suggestions aligned with Go best practices and idiomatic Go.", - ["typo3-development"] = "You are a coding assistant specialized in TYPO3 development. You have access to the project’s context and the user’s instructions. Your answers should focus on TYPO3 coding guidelines, extension development best practices, and TSconfig or TypoScript recommendations. When returning modifications, follow the YAML structure given.", - ["basic-prompt"] = "You are a coding assistant who receives a project's context and user instructions. The user will provide a prompt, and you will guide them through a workflow:\n1. First, ask them which files from the project are needed to understand and perform the coding task.\n2. If more context or information is needed, ask the user (outside of the YAML) to provide that.\n3. When all necessary information is gathered, provide the final YAML with the project's name and a list of files to be created or modified.\n\nThe final YAML must have a top-level key named 'project_name' that matches the project's configured name, and a top-level key named 'files', which is a list of file changes. Each element in 'files' must be a mapping with:\n- 'path' for the file path relative to the project’s root directory.\n- either 'content' with a multiline string for new content, or 'delete: true' if the file should be deleted.\n\nIf more context is needed at any point before providing the final YAML, request it outside of the YAML." -} - function M.load() local path = get_config_path() local fd = uv.fs_open(path, "r", 438) @@ -54,7 +47,8 @@ function M.load() default_prompt_blocks = {}, token_limit = 128000, project_name = "", - debug = false + debug = false, + initial_files = {} } if fd then @@ -82,14 +76,15 @@ function M.load() if type(result.debug) == "boolean" then config.debug = result.debug end + if type(result.initial_files) == "table" then + config.initial_files = result.initial_files + end end end else - -- Default fallback configuration config.initial_prompt = "You are a coding assistant who receives a project's context and user instructions. You will guide the user through a workflow:\n1. First, ask the user which files are needed from the project to understand and perform the coding task.\n2. If more information or context is needed, ask for it outside of the YAML.\n3. Once all information is obtained, return the final YAML with the project_name and the files to be created/modified or deleted.\n\nThe final YAML must have:\nproject_name: \nfiles:\n - path: \"relative/path/to/file\"\n content: |\n \n - path: \"relative/path/to/other_file\"\n delete: true\n\nIf more context is needed at any step before providing the final YAML, request it outside of the YAML." end - -- If default_prompt_blocks are specified, concatenate all matching prompts if type(config.default_prompt_blocks) == "table" and #config.default_prompt_blocks > 0 then local merged_prompt = {} for _, block_name in ipairs(config.default_prompt_blocks) do diff --git a/lua/chatgpt_nvim/init.lua b/lua/chatgpt_nvim/init.lua index 32e6238..c62878c 100644 --- a/lua/chatgpt_nvim/init.lua +++ b/lua/chatgpt_nvim/init.lua @@ -30,7 +30,6 @@ local function parse_response(raw) return data end --- A simple token estimation function (approx. 4 chars per token) local function estimate_tokens(text) local approx_chars_per_token = 4 local length = #text @@ -43,8 +42,25 @@ local function is_subpath(root, path) return target_abs:sub(1, #root_abs) == root_abs end --- We'll store requested files in a global variable for now, though a more robust solution might be needed. -local requested_files = {} +local function read_file(path) + local fd = vim.loop.fs_open(path, "r", 438) + if not fd then + return nil + end + local stat = vim.loop.fs_fstat(fd) + if not stat then + vim.loop.fs_close(fd) + return nil + end + local data = vim.loop.fs_read(fd, stat.size, 0) + vim.loop.fs_close(fd) + return data +end + +local function is_directory(path) + local stat = vim.loop.fs_stat(path) + return stat and stat.type == "directory" +end function M.run_chatgpt_command() local conf = config.load() @@ -60,11 +76,38 @@ function M.run_chatgpt_command() local dirs = conf.directories or {"."} local project_structure = context.get_project_structure(dirs) + local initial_files = conf.initial_files or {} + local included_sections = {} + + if #initial_files > 0 then + table.insert(included_sections, "\n\nIncluded files and directories (pre-selected):\n") + local root = vim.fn.getcwd() + for _, item in ipairs(initial_files) do + local full_path = root .. "/" .. item + if is_directory(full_path) then + local dir_files = context.get_project_files({item}) + for _, f in ipairs(dir_files) do + local path = root .. "/" .. f + local data = read_file(path) + if data then + table.insert(included_sections, "\nFile: `" .. f .. "`\n```\n" .. data .. "\n```\n") + end + end + else + local data = read_file(full_path) + if data then + table.insert(included_sections, "\nFile: `" .. item .. "`\n```\n" .. data .. "\n```\n") + end + end + end + end + local initial_sections = { conf.initial_prompt .. "\n" .. user_input, "\n\nProject name: " .. (conf.project_name or "") .. "\n", "\n\nProject Structure:\n", project_structure, + table.concat(included_sections, "\n"), "\n\nPlease respond with a YAML listing which files you need from the project. For example:\n", "project_name: " .. (conf.project_name or "") .. "\nfiles:\n - path: \"relative/path/to/file\"\n\n" } @@ -157,9 +200,7 @@ function M.run_chatgpt_paste_command() return else - -- Not final: the model is requesting these files. We must gather and send them. local dirs = conf.directories or {"."} - local all_files = context.get_project_files(dirs) local requested_paths = {} for _, fileinfo in ipairs(data.files) do if fileinfo.path then @@ -171,19 +212,10 @@ function M.run_chatgpt_paste_command() local root = vim.fn.getcwd() for _, f in ipairs(requested_paths) do - local full_path = root .. "/" .. f - local fd = vim.loop.fs_open(full_path, "r", 438) - if fd then - local stat = vim.loop.fs_fstat(fd) - if stat then - local content = vim.loop.fs_read(fd, stat.size, 0) - vim.loop.fs_close(fd) - if content then - table.insert(file_sections, "\nFile: `" .. f .. "`\n```\n" .. content .. "\n```\n") - end - else - vim.loop.fs_close(fd) - end + local path = root .. "/" .. f + local content = read_file(path) + if content then + table.insert(file_sections, "\nFile: `" .. f .. "`\n```\n" .. content .. "\n```\n") else table.insert(file_sections, "\nFile: `" .. f .. "`\n```\n(Could not read file)\n```\n") end