Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

REWRITE #10

Merged
merged 12 commits into from
Aug 22, 2024
54 changes: 52 additions & 2 deletions CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,57 @@ For developing any kind of plugin using `lazy.nvim` you basically just need to d
'chottolabs/kznllm.nvim',
dev = true,
dir = '$HOME/.config/nvim/plugins/kznllm.nvim',
dependencies = { 'nvim-lua/plenary.nvim' },
...
dependencies = {
{ 'nvim-lua/plenary.nvim' },
{ 'stevearc/dressing.nvim' },
},

-- this points to whatever you specified as the local `dir = path` in above
kznllm.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'

spec.SELECTED_MODEL = { name = 'hermes-3-llama-3.1-405b-fp8' }
spec.API_KEY_NAME = 'LAMBDA_API_KEY'
spec.URL = 'https://api.lambdalabs.com/v1/chat/completions'

local function invoke_llm()
kznllm.invoke_llm({
-- add more user/assistant stages to this and just supply a path to your custom template directories
-- every prompt template gets sent the same table of args for simplicity sake, add custom args as needed
{ role = 'system', prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.FILL_MODE_SYSTEM_PROMPT },
{ role = 'user', prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.FILL_MODE_USER_PROMPT },
}, spec.make_job)
end

-- add a new keymap with a new behavior
vim.keymap.set({ 'n', 'v' }, '<leader>k', invoke_llm, { desc = 'Send current selection to LLM invoke_llm' })
},
```

If you go into `init.lua` and focus on `invoke_llm` function there's literally this one table that gets pass to all prompt templates, write whatever logic you want in your templates.

```lua
local prompt_args = {
current_buffer_path = current_buffer_path,
current_buffer_context = current_buffer_context,
current_buffer_filetype = current_buffer_filetype,
visual_selection = visual_selection,
user_query = input,
replace = replace_mode,
}
```

The "no visual selection mode" is really just a "non replace" mode controlled by `local replace_mode = not (mode == 'n')`.

If you look at the system prompts, it's literally just defining all the logic in the same template, you can add whatever arguments you want in this to suit your use case:

```j2
{%- if replace -%}
You should replace the code that you are sent, only following the comments. Do not talk at all. Only output valid code. Do not provide any backticks that surround the code. Never ever output backticks like this ```. Any comment that is asking you for something should be removed after you satisfy them. Other comments should left alone. Do not output backticks
{%- else -%}
You are a Senior Engineer at a Fortune 500 Company. You will be provided with code samples, academic papers, and documentation as supporting context to assist you in answering user queries about coding. Your task is to analyze this information and use it to provide accurate, helpful responses to the user's coding-related questions.
{%- endif -%}
```

An interesting thing you might consider is implementing a "project-scoped" template directory that can look for documentation files and pipe it into the args (you can do this with the actual text or the file path `{% include <absolute_file_path> %}`). `minijinja-cli` makes this kind of stuff super easy to do.

_Note: Can't be bothered to read about how the event loop works in nvim so it gets weird when I'm using "vim.ui" and other async APIs._
216 changes: 88 additions & 128 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,159 +1,119 @@
Based on [dingllm.nvim](https://github.com/yacineMTB/dingllm.nvim) - but it'll probably diverge quite a bit from the original state.
Based on [dingllm.nvim](https://github.com/yacineMTB/dingllm.nvim) - but diverge quite a bit

- adds some docstring annotations
- refactored to respect some of the inherent coupling between neovim <> some LLM streaming spec.
- prompt user for additional context
- coupling between prompt templates <> models are expressed more neatly
- sensible defaults + simple method to override them
- uses jinja as templating engine

I recommend you fork the repo and make it work for you.

We're at a point in history where "AI-powered editor" is probably one of the best "hello world" projects for learning about custom Neovim plugins.

See [CONTRIBUTING](CONTRIBUTING.md) to understand the typical development workflow for Neovim plugins using `Lazy`.
- prompts user for additional context before filling
- structured to make the inherent coupling between neovim logic, LLM streaming spec, and model-specific templates more explicit
- uses jinja as templating engine for ensuring correctness in more complex prompts
- preset defaults + simple approach for overriding them
- free cursor movement during generation

> [!NOTE]
> This plugin depends on [fd](https://github.com/sharkdp/fd) and [minijinja-cli](https://github.com/mitsuhiko/minijinja) (`cargo install minijinja-cli`, but double-check) - way easier to compose prompts. You should have `fd` already from telescope

# How it works
The only supported command is `leader + k`, it does nothing more than fill in some LLM completion into the text buffer. It has two main behaviors:
1. If you made a visual selection, it will attempt to replace your selection with a valid code fragment.
2. If you make no visual selection, it can yap freely (or do something else specified by a good template).

**Buffer Mode** - like a long-form chat mode
- **Usage**: (1) make a visual selection (2) `leader + k` (3) type out your query/prompt (4) quit and save with `leader + q` or type `:q!` to quit without saving
- **Behavior**:
- (initial) opens up a buffer, copies in the prompt template + arguments, and then streams the answer out at the bottom.
- (debug / render input templates) if you hit `leader + d` while in the completion buffer, it will interrupt (if it is still writing) and open up a buffer with a debug template showing the rendered input context
- (quit) hit `leader + w` to interrupt + quit and save the buffer to `$HOME/.cache/nvim/kznllm/history` (`vim.fn.stdpath 'cache' .. '/kznllm/history'`) as `<timestamp>/output.xml` along with `args.json` and returns you back to the original buffer. You can also interrupt + quit without saving + delete the buffer from history using `leader + q`
- (search history) if you quit with `leader + w` the buffer stays listed and you can find it in open buffers again (e.g. using kickstart defaults it would be `space + space`), if you quit with `leader + q` it deletes the buffer and won't clutter up open buffers list
By default (in supported templates), it also pipes in the contents of your current buffer.

https://github.com/user-attachments/assets/89331af3-3c69-41e3-9293-83b4549a6025
---

**Project Mode** - same as buffer mode, but lets you retrieve any files using an `fd` call and formats it into multi-document context
- **Usage**: (1) make a visual selection (2) `leader + Kp` (3) input arguments to `fd` (4) navigate the picker and hit `Tab` to select the files you want in the context (5) same as buffer mode
It's easy to hack on and implement customize behaviors without actually understanding much about nvim plugins at all. I recommend you fork the repo and make it work for you.

> [!WARNING]
> experimental and mainly built for claude
See [CONTRIBUTING](CONTRIBUTING.md) to understand the typical development workflow for Neovim plugins using `Lazy` and some straightforward ways you can modify the plugin to suit your needs

https://github.com/user-attachments/assets/cfa01851-f2f5-42b5-b042-0bb1fc55e3f7
By keeping the plugin simple with some QOL features, you get **close-to-natty** coding experience because it can keep generating code while you are free to do whatever you want (almost) without getting in the way too much.

**Replace Mode** - basically infill specifically for coding
- **Usage**: (1) make a visual selection (2) `leader + Kr`
- **Behavior**: replaces current selection and rewrites the selection based on context provdied by comments + fixing any errors
https://github.com/user-attachments/assets/932fa67f-0332-4799-b467-ecaeea54c3d1

https://github.com/chottolabs/kznllm.nvim/assets/171991982/39da67df-1ebc-4866-b563-f6b30d393162
_editing code while it generates when 405b is too slow_

## Usage
## Configuration

Make your API keys available via environment variables
```
export LAMBDA_API_KEY=secret_...
export ANTHROPIC_API_KEY=sk-...
export GROQ_API_KEY=gsk_...
```

for lambda labs

```lua
{
'chottolabs/kznllm.nvim',
dependencies = {
{ 'nvim-lua/plenary.nvim' },
{ 'stevearc/dressing.nvim' }, -- optional
},
config = function(self)
local kznllm = require 'kznllm'
local provider = require 'kznllm.specs.openai'

-- falls back to `vim.fn.stdpath 'data' .. '/lazy/kznllm/templates'` when the plugin is not locally installed
kznllm.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'

provider.SELECTED_MODEL = { name = 'hermes-3-llama-3.1-405b-fp8' }
provider.API_KEY_NAME = 'LAMBDA_API_KEY'
provider.URL = 'https://api.lambdalabs.com/v1/chat/completions'

local function llm_fill()
kznllm.invoke_llm({
-- the first template must be for the system prompt when using anthropic
{ role = 'system', prompt_template = provider.PROMPT_TEMPLATES.NOUS_RESEARCH.FILL_MODE_SYSTEM_PROMPT },
{ role = 'user', prompt_template = provider.PROMPT_TEMPLATES.NOUS_RESEARCH.FILL_MODE_USER_PROMPT },
}, provider.make_job)
end

vim.keymap.set({ 'n', 'v' }, '<leader>k', llm_fill, { desc = 'Send current selection to LLM llm_fill' })
end,
},
```

```lua
{
'chottolabs/kznllm.nvim',
dependencies = { 'nvim-lua/plenary.nvim' },
config = function(self)
local kznllm = require 'kznllm'
local utils = require 'kznllm.utils'
local spec = require 'kznllm.specs.anthropic'

utils.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates'

local function llm_buffer()
kznllm.invoke_llm_buffer_mode({
system_prompt_template = spec.PROMPT_TEMPLATES.BUFFER_MODE_SYSTEM_PROMPT,
user_prompt_template = spec.PROMPT_TEMPLATES.BUFFER_MODE_USER_PROMPT,
}, spec.make_job)
end

local function llm_project()
kznllm.invoke_llm_project_mode({
system_prompt_template = spec.PROMPT_TEMPLATES.PROJECT_MODE_SYSTEM_PROMPT,
user_prompt_template = spec.PROMPT_TEMPLATES.PROJECT_MODE_USER_PROMPT,
}, spec.make_job)
end

local function llm_replace()
kznllm.invoke_llm_replace_mode({
system_prompt_template = spec.PROMPT_TEMPLATES.REPLACE_MODE_SYSTEM_PROMPT,
user_prompt_template = spec.PROMPT_TEMPLATES.REPLACE_MODE_USER_PROMPT,
}, spec.make_job)
end

vim.keymap.set({ 'n', 'v' }, '<leader>k', llm_buffer, { desc = 'Send current selection to LLM llm_buffer' })
vim.keymap.set({ 'n', 'v' }, '<leader>Kr', llm_replace, { desc = 'Send current selection to LLM llm_replace' })
vim.keymap.set({ 'n', 'v' }, '<leader>Kp', llm_project, { desc = 'Send current selection to LLM llm_project' })
end,
},
local kznllm = require 'kznllm'
local spec = require 'kznllm.specs.anthropic'

kznllm.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'

local function llm_fill()
kznllm.invoke_llm({
{ role = 'system', prompt_template = provider.PROMPT_TEMPLATES.FILL_MODE_SYSTEM_PROMPT },
{ role = 'user', prompt_template = provider.PROMPT_TEMPLATES.FILL_MODE_USER_PROMPT },
}, spec.make_job)
end

vim.keymap.set({ 'n', 'v' }, '<leader>k', llm_fill, { desc = 'Send current selection to LLM llm_fill' })
```

for groq
for groq (default)
```lua
{
'chottolabs/kznllm.nvim',
dependencies = { 'nvim-lua/plenary.nvim' },
config = function(self)
local kznllm = require 'kznllm'
local utils = require 'kznllm.utils'
local spec = require 'kznllm.specs.openai'

...
end,
},
local kznllm = require 'kznllm'
local utils = require 'kznllm.utils'
local spec = require 'kznllm.specs.openai'

kznllm.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'
-- fallsback to a preset default model configuration

local function llm_fill()
kznllm.invoke_llm({
{ role = 'system', prompt_template = provider.PROMPT_TEMPLATES.GROQ.FILL_MODE_SYSTEM_PROMPT },
{ role = 'user', prompt_template = provider.PROMPT_TEMPLATES.GROQ.FILL_MODE_USER_PROMPT },
}, spec.make_job)
end
...
```

for local openai server
(e.g. `vllm serve` w/ `--api-key <token>` and `--served-model-name meta-llama/Meta-Llama-3.1-8B-Instruct`) set `VLLM_API_KEY=<token>`
```lua
{
'chottolabs/kznllm.nvim',
dependencies = { 'nvim-lua/plenary.nvim' },
config = function(self)
local kznllm = require 'kznllm'
local utils = require 'kznllm.utils'
local spec = require 'kznllm.specs.openai'

spec.SELECTED_MODEL = { name = 'meta-llama/Meta-Llama-3.1-8B-Instruct', max_tokens = 8192 }
spec.URL = 'http://research.local:8000/v1/chat/completions'
spec.API_KEY_NAME = 'VLLM_API_KEY'

...
end,
},
```
local kznllm = require 'kznllm'
local spec = require 'kznllm.specs.openai'

for lambda labs
```lua
{
'chottolabs/kznllm.nvim',
dependencies = { 'nvim-lua/plenary.nvim' },
config = function(self)
local kznllm = require 'kznllm'
local utils = require 'kznllm.utils'
local spec = require 'kznllm.specs.openai'

spec.SELECTED_MODEL = { name = 'hermes-3-llama-3.1-405b-fp8' }
spec.API_KEY_NAME = 'LAMBDA_LABS_API_KEY'
spec.URL = 'https://api.lambdalabs.com/v1/chat/completions'

utils.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'

local function llm_buffer()
kznllm.invoke_llm_buffer_mode({
system_prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.BUFFER_MODE_SYSTEM_PROMPT,
user_prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.BUFFER_MODE_USER_PROMPT,
}, spec.make_job)
end

local function llm_project()
kznllm.invoke_llm_project_mode({
system_prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.PROJECT_MODE_SYSTEM_PROMPT,
user_prompt_template = spec.PROMPT_TEMPLATES.NOUS_RESEARCH.PROJECT_MODE_USER_PROMPT,
}, spec.make_job)
end
...
end,
},
kznllm.TEMPLATE_DIRECTORY = vim.fn.expand(self.dir) .. '/templates/'

spec.SELECTED_MODEL = { name = 'meta-llama/Meta-Llama-3.1-8B-Instruct', max_tokens = 8192 }
spec.API_KEY_NAME = 'VLLM_API_KEY'
spec.URL = 'http://research.local:8000/v1/chat/completions'
...
```

Loading