-
Notifications
You must be signed in to change notification settings - Fork 394
Custom providers
yetone edited this page Nov 16, 2024
·
9 revisions
To add support for custom providers, one add AvanteProvider spec into opts.vendors:
{
provider = "my-custom-provider", -- You can then change this provider here
vendors = {
["my-custom-provider"] = {...}
},
}
A custom provider should following the following spec:
---@type AvanteProvider
["my-custom-provider"] = {
endpoint = "https://api.openai.com/v1/chat/completions", -- The full endpoint of the provider
model = "gpt-4o", -- The model name to use with this provider
api_key_name = "OPENAI_API_KEY", -- The name of the environment variable that contains the API key
--- This function below will be used to parse in cURL arguments.
--- It takes in the provider options as the first argument, followed by code_opts retrieved from given buffer.
--- This code_opts include:
--- - question: Input from the users
--- - code_lang: the language of given code buffer
--- - code_content: content of code buffer
--- - selected_code_content: (optional) If given code content is selected in visual mode as context.
---@type fun(opts: AvanteProvider, code_opts: AvantePromptOptions): AvanteCurlOutput
parse_curl_args = function(opts, code_opts) end
--- This function will be used to parse incoming SSE stream
--- It takes in the data stream as the first argument, followed by SSE event state, and opts
--- retrieved from given buffer.
--- This opts include:
--- - on_chunk: (fun(chunk: string): any) this is invoked on parsing correct delta chunk
--- - on_complete: (fun(err: string|nil): any) this is invoked on either complete call or error chunk
---@type fun(data_stream: string, event_state: string, opts: ResponseParser): nil
parse_response = function(data_stream, event_state, opts) end
--- The following function SHOULD only be used when providers doesn't follow SSE spec [ADVANCED]
--- this is mutually exclusive with parse_response_data
---@type fun(data: string, handler_opts: AvanteHandlerOptions): nil
parse_stream_data = function(data, handler_opts) end
}You can use the __inherited_from field to inherit any of our built-in providers, so you can easily add an OpenAI Compatible vendor, for example:
provider = "qianwen",
vendors = {
qianwen = {
__inherited_from = "openai",
api_key_name = "DASHSCOPE_API_KEY",
endpoint = "https://dashscope.aliyuncs.com/compatible-mode/v1",
model = "qwen-coder-plus-latest",
},
},provider = "ollama",
vendors = {
ollama = {
__inherited_from = "openai",
api_key_name = "",
endpoint = "http://127.0.0.1:11434/v1",
model = "codegemma",
},
},See https://github.com/yetone/avante.nvim/pull/810/files
local M = {}
M.role_map = {
user = "human",
assistant = "assistant",
system = "system",
}
M.parse_messages = function(opts)
local messages = {
{ role = "system", content = opts.system_prompt },
}
vim
.iter(opts.messages)
:each(function(msg) table.insert(messages, { speaker = M.role_map[msg.role], text = msg.content }) end)
return messages
end
M.parse_response = function(data_stream, event_state, opts)
if event_state == "done" then
opts.on_complete()
return
end
if data_stream == nil or data_stream == "" then return end
local json = vim.json.decode(data_stream)
local delta = json.deltaText
local stopReason = json.stopReason
if stopReason == "end_turn" then return end
opts.on_chunk(delta)
end
---@type AvanteProvider
["my-custom-provider"] = {
endpoint = "https://sourcegraph.com",
model = "anthropic::2024-10-22::claude-3-5-sonnet-latest",
api_key_name = "SRC_ACCESS_TOKEN",
--- This function below will be used to parse in cURL arguments.
--- It takes in the provider options as the first argument, followed by code_opts retrieved from given buffer.
--- This code_opts include:
--- - question: Input from the users
--- - code_lang: the language of given code buffer
--- - code_content: content of code buffer
--- - selected_code_content: (optional) If given code content is selected in visual mode as context.
---@type fun(opts: AvanteProvider, code_opts: AvantePromptOptions): AvanteCurlOutput
parse_curl_args = function(opts, code_opts)
local headers = {
["Content-Type"] = "application/json",
["Authorization"] = "token " .. os.getenv(opts.api_key_name),
}
return {
url = opts.endpoint .. "/.api/completions/stream?api-version=2&client-name=web&client-version=0.0.1",
timeout = base.timeout,
insecure = false,
headers = headers,
body = vim.tbl_deep_extend("force", {
model = opts.model,
temperature = 0,
topK = -1,
topP = -1,
maxTokensToSample = 4000,
stream = true,
messages = M.parse_messages(code_opts),
}, {}),
}
end,
---@type fun(data_stream: string, event_state: string, opts: ResponseParser): nil
parse_response = function(data_stream, event_state, opts) M.parse_response(data_stream, event_state, opts) end
}vendors = {
---@type AvanteProvider
perplexity = {
endpoint = "https://api.perplexity.ai/chat/completions",
model = "llama-3.1-sonar-large-128k-online",
api_key_name = "cmd:bw get notes perplexity-api-key",
parse_curl_args = function(opts, code_opts)
return {
url = opts.endpoint,
headers = {
["Accept"] = "application/json",
["Content-Type"] = "application/json",
["Authorization"] = "Bearer " .. os.getenv(opts.api_key_name),
},
body = {
model = opts.model,
messages = { -- you can make your own message, but this is very advanced
{ role = "system", content = code_opts.system_prompt },
{ role = "user", content = require("avante.providers.openai").get_user_message(code_opts) },
},
temperature = 0,
max_tokens = 8192,
stream = true, -- this will be set by default.
},
}
end,
-- The below function is used if the vendors has specific SSE spec that is not claude or openai.
parse_response = function(data_stream, event_state, opts)
require("avante.providers").openai.parse_response(data_stream, event_state, opts)
end,
},
---@type AvanteProvider
groq = {
endpoint = "https://api.groq.com/openai/v1/chat/completions",
model = "llama-3.1-70b-versatile",
api_key_name = "GROQ_API_KEY",
parse_curl_args = function(opts, code_opts)
return {
url = opts.endpoint,
headers = {
["Accept"] = "application/json",
["Content-Type"] = "application/json",
["Authorization"] = "Bearer " .. os.getenv(opts.api_key_name),
},
body = {
model = opts.model,
messages = { -- you can make your own message, but this is very advanced
{ role = "system", content = code_opts.system_prompt },
{ role = "user", content = require("avante.providers.openai").get_user_message(code_opts) },
},
temperature = 0,
max_tokens = 4096,
stream = true, -- this will be set by default.
},
}
end,
parse_response = function(data_stream, event_state, opts)
require("avante.providers").openai.parse_response(data_stream, event_state, opts)
end,
},
---@type AvanteProvider
deepseek = {
endpoint = "https://api.deepseek.com/chat/completions",
model = "deepseek-coder",
api_key_name = "DEEPSEEK_API_KEY",
parse_curl_args = function(opts, code_opts)
return {
url = opts.endpoint,
headers = {
["Accept"] = "application/json",
["Content-Type"] = "application/json",
["Authorization"] = "Bearer " .. os.getenv(opts.api_key_name),
},
body = {
model = opts.model,
messages = { -- you can make your own message, but this is very advanced
{ role = "system", content = code_opts.system_prompt },
{ role = "user", content = require("avante.providers.openai").get_user_message(code_opts) },
},
temperature = 0,
max_tokens = 4096,
stream = true, -- this will be set by default.
},
}
end,
parse_response = function(data_stream, event_state, opts)
require("avante.providers").openai.parse_response(data_stream, event_state, opts)
end,
},
}If certain providers don't follow SSE streaming spec, you might want to implement parse_stream_data for your custom providers.
See parse_and_call implementation for more information.