Feat: model selection based on API result

- add API calls to ask for available models
- fix new state: now stores chat and command agent separately

Author: frankroeder

lua/parrot/chat_handler.lua

@@ -656,6 +656,7 @@ function ChatHandler:_chat_respond(params)
   end
 
   local query_prov = self:get_provider(true)
+  query_prov:set_model(model_obj.name)
 
   local llm_prefix = self.options.llm_prefix
   local llm_suffix = "[{{llm}}]"
@@ -761,8 +762,7 @@ function ChatHandler:_chat_respond(params)
         -- prepare invisible buffer for the model to write to
         local topic_buf = vim.api.nvim_create_buf(false, true)
         local topic_handler = chatutils.create_handler(self.queries, topic_buf, nil, 0, false, "", false)
-
-        topic_prov:check({ model = self.providers[topic_prov.name].topic_model })
+        topic_prov:set_model(self.providers[topic_prov.name].topic.model)
 
         local topic_spinner = nil
         if self.options.enable_spinner then
@@ -927,7 +927,7 @@ function ChatHandler:switch_provider(selected_prov, is_chat)
   end
 
   if self.providers[selected_prov] then
-    self.set_provider(selected_prov, is_chat)
+    self:set_provider(selected_prov, is_chat)
     logger.info("Switched to provider: " .. selected_prov)
     return
   else
@@ -949,7 +949,7 @@ function ChatHandler:provider(params)
       prompt = "Provider selection ❯",
       fzf_opts = self.options.fzf_lua_opts,
       complete = function(selection)
-        self:switch_provider(selection[1])
+        self:switch_provider(selection[1], is_chat)
       end,
     })
   else
@@ -966,7 +966,6 @@ function ChatHandler:switch_model(is_chat, selected_model, prov)
     logger.warning("Empty model selection")
     return
   end
-  prov:check(selected_model)
   if is_chat then
     self.state:set_model(prov.name, selected_model, "chat")
     logger.info("Chat model: " .. selected_model)
@@ -1281,6 +1280,7 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template)
     end
 
     -- call the model and write the response
+		prov:set_model(model_obj.name)
 
     local spinner = nil
     if self.options.enable_spinner then

lua/parrot/config.lua

@@ -326,9 +326,9 @@ function M.setup(opts)
 
   M.available_providers = vim.tbl_keys(M.providers)
 
-  available_models = {}
+  local available_models = {}
   for _, prov_name in ipairs(M.available_providers) do
-    local _prov = init_provider(prov_name, "", "")
+    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()
   end
   M.available_models = available_models

lua/parrot/provider/anthropic.lua

@@ -4,13 +4,6 @@ local utils = require("parrot.utils")
 local Anthropic = {}
 Anthropic.__index = Anthropic
 
-local available_model_set = {
-  ["claude-3-5-sonnet-20240620"] = true,
-  ["claude-3-opus-20240229"] = true,
-  ["claude-3-sonnet-20240229"] = true,
-  ["claude-3-haiku-20240307"] = true,
-}
-
 -- https://docs.anthropic.com/en/api/messages
 local available_api_parameters = {
   -- required
@@ -100,10 +93,6 @@ function Anthropic:process_onexit(res)
   end
 end
 
-function Anthropic:check(model)
-  return available_model_set[model]
-end
-
 function Anthropic:get_available_models()
   return {
     "claude-3-5-sonnet-20240620",

lua/parrot/provider/gemini.lua

@@ -4,11 +4,6 @@ local utils = require("parrot.utils")
 local Gemini = {}
 Gemini.__index = Gemini
 
-local available_model_set = {
-  ["gemini-1.5-flash"] = true,
-  ["gemini-1.5-pro"] = true,
-}
-
 -- https://ai.google.dev/gemini-api/docs/models/generative-models#model_parameters
 local available_api_parameters = {
   ["contents"] = true,
@@ -120,14 +115,11 @@ function Gemini:process_onexit(res)
   end
 end
 
-function Gemini:check(model)
-  return available_model_set[model]
-end
-
 function Gemini:get_available_models()
   return {
     "gemini-1.5-flash",
     "gemini-1.5-pro",
+    "gemini-1.0-pro",
   }
 end
 

lua/parrot/provider/groq.lua

@@ -1,23 +1,10 @@
 local logger = require("parrot.logger")
 local utils = require("parrot.utils")
+local Job = require("plenary.job")
 
 local Groq = {}
 Groq.__index = Groq
 
-local available_model_set = {
-  ["llama-3.1-405b-reasoning"] = true,
-  ["llama-3.1-70b-versatile"] = true,
-  ["llama-3.1-8b-instant"] = true,
-  ["llama3-groq-70b-8192-tool-use-preview"] = true,
-  ["llama3-groq-8b-8192-tool-use-preview"] = true,
-  ["llama-guard-3-8b"] = true,
-  ["llama3-70b-8192"] = true,
-  ["llama3-8b-8192"] = true,
-  ["mixtral-8x7b-32768"] = true,
-  ["gemma-7b-it"] = true,
-  ["gemma2-9b-it"] = true,
-}
-
 -- https://console.groq.com/docs/api-reference#chat-create
 local available_api_parameters = {
   -- required
@@ -100,18 +87,33 @@ function Groq:process_stdout(response)
 end
 
 function Groq:process_onexit(res)
-  -- '{"error":{"message":"Invalid API Key","type":"invalid_request_error","code":"invalid_api_key"}}'
   local success, parsed = pcall(vim.json.decode, res)
   if success and parsed.error then
     logger.error("Groq - message: " .. parsed.error.message)
   end
 end
 
-function Groq:check(model)
-  return available_model_set[model]
-end
-
 function Groq:get_available_models()
+  if self:verify() then
+    Job:new({
+      command = "curl",
+      args = {
+        "https://api.groq.com/openai/v1/models",
+        "-H",
+        "Authorization: Bearer " .. self.api_key,
+        "-H",
+        "Content-Type: application/json",
+      },
+      on_exit = function(job)
+        local parsed_response = utils.parse_raw_response(job:result())
+        local ids = {}
+        for _, item in ipairs(vim.json.decode(parsed_response).data) do
+          table.insert(ids, item.id)
+        end
+        return ids
+      end,
+    }):start()
+  end
   return {
     "llama-3.1-405b-reasoning",
     "llama-3.1-70b-versatile",

lua/parrot/provider/mistral.lua

@@ -4,17 +4,6 @@ local utils = require("parrot.utils")
 local Mistral = {}
 Mistral.__index = Mistral
 
-local available_model_set = {
-  ["codestral-latest"] = true,
-  ["mistral-tiny"] = true,
-  ["mistral-small-latest"] = true,
-  ["mistral-medium-latest"] = true,
-  ["mistral-large-latest"] = true,
-  ["open-mistral-7b"] = true,
-  ["open-mixtral-8x7b"] = true,
-  ["open-mixtral-8x22b"] = true,
-}
-
 -- https://docs.mistral.ai/api/#operation/createChatCompletion
 local available_api_parameters = {
   -- required
@@ -91,10 +80,6 @@ function Mistral:process_onexit(res)
   end
 end
 
-function Mistral:check(model)
-  return available_model_set[model]
-end
-
 function Mistral:get_available_models()
   return {
     "codestral-latest",

lua/parrot/provider/ollama.lua

@@ -77,57 +77,28 @@ function Ollama:process_onexit(res)
   end
 end
 
-function Ollama:check(model)
-  if not self.ollama_installed then
-    logger.warning("ollama not found.")
-    return false
-  end
-
-  local handle = io.popen("ollama list")
-  local result = handle:read("*a")
-  handle:close()
+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()
 
-  local found_match = false
-  for line in result:gmatch("[^\r\n]+") do
-    if string.match(line, model) then
-      found_match = true
-      break
-    end
+  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
-
-  if not found_match then
-    if not pcall(require, "plenary") then
-      logger.error("Plenary not installed. Please install nvim-lua/plenary.nvim to use this plugin.")
-      return false
-    end
-    local confirm = vim.fn.confirm("ollama model " .. model .. " not found. Download now?", "&Yes\n&No", 1)
-    if confirm == 1 then
-      local job = Job:new({
-        command = "ollama",
-        args = { "pull", model },
-        on_exit = function(_, return_val)
-          logger.info("Download finished with exit code: " .. return_val)
-        end,
-        on_stderr = function(j, data)
-          print("Downloading, please wait: " .. data)
-          if j ~= nil then
-            logger.error(vim.inspect(j:result()))
-          end
-        end,
-      })
-      job:start()
-      return true
+  local names = {}
+  if parsed_data.models then
+    for _, model in ipairs(parsed_data.models) do
+      table.insert(names, model.name)
     end
   else
-    return true
+    return { "No model found, please download" }
   end
-  return false
-end
-
-function Ollama:get_available_models()
-  return {
-    "llama3",
-  }
+  return names
 end
 
 return Ollama

lua/parrot/provider/openai.lua

@@ -1,29 +1,10 @@
 local logger = require("parrot.logger")
 local utils = require("parrot.utils")
+local Job = require("plenary.job")
 
 local OpenAI = {}
 OpenAI.__index = OpenAI
 
-local available_model_set = {
-  ["gpt-3.5-turbo"] = true,
-  ["gpt-3.5-turbo-0125"] = true,
-  ["gpt-3.5-turbo-1106"] = true,
-  ["gpt-3.5-turbo-16k"] = true,
-  ["gpt-3.5-turbo-instruct"] = true,
-  ["gpt-3.5-turbo-instruct-0914"] = true,
-  ["gpt-4"] = true,
-  ["gpt-4-0125-preview"] = true,
-  ["gpt-4-0613"] = true,
-  ["gpt-4-1106-preview"] = true,
-  ["gpt-4-turbo"] = true,
-  ["gpt-4-turbo-2024-04-09"] = true,
-  ["gpt-4-turbo-preview"] = true,
-  ["gpt-4o"] = true,
-  ["gpt-4o-2024-05-13"] = true,
-  ["gpt-4o-mini"] = true,
-  ["gpt-4o-mini-2024-07-18"] = true,
-}
-
 -- https://platform.openai.com/docs/api-reference/chat/create
 local available_api_parameters = {
   -- required
@@ -110,32 +91,27 @@ function OpenAI:process_onexit(res)
   end
 end
 
-function OpenAI:check(model)
-  return available_model_set[model]
-end
-
 function OpenAI:get_available_models()
   -- curl https://api.openai.com/v1/models -H "Authorization: Bearer $OPENAI_API_KEY"
-  -- local Job = require("plenary.job")
-  -- self:verify()
-  -- Job:new({
-  --   command = "curl",
-  --   args = {
-  --     "https://api.openai.com/v1/models",
-  --     "-H", "Authorization: Bearer " .. self.api_key,
-  --   },
-  --   on_exit = function(job)
-  --      local parsed_response = utils.parse_raw_response(job:result())
-  -- 		print("JSON DEOCDE", vim.inspect(vim.json.decode(parsed_response)))
-  --      local ids = {}
-  --      for _, item in ipairs(vim.json.decode(parsed_response).data) do
-  --        table.insert(ids, item.id)
-  --      end
-  --      print("IDS", vim.inspect(ids))
-  --      return ids
-  --   end,
-  -- }):start()
 
+  if self:verify() then
+    Job:new({
+      command = "curl",
+      args = {
+        "https://api.openai.com/v1/models",
+        "-H",
+        "Authorization: Bearer " .. self.api_key,
+      },
+      on_exit = function(job)
+        local parsed_response = utils.parse_raw_response(job:result())
+        local ids = {}
+        for _, item in ipairs(vim.json.decode(parsed_response).data) do
+          table.insert(ids, item.id)
+        end
+        return ids
+      end,
+    }):start()
+  end
   return {
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-0125",

lua/parrot/provider/perplexity.lua

@@ -4,17 +4,6 @@ local utils = require("parrot.utils")
 local Perplexity = {}
 Perplexity.__index = Perplexity
 
-local available_model_set = {
-  -- adjust the new data
-  ["llama-3.1-sonar-small-128k-chat"] = true,
-  ["llama-3.1-sonar-small-128k-online"] = true,
-  ["llama-3.1-sonar-large-128k-chat"] = true,
-  ["llama-3.1-sonar-large-128k-online"] = true,
-  ["llama-3.1-8b-instruct"] = true,
-  ["llama-3.1-70b-instruct"] = true,
-  ["mixtral-8x7b-instruct"] = true,
-}
-
 -- https://docs.perplexity.ai/reference/post_chat_completions
 local available_api_parameters = {
   -- required
@@ -95,10 +84,6 @@ function Perplexity:process_onexit(res)
   end
 end
 
-function Perplexity:check(model)
-  return available_model_set[model]
-end
-
 function Perplexity:get_available_models()
   return {
     "llama-3.1-sonar-small-128k-chat",