109 lines
2.8 KiB
Lua
109 lines
2.8 KiB
Lua
local Utils = require("avante.utils")
|
|
local P = require("avante.providers")
|
|
local Clipboard = require("avante.clipboard")
|
|
|
|
---@class AvanteProviderFunctor
|
|
local M = {}
|
|
|
|
M.api_key_name = "GEMINI_API_KEY"
|
|
M.role_map = {
|
|
user = "user",
|
|
assistant = "model",
|
|
}
|
|
-- M.tokenizer_id = "google/gemma-2b"
|
|
|
|
M.parse_messages = function(opts)
|
|
local contents = {}
|
|
local prev_role = nil
|
|
|
|
vim.iter(opts.messages):each(function(message)
|
|
local role = message.role
|
|
if role == prev_role then
|
|
if role == M.role_map["user"] then
|
|
table.insert(
|
|
contents,
|
|
{ role = M.role_map["assistant"], parts = {
|
|
{ text = "Ok, I understand." },
|
|
} }
|
|
)
|
|
else
|
|
table.insert(contents, { role = M.role_map["user"], parts = {
|
|
{ text = "Ok" },
|
|
} })
|
|
end
|
|
end
|
|
prev_role = role
|
|
table.insert(contents, { role = M.role_map[role] or role, parts = {
|
|
{ text = message.content },
|
|
} })
|
|
end)
|
|
|
|
if Clipboard.support_paste_image() and opts.image_paths then
|
|
for _, image_path in ipairs(opts.image_paths) do
|
|
local image_data = {
|
|
inline_data = {
|
|
mime_type = "image/png",
|
|
data = Clipboard.get_base64_content(image_path),
|
|
},
|
|
}
|
|
|
|
table.insert(contents[#contents].parts, image_data)
|
|
end
|
|
end
|
|
|
|
return {
|
|
systemInstruction = {
|
|
role = "user",
|
|
parts = {
|
|
{
|
|
text = opts.system_prompt,
|
|
},
|
|
},
|
|
},
|
|
contents = contents,
|
|
}
|
|
end
|
|
|
|
M.parse_response = function(ctx, data_stream, _, opts)
|
|
local ok, json = pcall(vim.json.decode, data_stream)
|
|
if not ok then opts.on_stop({ reason = "error", error = json }) end
|
|
if json.candidates then
|
|
if #json.candidates > 0 then
|
|
if json.candidates[1].finishReason and json.candidates[1].finishReason == "STOP" then
|
|
opts.on_chunk(json.candidates[1].content.parts[1].text)
|
|
opts.on_stop({ reason = "complete" })
|
|
else
|
|
opts.on_chunk(json.candidates[1].content.parts[1].text)
|
|
end
|
|
else
|
|
opts.on_stop({ reason = "complete" })
|
|
end
|
|
end
|
|
end
|
|
|
|
M.parse_curl_args = function(provider, code_opts)
|
|
local base, body_opts = P.parse_config(provider)
|
|
|
|
body_opts = vim.tbl_deep_extend("force", body_opts, {
|
|
generationConfig = {
|
|
temperature = body_opts.temperature,
|
|
maxOutputTokens = body_opts.max_tokens,
|
|
},
|
|
})
|
|
body_opts.temperature = nil
|
|
body_opts.max_tokens = nil
|
|
|
|
return {
|
|
url = Utils.url_join(
|
|
base.endpoint,
|
|
base.model .. ":streamGenerateContent?alt=sse&key=" .. provider.parse_api_key()
|
|
),
|
|
proxy = base.proxy,
|
|
insecure = base.allow_insecure,
|
|
headers = { ["Content-Type"] = "application/json" },
|
|
body = vim.tbl_deep_extend("force", {}, M.parse_messages(code_opts), body_opts),
|
|
}
|
|
end
|
|
|
|
return M
|