From 334bdac260165273da7c87625b0e17854bd4ebeb Mon Sep 17 00:00:00 2001 From: Frank Roeder Date: Thu, 8 Aug 2024 08:44:30 +0200 Subject: [PATCH] fix: do not fetch models from API on startup - only do it in later invoked provider selections - fix health check - add check if ollama is installed --- lua/parrot/chat_handler.lua | 5 ++-- lua/parrot/config.lua | 3 +- lua/parrot/health.lua | 47 +------------------------------ lua/parrot/provider/anthropic.lua | 2 +- lua/parrot/provider/groq.lua | 4 +-- lua/parrot/provider/ollama.lua | 39 +++++++++++++------------ lua/parrot/provider/openai.lua | 5 ++-- 7 files changed, 32 insertions(+), 73 deletions(-) diff --git a/lua/parrot/chat_handler.lua b/lua/parrot/chat_handler.lua index 1ce5868..879e84c 100644 --- a/lua/parrot/chat_handler.lua +++ b/lua/parrot/chat_handler.lua @@ -983,11 +983,12 @@ function ChatHandler:model(params) local prov = self:get_provider(is_chat) local model_name = string.gsub(params.args, "^%s*(.-)%s*$", "%1") local has_fzf, fzf_lua = pcall(require, "fzf-lua") + local fetch_online = true if model_name ~= "" then self:switch_model(is_chat, model_name, prov) elseif has_fzf then - fzf_lua.fzf_exec(prov:get_available_models(), { + fzf_lua.fzf_exec(prov:get_available_models(fetch_online), { prompt = "Model selection ❯", fzf_opts = self.options.fzf_lua_opts, complete = function(selection) @@ -1000,7 +1001,7 @@ function ChatHandler:model(params) end, }) else - vim.ui.select(prov:get_available_models(), { + vim.ui.select(prov:get_available_models(fetch_online), { prompt = "Select your model:", }, function(selected_model) self:switch_model(is_chat, selected_model, prov) diff --git a/lua/parrot/config.lua b/lua/parrot/config.lua index 6cf2958..7114b43 100644 --- a/lua/parrot/config.lua +++ b/lua/parrot/config.lua @@ -329,7 +329,8 @@ function M.setup(opts) local available_models = {} for _, prov_name in ipairs(M.available_providers) do local _prov = init_provider(prov_name, M.providers[prov_name].endpoint, M.providers[prov_name].api_key) - available_models[prov_name] = _prov:get_available_models() + -- do not make an API call on startup + available_models[prov_name] = _prov:get_available_models(false) end M.available_models = available_models diff --git a/lua/parrot/health.lua b/lua/parrot/health.lua index 9eab5de..8405ac7 100644 --- a/lua/parrot/health.lua +++ b/lua/parrot/health.lua @@ -1,47 +1,5 @@ local M = {} -local check_provider = function(parrot, prov_key) - local providers = parrot.providers - - if prov_key ~= "ollama" then - local api_key = providers[prov_key].api_key - if api_key then - vim.health.ok(prov_key .. " api_key is set") - else - vim.health.error( - "require('parrot').setup({provider {.." - .. prov_key - .. "..: {api_key: ???}}) is not set: " - .. vim.inspect(api_key) - ) - end - end - - local endpoint = providers[prov_key].endpoint - if endpoint and string.match(endpoint, "%S") then - vim.health.ok(prov_key .. " endpoint is set") - else - vim.health.error( - "require('parrot').setup({provider {.." - .. prov_key - .. "..: {endpoint: ???}}) is not set: " - .. vim.inspect(endpoint) - ) - end - - local topic_prompt = providers[prov_key].topic_prompt - if topic_prompt and string.match(topic_prompt, "%S") then - vim.health.ok(prov_key .. " topic_prompt is set") - else - vim.health.error( - "require('parrot').setup({provider {.." - .. prov_key - .. "..: {topic_prompt: ???}}) is not set: " - .. vim.inspect(topic_prompt) - ) - end -end - function M.check() vim.health.start("parrot.nvim checks") @@ -63,14 +21,11 @@ function M.check() else vim.health.ok("require('parrot') succeeded") - if parrot._setup_called then + if parrot.did_setup then vim.health.ok("require('parrot').setup() has been called") else vim.health.error("require('parrot').setup() has not been called") end - for _, name in ipairs(parrot._available_providers) do - check_provider(parrot, name) - end end for _, name in ipairs({ "curl", "grep", "rg", "ln" }) do diff --git a/lua/parrot/provider/anthropic.lua b/lua/parrot/provider/anthropic.lua index 0257cd4..181866a 100644 --- a/lua/parrot/provider/anthropic.lua +++ b/lua/parrot/provider/anthropic.lua @@ -39,7 +39,7 @@ function Anthropic:preprocess_payload(payload) if payload.messages[1].role == "system" then local system_prompt = payload.messages[1].content -- remove the first message that serves as the system prompt as anthropic - -- expects the system prompt to be part of the curl request and not the messages + -- expects the system prompt to be part of the curl request body and not the messages table.remove(payload.messages, 1) payload.system = system_prompt end diff --git a/lua/parrot/provider/groq.lua b/lua/parrot/provider/groq.lua index b0c7252..1d42d9f 100644 --- a/lua/parrot/provider/groq.lua +++ b/lua/parrot/provider/groq.lua @@ -93,8 +93,8 @@ function Groq:process_onexit(res) end end -function Groq:get_available_models() - if self:verify() then +function Groq:get_available_models(online) + if online and self:verify() then Job:new({ command = "curl", args = { diff --git a/lua/parrot/provider/ollama.lua b/lua/parrot/provider/ollama.lua index 8456f98..ad6c80a 100644 --- a/lua/parrot/provider/ollama.lua +++ b/lua/parrot/provider/ollama.lua @@ -79,27 +79,30 @@ end function Ollama:get_available_models() -- curl https://api.openai.com/v1/models -H "Authorization: Bearer $OPENAI_API_KEY" - local job = Job:new({ - command = "curl", - args = { "-H", "Content-Type: application/json", "http://localhost:11434/api/tags" }, - }):sync() + if vim.fn.executable("ollama") then + local job = Job:new({ + command = "curl", + args = { "-H", "Content-Type: application/json", "http://localhost:11434/api/tags" }, + }):sync() - local parsed_response = utils.parse_raw_response(job) - local success, parsed_data = pcall(vim.json.decode, parsed_response) - if not success then - logger.error("Error parsing JSON:" .. vim.inspect(parsed_data)) - return {} - end - local names = {} - if parsed_data.models then - for _, model in ipairs(parsed_data.models) do - table.insert(names, model.name) + local parsed_response = utils.parse_raw_response(job) + local success, parsed_data = pcall(vim.json.decode, parsed_response) + if not success then + logger.error("Error parsing JSON:" .. vim.inspect(parsed_data)) + return {} + end + local names = {} + if parsed_data.models then + for _, model in ipairs(parsed_data.models) do + table.insert(names, model.name) + end + else + logger.error("No models found. Please use 'ollama pull' to download one.") + return {} end - else - logger.error("No models found. Please use 'ollama pull' to download one.") - return {} + return names end - return names + return {} end return Ollama diff --git a/lua/parrot/provider/openai.lua b/lua/parrot/provider/openai.lua index 8efe2d4..9e1c5ba 100644 --- a/lua/parrot/provider/openai.lua +++ b/lua/parrot/provider/openai.lua @@ -91,10 +91,9 @@ function OpenAI:process_onexit(res) end end -function OpenAI:get_available_models() +function OpenAI:get_available_models(online) -- curl https://api.openai.com/v1/models -H "Authorization: Bearer $OPENAI_API_KEY" - - if self:verify() then + if online and self:verify() then Job:new({ command = "curl", args = {