fix: model election algorithm

override -> command options -> defaults

Signed-off-by: blob42 <contact@blob42.xyz>

Author: blob42

lua/codegpt/commands.lua

@@ -2,20 +2,20 @@ local Utils = require("codegpt.utils")
 local Ui = require("codegpt.ui")
 local Providers = require("codegpt.providers")
 local Api = require("codegpt.api")
-local config = require("codegpt.config")
+local Config = require("codegpt.config")
 local models = require("codegpt.models")
 
 local M = {}
 
 local text_popup_stream = function(stream, bufnr, start_row, start_col, end_row, end_col)
-	local popup_filetype = config.opts.ui.text_popup_filetype
+	local popup_filetype = Config.opts.ui.text_popup_filetype
 	Ui.popup_stream(stream, popup_filetype, bufnr, start_row, start_col, end_row, end_col)
 end
 
 M.CallbackTypes = {
 	["text_popup_stream"] = text_popup_stream,
 	["text_popup"] = function(lines, bufnr, start_row, start_col, end_row, end_col)
-		local popup_filetype = config.opts.ui.text_popup_filetype
+		local popup_filetype = Config.opts.ui.text_popup_filetype
 		Ui.popup(lines, popup_filetype, bufnr, start_row, start_col, end_row, end_col)
 	end,
 	["code_popup"] = function(lines, bufnr, start_row, start_col, end_row, end_col)
@@ -45,15 +45,21 @@ M.CallbackTypes = {
 ---@return table opts parsed options
 ---@return boolean is_stream streaming enabled
 local function get_cmd_opts(cmd)
-	local opts = config.opts.commands[cmd]
-	local cmd_defaults = config.opts.global_defaults
+	local opts = Config.opts.commands[cmd]
+	-- print(vim.inspect(opts))
+	local cmd_defaults = Config.opts.global_defaults
 	local is_stream = false
 
 	-- print(vim.inspect(cmd))
 	-- print(vim.inspect(config.opts.commands))
 	-- print(vim.inspect(opts))
 
-	local _, model = models.get_model()
+	local model
+	if opts.model then
+		_, model = models.get_model_by_name(opts.model)
+	else
+		_, model = models.get_model()
+	end
 
 	---@type codegpt.CommandOpts
 	--- options priority heighest->lowest: cmd options, model options, global
@@ -63,7 +69,7 @@ local function get_cmd_opts(cmd)
 		opts.callback = opts.callback_type
 	else
 		if
-			(config.opts.ui.stream_output and opts.callback_type == "text_popup")
+			(Config.opts.ui.stream_output and opts.callback_type == "text_popup")
 			or opts.callback_type == "test_popup_stream"
 		then
 			opts.callback = text_popup_stream
@@ -72,6 +78,8 @@ local function get_cmd_opts(cmd)
 			opts.callback = M.CallbackTypes[opts.callback_type]
 		end
 	end
+	-- print(vim.inspect(opts))
+	-- error(1)
 
 	return opts, is_stream
 end

lua/codegpt/config.lua

@@ -42,6 +42,7 @@ local M = {}
 ---@field temperature? number Custom temperature for this command
 ---@field max_tokens? number Custom max_tokens for this command
 ---@field append_string? string String to append to prompt -- ex: /no_think
+---@field model? string Model to always use with this command
 
 ---@type { [string]: codegpt.CommandOpts }
 local default_commands = {

lua/codegpt/models.lua

@@ -1,7 +1,38 @@
-local config = require("codegpt.config")
-local providers = require("codegpt.providers")
+local Config = require("codegpt.config")
+local Providers = require("codegpt.providers")
 local M = {}
 
+function M.get_model_by_name(name)
+	---@type codegpt.Model
+
+	local provider_name = vim.fn.tolower(Config.opts.connection.api_provider)
+	local provider_config = Config.opts.models[provider_name]
+
+	if type(provider_config) == "string" and #provider_config > 0 then
+		return provider_config, nil
+	end
+
+	assert(type(provider_config) == "table")
+
+	if provider_config == nil then
+		error("no models defined for " .. provider_name)
+	end
+
+	local selected = name
+	local result = provider_config[selected]
+
+	if result == nil then
+		for model_name, model in pairs(provider_config) do
+			if model.alias == selected then
+				result = model
+				selected = model_name
+				break
+			end
+		end
+	end
+
+	return selected, result
+end
 --- default model selection order from highest to lowest priority
 --- 1. global model_override (manual selection, always temporary for an nvim session)
 --- 2. provider default model
@@ -12,10 +43,10 @@ function M.get_model()
 	---@type codegpt.Model
 	local result
 
-	local provider_name = vim.fn.tolower(config.opts.connection.api_provider)
+	local provider_name = vim.fn.tolower(Config.opts.connection.api_provider)
 	-- local selected = config.model_override or config.opts.models[provider].default or config.opts.models.default
 	local selected
-	local provider_config = config.opts.models[provider_name]
+	local provider_config = Config.opts.models[provider_name]
 
 	-- of provider config is a string, then it must be just a model name
 	if type(provider_config) == "string" and #provider_config > 0 then
@@ -27,7 +58,7 @@ function M.get_model()
 		error("no models defined for " .. provider_name)
 	end
 
-	selected = config.model_override or provider_config.default or config.opts.models.default
+	selected = Config.model_override or provider_config.default or Config.opts.models.default
 	assert(type(selected) == "string")
 
 	result = provider_config[selected]
@@ -52,8 +83,8 @@ end
 function M.get_model_by_alias(alias)
 	assert(alias and #alias > 0)
 
-	local provider_name = vim.fn.tolower(config.opts.connection.api_provider)
-	local provider_config = config.opts.models[provider_name]
+	local provider_name = vim.fn.tolower(Config.opts.connection.api_provider)
+	local provider_config = Config.opts.models[provider_name]
 	if provider_config == nil then
 		error("no models defined for " .. provider_name)
 	end
@@ -68,7 +99,7 @@ end
 
 ---@return table[]
 function M.get_remote_models()
-	local models = providers.get_provider().get_models()
+	local models = Providers.get_provider().get_models()
 	if models ~= nil then
 		models = vim.tbl_map(function(remote)
 			remote.model_source = "remote"
@@ -83,7 +114,7 @@ end
 ---@param provider string
 ---@return table[] models list of locally defined models
 function M.get_local_models(provider)
-	local provider_config = config.opts.models[provider]
+	local provider_config = Config.opts.models[provider]
 
 	-- models defined by name only are skipped since they must be a remote one
 	if type(provider_config) == "string" then
@@ -107,12 +138,12 @@ function M.get_local_models(provider)
 end
 
 --- List available models
-function M.list_models()
+function M.select_model()
 	local remote_models = M.get_remote_models()
 	local models = vim.tbl_extend("force", {}, remote_models)
 
 	-- get local defined models
-	local used_provider = vim.fn.tolower(config.opts.connection.api_provider)
+	local used_provider = vim.fn.tolower(Config.opts.connection.api_provider)
 	local local_models = M.get_local_models(used_provider)
 	models = vim.tbl_extend("force", models, local_models)
 
@@ -137,28 +168,26 @@ function M.list_models()
 	}, function(choice)
 		if choice ~= nil then
 			if choice.name ~= nil and #choice.name > 0 then
-				print("selected <" .. choice.name .. "> (" .. choice.model_source .. " defined)")
+				Config.model_override = choice.name
+				print("model override = <" .. choice.name .. "> (" .. choice.model_source .. " defined)")
 			end
 		end
 	end)
 end
 
-function M.select_model()
-	local models = providers.get_provider().get_models()
-	if models == nil then
-		error("querying models")
+---@param cmd_opts table
+---@return string name
+---@return table? model
+function M.get_model_for_cmdopts(cmd_opts)
+	local model_name, model
+	if cmd_opts.model ~= nil and Config.model_override == nil then
+		model_name, model = M.get_model_by_name(cmd_opts.model)
+	else
+		model_name, model = M.get_model()
 	end
-	vim.ui.select(models, {
-		prompt = "ollama: available models",
-		format_item = function(item)
-			return item.name
-		end,
-	}, function(choice)
-		if choice ~= nil and #choice.name > 0 then
-			config.model_override = choice.name
-			print("model override = <" .. choice.name .. ">")
-		end
-	end)
+	assert(model_name and #model_name > 0, "undefined model")
+
+	return model_name, model
 end
 
 return M

lua/codegpt/providers/ollama.lua

@@ -46,14 +46,11 @@ end
 ---@param is_stream? boolean
 function M.make_request(command, cmd_opts, command_args, text_selection, is_stream)
 	local models = require("codegpt.models")
-
 	local messages = generate_messages(command, cmd_opts, command_args, text_selection)
 
 	-- max # of tokens to generate
 	local max_tokens = get_max_tokens(cmd_opts.max_tokens, messages)
-
-	local model_name, model = models.get_model()
-	assert(model_name and #model_name > 0, "undefined model")
+	local model_name, model = models.get_model_for_cmdopts(cmd_opts)
 
 	local model_opts = {}
 

lua/codegpt/providers/openai.lua

@@ -58,8 +58,7 @@ function M.make_request(command, cmd_opts, command_args, text_selection, is_stre
 		max_tokens = get_max_tokens(cmd_opts.max_tokens, messages)
 	end
 
-	local model_name, model = models.get_model()
-	assert(model_name and #model_name > 0, "undefined model")
+	local model_name, model = models.get_model_for_cmdopts(cmd_opts)
 
 	local request = {
 		temperature = cmd_opts.temperature,

lua/codegpt/ui.lua

@@ -151,7 +151,6 @@ function M.popup_stream(stream, filetype, bufnr, start_row, start_col, end_row,
 	if stream == nil and #buffer > 0 then
 		table.insert(lines, buffer)
 		buffer = ""
-		print("trailing buffer !")
 		streaming = false
 	elseif stream == nil then
 		streaming = false

tests/command_spec.lua

@@ -87,4 +87,28 @@ describe("cmd opts", function()
 			assert(opts.max_tokens == 4242)
 		end)
 	end)
+
+	it("should prioritize command model ", function()
+		codegpt.setup({
+			commands = {
+				foocmd = {
+					model = "foomodel",
+				},
+			},
+			models = {
+				openai = {
+					default = "barmodel",
+					barmodel = {
+						alias = "llamabar",
+					},
+					foomodel = {
+						alias = "llamafoo",
+					},
+				},
+			},
+		})
+		local cmds = require("codegpt.commands")
+		local opts = cmds.get_cmd_opts("foocmd")
+		assert(opts.model == "foomodel")
+	end)
 end)