Skip to content

Commit

Permalink
add prefix caching for deepseek
Browse files Browse the repository at this point in the history
  • Loading branch information
chottolabs committed Sep 8, 2024
1 parent 116eea6 commit 5e67b9f
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion lua/kznllm/presets.lua
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,35 @@ local function make_data_for_openai_chat(prompt_args, opts)
return data
end

local function make_data_for_deepseek_chat(prompt_args, opts)
local messages = {
{
role = 'system',
content = kznllm.make_prompt_from_template(opts.template_directory / 'nous_research/fill_mode_system_prompt.xml.jinja', prompt_args),
},
{
role = 'user',
content = kznllm.make_prompt_from_template(opts.template_directory / 'nous_research/fill_mode_user_prompt.xml.jinja', prompt_args),
},
}

local data = {
messages = messages,
model = opts.model,
stream = true,
}
if M.PROMPT_ARGS_STATE.replace and opts.prefill and opts.stop_param then
table.insert(messages, {
role = 'assistant',
content = opts.prefill .. prompt_args.current_buffer_filetype .. '\n',
prefix = true,
})
end
data = vim.tbl_extend('keep', data, opts.data_params, opts.stop_param or {})

return data
end

---Example implementation of a `make_data_fn` compatible with `kznllm.invoke_llm` for anthropic spec
---@param prompt_args any
---@param opts any
Expand Down Expand Up @@ -308,14 +337,16 @@ local presets = {
{
id = 'chat-model',
provider = 'deepseek',
make_data_fn = make_data_for_openai_chat,
make_data_fn = make_data_for_deepseek_chat,
debug_fn = openai_debug_fn,
opts = {
model = 'deepseek-coder',
data_params = {
max_tokens = 8192,
temperature = 0,
},
stop_param = { stop = { '```' } },
prefill = '```',
debug_fn = openai_debug_fn,
base_url = 'https://api.deepseek.com',
endpoint = '/beta/v1/chat/completions',
Expand Down

0 comments on commit 5e67b9f

Please sign in to comment.