122 lines
3.1 KiB
Lua
122 lines
3.1 KiB
Lua
|
local Utils = require("avante.utils")
|
||
|
|
local P = require("avante.providers")
|
||
|
|
|
||
|
|
---@alias CohereFinishReason "COMPLETE" | "LENGTH" | "ERROR"
|
||
|
|
---
|
||
|
|
---@class CohereChatStreamResponse
|
||
|
|
---@field event_type "stream-start" | "text-generation" | "stream-end"
|
||
|
|
---@field is_finished boolean
|
||
|
|
---
|
||
|
|
---@class CohereTextGenerationResponse: CohereChatStreamResponse
|
||
|
|
---@field text string
|
||
|
|
---
|
||
|
|
---@class CohereStreamEndResponse: CohereChatStreamResponse
|
||
|
|
---@field response CohereChatResponse
|
||
|
|
---@field finish_reason CohereFinishReason
|
||
|
|
---
|
||
|
|
---@class CohereChatResponse
|
||
|
|
---@field text string
|
||
|
|
---@field generation_id string
|
||
|
|
---@field chat_history CohereMessage[]
|
||
|
|
---@field finish_reason CohereFinishReason
|
||
|
|
---@field meta {api_version: {version: integer}, billed_units: {input_tokens: integer, output_tokens: integer}, tokens: {input_tokens: integer, output_tokens: integer}}
|
||
|
|
---
|
||
|
|
---@class CohereMessage
|
||
|
|
---@field role? "USER" | "SYSTEM" | "CHATBOT"
|
||
|
|
---@field message string
|
||
|
|
---
|
||
|
|
---@class AvanteProviderFunctor
|
||
|
|
local M = {}
|
||
|
|
|
||
|
|
M.api_key_name = "CO_API_KEY"
|
||
|
|
|
||
|
|
M.has = function()
|
||
|
|
return os.getenv(M.api_key_name) and true or false
|
||
|
|
end
|
||
|
|
|
||
|
|
M.parse_message = function(opts)
|
||
|
|
local user_prompt = opts.base_prompt
|
||
|
|
.. "\n\nCODE:\n"
|
||
|
|
.. "```"
|
||
|
|
.. opts.code_lang
|
||
|
|
.. "\n"
|
||
|
|
.. opts.code_content
|
||
|
|
.. "\n```"
|
||
|
|
.. "\n\nQUESTION:\n"
|
||
|
|
.. opts.question
|
||
|
|
|
||
|
|
if opts.selected_code_content ~= nil then
|
||
|
|
user_prompt = opts.base_prompt
|
||
|
|
.. "\n\nCODE CONTEXT:\n"
|
||
|
|
.. "```"
|
||
|
|
.. opts.code_lang
|
||
|
|
.. "\n"
|
||
|
|
.. opts.code_content
|
||
|
|
.. "\n```"
|
||
|
|
.. "\n\nCODE:\n"
|
||
|
|
.. "```"
|
||
|
|
.. opts.code_lang
|
||
|
|
.. "\n"
|
||
|
|
.. opts.selected_code_content
|
||
|
|
.. "\n```"
|
||
|
|
.. "\n\nQUESTION:\n"
|
||
|
|
.. opts.question
|
||
|
|
end
|
||
|
|
|
||
|
|
return {
|
||
|
|
preamble = opts.system_prompt,
|
||
|
|
message = user_prompt,
|
||
|
|
}
|
||
|
|
end
|
||
|
|
|
||
|
|
M.parse_stream_data = function(data, opts)
|
||
|
|
---@type CohereChatStreamResponse
|
||
|
|
local json = vim.json.decode(data)
|
||
|
|
if json.is_finished then
|
||
|
|
opts.on_complete(nil)
|
||
|
|
return
|
||
|
|
end
|
||
|
|
if json.event_type ~= nil then
|
||
|
|
---@cast json CohereStreamEndResponse
|
||
|
|
if json.event_type == "stream-end" and json.finish_reason == "COMPLETE" then
|
||
|
|
opts.on_complete(nil)
|
||
|
|
return
|
||
|
|
end
|
||
|
|
---@cast json CohereTextGenerationResponse
|
||
|
|
if json.event_type == "text-generation" then
|
||
|
|
opts.on_chunk(json.text)
|
||
|
|
end
|
||
|
|
end
|
||
|
|
end
|
||
|
|
|
||
|
|
M.parse_curl_args = function(provider, code_opts)
|
||
|
|
local base, body_opts = P.parse_config(provider)
|
||
|
|
|
||
|
|
local headers = {
|
||
|
|
["Accept"] = "application/json",
|
||
|
|
["Content-Type"] = "application/json",
|
||
|
|
["X-Client-Name"] = "avante.nvim/Neovim/"
|
||
|
|
.. vim.version().major
|
||
|
|
.. "."
|
||
|
|
.. vim.version().minor
|
||
|
|
.. "."
|
||
|
|
.. vim.version().patch,
|
||
|
|
}
|
||
|
|
if not P.env.is_local("openai") then
|
||
|
|
headers["Authorization"] = "Bearer " .. os.getenv(base.api_key_name or M.api_key_name)
|
||
|
|
end
|
||
|
|
|
||
|
|
return {
|
||
|
|
url = Utils.trim(base.endpoint, { suffix = "/" }) .. "/v1/chat",
|
||
|
|
proxy = base.proxy,
|
||
|
|
insecure = base.allow_insecure,
|
||
|
|
headers = headers,
|
||
|
|
body = vim.tbl_deep_extend("force", {
|
||
|
|
model = base.model,
|
||
|
|
stream = true,
|
||
|
|
}, M.parse_message(code_opts), body_opts),
|
||
|
|
}
|
||
|
|
end
|
||
|
|
|
||
|
|
return M
|