feat(llm): cohere support (#167)
should be good set of defaults now, one in US, one in canada, and microsoft :/ Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
This commit is contained in:
121
lua/avante/providers/cohere.lua
Normal file
121
lua/avante/providers/cohere.lua
Normal file
@@ -0,0 +1,121 @@
|
||||
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
|
||||
Reference in New Issue
Block a user