557 lines
19 KiB
Lua
557 lines
19 KiB
Lua
local Utils = require("avante.utils")
|
|
local Config = require("avante.config")
|
|
local Clipboard = require("avante.clipboard")
|
|
local Providers = require("avante.providers")
|
|
local HistoryMessage = require("avante.history_message")
|
|
local XMLParser = require("avante.libs.xmlparser")
|
|
local JsonParser = require("avante.libs.jsonparser")
|
|
local Prompts = require("avante.utils.prompts")
|
|
local LlmTools = require("avante.llm_tools")
|
|
|
|
---@class AvanteProviderFunctor
|
|
local M = {}
|
|
|
|
M.api_key_name = "OPENAI_API_KEY"
|
|
|
|
M.role_map = {
|
|
user = "user",
|
|
assistant = "assistant",
|
|
}
|
|
|
|
function M:is_disable_stream() return false end
|
|
|
|
---@param tool AvanteLLMTool
|
|
---@return AvanteOpenAITool
|
|
function M:transform_tool(tool)
|
|
local input_schema_properties, required = Utils.llm_tool_param_fields_to_json_schema(tool.param.fields)
|
|
---@type AvanteOpenAIToolFunctionParameters
|
|
local parameters = nil
|
|
if not vim.tbl_isempty(input_schema_properties) then
|
|
parameters = {
|
|
type = "object",
|
|
properties = input_schema_properties,
|
|
required = required,
|
|
additionalProperties = false,
|
|
}
|
|
end
|
|
---@type AvanteOpenAITool
|
|
local res = {
|
|
type = "function",
|
|
["function"] = {
|
|
name = tool.name,
|
|
description = tool.get_description and tool.get_description() or tool.description,
|
|
parameters = parameters,
|
|
},
|
|
}
|
|
return res
|
|
end
|
|
|
|
function M.is_openrouter(url) return url:match("^https://openrouter%.ai/") end
|
|
|
|
---@param opts AvantePromptOptions
|
|
function M.get_user_message(opts)
|
|
vim.deprecate("get_user_message", "parse_messages", "0.1.0", "avante.nvim")
|
|
return table.concat(
|
|
vim
|
|
.iter(opts.messages)
|
|
:filter(function(_, value) return value == nil or value.role ~= "user" end)
|
|
:fold({}, function(acc, value)
|
|
acc = vim.list_extend({}, acc)
|
|
acc = vim.list_extend(acc, { value.content })
|
|
return acc
|
|
end),
|
|
"\n"
|
|
)
|
|
end
|
|
|
|
function M.is_reasoning_model(model) return model and string.match(model, "^o%d+") ~= nil end
|
|
|
|
function M.set_allowed_params(provider_conf, request_body)
|
|
if M.is_reasoning_model(provider_conf.model) then
|
|
request_body.temperature = 1
|
|
else
|
|
request_body.reasoning_effort = nil
|
|
end
|
|
-- If max_tokens is set in config, unset max_completion_tokens
|
|
if request_body.max_tokens then request_body.max_completion_tokens = nil end
|
|
end
|
|
|
|
function M:parse_messages(opts)
|
|
local messages = {}
|
|
local provider_conf, _ = Providers.parse_config(self)
|
|
|
|
local use_ReAct_prompt = provider_conf.use_ReAct_prompt == true
|
|
local system_prompt = opts.system_prompt
|
|
|
|
if use_ReAct_prompt then system_prompt = Prompts.get_ReAct_system_prompt(provider_conf, opts) end
|
|
|
|
if self.is_reasoning_model(provider_conf.model) then
|
|
table.insert(messages, { role = "developer", content = system_prompt })
|
|
else
|
|
table.insert(messages, { role = "system", content = system_prompt })
|
|
end
|
|
|
|
local has_tool_use = false
|
|
|
|
vim.iter(opts.messages):each(function(msg)
|
|
if type(msg.content) == "string" then
|
|
table.insert(messages, { role = self.role_map[msg.role], content = msg.content })
|
|
elseif type(msg.content) == "table" then
|
|
local content = {}
|
|
local tool_calls = {}
|
|
local tool_results = {}
|
|
for _, item in ipairs(msg.content) do
|
|
if type(item) == "string" then
|
|
table.insert(content, { type = "text", text = item })
|
|
elseif item.type == "text" then
|
|
table.insert(content, { type = "text", text = item.text })
|
|
elseif item.type == "image" then
|
|
table.insert(content, {
|
|
type = "image_url",
|
|
image_url = {
|
|
url = "data:" .. item.source.media_type .. ";" .. item.source.type .. "," .. item.source.data,
|
|
},
|
|
})
|
|
elseif item.type == "tool_use" and not use_ReAct_prompt then
|
|
has_tool_use = true
|
|
table.insert(tool_calls, {
|
|
id = item.id,
|
|
type = "function",
|
|
["function"] = { name = item.name, arguments = vim.json.encode(item.input) },
|
|
})
|
|
elseif item.type == "tool_result" and has_tool_use and not use_ReAct_prompt then
|
|
table.insert(
|
|
tool_results,
|
|
{ tool_call_id = item.tool_use_id, content = item.is_error and "Error: " .. item.content or item.content }
|
|
)
|
|
end
|
|
end
|
|
if not provider_conf.disable_tools and use_ReAct_prompt then
|
|
if msg.content[1].type == "tool_result" then
|
|
local tool_use = nil
|
|
for _, msg_ in ipairs(opts.messages) do
|
|
if type(msg_.content) == "table" and #msg_.content > 0 then
|
|
if msg_.content[1].type == "tool_use" and msg_.content[1].id == msg.content[1].tool_use_id then
|
|
tool_use = msg_
|
|
break
|
|
end
|
|
end
|
|
end
|
|
if tool_use then
|
|
msg.role = "user"
|
|
table.insert(content, {
|
|
type = "text",
|
|
text = "["
|
|
.. tool_use.content[1].name
|
|
.. " for '"
|
|
.. (tool_use.content[1].input.path or tool_use.content[1].input.rel_path or "")
|
|
.. "'] Result:",
|
|
})
|
|
table.insert(content, {
|
|
type = "text",
|
|
text = msg.content[1].content,
|
|
})
|
|
end
|
|
end
|
|
end
|
|
if #content > 0 then table.insert(messages, { role = self.role_map[msg.role], content = content }) end
|
|
if not provider_conf.disable_tools and not use_ReAct_prompt then
|
|
if #tool_calls > 0 then
|
|
local last_message = messages[#messages]
|
|
if last_message and last_message.role == self.role_map["assistant"] and last_message.tool_calls then
|
|
last_message.tool_calls = vim.list_extend(last_message.tool_calls, tool_calls)
|
|
else
|
|
table.insert(messages, { role = self.role_map["assistant"], tool_calls = tool_calls })
|
|
end
|
|
end
|
|
if #tool_results > 0 then
|
|
for _, tool_result in ipairs(tool_results) do
|
|
table.insert(
|
|
messages,
|
|
{ role = "tool", tool_call_id = tool_result.tool_call_id, content = tool_result.content or "" }
|
|
)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end)
|
|
|
|
if Config.behaviour.support_paste_from_clipboard and opts.image_paths and #opts.image_paths > 0 then
|
|
local message_content = messages[#messages].content
|
|
if type(message_content) ~= "table" or message_content[1] == nil then
|
|
message_content = { { type = "text", text = message_content } }
|
|
end
|
|
for _, image_path in ipairs(opts.image_paths) do
|
|
table.insert(message_content, {
|
|
type = "image_url",
|
|
image_url = {
|
|
url = "data:image/png;base64," .. Clipboard.get_base64_content(image_path),
|
|
},
|
|
})
|
|
end
|
|
messages[#messages].content = message_content
|
|
end
|
|
|
|
local final_messages = {}
|
|
local prev_role = nil
|
|
|
|
vim.iter(messages):each(function(message)
|
|
local role = message.role
|
|
if role == prev_role and role ~= "tool" then
|
|
if role == self.role_map["assistant"] then
|
|
table.insert(final_messages, { role = self.role_map["user"], content = "Ok" })
|
|
else
|
|
table.insert(final_messages, { role = self.role_map["assistant"], content = "Ok, I understand." })
|
|
end
|
|
end
|
|
prev_role = role
|
|
table.insert(final_messages, message)
|
|
end)
|
|
|
|
return final_messages
|
|
end
|
|
|
|
function M:finish_pending_messages(ctx, opts)
|
|
if ctx.content ~= nil and ctx.content ~= "" then self:add_text_message(ctx, "", "generated", opts) end
|
|
if ctx.tool_use_list then
|
|
for _, tool_use in ipairs(ctx.tool_use_list) do
|
|
if tool_use.state == "generating" then self:add_tool_use_message(tool_use, "generated", opts) end
|
|
end
|
|
end
|
|
end
|
|
|
|
local llm_tool_names = nil
|
|
|
|
function M:add_text_message(ctx, text, state, opts)
|
|
if llm_tool_names == nil then llm_tool_names = LlmTools.get_tool_names() end
|
|
if ctx.content == nil then ctx.content = "" end
|
|
ctx.content = ctx.content .. text
|
|
local content = ctx.content
|
|
:gsub("<tool_code>", "")
|
|
:gsub("</tool_code>", "")
|
|
:gsub("<tool_call>", "")
|
|
:gsub("</tool_call>", "")
|
|
:gsub("<tool_use>", "")
|
|
:gsub("</tool_use>", "")
|
|
ctx.content = content
|
|
local msg = HistoryMessage:new({
|
|
role = "assistant",
|
|
content = ctx.content,
|
|
}, {
|
|
state = state,
|
|
uuid = ctx.content_uuid,
|
|
})
|
|
ctx.content_uuid = msg.uuid
|
|
local msgs = { msg }
|
|
local xml_content = ctx.content
|
|
local xml_lines = vim.split(xml_content, "\n")
|
|
local cleaned_xml_lines = {}
|
|
local prev_tool_name = nil
|
|
for _, line in ipairs(xml_lines) do
|
|
if line:match("<tool_name>") then
|
|
local tool_name = line:match("<tool_name>(.*)</tool_name>")
|
|
if tool_name then prev_tool_name = tool_name end
|
|
elseif line:match("<parameters>") then
|
|
if prev_tool_name then table.insert(cleaned_xml_lines, "<" .. prev_tool_name .. ">") end
|
|
goto continue
|
|
elseif line:match("</parameters>") then
|
|
if prev_tool_name then table.insert(cleaned_xml_lines, "</" .. prev_tool_name .. ">") end
|
|
goto continue
|
|
end
|
|
table.insert(cleaned_xml_lines, line)
|
|
::continue::
|
|
end
|
|
local cleaned_xml_content = table.concat(cleaned_xml_lines, "\n")
|
|
local stream_parser = XMLParser.createStreamParser()
|
|
stream_parser:addData(cleaned_xml_content)
|
|
local has_tool_use = false
|
|
local xml = stream_parser:getAllElements()
|
|
if xml then
|
|
local new_content_list = {}
|
|
local xml_md_openned = false
|
|
for idx, item in ipairs(xml) do
|
|
if item._name == "_text" then
|
|
local cleaned_lines = {}
|
|
local lines = vim.split(item._text, "\n")
|
|
for _, line in ipairs(lines) do
|
|
if line:match("^```xml") or line:match("^```tool_code") or line:match("^```tool_use") then
|
|
xml_md_openned = true
|
|
elseif line:match("^```$") then
|
|
if xml_md_openned then
|
|
xml_md_openned = false
|
|
else
|
|
table.insert(cleaned_lines, line)
|
|
end
|
|
else
|
|
table.insert(cleaned_lines, line)
|
|
end
|
|
end
|
|
table.insert(new_content_list, table.concat(cleaned_lines, "\n"))
|
|
goto continue
|
|
end
|
|
if not vim.tbl_contains(llm_tool_names, item._name) then goto continue end
|
|
local ok, input = pcall(vim.json.decode, item._text)
|
|
if not ok then input = {} end
|
|
if not ok and item.children and #item.children > 0 then
|
|
for _, item_ in ipairs(item.children) do
|
|
local ok_, input_ = pcall(vim.json.decode, item_._text)
|
|
if ok_ and input_ then
|
|
input[item_._name] = input_
|
|
else
|
|
input[item_._name] = item_._text
|
|
end
|
|
end
|
|
end
|
|
if next(input) ~= nil then
|
|
local msg_uuid = ctx.content_uuid .. "-" .. idx
|
|
local tool_use_id = msg_uuid
|
|
local msg_ = HistoryMessage:new({
|
|
role = "assistant",
|
|
content = {
|
|
{
|
|
type = "tool_use",
|
|
name = item._name,
|
|
id = tool_use_id,
|
|
input = input,
|
|
},
|
|
},
|
|
}, {
|
|
state = state,
|
|
uuid = msg_uuid,
|
|
})
|
|
msgs[#msgs + 1] = msg_
|
|
ctx.tool_use_list = ctx.tool_use_list or {}
|
|
ctx.tool_use_list[#ctx.tool_use_list + 1] = {
|
|
id = tool_use_id,
|
|
name = item._name,
|
|
input_json = input,
|
|
}
|
|
has_tool_use = true
|
|
end
|
|
if #new_content_list > 0 then msg.message.content = table.concat(new_content_list, "\n") end
|
|
::continue::
|
|
end
|
|
end
|
|
if opts.on_messages_add then opts.on_messages_add(msgs) end
|
|
if has_tool_use and state == "generating" then opts.on_stop({ reason = "tool_use", streaming_tool_use = true }) end
|
|
end
|
|
|
|
function M:add_thinking_message(ctx, text, state, opts)
|
|
if ctx.reasonging_content == nil then ctx.reasonging_content = "" end
|
|
ctx.reasonging_content = ctx.reasonging_content .. text
|
|
local msg = HistoryMessage:new({
|
|
role = "assistant",
|
|
content = {
|
|
{
|
|
type = "thinking",
|
|
thinking = ctx.reasonging_content,
|
|
signature = "",
|
|
},
|
|
},
|
|
}, {
|
|
state = state,
|
|
uuid = ctx.reasonging_content_uuid,
|
|
})
|
|
ctx.reasonging_content_uuid = msg.uuid
|
|
if opts.on_messages_add then opts.on_messages_add({ msg }) end
|
|
end
|
|
|
|
function M:add_tool_use_message(tool_use, state, opts)
|
|
local jsn = JsonParser.parse(tool_use.input_json)
|
|
local msg = HistoryMessage:new({
|
|
role = "assistant",
|
|
content = {
|
|
{
|
|
type = "tool_use",
|
|
name = tool_use.name,
|
|
id = tool_use.id,
|
|
input = jsn or {},
|
|
},
|
|
},
|
|
}, {
|
|
state = state,
|
|
uuid = tool_use.uuid,
|
|
})
|
|
tool_use.uuid = msg.uuid
|
|
tool_use.state = state
|
|
if opts.on_messages_add then opts.on_messages_add({ msg }) end
|
|
if state == "generating" then opts.on_stop({ reason = "tool_use", streaming_tool_use = true }) end
|
|
end
|
|
|
|
function M:parse_response(ctx, data_stream, _, opts)
|
|
local orig_on_stop = opts.on_stop
|
|
local stopped = false
|
|
---@param stop_opts AvanteLLMStopCallbackOptions
|
|
opts.on_stop = function(stop_opts)
|
|
if stop_opts and not stop_opts.streaming_tool_use then
|
|
if stopped then return end
|
|
stopped = true
|
|
end
|
|
return orig_on_stop(stop_opts)
|
|
end
|
|
if data_stream:match('"%[DONE%]":') then
|
|
self:finish_pending_messages(ctx, opts)
|
|
if ctx.tool_use_list and #ctx.tool_use_list > 0 then
|
|
opts.on_stop({ reason = "tool_use" })
|
|
else
|
|
opts.on_stop({ reason = "complete" })
|
|
end
|
|
return
|
|
end
|
|
if data_stream == "[DONE]" then
|
|
opts.on_stop({ reason = "complete" })
|
|
return
|
|
end
|
|
local jsn = vim.json.decode(data_stream)
|
|
---@cast jsn AvanteOpenAIChatResponse
|
|
if not jsn.choices then return end
|
|
local choice = jsn.choices[1]
|
|
if not choice then return end
|
|
local delta = choice.delta
|
|
if not delta then
|
|
local provider_conf = Providers.parse_config(self)
|
|
if provider_conf.model:match("o1") then delta = choice.message end
|
|
end
|
|
if not delta then return end
|
|
if delta.reasoning_content and delta.reasoning_content ~= vim.NIL and delta.reasoning_content ~= "" then
|
|
if ctx.returned_think_start_tag == nil or not ctx.returned_think_start_tag then
|
|
ctx.returned_think_start_tag = true
|
|
if opts.on_chunk then opts.on_chunk("<think>\n") end
|
|
end
|
|
ctx.last_think_content = delta.reasoning_content
|
|
self:add_thinking_message(ctx, delta.reasoning_content, "generating", opts)
|
|
if opts.on_chunk then opts.on_chunk(delta.reasoning_content) end
|
|
elseif delta.reasoning and delta.reasoning ~= vim.NIL then
|
|
if ctx.returned_think_start_tag == nil or not ctx.returned_think_start_tag then
|
|
ctx.returned_think_start_tag = true
|
|
if opts.on_chunk then opts.on_chunk("<think>\n") end
|
|
end
|
|
ctx.last_think_content = delta.reasoning
|
|
self:add_thinking_message(ctx, delta.reasoning, "generating", opts)
|
|
if opts.on_chunk then opts.on_chunk(delta.reasoning) end
|
|
elseif delta.tool_calls and delta.tool_calls ~= vim.NIL then
|
|
local choice_index = choice.index or 0
|
|
for idx, tool_call in ipairs(delta.tool_calls) do
|
|
--- In Gemini's so-called OpenAI Compatible API, tool_call.index is nil, which is quite absurd! Therefore, a compatibility fix is needed here.
|
|
if tool_call.index == nil then tool_call.index = choice_index + idx - 1 end
|
|
if not ctx.tool_use_list then ctx.tool_use_list = {} end
|
|
if not ctx.tool_use_list[tool_call.index + 1] then
|
|
if tool_call.index > 0 and ctx.tool_use_list[tool_call.index] then
|
|
local prev_tool_use = ctx.tool_use_list[tool_call.index]
|
|
self:add_tool_use_message(prev_tool_use, "generated", opts)
|
|
end
|
|
local tool_use = {
|
|
name = tool_call["function"].name,
|
|
id = tool_call.id,
|
|
input_json = type(tool_call["function"].arguments) == "string" and tool_call["function"].arguments or "",
|
|
}
|
|
ctx.tool_use_list[tool_call.index + 1] = tool_use
|
|
self:add_tool_use_message(tool_use, "generating", opts)
|
|
else
|
|
local tool_use = ctx.tool_use_list[tool_call.index + 1]
|
|
tool_use.input_json = tool_use.input_json .. tool_call["function"].arguments
|
|
self:add_tool_use_message(tool_use, "generating", opts)
|
|
end
|
|
end
|
|
elseif delta.content then
|
|
if
|
|
ctx.returned_think_start_tag ~= nil and (ctx.returned_think_end_tag == nil or not ctx.returned_think_end_tag)
|
|
then
|
|
ctx.returned_think_end_tag = true
|
|
if opts.on_chunk then
|
|
if ctx.last_think_content and ctx.last_think_content ~= vim.NIL and ctx.last_think_content:sub(-1) ~= "\n" then
|
|
opts.on_chunk("\n</think>\n")
|
|
else
|
|
opts.on_chunk("</think>\n")
|
|
end
|
|
end
|
|
self:add_thinking_message(ctx, "", "generated", opts)
|
|
end
|
|
if delta.content ~= vim.NIL then
|
|
if opts.on_chunk then opts.on_chunk(delta.content) end
|
|
self:add_text_message(ctx, delta.content, "generating", opts)
|
|
end
|
|
end
|
|
if choice.finish_reason == "stop" or choice.finish_reason == "eos_token" or choice.finish_reason == "length" then
|
|
self:finish_pending_messages(ctx, opts)
|
|
if ctx.tool_use_list and #ctx.tool_use_list > 0 then
|
|
opts.on_stop({ reason = "tool_use", usage = jsn.usage })
|
|
else
|
|
opts.on_stop({ reason = "complete", usage = jsn.usage })
|
|
end
|
|
end
|
|
if choice.finish_reason == "tool_calls" then
|
|
self:finish_pending_messages(ctx, opts)
|
|
opts.on_stop({
|
|
reason = "tool_use",
|
|
usage = jsn.usage,
|
|
})
|
|
end
|
|
end
|
|
|
|
function M:parse_response_without_stream(data, _, opts)
|
|
---@type AvanteOpenAIChatResponse
|
|
local json = vim.json.decode(data)
|
|
if json.choices and json.choices[1] then
|
|
local choice = json.choices[1]
|
|
if choice.message and choice.message.content then
|
|
if opts.on_chunk then opts.on_chunk(choice.message.content) end
|
|
self:add_text_message({}, choice.message.content, "generated", opts)
|
|
vim.schedule(function() opts.on_stop({ reason = "complete" }) end)
|
|
end
|
|
end
|
|
end
|
|
|
|
function M:parse_curl_args(prompt_opts)
|
|
local provider_conf, request_body = Providers.parse_config(self)
|
|
local disable_tools = provider_conf.disable_tools or false
|
|
|
|
local headers = {
|
|
["Content-Type"] = "application/json",
|
|
}
|
|
|
|
if Providers.env.require_api_key(provider_conf) then
|
|
local api_key = self.parse_api_key()
|
|
if api_key == nil then
|
|
error(Config.provider .. " API key is not set, please set it in your environment variable or config file")
|
|
end
|
|
headers["Authorization"] = "Bearer " .. api_key
|
|
end
|
|
|
|
if M.is_openrouter(provider_conf.endpoint) then
|
|
headers["HTTP-Referer"] = "https://github.com/yetone/avante.nvim"
|
|
headers["X-Title"] = "Avante.nvim"
|
|
request_body.include_reasoning = true
|
|
end
|
|
|
|
self.set_allowed_params(provider_conf, request_body)
|
|
|
|
local use_ReAct_prompt = provider_conf.use_ReAct_prompt == true
|
|
|
|
local tools = nil
|
|
if not disable_tools and prompt_opts.tools and not use_ReAct_prompt then
|
|
tools = {}
|
|
for _, tool in ipairs(prompt_opts.tools) do
|
|
table.insert(tools, self:transform_tool(tool))
|
|
end
|
|
end
|
|
|
|
Utils.debug("endpoint", provider_conf.endpoint)
|
|
Utils.debug("model", provider_conf.model)
|
|
|
|
return {
|
|
url = Utils.url_join(provider_conf.endpoint, "/chat/completions"),
|
|
proxy = provider_conf.proxy,
|
|
insecure = provider_conf.allow_insecure,
|
|
headers = Utils.tbl_override(headers, self.extra_headers),
|
|
body = vim.tbl_deep_extend("force", {
|
|
model = provider_conf.model,
|
|
messages = self:parse_messages(prompt_opts),
|
|
stream = true,
|
|
tools = tools,
|
|
}, request_body),
|
|
}
|
|
end
|
|
|
|
return M
|